From 12c5be4690b23375468af16b00790c106232f539 Mon Sep 17 00:00:00 2001 From: Vinny Senthil Date: Mon, 3 May 2021 08:36:32 -0600 Subject: [PATCH] feat: Add export model (#353) --- google/cloud/aiplatform/models.py | 162 +++++++++++++++ tests/unit/aiplatform/test_models.py | 283 +++++++++++++++++++++++++-- 2 files changed, 430 insertions(+), 15 deletions(-) diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index ea8b154a20..cecc992644 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -17,6 +17,7 @@ import proto from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple, Union +from google.api_core import operation from google.auth import credentials as auth_credentials from google.cloud.aiplatform import base @@ -35,9 +36,11 @@ endpoint_v1 as gca_endpoint_v1, endpoint_v1beta1 as gca_endpoint_v1beta1, explanation_v1beta1 as gca_explanation_v1beta1, + io as gca_io_compat, machine_resources as gca_machine_resources_compat, machine_resources_v1beta1 as gca_machine_resources_v1beta1, model as gca_model_compat, + model_service as gca_model_service_compat, model_v1beta1 as gca_model_v1beta1, env_var as gca_env_var_compat, env_var_v1beta1 as gca_env_var_v1beta1, @@ -1217,6 +1220,26 @@ def description(self): """Description of the model.""" return self._gca_resource.description + @property + def supported_export_formats( + self, + ) -> Dict[str, List[gca_model_compat.Model.ExportFormat.ExportableContent]]: + """The formats and content types in which this Model may be exported. + If empty, this Model is not available for export. + + For example, if this model can be exported as a Tensorflow SavedModel and + have the artifacts written to Cloud Storage, the expected value would be: + + {'tf-saved-model': []} + """ + return { + export_format.id: [ + gca_model_compat.Model.ExportFormat.ExportableContent(content) + for content in export_format.exportable_contents + ] + for export_format in self._gca_resource.supported_export_formats + } + def __init__( self, model_name: str, @@ -2030,3 +2053,142 @@ def list( location=location, credentials=credentials, ) + + @base.optional_sync() + def _wait_on_export(self, operation_future: operation.Operation, sync=True) -> None: + operation_future.result() + + def export_model( + self, + export_format_id: str, + artifact_destination: Optional[str] = None, + image_destination: Optional[str] = None, + sync: bool = True, + ) -> Dict[str, str]: + """Exports a trained, exportable Model to a location specified by the user. + A Model is considered to be exportable if it has at least one `supported_export_formats`. + Either `artifact_destination` or `image_destination` must be provided. + + Usage: + my_model.export( + export_format_id='tf-saved-model' + artifact_destination='gs://my-bucket/models/' + ) + + or + + my_model.export( + export_format_id='custom-model' + image_destination='us-central1-docker.pkg.dev/projectId/repo/image' + ) + + Args: + export_format_id (str): + Required. The ID of the format in which the Model must be exported. + The list of export formats that this Model supports can be found + by calling `Model.supported_export_formats`. + artifact_destination (str): + The Cloud Storage location where the Model artifact is to be + written to. Under the directory given as the destination a + new one with name + "``model-export--``", + where timestamp is in YYYY-MM-DDThh:mm:ss.sssZ ISO-8601 + format, will be created. Inside, the Model and any of its + supporting files will be written. + + This field should only be set when, in [Model.supported_export_formats], + the value for the key given in `export_format_id` contains ``ARTIFACT``. + image_destination (str): + The Google Container Registry or Artifact Registry URI where + the Model container image will be copied to. Accepted forms: + + - Google Container Registry path. For example: + ``gcr.io/projectId/imageName:tag``. + + - Artifact Registry path. For example: + ``us-central1-docker.pkg.dev/projectId/repoName/imageName:tag``. + + This field should only be set when, in [Model.supported_export_formats], + the value for the key given in `export_format_id` contains ``IMAGE``. + sync (bool): + Whether to execute this export synchronously. If False, this method + will be executed in concurrent Future and any downstream object will + be immediately returned and synced when the Future has completed. + Returns: + output_info (Dict[str, str]): + Details of the completed export with output destination paths to + the artifacts or container image. + Raises: + ValueError if model does not support exporting. + + ValueError if invalid arguments or export formats are provided. + """ + + # Model does not support exporting + if not self.supported_export_formats: + raise ValueError(f"The model `{self.resource_name}` is not exportable.") + + # No destination provided + if not any((artifact_destination, image_destination)): + raise ValueError( + "Please provide an `artifact_destination` or `image_destination`." + ) + + export_format_id = export_format_id.lower() + + # Unsupported export type + if export_format_id not in self.supported_export_formats: + raise ValueError( + f"'{export_format_id}' is not a supported export format for this model. " + f"Choose one of the following: {self.supported_export_formats}" + ) + + content_types = gca_model_compat.Model.ExportFormat.ExportableContent + supported_content_types = self.supported_export_formats[export_format_id] + + if ( + artifact_destination + and content_types.ARTIFACT not in supported_content_types + ): + raise ValueError( + "This model can not be exported as an artifact in '{export_format_id}' format. " + "Try exporting as a container image by passing the `image_destination` argument." + ) + + if image_destination and content_types.IMAGE not in supported_content_types: + raise ValueError( + "This model can not be exported as a container image in '{export_format_id}' format. " + "Try exporting the model artifacts by passing a `artifact_destination` argument." + ) + + # Construct request payload + output_config = gca_model_service_compat.ExportModelRequest.OutputConfig( + export_format_id=export_format_id + ) + + if artifact_destination: + output_config.artifact_destination = gca_io_compat.GcsDestination( + output_uri_prefix=artifact_destination + ) + + if image_destination: + output_config.image_destination = gca_io_compat.ContainerRegistryDestination( + output_uri=image_destination + ) + + _LOGGER.log_action_start_against_resource("Exporting", "model", self) + + operation_future = self.api_client.export_model( + name=self.resource_name, output_config=output_config + ) + + _LOGGER.log_action_started_against_resource_with_lro( + "Export", "model", self.__class__, operation_future + ) + + # Block before returning + self._wait_on_export(operation_future=operation_future, sync=sync) + + _LOGGER.log_action_completed_against_resource("model", "exported", self) + + return json_format.MessageToDict(operation_future.metadata.output_info._pb) diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index e76c1451f7..ad84fde65b 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/tests/unit/aiplatform/test_models.py @@ -157,6 +157,39 @@ ) _TEST_OUTPUT_DIR = "gs://my-output-bucket" +_TEST_CONTAINER_REGISTRY_DESTINATION = ( + "us-central1-docker.pkg.dev/projectId/repoName/imageName" +) + +_TEST_EXPORT_FORMAT_ID_IMAGE = "custom-trained" +_TEST_EXPORT_FORMAT_ID_ARTIFACT = "tf-saved-model" + +_TEST_SUPPORTED_EXPORT_FORMATS_IMAGE = [ + gca_model.Model.ExportFormat( + id=_TEST_EXPORT_FORMAT_ID_IMAGE, + exportable_contents=[gca_model.Model.ExportFormat.ExportableContent.IMAGE], + ) +] + +_TEST_SUPPORTED_EXPORT_FORMATS_ARTIFACT = [ + gca_model.Model.ExportFormat( + id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + exportable_contents=[gca_model.Model.ExportFormat.ExportableContent.ARTIFACT], + ) +] + +_TEST_SUPPORTED_EXPORT_FORMATS_BOTH = [ + gca_model.Model.ExportFormat( + id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + exportable_contents=[ + gca_model.Model.ExportFormat.ExportableContent.ARTIFACT, + gca_model.Model.ExportFormat.ExportableContent.IMAGE, + ], + ) +] + +_TEST_SUPPORTED_EXPORT_FORMATS_UNSUPPORTED = [] +_TEST_CONTAINER_REGISTRY_DESTINATION @pytest.fixture @@ -219,6 +252,58 @@ def get_model_with_custom_project_mock(): yield get_model_mock +@pytest.fixture +def get_model_with_supported_export_formats_image(): + with mock.patch.object( + model_service_client.ModelServiceClient, "get_model" + ) as get_model_mock: + get_model_mock.return_value = gca_model.Model( + display_name=_TEST_MODEL_NAME, + name=_TEST_MODEL_RESOURCE_NAME, + supported_export_formats=_TEST_SUPPORTED_EXPORT_FORMATS_IMAGE, + ) + yield get_model_mock + + +@pytest.fixture +def get_model_with_supported_export_formats_artifact(): + with mock.patch.object( + model_service_client.ModelServiceClient, "get_model" + ) as get_model_mock: + get_model_mock.return_value = gca_model.Model( + display_name=_TEST_MODEL_NAME, + name=_TEST_MODEL_RESOURCE_NAME, + supported_export_formats=_TEST_SUPPORTED_EXPORT_FORMATS_ARTIFACT, + ) + yield get_model_mock + + +@pytest.fixture +def get_model_with_both_supported_export_formats(): + with mock.patch.object( + model_service_client.ModelServiceClient, "get_model" + ) as get_model_mock: + get_model_mock.return_value = gca_model.Model( + display_name=_TEST_MODEL_NAME, + name=_TEST_MODEL_RESOURCE_NAME, + supported_export_formats=_TEST_SUPPORTED_EXPORT_FORMATS_BOTH, + ) + yield get_model_mock + + +@pytest.fixture +def get_model_with_unsupported_export_formats(): + with mock.patch.object( + model_service_client.ModelServiceClient, "get_model" + ) as get_model_mock: + get_model_mock.return_value = gca_model.Model( + display_name=_TEST_MODEL_NAME, + name=_TEST_MODEL_RESOURCE_NAME, + supported_export_formats=_TEST_SUPPORTED_EXPORT_FORMATS_UNSUPPORTED, + ) + yield get_model_mock + + @pytest.fixture def upload_model_mock(): with mock.patch.object( @@ -271,6 +356,22 @@ def upload_model_with_custom_location_mock(): yield upload_model_mock +@pytest.fixture +def export_model_mock(): + with mock.patch.object( + model_service_client.ModelServiceClient, "export_model" + ) as export_model_mock: + export_model_lro_mock = mock.Mock(ga_operation.Operation) + export_model_lro_mock.metadata = gca_model_service.ExportModelOperationMetadata( + output_info=gca_model_service.ExportModelOperationMetadata.OutputInfo( + artifact_output_uri=_TEST_OUTPUT_DIR + ) + ) + export_model_lro_mock.result.return_value = None + export_model_mock.return_value = export_model_lro_mock + yield export_model_mock + + @pytest.fixture def delete_model_mock(): with mock.patch.object( @@ -406,7 +507,6 @@ def test_constructor_create_client_with_custom_location(self, create_client_mock def test_constructor_creates_client_with_custom_credentials( self, create_client_mock ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) creds = auth_credentials.AnonymousCredentials() models.Model(_TEST_ID, credentials=creds) create_client_mock.assert_called_once_with( @@ -417,12 +517,10 @@ def test_constructor_creates_client_with_custom_credentials( ) def test_constructor_gets_model(self, get_model_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) models.Model(_TEST_ID) get_model_mock.assert_called_once_with(name=_TEST_MODEL_RESOURCE_NAME) def test_constructor_gets_model_with_custom_project(self, get_model_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) models.Model(_TEST_ID, project=_TEST_PROJECT_2) test_model_resource_name = model_service_client.ModelServiceClient.model_path( _TEST_PROJECT_2, _TEST_LOCATION, _TEST_ID @@ -430,7 +528,6 @@ def test_constructor_gets_model_with_custom_project(self, get_model_mock): get_model_mock.assert_called_once_with(name=test_model_resource_name) def test_constructor_gets_model_with_custom_location(self, get_model_mock): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) models.Model(_TEST_ID, location=_TEST_LOCATION_2) test_model_resource_name = model_service_client.ModelServiceClient.model_path( _TEST_PROJECT, _TEST_LOCATION_2, _TEST_ID @@ -442,7 +539,6 @@ def test_upload_uploads_and_gets_model( self, upload_model_mock, get_model_mock, sync ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) my_model = models.Model.upload( display_name=_TEST_MODEL_NAME, serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE, @@ -489,8 +585,6 @@ def test_upload_uploads_and_gets_model_with_all_args( self, upload_model_with_explanations_mock, get_model_mock, sync ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - my_model = models.Model.upload( display_name=_TEST_MODEL_NAME, artifact_uri=_TEST_ARTIFACT_URI, @@ -564,8 +658,6 @@ def test_upload_uploads_and_gets_model_with_custom_project( sync, ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - test_model_resource_name = model_service_client.ModelServiceClient.model_path( _TEST_PROJECT_2, _TEST_LOCATION, _TEST_ID ) @@ -612,7 +704,6 @@ def test_upload_uploads_and_gets_model_with_custom_location( get_model_with_custom_location_mock, sync, ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_model_resource_name = model_service_client.ModelServiceClient.model_path( _TEST_PROJECT, _TEST_LOCATION_2, _TEST_ID ) @@ -751,7 +842,6 @@ def test_deploy_no_endpoint_dedicated_resources(self, deploy_model_mock, sync): def test_deploy_no_endpoint_with_explanations( self, deploy_model_with_explanations_mock, sync ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_model = models.Model(_TEST_ID) test_endpoint = test_model.deploy( machine_type=_TEST_MACHINE_TYPE, @@ -943,7 +1033,6 @@ def test_batch_predict_gcs_source_bq_dest( def test_batch_predict_with_all_args( self, create_batch_prediction_job_with_explanations_mock, sync ): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_model = models.Model(_TEST_ID) creds = auth_credentials.AnonymousCredentials() @@ -1102,7 +1191,6 @@ def test_delete_model(self, delete_model_mock, sync): @pytest.mark.usefixtures("get_model_mock") def test_print_model(self): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_model = models.Model(_TEST_ID) assert ( repr(test_model) @@ -1111,7 +1199,6 @@ def test_print_model(self): @pytest.mark.usefixtures("get_model_mock") def test_print_model_if_waiting(self): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_model = models.Model(_TEST_ID) test_model._gca_resource = None test_model._latest_future = futures.Future() @@ -1122,7 +1209,6 @@ def test_print_model_if_waiting(self): @pytest.mark.usefixtures("get_model_mock") def test_print_model_if_exception(self): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) test_model = models.Model(_TEST_ID) test_model._gca_resource = None mock_exception = Exception("mock exception") @@ -1131,3 +1217,170 @@ def test_print_model_if_exception(self): repr(test_model) == f"{object.__repr__(test_model)} failed with {str(mock_exception)}" ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_supported_export_formats_artifact") + def test_export_model_as_artifact(self, export_model_mock, sync): + test_model = models.Model(_TEST_ID) + + if not sync: + test_model.wait() + + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + artifact_destination=_TEST_OUTPUT_DIR, + ) + + expected_output_config = gca_model_service.ExportModelRequest.OutputConfig( + export_format_id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + artifact_destination=gca_io.GcsDestination( + output_uri_prefix=_TEST_OUTPUT_DIR + ), + ) + + export_model_mock.assert_called_once_with( + name=f"{_TEST_PARENT}/models/{_TEST_ID}", + output_config=expected_output_config, + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_supported_export_formats_image") + def test_export_model_as_image(self, export_model_mock, sync): + test_model = models.Model(_TEST_ID) + + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_IMAGE, + image_destination=_TEST_CONTAINER_REGISTRY_DESTINATION, + ) + + if not sync: + test_model.wait() + + expected_output_config = gca_model_service.ExportModelRequest.OutputConfig( + export_format_id=_TEST_EXPORT_FORMAT_ID_IMAGE, + image_destination=gca_io.ContainerRegistryDestination( + output_uri=_TEST_CONTAINER_REGISTRY_DESTINATION + ), + ) + + export_model_mock.assert_called_once_with( + name=f"{_TEST_PARENT}/models/{_TEST_ID}", + output_config=expected_output_config, + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_both_supported_export_formats") + def test_export_model_as_both_formats(self, export_model_mock, sync): + """Exports a 'tf-saved-model' as both an artifact and an image""" + + test_model = models.Model(_TEST_ID) + + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + image_destination=_TEST_CONTAINER_REGISTRY_DESTINATION, + artifact_destination=_TEST_OUTPUT_DIR, + ) + + if not sync: + test_model.wait() + + expected_output_config = gca_model_service.ExportModelRequest.OutputConfig( + export_format_id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + image_destination=gca_io.ContainerRegistryDestination( + output_uri=_TEST_CONTAINER_REGISTRY_DESTINATION + ), + artifact_destination=gca_io.GcsDestination( + output_uri_prefix=_TEST_OUTPUT_DIR + ), + ) + + export_model_mock.assert_called_once_with( + name=f"{_TEST_PARENT}/models/{_TEST_ID}", + output_config=expected_output_config, + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_unsupported_export_formats") + def test_export_model_not_supported(self, export_model_mock, sync): + test_model = models.Model(_TEST_ID) + + with pytest.raises(ValueError) as e: + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_IMAGE, + image_destination=_TEST_CONTAINER_REGISTRY_DESTINATION, + ) + + if not sync: + test_model.wait() + + assert e.match( + regexp=f"The model `{_TEST_PARENT}/models/{_TEST_ID}` is not exportable." + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_supported_export_formats_image") + def test_export_model_as_image_with_invalid_args(self, export_model_mock, sync): + + # Passing an artifact destination on an image-only Model + with pytest.raises(ValueError) as dest_type_err: + test_model = models.Model(_TEST_ID) + + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_IMAGE, + artifact_destination=_TEST_OUTPUT_DIR, + sync=sync, + ) + + if not sync: + test_model.wait() + + # Passing no destination type + with pytest.raises(ValueError) as no_dest_err: + test_model = models.Model(_TEST_ID) + + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_IMAGE, sync=sync, + ) + + if not sync: + test_model.wait() + + # Passing an invalid export format ID + with pytest.raises(ValueError) as format_err: + test_model = models.Model(_TEST_ID) + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + image_destination=_TEST_CONTAINER_REGISTRY_DESTINATION, + sync=sync, + ) + + if not sync: + test_model.wait() + + assert dest_type_err.match( + regexp=r"This model can not be exported as an artifact." + ) + assert no_dest_err.match(regexp=r"Please provide an") + assert format_err.match( + regexp=f"'{_TEST_EXPORT_FORMAT_ID_ARTIFACT}' is not a supported export format" + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_supported_export_formats_artifact") + def test_export_model_as_artifact_with_invalid_args(self, export_model_mock, sync): + test_model = models.Model(_TEST_ID) + + # Passing an image destination on an artifact-only Model + with pytest.raises(ValueError) as e: + test_model.export_model( + export_format_id=_TEST_EXPORT_FORMAT_ID_ARTIFACT, + image_destination=_TEST_CONTAINER_REGISTRY_DESTINATION, + sync=sync, + ) + + if not sync: + test_model.wait() + + assert e.match( + regexp=r"This model can not be exported as a container image." + )