diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 732e6b9acf..2df7d7234e 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -819,8 +819,9 @@ def _sync_object_with_future_result( if value: setattr(self, attribute, value) + @classmethod def _construct_sdk_resource_from_gapic( - self, + cls, gapic_resource: proto.Message, project: Optional[str] = None, location: Optional[str] = None, @@ -846,7 +847,7 @@ def _construct_sdk_resource_from_gapic( VertexAiResourceNoun: An initialized SDK object that represents GAPIC type. """ - sdk_resource = self._empty_constructor( + sdk_resource = cls._empty_constructor( project=project, location=location, credentials=credentials ) sdk_resource._gca_resource = gapic_resource @@ -894,14 +895,14 @@ def _list( Returns: List[VertexAiResourceNoun] - A list of SDK resource objects """ - self = cls._empty_constructor( + resource = cls._empty_constructor( project=project, location=location, credentials=credentials ) # Fetch credentials once and re-use for all `_empty_constructor()` calls creds = initializer.global_config.credentials - resource_list_method = getattr(self.api_client, self._list_method) + resource_list_method = getattr(resource.api_client, resource._list_method) list_request = { "parent": initializer.global_config.common_location_path( @@ -916,7 +917,7 @@ def _list( resource_list = resource_list_method(request=list_request) or [] return [ - self._construct_sdk_resource_from_gapic( + cls._construct_sdk_resource_from_gapic( gapic_resource, project=project, location=location, credentials=creds ) for gapic_resource in resource_list diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index b287581431..56cad667cb 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -116,9 +116,9 @@ def __init__( resource_name=endpoint_name, ) self._gca_resource = self._get_gca_resource(resource_name=endpoint_name) + self._prediction_client = self._instantiate_prediction_client( - location=location or initializer.global_config.location, - credentials=credentials, + location=self.location, credentials=credentials, ) @property @@ -324,6 +324,46 @@ def _create( credentials=credentials, ) + @classmethod + def _construct_sdk_resource_from_gapic( + cls, + gapic_resource: proto.Message, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Endpoint": + """Given a GAPIC Endpoint object, return the SDK representation. + + Args: + gapic_resource (proto.Message): + A GAPIC representation of a Endpoint resource, usually + retrieved by a get_* or in a list_* API call. + project (str): + Optional. Project to construct Endpoint object from. If not set, + project set in aiplatform.init will be used. + location (str): + Optional. Location to construct Endpoint object from. If not set, + location set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to construct Endpoint. + Overrides credentials set in aiplatform.init. + + Returns: + Endpoint: + An initialized Endpoint resource. + """ + endpoint = cls._empty_constructor( + project=project, location=location, credentials=credentials + ) + + endpoint._gca_resource = gapic_resource + + endpoint._prediction_client = cls._instantiate_prediction_client( + location=endpoint.location, credentials=credentials, + ) + + return endpoint + @staticmethod def _allocate_traffic( traffic_split: Dict[str, int], traffic_percentage: int, diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index 03c3f38667..e9f7de971a 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/tests/unit/aiplatform/test_endpoints.py @@ -66,7 +66,6 @@ _TEST_LOCATION = "us-central1" _TEST_LOCATION_2 = "europe-west4" -_TEST_ENDPOINT_NAME = "test-endpoint" _TEST_DISPLAY_NAME = "test-display-name" _TEST_DISPLAY_NAME_2 = "test-display-name-2" _TEST_ID = "1028944691210842416" @@ -76,6 +75,9 @@ _TEST_ENDPOINT_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/{_TEST_ID}" ) +_TEST_ENDPOINT_NAME_ALT_LOCATION = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION_2}/endpoints/{_TEST_ID}" +) _TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" _TEST_MODEL_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_ID}" @@ -139,6 +141,9 @@ kms_key_name=_TEST_ENCRYPTION_KEY_NAME ) +_TEST_ENDPOINT_GAPIC = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, name=_TEST_ENDPOINT_NAME +) _TEST_ENDPOINT_LIST = [ gca_endpoint.Endpoint( @@ -170,6 +175,19 @@ def get_endpoint_mock(): yield get_endpoint_mock +@pytest.fixture +def get_endpoint_alt_location_mock(): + with mock.patch.object( + endpoint_service_client.EndpointServiceClient, "get_endpoint" + ) as get_endpoint_mock: + get_endpoint_mock.return_value = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, + name=_TEST_ENDPOINT_NAME_ALT_LOCATION, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + yield get_endpoint_mock + + @pytest.fixture def get_endpoint_with_models_mock(): with mock.patch.object( @@ -293,14 +311,16 @@ def list_endpoints_mock(): @pytest.fixture -def create_client_mock(): +def create_endpoint_client_mock(): with mock.patch.object( initializer.global_config, "create_client", autospec=True, - ) as create_client_mock: - create_client_mock.return_value = mock.Mock( + ) as create_endpoint_client_mock: + endpoint_client_mock = mock.Mock( spec=endpoint_service_client.EndpointServiceClient ) - yield create_client_mock + endpoint_client_mock.get_endpoint.return_value = _TEST_ENDPOINT_GAPIC + create_endpoint_client_mock.return_value = endpoint_client_mock + yield create_endpoint_client_mock @pytest.fixture @@ -340,14 +360,14 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - def test_constructor(self, create_client_mock): + def test_constructor(self, create_endpoint_client_mock): aiplatform.init( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, ) models.Endpoint(_TEST_ENDPOINT_NAME) - create_client_mock.assert_has_calls( + create_endpoint_client_mock.assert_has_calls( [ mock.call( client_class=utils.EndpointClientWithOverride, @@ -365,37 +385,47 @@ def test_constructor(self, create_client_mock): ) def test_constructor_with_endpoint_id(self, get_endpoint_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) models.Endpoint(_TEST_ID) get_endpoint_mock.assert_called_with(name=_TEST_ENDPOINT_NAME) def test_constructor_with_endpoint_name(self, get_endpoint_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) models.Endpoint(_TEST_ENDPOINT_NAME) get_endpoint_mock.assert_called_with(name=_TEST_ENDPOINT_NAME) def test_constructor_with_custom_project(self, get_endpoint_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) models.Endpoint(endpoint_name=_TEST_ID, project=_TEST_PROJECT_2) test_endpoint_resource_name = endpoint_service_client.EndpointServiceClient.endpoint_path( _TEST_PROJECT_2, _TEST_LOCATION, _TEST_ID ) get_endpoint_mock.assert_called_with(name=test_endpoint_resource_name) - def test_constructor_with_custom_location(self, get_endpoint_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + @pytest.mark.usefixtures("get_endpoint_mock") + def test_constructor_with_conflicting_location(self): + """Passing a full resource name with `_TEST_LOCATION` and providing `_TEST_LOCATION_2` as location""" + + with pytest.raises(RuntimeError) as err: + models.Endpoint( + endpoint_name=_TEST_ENDPOINT_NAME, location=_TEST_LOCATION_2 + ) + + assert err.match( + regexp=r"is provided, but different from the resource location" + ) + + def test_constructor_with_custom_location(self, get_endpoint_alt_location_mock): models.Endpoint(endpoint_name=_TEST_ID, location=_TEST_LOCATION_2) test_endpoint_resource_name = endpoint_service_client.EndpointServiceClient.endpoint_path( _TEST_PROJECT, _TEST_LOCATION_2, _TEST_ID ) - get_endpoint_mock.assert_called_with(name=test_endpoint_resource_name) + get_endpoint_alt_location_mock.assert_called_with( + name=test_endpoint_resource_name + ) - def test_constructor_with_custom_credentials(self, create_client_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + def test_constructor_with_custom_credentials(self, create_endpoint_client_mock): creds = auth_credentials.AnonymousCredentials() models.Endpoint(_TEST_ENDPOINT_NAME, credentials=creds) - create_client_mock.assert_has_calls( + create_endpoint_client_mock.assert_has_calls( [ mock.call( client_class=utils.EndpointClientWithOverride, @@ -440,7 +470,6 @@ def test_init_aiplatform_with_encryption_key_name_and_create_endpoint( @pytest.mark.usefixtures("get_endpoint_mock") @pytest.mark.parametrize("sync", [True, False]) def test_create(self, create_endpoint_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) my_endpoint = models.Endpoint.create( display_name=_TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, @@ -463,7 +492,6 @@ def test_create(self, create_endpoint_mock, sync): @pytest.mark.usefixtures("get_endpoint_mock") @pytest.mark.parametrize("sync", [True, False]) def test_create_with_description(self, create_endpoint_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) my_endpoint = models.Endpoint.create( display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, sync=sync ) @@ -480,7 +508,6 @@ def test_create_with_description(self, create_endpoint_mock, sync): @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(test_model, sync=sync) @@ -506,7 +533,6 @@ def test_deploy(self, deploy_model_mock, sync): @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_display_name(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy( @@ -535,7 +561,6 @@ def test_deploy_with_display_name(self, deploy_model_mock, sync): @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_traffic_80(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, traffic_percentage=80, sync=sync) @@ -547,7 +572,6 @@ def test_deploy_raise_error_traffic_80(self, sync): @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_traffic_120(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, traffic_percentage=120, sync=sync) @@ -556,7 +580,6 @@ def test_deploy_raise_error_traffic_120(self, sync): @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_traffic_negative(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, traffic_percentage=-18, sync=sync) @@ -565,7 +588,6 @@ def test_deploy_raise_error_traffic_negative(self, sync): @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_min_replica(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, min_replica_count=-1, sync=sync) @@ -574,7 +596,6 @@ def test_deploy_raise_error_min_replica(self, sync): @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_max_replica(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, max_replica_count=-2, sync=sync) @@ -583,7 +604,6 @@ def test_deploy_raise_error_max_replica(self, sync): @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_traffic_split(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, traffic_split={"a": 99}, sync=sync) @@ -591,7 +611,6 @@ def test_deploy_raise_error_traffic_split(self, sync): @pytest.mark.usefixtures("get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_traffic_percent(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) with mock.patch.object( endpoint_service_client.EndpointServiceClient, "get_endpoint" ) as get_endpoint_mock: @@ -624,7 +643,6 @@ def test_deploy_with_traffic_percent(self, deploy_model_mock, sync): @pytest.mark.usefixtures("get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_traffic_split(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) with mock.patch.object( endpoint_service_client.EndpointServiceClient, "get_endpoint" ) as get_endpoint_mock: @@ -660,7 +678,6 @@ def test_deploy_with_traffic_split(self, deploy_model_mock, sync): @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_dedicated_resources(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy( @@ -701,7 +718,6 @@ def test_deploy_with_dedicated_resources(self, deploy_model_mock, sync): @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_explanations(self, deploy_model_with_explanations_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy( @@ -746,7 +762,6 @@ def test_deploy_with_explanations(self, deploy_model_with_explanations_mock, syn @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_min_replica_count(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, min_replica_count=2, sync=sync) @@ -771,7 +786,6 @@ def test_deploy_with_min_replica_count(self, deploy_model_mock, sync): @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_max_replica_count(self, deploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_model = models.Model(_TEST_ID) test_endpoint.deploy(model=test_model, max_replica_count=2, sync=sync) @@ -860,7 +874,6 @@ def test_unallocate_traffic(self, model1, model2, model3, deployed_model): @pytest.mark.parametrize("sync", [True, False]) def test_undeploy(self, undeploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) with mock.patch.object( endpoint_service_client.EndpointServiceClient, "get_endpoint" ) as get_endpoint_mock: @@ -884,7 +897,6 @@ def test_undeploy(self, undeploy_model_mock, sync): @pytest.mark.parametrize("sync", [True, False]) def test_undeploy_with_traffic_split(self, undeploy_model_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) with mock.patch.object( endpoint_service_client.EndpointServiceClient, "get_endpoint" ) as get_endpoint_mock: @@ -914,7 +926,6 @@ def test_undeploy_with_traffic_split(self, undeploy_model_mock, sync): @pytest.mark.parametrize("sync", [True, False]) def test_undeploy_raise_error_traffic_split_total(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_endpoint.undeploy( deployed_model_id="model1", traffic_split={"model2": 99}, sync=sync @@ -924,7 +935,6 @@ def test_undeploy_raise_error_traffic_split_total(self, sync): @pytest.mark.parametrize("sync", [True, False]) def test_undeploy_raise_error_undeployed_model_traffic(self, sync): with pytest.raises(ValueError): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) test_endpoint.undeploy( deployed_model_id="model1", @@ -933,7 +943,6 @@ def test_undeploy_raise_error_undeployed_model_traffic(self, sync): ) def test_predict(self, get_endpoint_mock, predict_client_predict_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ID) test_prediction = test_endpoint.predict( @@ -952,7 +961,6 @@ def test_predict(self, get_endpoint_mock, predict_client_predict_mock): ) def test_explain(self, get_endpoint_mock, predict_client_explain_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_endpoint = models.Endpoint(_TEST_ID) test_prediction = test_endpoint.explain( @@ -978,7 +986,6 @@ def test_explain(self, get_endpoint_mock, predict_client_explain_mock): ) def test_list_models(self, get_endpoint_with_models_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) ept = aiplatform.Endpoint(_TEST_ID) my_models = ept.list_models() @@ -988,7 +995,6 @@ def test_list_models(self, get_endpoint_with_models_mock): @pytest.mark.usefixtures("get_endpoint_with_models_mock") @pytest.mark.parametrize("sync", [True, False]) def test_undeploy_all(self, sdk_private_undeploy_mock, sync): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) ept = aiplatform.Endpoint(_TEST_ID) ept.undeploy_all(sync=sync) @@ -1005,6 +1011,23 @@ def test_undeploy_all(self, sdk_private_undeploy_mock, sync): any_order=True, ) + @pytest.mark.usefixtures("list_endpoints_mock") + def test_list_endpoint_has_prediction_client(self): + """Test call to Endpoint.list() and ensure Endpoints have prediction client set""" + ep_list = aiplatform.Endpoint.list(order_by=_TEST_LIST_ORDER_BY_CREATE_TIME) + + assert ep_list # Ensure list is not empty + + # Confirm every Endpoint object in the list has a prediction client + assert all( + [ + isinstance( + e._prediction_client, aiplatform.utils.PredictionClientWithOverride + ) + for e in ep_list + ] + ) + def test_list_endpoint_order_by_time(self, list_endpoints_mock): """Test call to Endpoint.list() and ensure list is returned in descending order of create_time""" @@ -1050,7 +1073,6 @@ def test_list_endpoint_order_by_display_name(self, list_endpoints_mock): def test_delete_endpoint_without_force( self, sdk_undeploy_all_mock, delete_endpoint_mock, sync ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) ept = aiplatform.Endpoint(_TEST_ID) ept.delete(sync=sync) @@ -1068,7 +1090,6 @@ def test_delete_endpoint_without_force( def test_delete_endpoint_with_force( self, sdk_undeploy_all_mock, delete_endpoint_mock, sync ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) ept = aiplatform.Endpoint(_TEST_ID) ept.delete(force=True, sync=sync)