diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index ceb9287322..ab2b983dbe 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -21,6 +21,7 @@ import functools import inspect import logging +import re import sys import threading import time @@ -454,6 +455,24 @@ def _format_resource_name_method(self) -> str: # to use custom resource id validators per resource _resource_id_validator: Optional[Callable[[str], None]] = None + @staticmethod + def _revisioned_resource_id_validator( + resource_id: str, + ) -> None: + """Some revisioned resource names can have '@' in them + to separate the resource ID from the revision ID. + Thus, they need their own resource id validator. + See https://google.aip.dev/162 + + Args: + resource_id(str): A resource ID for a resource type that accepts revision syntax. + See https://google.aip.dev/162. + Raises: + ValueError: If a `resource_id` doesn't conform to appropriate revision syntax. + """ + if not re.compile(r"^[\w-]+@?[\w-]+$").match(resource_id): + raise ValueError(f"Resource {resource_id} is not a valid resource ID.") + def __init__( self, project: Optional[str] = None, diff --git a/google/cloud/aiplatform/jobs.py b/google/cloud/aiplatform/jobs.py index ab24afb171..c3e0459aac 100644 --- a/google/cloud/aiplatform/jobs.py +++ b/google/cloud/aiplatform/jobs.py @@ -392,6 +392,8 @@ def create( Required. A fully-qualified model resource name or model ID. Example: "projects/123/locations/us-central1/models/456" or "456" when project and location are initialized or passed. + May optionally contain a version ID or alias in + {model_name}@{version} form. Or an instance of aiplatform.Model. instances_format (str): @@ -564,6 +566,7 @@ def create( format_resource_name_method=aiplatform.Model._format_resource_name, project=project, location=location, + resource_id_validator=super()._revisioned_resource_id_validator, ) # Raise error if both or neither source URIs are provided @@ -713,7 +716,9 @@ def _create( Required. BatchPredictionJob without _gca_resource populated. model_or_model_name (Union[str, aiplatform.Model]): Required. Required. A fully-qualified model resource name or - an instance of aiplatform.Model. + an instance of aiplatform.Model. If a resource name, it may + optionally contain a version ID or alias in + {model_name}@{version} form. gca_batch_prediction_job (gca_bp_job.BatchPredictionJob): Required. a batch prediction job proto for creating a batch prediction job on Vertex AI. generate_explanation (bool): @@ -731,7 +736,6 @@ def _create( provided instances_format or predictions_format are not supported by Vertex AI. """ - # select v1beta1 if explain else use default v1 parent = initializer.global_config.common_location_path( project=empty_batch_prediction_job.project, @@ -741,7 +745,7 @@ def _create( model_resource_name = ( model_or_model_name if isinstance(model_or_model_name, str) - else model_or_model_name.resource_name + else model_or_model_name.versioned_resource_name ) gca_batch_prediction_job.model = model_resource_name diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index e6e108d45f..895d0e00c2 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -49,7 +49,7 @@ env_var as gca_env_var_compat, ) -from google.protobuf import field_mask_pb2, json_format +from google.protobuf import field_mask_pb2, json_format, timestamp_pb2 _DEFAULT_MACHINE_TYPE = "n1-standard-2" _DEPLOYING_MODEL_TRAFFIC_SPLIT_KEY = "0" @@ -67,6 +67,39 @@ ] +class VersionInfo(NamedTuple): + """VersionInfo class envelopes returned Model version information. + + Attributes: + version_id: + The version ID of the model. + create_time: + Timestamp when this Model version was uploaded into Vertex AI. + update_time: + Timestamp when this Model version was most recently updated. + model_display_name: + The user-defined name of the model this version belongs to. + model_resource_name: + The fully-qualified model resource name. + e.g. projects/{project}/locations/{location}/models/{model_display_name} + version_aliases: + User provided version aliases so that a model version can be referenced via + alias (i.e. projects/{project}/locations/{location}/models/{model_display_name}@{version_alias}). + Default is None. + version_description: + The description of this version. + Default is None. + """ + + version_id: str + version_create_time: timestamp_pb2.Timestamp + version_update_time: timestamp_pb2.Timestamp + model_display_name: str + model_resource_name: str + version_aliases: Optional[Sequence[str]] = None + version_description: Optional[str] = None + + class Prediction(NamedTuple): """Prediction class envelopes returned Model predictions and the Model id. @@ -78,6 +111,10 @@ class Prediction(NamedTuple): [PredictSchemata's][google.cloud.aiplatform.v1beta1.Model.predict_schemata] deployed_model_id: ID of the Endpoint's DeployedModel that served this prediction. + model_version_id: + ID of the DeployedModel's version that served this prediction. + model_resource_name: + The fully-qualified resource name of the model that served this prediction. explanations: The explanations of the Model's predictions. It has the same number of elements as instances to be explained. Default is None. @@ -85,6 +122,8 @@ class Prediction(NamedTuple): predictions: List[Dict[str, Any]] deployed_model_id: str + model_version_id: Optional[str] = None + model_resource_name: Optional[str] = None explanations: Optional[Sequence[gca_explanation_compat.Explanation]] = None @@ -1049,17 +1088,17 @@ def _deploy_call( ) deployed_model = gca_endpoint_compat.DeployedModel( - model=model.resource_name, + model=model.versioned_resource_name, display_name=deployed_model_display_name, service_account=service_account, ) supports_automatic_resources = ( - aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES + gca_model_compat.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES in model.supported_deployment_resources_types ) supports_dedicated_resources = ( - aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES + gca_model_compat.Model.DeploymentResourcesType.DEDICATED_RESOURCES in model.supported_deployment_resources_types ) provided_custom_machine_spec = ( @@ -1302,7 +1341,6 @@ def _instantiate_prediction_client( location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, ) -> utils.PredictionClientWithOverride: - """Helper method to instantiates prediction client with optional overrides for this endpoint. @@ -1471,6 +1509,8 @@ def predict( for item in prediction_response.predictions.pb ], deployed_model_id=prediction_response.deployed_model_id, + model_version_id=prediction_response.model_version_id, + model_resource_name=prediction_response.model, ) def explain( @@ -2408,12 +2448,81 @@ def container_spec(self) -> Optional[aiplatform.gapic.ModelContainerSpec]: self._assert_gca_resource_is_available() return getattr(self._gca_resource, "container_spec") + @property + def version_id(self) -> str: + """The version ID of the model. + A new version is committed when a new model version is uploaded or + trained under an existing model id. It is an auto-incrementing decimal + number in string representation.""" + self._assert_gca_resource_is_available() + return getattr(self._gca_resource, "version_id") + + @property + def version_aliases(self) -> Sequence[str]: + """User provided version aliases so that a model version can be referenced via + alias (i.e. projects/{project}/locations/{location}/models/{model_id}@{version_alias} + instead of auto-generated version id (i.e. + projects/{project}/locations/{location}/models/{model_id}@{version_id}). + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] to distinguish from + version_id. A default version alias will be created for the first version + of the model, and there must be exactly one default version alias for a model.""" + self._assert_gca_resource_is_available() + return getattr(self._gca_resource, "version_aliases") + + @property + def version_create_time(self) -> timestamp_pb2.Timestamp: + """Timestamp when this version was created.""" + self._assert_gca_resource_is_available() + return getattr(self._gca_resource, "version_create_time") + + @property + def version_update_time(self) -> timestamp_pb2.Timestamp: + """Timestamp when this version was updated.""" + self._assert_gca_resource_is_available() + return getattr(self._gca_resource, "version_update_time") + + @property + def version_description(self) -> str: + """The description of this version.""" + self._assert_gca_resource_is_available() + return getattr(self._gca_resource, "version_description") + + @property + def resource_name(self) -> str: + """Full qualified resource name, without any version ID.""" + self._assert_gca_resource_is_available() + return ModelRegistry._parse_versioned_name(self._gca_resource.name)[0] + + @property + def name(self) -> str: + """Name of this resource.""" + self._assert_gca_resource_is_available() + return ModelRegistry._parse_versioned_name(super().name)[0] + + @property + def versioned_resource_name(self) -> str: + """The fully-qualified resource name, including the version ID. For example, + projects/{project}/locations/{location}/models/{model_id}@{version_id} + """ + self._assert_gca_resource_is_available() + return ModelRegistry._get_versioned_name( + self.resource_name, + self.version_id, + ) + + @property + def versioning_registry(self) -> "ModelRegistry": + """The registry of model versions associated with this + Model instance.""" + return self._registry + def __init__( self, model_name: str, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, + version: Optional[str] = None, ): """Retrieves the model resource and instantiates its representation. @@ -2422,6 +2531,8 @@ def __init__( Required. A fully-qualified model resource name or model ID. Example: "projects/123/locations/us-central1/models/456" or "456" when project and location are initialized or passed. + May optionally contain a version ID or version alias in + {model_name}@{version} form. See version arg. project (str): Optional project to retrieve model from. If not set, project set in aiplatform.init will be used. @@ -2431,7 +2542,24 @@ def __init__( credentials: Optional[auth_credentials.Credentials]=None, Custom credentials to use to upload this model. If not set, credentials set in aiplatform.init will be used. + version (str): + Optional. Version ID or version alias. + When set, the specified model version will be targeted + unless overridden in method calls. + When not set, the model with the "default" alias will + be targeted unless overridden in method calls. + No behavior change if only one version of a model exists. + Raises: + ValueError: If `version` is passed alongside a model_name referencing a different version. """ + # If the version was passed in model_name, parse it + model_name, parsed_version = ModelRegistry._parse_versioned_name(model_name) + if parsed_version: + if version and version != parsed_version: + raise ValueError( + f"A version of {version} was passed that conflicts with the version of {parsed_version} in the model_name." + ) + version = parsed_version super().__init__( project=project, @@ -2439,7 +2567,21 @@ def __init__( credentials=credentials, resource_name=model_name, ) - self._gca_resource = self._get_gca_resource(resource_name=model_name) + + # Model versions can include @{version} in the resource name. + self._resource_id_validator = super()._revisioned_resource_id_validator + + # Create a versioned model_name, if it exists, for getting the GCA model + versioned_model_name = ModelRegistry._get_versioned_name(model_name, version) + self._gca_resource = self._get_gca_resource(resource_name=versioned_model_name) + + # Create ModelRegistry with the unversioned resource name + self._registry = ModelRegistry( + self.resource_name, + location=location, + project=project, + credentials=credentials, + ) def update( self, @@ -2488,6 +2630,10 @@ def update( update_mask: List[str] = [] + # Updates to base model properties cannot occur if a versioned model is passed. + # Use the unversioned model resource name. + copied_model_proto.name = self.resource_name + if display_name: utils.validate_display_name(display_name) @@ -2521,6 +2667,11 @@ def upload( serving_container_image_uri: str, *, artifact_uri: Optional[str] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: bool = True, + version_aliases: Optional[Sequence[str]] = None, + version_description: Optional[str] = None, serving_container_predict_route: Optional[str] = None, serving_container_health_route: Optional[str] = None, description: Optional[str] = None, @@ -2560,6 +2711,34 @@ def upload( Optional. The path to the directory containing the Model artifact and any of its supporting files. Leave blank for custom container prediction. Not present for AutoML Models. + model_id (str): + Optional. The ID to use for the uploaded Model, which will + become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model that the + newly-uploaded model will be a version of. + + Only set this field when uploading a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + this model without a version specified will use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the newly-uploaded model version will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + version_aliases (Sequence[str]): + Optional. User provided version aliases so that a model version + can be referenced via alias instead of auto-generated version ID. + A default version alias will be created for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + version_description (str): + Optional. The description of the model version being uploaded. serving_container_predict_route (str): Optional. An HTTP path to send prediction requests to the container, and which must be supported by it. If not specified a default HTTP path will @@ -2746,9 +2925,19 @@ def upload( encryption_spec_key_name=encryption_spec_key_name, ) + parent_model = ModelRegistry._get_true_version_parent( + location=location, project=project, parent_model=parent_model + ) + + version_aliases = ModelRegistry._get_true_alias_list( + version_aliases=version_aliases, is_default_version=is_default_version + ) + managed_model = gca_model_compat.Model( display_name=display_name, description=description, + version_aliases=version_aliases, + version_description=version_description, container_spec=container_spec, predict_schemata=model_predict_schemata, labels=labels, @@ -2795,9 +2984,15 @@ def upload( explanation_spec.parameters = explanation_parameters managed_model.explanation_spec = explanation_spec - lro = api_client.upload_model( + request = gca_model_service_compat.UploadModelRequest( parent=initializer.global_config.common_location_path(project, location), model=managed_model, + parent_model=parent_model, + model_id=model_id, + ) + + lro = api_client.upload_model( + request=request, timeout=upload_request_timeout, ) @@ -2805,7 +3000,9 @@ def upload( model_upload_response = lro.result() - this_model = cls(model_upload_response.model) + this_model = cls( + model_upload_response.model, version=model_upload_response.model_version_id + ) _LOGGER.log_create_complete(cls, this_model._gca_resource, "model") @@ -3434,6 +3631,51 @@ def list( credentials=credentials, ) + @classmethod + def _construct_sdk_resource_from_gapic( + cls, + gapic_resource: gca_model_compat.Model, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Model": + """Override base._construct_sdk_resource_from_gapic to allow for setting + a ModelRegistry and resource_id_validator. + + Args: + gapic_resource (gca_model_compat.Model): + A GAPIC representation of a Model resource. + project (str): + Optional. Project to construct SDK object from. If not set, + project set in aiplatform.init will be used. + location (str): + Optional. Location to construct SDK object from. If not set, + location set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to construct SDK object. + Overrides credentials set in aiplatform.init. + + Returns: + Model: + An initialized SDK Model object that represents the Model GAPIC type. + """ + sdk_resource = super()._construct_sdk_resource_from_gapic( + gapic_resource=gapic_resource, + project=project, + location=location, + credentials=credentials, + ) + sdk_resource._resource_id_validator = super()._revisioned_resource_id_validator + + sdk_resource._registry = ModelRegistry( + sdk_resource.resource_name, + location=location, + project=project, + credentials=credentials, + ) + + return sdk_resource + @base.optional_sync() def _wait_on_export(self, operation_future: operation.Operation, sync=True) -> None: operation_future.result() @@ -3561,8 +3803,10 @@ def export_model( _LOGGER.log_action_start_against_resource("Exporting", "model", self) + model_name = self.versioned_resource_name + operation_future = self.api_client.export_model( - name=self.resource_name, output_config=output_config + name=model_name, output_config=output_config ) _LOGGER.log_action_started_against_resource_with_lro( @@ -3584,6 +3828,11 @@ def upload_xgboost_model_file( xgboost_version: Optional[str] = None, display_name: Optional[str] = None, description: Optional[str] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + version_aliases: Optional[Sequence[str]] = None, + version_description: Optional[str] = None, instance_schema_uri: Optional[str] = None, parameters_schema_uri: Optional[str] = None, prediction_schema_uri: Optional[str] = None, @@ -3620,6 +3869,34 @@ def upload_xgboost_model_file( characters long and can be consist of any UTF-8 characters. description (str): The description of the model. + model_id (str): + Optional. The ID to use for the uploaded Model, which will + become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model that the + newly-uploaded model will be a version of. + + Only set this field when uploading a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + this model without a version specified will use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the newly-uploaded model version will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + version_aliases (Sequence[str]): + Optional. User provided version aliases so that a model version + can be referenced via alias instead of auto-generated version ID. + A default version alias will be created for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + version_description (str): + Optional. The description of the model version being uploaded. instance_schema_uri (str): Optional. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which @@ -3766,6 +4043,11 @@ def upload_xgboost_model_file( artifact_uri=prepared_model_dir, display_name=display_name, description=description, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + version_aliases=version_aliases, + version_description=version_description, instance_schema_uri=instance_schema_uri, parameters_schema_uri=parameters_schema_uri, prediction_schema_uri=prediction_schema_uri, @@ -3789,6 +4071,11 @@ def upload_scikit_learn_model_file( sklearn_version: Optional[str] = None, display_name: Optional[str] = None, description: Optional[str] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + version_aliases: Optional[Sequence[str]] = None, + version_description: Optional[str] = None, instance_schema_uri: Optional[str] = None, parameters_schema_uri: Optional[str] = None, prediction_schema_uri: Optional[str] = None, @@ -3826,6 +4113,34 @@ def upload_scikit_learn_model_file( characters long and can be consist of any UTF-8 characters. description (str): The description of the model. + model_id (str): + Optional. The ID to use for the uploaded Model, which will + become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model that the + newly-uploaded model will be a version of. + + Only set this field when uploading a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + this model without a version specified will use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the newly-uploaded model version will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + version_aliases (Sequence[str]): + Optional. User provided version aliases so that a model version + can be referenced via alias instead of auto-generated version ID. + A default version alias will be created for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + version_description (str): + Optional. The description of the model version being uploaded. instance_schema_uri (str): Optional. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which @@ -3975,6 +4290,11 @@ def upload_scikit_learn_model_file( artifact_uri=prepared_model_dir, display_name=display_name, description=description, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + version_aliases=version_aliases, + version_description=version_description, instance_schema_uri=instance_schema_uri, parameters_schema_uri=parameters_schema_uri, prediction_schema_uri=prediction_schema_uri, @@ -3998,6 +4318,11 @@ def upload_tensorflow_saved_model( use_gpu: bool = False, display_name: Optional[str] = None, description: Optional[str] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + version_aliases: Optional[Sequence[str]] = None, + version_description: Optional[str] = None, instance_schema_uri: Optional[str] = None, parameters_schema_uri: Optional[str] = None, prediction_schema_uri: Optional[str] = None, @@ -4037,6 +4362,34 @@ def upload_tensorflow_saved_model( characters long and can be consist of any UTF-8 characters. description (str): The description of the model. + model_id (str): + Optional. The ID to use for the uploaded Model, which will + become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model that the + newly-uploaded model will be a version of. + + Only set this field when uploading a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + this model without a version specified will use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the newly-uploaded model version will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + version_aliases (Sequence[str]): + Optional. User provided version aliases so that a model version + can be referenced via alias instead of auto-generated version ID. + A default version alias will be created for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + version_description (str): + Optional. The description of the model version being uploaded. instance_schema_uri (str): Optional. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which @@ -4154,6 +4507,11 @@ def upload_tensorflow_saved_model( artifact_uri=saved_model_dir, display_name=display_name, description=description, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + version_aliases=version_aliases, + version_description=version_description, instance_schema_uri=instance_schema_uri, parameters_schema_uri=parameters_schema_uri, prediction_schema_uri=prediction_schema_uri, @@ -4243,3 +4601,311 @@ def get_model_evaluation( evaluation_name=evaluation_resource_name, credentials=self.credentials, ) + + +# TODO (b/232546878): Async support +class ModelRegistry: + def __init__( + self, + model: Union[Model, str], + location: Optional[str] = None, + project: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Creates a ModelRegistry instance for version management of a registered model. + + Args: + model (Union[Model, str]): + Required. One of the following: + 1. A Model instance + 2. A fully-qualified model resource name + 3. A model ID. A location and project must be provided. + location (str): + Optional. The model location. Used when passing a model name as model. + If not set, project set in aiplatform.init will be used. + project (str): + Optional. The model project. Used when passing a model name as model. + If not set, project set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use with model access. If not set, + credentials set in aiplatform.init will be used. + """ + + if isinstance(model, Model): + self.model_resource_name = model.resource_name + else: + self.model_resource_name = utils.full_resource_name( + resource_name=model, + resource_noun="models", + parse_resource_name_method=Model._parse_resource_name, + format_resource_name_method=Model._format_resource_name, + project=project, + location=location, + resource_id_validator=base.VertexAiResourceNoun._revisioned_resource_id_validator, + ) + + self.credentials = credentials or ( + model.credentials + if isinstance(model, Model) + else initializer.global_config.credentials + ) + self.client = Model._instantiate_client(location, self.credentials) + + def get_model( + self, + version: Optional[str] = None, + ) -> Model: + """Gets a registered model with optional version. + + Args: + version (str): + Optional. A model version ID or alias to target. + Defaults to the model with the "default" alias. + + Returns: + Model: An instance of a Model from this ModelRegistry. + """ + return Model( + self.model_resource_name, version=version, credentials=self.credentials + ) + + def list_versions( + self, + ) -> List[VersionInfo]: + """Lists the versions and version info of a model. + + Returns: + List[VersionInfo]: + A list of VersionInfo, each containing + info about specific model versions. + """ + + _LOGGER.info(f"Getting versions for {self.model_resource_name}") + + page_result = self.client.list_model_versions( + name=self.model_resource_name, + ) + + versions = [ + VersionInfo( + version_id=model.version_id, + version_create_time=model.version_create_time, + version_update_time=model.version_update_time, + model_display_name=model.display_name, + model_resource_name=self._parse_versioned_name(model.name)[0], + version_aliases=model.version_aliases, + version_description=model.version_description, + ) + for model in page_result + ] + + return versions + + def get_version_info( + self, + version: str, + ) -> VersionInfo: + """Gets information about a specific model version. + + Args: + version (str): Required. The model version to obtain info for. + + Returns: + VersionInfo: Contains info about the model version. + """ + + _LOGGER.info(f"Getting version {version} info for {self.model_resource_name}") + + model = self.client.get_model( + name=self._get_versioned_name(self.model_resource_name, version), + ) + + return VersionInfo( + version_id=model.version_id, + version_create_time=model.version_create_time, + version_update_time=model.version_update_time, + model_display_name=model.display_name, + model_resource_name=self._parse_versioned_name(model.name)[0], + version_aliases=model.version_aliases, + version_description=model.version_description, + ) + + def delete_version( + self, + version: str, + ) -> None: + """Deletes a model version from the registry. + + Cannot delete a version if it is the last remaining version. + Use Model.delete() in that case. + + Args: + version (str): Required. The model version ID or alias to delete. + """ + + lro = self.client.delete_model_version( + name=self._get_versioned_name(self.model_resource_name, version), + ) + + _LOGGER.info(f"Deleting version {version} for {self.model_resource_name}") + + lro.result() + + _LOGGER.info(f"Deleted version {version} for {self.model_resource_name}") + + def add_version_aliases( + self, + new_aliases: List[str], + version: str, + ) -> None: + """Adds version alias(es) to a model version. + + Args: + new_aliases (List[str]): Required. The alias(es) to add to a model version. + version (str): Required. The version ID to receive the new alias(es). + """ + + self._merge_version_aliases( + version_aliases=new_aliases, + version=version, + ) + + def remove_version_aliases( + self, + target_aliases: List[str], + version: str, + ) -> None: + """Removes version alias(es) from a model version. + + Args: + target_aliases (List[str]): Required. The alias(es) to remove from a model version. + version (str): Required. The version ID to be stripped of the target alias(es). + """ + + self._merge_version_aliases( + version_aliases=[f"-{alias}" for alias in target_aliases], + version=version, + ) + + def _merge_version_aliases( + self, + version_aliases: List[str], + version: str, + ) -> None: + """Merges a list of version aliases with a model's existing alias list. + + Args: + version_aliases (List[str]): Required. The version alias change list. + version (str): Required. The version ID to have its alias list changed. + """ + + _LOGGER.info(f"Merging version aliases for {self.model_resource_name}") + + self.client.merge_version_aliases( + name=self._get_versioned_name(self.model_resource_name, version), + version_aliases=version_aliases, + ) + + _LOGGER.info( + f"Completed merging version aliases for {self.model_resource_name}" + ) + + @staticmethod + def _get_versioned_name( + resource_name: str, + version: Optional[str] = None, + ) -> str: + """Creates a versioned form of a model resource name. + + Args: + resource_name (str): Required. A fully-qualified resource name or resource ID. + version (str): Optional. The version or alias of the resource. + + Returns: + versioned_name (str): The versioned resource name in revisioned format. + """ + if version: + return f"{resource_name}@{version}" + return resource_name + + @staticmethod + def _parse_versioned_name( + model_name: str, + ) -> Tuple[str, Optional[str]]: + """Return a model name and, if included in the model name, a model version. + + Args: + model_name (str): Required. A fully-qualified model name or model ID, + optionally with an included version. + + Returns: + parsed_version_name (Tuple[str, Optional[str]]): + A tuple containing the model name or ID as the first element, + and the model version as the second element, if present in `model_name`. + + Raises: + ValueError: If the `model_name` is invalid and contains too many '@' symbols. + """ + if "@" not in model_name: + return model_name, None + elif model_name.count("@") > 1: + raise ValueError( + f"Received an invalid model_name with too many `@`s: {model_name}" + ) + else: + return model_name.split("@") + + @staticmethod + def _get_true_version_parent( + parent_model: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + ) -> Optional[str]: + """Gets the true `parent_model` with full resource name. + + Args: + parent_model (str): Optional. A fully-qualified resource name or resource ID + of the model that would be the parent of another model. + project (str): Optional. The project of `parent_model`, if not included in `parent_model`. + location (str): Optional. The location of `parent_model`, if not included in `parent_model`. + + Returns: + true_parent_model (str): + Optional. The true resource name of the parent model, if one should exist. + """ + if parent_model: + existing_resource = utils.full_resource_name( + resource_name=parent_model, + resource_noun="models", + parse_resource_name_method=Model._parse_resource_name, + format_resource_name_method=Model._format_resource_name, + project=project, + location=location, + ) + parent_model = existing_resource + return parent_model + + @staticmethod + def _get_true_alias_list( + version_aliases: Optional[Sequence[str]] = None, + is_default_version: bool = True, + ) -> Optional[Sequence[str]]: + """Gets the true `version_aliases` list based on `is_default_version`. + + Args: + version_aliases (Sequence[str]): Optional. The user-provided list of model aliases. + is_default_version (bool): + Optional. When set, includes the "default" alias in `version_aliases`. + Defaults to True. + + Returns: + true_alias_list (Sequence[str]): + Optional: The true alias list, should one exist, + containing "default" if specified. + """ + if is_default_version: + if version_aliases and "default" not in version_aliases: + version_aliases.append("default") + elif not version_aliases: + version_aliases = ["default"] + return version_aliases diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py index b2a93d952c..5d36cd97fa 100644 --- a/google/cloud/aiplatform/training_jobs.py +++ b/google/cloud/aiplatform/training_jobs.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2020 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -575,6 +575,11 @@ def _run_job( timestamp_split_column_name: Optional[str] = None, annotation_schema_uri: Optional[str] = None, model: Optional[gca_model.Model] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, gcs_destination_uri_prefix: Optional[str] = None, bigquery_destination: Optional[str] = None, create_request_timeout: Optional[float] = None, @@ -698,6 +703,36 @@ def _run_job( resource ``name`` is populated. The Model is always uploaded into the Project and Location in which this pipeline is. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. gcs_destination_uri_prefix (str): Optional. The Google Cloud Storage location. @@ -743,12 +778,27 @@ def _run_job( bigquery_destination=bigquery_destination, ) + parent_model = models.ModelRegistry._get_true_version_parent( + parent_model=parent_model, + project=self.project, + location=self.location, + ) + + if model: + model.version_aliases = models.ModelRegistry._get_true_alias_list( + version_aliases=model_version_aliases, + is_default_version=is_default_version, + ) + model.version_description = model_version_description + # create training pipeline training_pipeline = gca_training_pipeline.TrainingPipeline( display_name=self._display_name, training_task_definition=training_task_definition, training_task_inputs=training_task_inputs, model_to_upload=model, + model_id=model_id, + parent_model=parent_model, input_data_config=input_data_config, labels=self._labels, encryption_spec=self._training_encryption_spec, @@ -864,7 +914,10 @@ def _get_model(self) -> Optional[models.Model]: return None if self._gca_resource.model_to_upload.name: - return models.Model(model_name=self._gca_resource.model_to_upload.name) + return models.Model( + model_name=self._gca_resource.model_to_upload.name, + version=self._gca_resource.model_to_upload.version_id, + ) def _wait_callback(self): """Callback performs custom logging during _block_until_complete. Override in subclass.""" @@ -1726,6 +1779,11 @@ def run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, additional_experiments: Optional[List[str]] = None, hierarchy_group_columns: Optional[List[str]] = None, hierarchy_group_total_weight: Optional[float] = None, @@ -1894,6 +1952,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. additional_experiments (List[str]): Optional. Additional experiment flags for the time series forcasting training. create_request_timeout (float): @@ -1998,6 +2086,11 @@ def run( validation_options=validation_options, model_display_name=model_display_name, model_labels=model_labels, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, hierarchy_group_columns=hierarchy_group_columns, hierarchy_group_total_weight=hierarchy_group_total_weight, hierarchy_temporal_total_weight=hierarchy_temporal_total_weight, @@ -2038,6 +2131,11 @@ def _run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, hierarchy_group_columns: Optional[List[str]] = None, hierarchy_group_total_weight: Optional[float] = None, hierarchy_temporal_total_weight: Optional[float] = None, @@ -2214,6 +2312,36 @@ def _run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. hierarchy_group_columns (List[str]): Optional. A list of time series attribute column names that define the time series hierarchy. Only one level of hierarchy is @@ -2364,6 +2492,11 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, model=model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, create_request_timeout=create_request_timeout, ) @@ -2683,6 +2816,11 @@ def run( annotation_schema_uri: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, base_output_dir: Optional[str] = None, service_account: Optional[str] = None, network: Optional[str] = None, @@ -2809,6 +2947,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. base_output_dir (str): GCS output directory of job. If not provided a timestamped directory in the staging directory will be used. @@ -2990,6 +3158,11 @@ def run( annotation_schema_uri=annotation_schema_uri, worker_pool_specs=worker_pool_specs, managed_model=managed_model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, args=args, environment_variables=environment_variables, base_output_dir=base_output_dir, @@ -3030,6 +3203,11 @@ def _run( annotation_schema_uri: Optional[str], worker_pool_specs: worker_spec_utils._DistributedTrainingSpec, managed_model: Optional[gca_model.Model] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, environment_variables: Optional[Dict[str, str]] = None, base_output_dir: Optional[str] = None, @@ -3073,6 +3251,36 @@ def _run( Worker pools pecs required to run job. managed_model (gca_model.Model): Model proto if this script produces a Managed Model. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. args (List[Unions[str, int, float]]): Command line arguments to be passed to the Python script. environment_variables (Dict[str, str]): @@ -3268,6 +3476,11 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, model=managed_model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, gcs_destination_uri_prefix=base_output_dir, bigquery_destination=bigquery_destination, create_request_timeout=create_request_timeout, @@ -3520,6 +3733,11 @@ def run( annotation_schema_uri: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, base_output_dir: Optional[str] = None, service_account: Optional[str] = None, network: Optional[str] = None, @@ -3639,6 +3857,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. base_output_dir (str): GCS output directory of job. If not provided a timestamped directory in the staging directory will be used. @@ -3819,6 +4067,11 @@ def run( annotation_schema_uri=annotation_schema_uri, worker_pool_specs=worker_pool_specs, managed_model=managed_model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, args=args, environment_variables=environment_variables, base_output_dir=base_output_dir, @@ -3858,6 +4111,11 @@ def _run( annotation_schema_uri: Optional[str], worker_pool_specs: worker_spec_utils._DistributedTrainingSpec, managed_model: Optional[gca_model.Model] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, environment_variables: Optional[Dict[str, str]] = None, base_output_dir: Optional[str] = None, @@ -3898,6 +4156,36 @@ def _run( Worker pools pecs required to run job. managed_model (gca_model.Model): Model proto if this script produces a Managed Model. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. args (List[Unions[str, int, float]]): Command line arguments to be passed to the Python script. environment_variables (Dict[str, str]): @@ -4086,6 +4374,11 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, model=managed_model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, gcs_destination_uri_prefix=base_output_dir, bigquery_destination=bigquery_destination, create_request_timeout=create_request_timeout, @@ -4293,6 +4586,11 @@ def run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, disable_early_stopping: bool = False, export_evaluated_data_items: bool = False, export_evaluated_data_items_bigquery_destination_uri: Optional[str] = None, @@ -4398,6 +4696,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. disable_early_stopping (bool): Required. If true, the entire budget is used. This disables the early stopping feature. By default, the early stopping feature is enabled, which means @@ -4466,6 +4794,11 @@ def run( budget_milli_node_hours=budget_milli_node_hours, model_display_name=model_display_name, model_labels=model_labels, + model_id=model_id, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, + parent_model=parent_model, + is_default_version=is_default_version, disable_early_stopping=disable_early_stopping, export_evaluated_data_items=export_evaluated_data_items, export_evaluated_data_items_bigquery_destination_uri=export_evaluated_data_items_bigquery_destination_uri, @@ -4488,6 +4821,11 @@ def _run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, disable_early_stopping: bool = False, export_evaluated_data_items: bool = False, export_evaluated_data_items_bigquery_destination_uri: Optional[str] = None, @@ -4592,6 +4930,36 @@ def _run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. disable_early_stopping (bool): Required. If true, the entire budget is used. This disables the early stopping feature. By default, the early stopping feature is enabled, which means @@ -4698,6 +5066,11 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, model=model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, create_request_timeout=create_request_timeout, ) @@ -4789,6 +5162,11 @@ def run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, additional_experiments: Optional[List[str]] = None, hierarchy_group_columns: Optional[List[str]] = None, hierarchy_group_total_weight: Optional[float] = None, @@ -4827,6 +5205,11 @@ def run( validation_options=validation_options, model_display_name=model_display_name, model_labels=model_labels, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, additional_experiments=additional_experiments, hierarchy_group_columns=hierarchy_group_columns, hierarchy_group_total_weight=hierarchy_group_total_weight, @@ -4875,6 +5258,11 @@ def run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, additional_experiments: Optional[List[str]] = None, hierarchy_group_columns: Optional[List[str]] = None, hierarchy_group_total_weight: Optional[float] = None, @@ -4902,6 +5290,11 @@ def run( test_fraction_split=test_fraction_split, predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, weight_column=weight_column, time_series_attribute_columns=time_series_attribute_columns, context_window=context_window, @@ -5103,6 +5496,11 @@ def run( budget_milli_node_hours: Optional[int] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, disable_early_stopping: bool = False, sync: bool = True, create_request_timeout: Optional[float] = None, @@ -5201,6 +5599,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. disable_early_stopping: bool = False Required. If true, the entire budget is used. This disables the early stopping feature. By default, the early stopping feature is enabled, which means @@ -5244,6 +5672,11 @@ def run( budget_milli_node_hours=budget_milli_node_hours, model_display_name=model_display_name, model_labels=model_labels, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, disable_early_stopping=disable_early_stopping, sync=sync, create_request_timeout=create_request_timeout, @@ -5263,6 +5696,11 @@ def _run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, disable_early_stopping: bool = False, sync: bool = True, create_request_timeout: Optional[float] = None, @@ -5303,6 +5741,36 @@ def _run( Otherwise, the new model will be trained from scratch. The `base` model must be in the same Project and Location as the new Model to train, and have the same model_type. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. training_fraction_split (float): Optional. The fraction of the input data that is to be used to train the Model. This is ignored if Dataset is not provided. @@ -5421,6 +5889,11 @@ def _run( validation_filter_split=validation_filter_split, test_filter_split=test_filter_split, model=model_tbt, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, create_request_timeout=create_request_timeout, ) @@ -5687,6 +6160,11 @@ def run( annotation_schema_uri: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, base_output_dir: Optional[str] = None, service_account: Optional[str] = None, network: Optional[str] = None, @@ -5806,6 +6284,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. base_output_dir (str): GCS output directory of job. If not provided a timestamped directory in the staging directory will be used. @@ -5981,6 +6489,11 @@ def run( annotation_schema_uri=annotation_schema_uri, worker_pool_specs=worker_pool_specs, managed_model=managed_model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, args=args, environment_variables=environment_variables, base_output_dir=base_output_dir, @@ -6020,6 +6533,11 @@ def _run( annotation_schema_uri: Optional[str], worker_pool_specs: worker_spec_utils._DistributedTrainingSpec, managed_model: Optional[gca_model.Model] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, args: Optional[List[Union[str, float, int]]] = None, environment_variables: Optional[Dict[str, str]] = None, base_output_dir: Optional[str] = None, @@ -6061,6 +6579,36 @@ def _run( Worker pools pecs required to run job. managed_model (gca_model.Model): Model proto if this script produces a Managed Model. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. args (List[Unions[str, int, float]]): Command line arguments to be passed to the Python script. environment_variables (Dict[str, str]): @@ -6235,6 +6783,11 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, model=managed_model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, gcs_destination_uri_prefix=base_output_dir, bigquery_destination=bigquery_destination, create_request_timeout=create_request_timeout, @@ -6388,6 +6941,11 @@ def run( test_filter_split: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, sync: bool = True, create_request_timeout: Optional[float] = None, ) -> models.Model: @@ -6453,6 +7011,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. sync: bool = True Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -6486,6 +7074,11 @@ def run( test_filter_split=test_filter_split, model_display_name=model_display_name, model_labels=model_labels, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, sync=sync, create_request_timeout=create_request_timeout, ) @@ -6500,6 +7093,11 @@ def _run( test_filter_split: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, sync: bool = True, create_request_timeout: Optional[float] = None, ) -> models.Model: @@ -6567,6 +7165,36 @@ def _run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -6610,6 +7238,11 @@ def _run( validation_filter_split=validation_filter_split, test_filter_split=test_filter_split, model=model_tbt, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, create_request_timeout=create_request_timeout, ) @@ -6777,6 +7410,11 @@ def run( test_filter_split: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, sync: bool = True, create_request_timeout: Optional[float] = None, ) -> models.Model: @@ -6854,6 +7492,36 @@ def run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -6888,6 +7556,11 @@ def run( test_filter_split=test_filter_split, model_display_name=model_display_name, model_labels=model_labels, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, sync=sync, create_request_timeout=create_request_timeout, ) @@ -6904,6 +7577,11 @@ def _run( test_filter_split: Optional[str] = None, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + model_id: Optional[str] = None, + parent_model: Optional[str] = None, + is_default_version: Optional[bool] = True, + model_version_aliases: Optional[Sequence[str]] = None, + model_version_description: Optional[str] = None, sync: bool = True, create_request_timeout: Optional[float] = None, ) -> models.Model: @@ -6983,6 +7661,36 @@ def _run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + model_id (str): + Optional. The ID to use for the Model produced by this job, + which will become the final component of the model resource name. + This value may be up to 63 characters, and valid characters + are `[a-z0-9_-]`. The first character cannot be a number or hyphen. + parent_model (str): + Optional. The resource name or model ID of an existing model. + The new model uploaded by this job will be a version of `parent_model`. + + Only set this field when training a new version of an existing model. + is_default_version (bool): + Optional. When set to True, the newly uploaded model version will + automatically have alias "default" included. Subsequent uses of + the model produced by this job without a version specified will + use this "default" version. + + When set to False, the "default" alias will not be moved. + Actions targeting the model version produced by this job will need + to specifically reference this version by ID or alias. + + New model uploads, i.e. version 1, will always be "default" aliased. + model_version_aliases (Sequence[str]): + Optional. User provided version aliases so that the model version + uploaded by this job can be referenced via alias instead of + auto-generated version ID. A default version alias will be created + for the first version of the model. + + The format is [a-z][a-zA-Z0-9-]{0,126}[a-z0-9] + model_version_description (str): + Optional. The description of the model version being uploaded by this job. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -7012,6 +7720,11 @@ def _run( validation_filter_split=validation_filter_split, test_filter_split=test_filter_split, model=model, + model_id=model_id, + parent_model=parent_model, + is_default_version=is_default_version, + model_version_aliases=model_version_aliases, + model_version_description=model_version_description, create_request_timeout=create_request_timeout, ) diff --git a/tests/system/aiplatform/test_model_version_management.py b/tests/system/aiplatform/test_model_version_management.py new file mode 100644 index 0000000000..2ac7021c98 --- /dev/null +++ b/tests/system/aiplatform/test_model_version_management.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import tempfile +import uuid + +import pytest + +from google.cloud import aiplatform +from google.cloud import storage +from google.cloud.aiplatform.models import ModelRegistry + +from tests.system.aiplatform import e2e_base +from tests.system.aiplatform import test_model_upload + + +@pytest.mark.usefixtures("delete_staging_bucket", "tear_down_resources") +class TestVersionManagement(e2e_base.TestEndToEnd): + + _temp_prefix = "temp_vertex_sdk_e2e_model_upload_test" + + def test_upload_deploy_manage_versioned_model(self, shared_state): + """Upload XGBoost model from local file and deploy it for prediction. Additionally, update model name, description and labels""" + + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + ) + + storage_client = storage.Client(project=e2e_base._PROJECT) + model_blob = storage.Blob.from_string( + uri=test_model_upload._XGBOOST_MODEL_URI, client=storage_client + ) + model_path = tempfile.mktemp() + ".my_model.xgb" + model_blob.download_to_filename(filename=model_path) + + model_id = "my_model_id" + uuid.uuid4().hex + version_description = "My description" + version_aliases = ["system-test-model", "testing"] + + model = aiplatform.Model.upload_xgboost_model_file( + model_file_path=model_path, + version_aliases=version_aliases, + model_id=model_id, + version_description=version_description, + ) + shared_state["resources"] = [model] + + staging_bucket = storage.Blob.from_string( + uri=model.uri, client=storage_client + ).bucket + # Checking that the bucket is auto-generated + assert "-vertex-staging-" in staging_bucket.name + + shared_state["bucket"] = staging_bucket + + assert model.version_description == version_description + assert model.version_aliases == version_aliases + assert "default" in model.version_aliases + + model2 = aiplatform.Model.upload_xgboost_model_file( + model_file_path=model_path, parent_model=model_id, is_default_version=False + ) + shared_state["resources"].append(model2) + + assert model2.version_id == "2" + assert model2.resource_name == model.resource_name + assert model2.version_aliases == [] + + # Test that VersionInfo properties are correct. + model_info = model2.versioning_registry.get_version_info("testing") + version_list = model2.versioning_registry.list_versions() + assert len(version_list) == 2 + list_info = version_list[0] + assert model_info.version_id == list_info.version_id == model.version_id + assert ( + model_info.version_aliases + == list_info.version_aliases + == model.version_aliases + ) + assert ( + model_info.version_description + == list_info.version_description + == model.version_description + ) + assert ( + model_info.model_display_name + == list_info.model_display_name + == model.display_name + ) + assert ( + model_info.version_update_time + == list_info.version_update_time + == model.version_update_time + ) + + # Test that get_model yields a new instance of `model` + model_clone = model2.versioning_registry.get_model() + assert model.resource_name == model_clone.resource_name + assert model.version_id == model_clone.version_id + assert model.name == model_clone.name + + # Test add and removal of aliases + registry = ModelRegistry(model) + registry.add_version_aliases(["new-alias"], "default") + registry.remove_version_aliases(["testing"], "new-alias") + model = registry.get_model("new-alias") + assert "testing" not in model.version_aliases + + # Test deletion of a model version + registry.delete_version("2") + versions = registry.list_versions() + assert "2" not in [version.version_id for version in versions] diff --git a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py index 236275237a..43be93c003 100644 --- a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py @@ -229,7 +229,7 @@ def mock_model_service_get(): with mock.patch.object( model_service_client.ModelServiceClient, "get_model" ) as mock_get_model: - mock_get_model.return_value = gca_model.Model() + mock_get_model.return_value = gca_model.Model(name=_TEST_MODEL_NAME) yield mock_get_model @@ -341,7 +341,9 @@ def test_run_call_pipeline_service_create( model_from_job.wait() true_managed_model = gca_model.Model( - display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -447,7 +449,9 @@ def test_run_call_pipeline_service_create_with_timeout( model_from_job.wait() true_managed_model = gca_model.Model( - display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -538,6 +542,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( true_managed_model = gca_model.Model( display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -623,7 +628,10 @@ def test_run_call_pipeline_if_set_additional_experiments( model_from_job.wait() # Test that if defaults to the job display name - true_managed_model = gca_model.Model(display_name=_TEST_DISPLAY_NAME) + true_managed_model = gca_model.Model( + display_name=_TEST_DISPLAY_NAME, + version_aliases=["default"], + ) true_input_data_config = gca_training_pipeline.InputDataConfig( dataset_id=mock_dataset_time_series.name, @@ -910,6 +918,7 @@ def test_splits_fraction( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1017,6 +1026,7 @@ def test_splits_timestamp( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1116,6 +1126,7 @@ def test_splits_predefined( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1211,6 +1222,7 @@ def test_splits_default( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( diff --git a/tests/unit/aiplatform/test_automl_image_training_jobs.py b/tests/unit/aiplatform/test_automl_image_training_jobs.py index 563f16613a..861cf41d36 100644 --- a/tests/unit/aiplatform/test_automl_image_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_image_training_jobs.py @@ -178,7 +178,7 @@ def mock_model_service_get(): with mock.patch.object( model_service_client.ModelServiceClient, "get_model" ) as mock_get_model: - mock_get_model.return_value = gca_model.Model() + mock_get_model.return_value = gca_model.Model(name=_TEST_MODEL_NAME) yield mock_get_model @@ -318,6 +318,7 @@ def test_run_call_pipeline_service_create( labels=mock_model._gca_resource.labels, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -402,6 +403,7 @@ def test_run_call_pipeline_service_create_with_timeout( labels=mock_model._gca_resource.labels, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -460,6 +462,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -645,6 +648,7 @@ def test_splits_fraction( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -718,6 +722,7 @@ def test_splits_filter( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -782,6 +787,7 @@ def test_splits_default( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( diff --git a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py index 7a3d3340fc..9464ed2e75 100644 --- a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py @@ -258,7 +258,7 @@ def mock_model_service_get(): with mock.patch.object( model_service_client.ModelServiceClient, "get_model" ) as mock_get_model: - mock_get_model.return_value = gca_model.Model() + mock_get_model.return_value = gca_model.Model(name=_TEST_MODEL_NAME) yield mock_get_model @@ -374,6 +374,7 @@ def test_run_call_pipeline_service_create( display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -459,6 +460,7 @@ def test_run_call_pipeline_service_create_with_timeout( display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -531,6 +533,7 @@ def test_run_call_pipeline_service_create_with_export_eval_data_items( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -612,6 +615,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -682,6 +686,7 @@ def test_run_call_pipeline_service_create_if_no_column_transformations( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -753,6 +758,7 @@ def test_run_call_pipeline_service_create_if_set_additional_experiments( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -817,7 +823,10 @@ def test_run_call_pipeline_service_create_with_column_specs( if not sync: model_from_job.wait() - true_managed_model = gca_model.Model(display_name=_TEST_MODEL_DISPLAY_NAME) + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + version_aliases=["default"], + ) true_input_data_config = gca_training_pipeline.InputDataConfig( dataset_id=mock_dataset_tabular_alternative.name, @@ -927,7 +936,10 @@ def test_run_call_pipeline_service_create_with_column_specs_not_auto( if not sync: model_from_job.wait() - true_managed_model = gca_model.Model(display_name=_TEST_MODEL_DISPLAY_NAME) + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + version_aliases=["default"], + ) true_input_data_config = gca_training_pipeline.InputDataConfig( dataset_id=mock_dataset_tabular_alternative.name, @@ -1254,6 +1266,7 @@ def test_splits_fraction( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1333,6 +1346,7 @@ def test_splits_timestamp( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1406,6 +1420,7 @@ def test_splits_predefined( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1474,6 +1489,7 @@ def test_splits_default( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( diff --git a/tests/unit/aiplatform/test_automl_text_training_jobs.py b/tests/unit/aiplatform/test_automl_text_training_jobs.py index d424cdd654..deb9e72aba 100644 --- a/tests/unit/aiplatform/test_automl_text_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_text_training_jobs.py @@ -166,7 +166,7 @@ def mock_model_service_get(): with mock.patch.object( model_service_client.ModelServiceClient, "get_model" ) as mock_get_model: - mock_get_model.return_value = gca_model.Model() + mock_get_model.return_value = gca_model.Model(name=_TEST_MODEL_NAME) yield mock_get_model @@ -314,6 +314,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -383,6 +384,7 @@ def test_run_call_pipeline_service_create_classification( display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS, encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -463,6 +465,7 @@ def test_run_call_pipeline_service_create_classification_with_timeout( display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS, encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -530,6 +533,7 @@ def test_run_call_pipeline_service_create_extraction( true_managed_model = gca_model.Model( display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -604,7 +608,9 @@ def test_run_call_pipeline_service_create_sentiment( ) true_managed_model = gca_model.Model( - display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -671,6 +677,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( true_managed_model = gca_model.Model( display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -843,6 +850,7 @@ def test_splits_fraction( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -916,6 +924,7 @@ def test_splits_filter( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -980,6 +989,7 @@ def test_splits_default( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( diff --git a/tests/unit/aiplatform/test_automl_video_training_jobs.py b/tests/unit/aiplatform/test_automl_video_training_jobs.py index 94ed81e5b1..8681c4b7f8 100644 --- a/tests/unit/aiplatform/test_automl_video_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_video_training_jobs.py @@ -163,7 +163,7 @@ def mock_model_service_get(): with mock.patch.object( model_service_client.ModelServiceClient, "get_model" ) as mock_get_model: - mock_get_model.return_value = gca_model.Model() + mock_get_model.return_value = gca_model.Model(name=_TEST_MODEL_NAME) yield mock_get_model @@ -281,6 +281,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -361,6 +362,7 @@ def test_splits_fraction( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -433,6 +435,7 @@ def test_splits_filter( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -497,6 +500,7 @@ def test_splits_default( display_name=_TEST_MODEL_DISPLAY_NAME, description=mock_model._gca_resource.description, encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -567,6 +571,7 @@ def test_run_call_pipeline_service_create( labels=_TEST_MODEL_LABELS, description=mock_model._gca_resource.description, encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -648,6 +653,7 @@ def test_run_call_pipeline_service_create_with_timeout( labels=_TEST_MODEL_LABELS, description=mock_model._gca_resource.description, encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -711,6 +717,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( true_managed_model = gca_model.Model( display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( diff --git a/tests/unit/aiplatform/test_end_to_end.py b/tests/unit/aiplatform/test_end_to_end.py index 5f42b7628d..0b9d5c0345 100644 --- a/tests/unit/aiplatform/test_end_to_end.py +++ b/tests/unit/aiplatform/test_end_to_end.py @@ -174,6 +174,8 @@ def test_dataset_create_to_model_predict( true_prediction = models.Prediction( predictions=test_endpoints._TEST_PREDICTION, deployed_model_id=test_endpoints._TEST_ID, + model_resource_name=model_from_job.resource_name, + model_version_id=model_from_job.version_id, ) assert true_prediction == test_prediction @@ -255,6 +257,7 @@ def test_dataset_create_to_model_predict( true_managed_model = gca_model.Model( display_name=test_training_jobs._TEST_MODEL_DISPLAY_NAME, container_spec=true_container_spec, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -453,6 +456,7 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( display_name=test_training_jobs._TEST_MODEL_DISPLAY_NAME, container_spec=true_container_spec, encryption_spec=_TEST_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index e14c20ece6..07696464dc 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/tests/unit/aiplatform/test_endpoints.py @@ -41,6 +41,7 @@ endpoint_service_client, prediction_service_client, ) + from google.cloud.aiplatform.compat.types import ( endpoint as gca_endpoint, model as gca_model, @@ -76,6 +77,9 @@ _TEST_MODEL_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_ID}" ) + +_TEST_VERSION_ID = "1" + _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_ID}" _TEST_MODEL_ID = "1028944691210842416" @@ -419,7 +423,9 @@ def predict_client_predict_mock(): prediction_service_client.PredictionServiceClient, "predict" ) as predict_mock: predict_mock.return_value = gca_prediction_service.PredictResponse( - deployed_model_id=_TEST_MODEL_ID + deployed_model_id=_TEST_MODEL_ID, + model_version_id=_TEST_VERSION_ID, + model=_TEST_MODEL_NAME, ) predict_mock.return_value.predictions.extend(_TEST_PREDICTION) yield predict_mock @@ -1689,7 +1695,10 @@ def test_predict(self, predict_client_predict_mock): ) true_prediction = models.Prediction( - predictions=_TEST_PREDICTION, deployed_model_id=_TEST_ID + predictions=_TEST_PREDICTION, + deployed_model_id=_TEST_ID, + model_version_id=_TEST_VERSION_ID, + model_resource_name=_TEST_MODEL_NAME, ) assert true_prediction == test_prediction diff --git a/tests/unit/aiplatform/test_jobs.py b/tests/unit/aiplatform/test_jobs.py index e69378f55b..ca1d68fe35 100644 --- a/tests/unit/aiplatform/test_jobs.py +++ b/tests/unit/aiplatform/test_jobs.py @@ -66,6 +66,10 @@ _TEST_MODEL_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_ALT_ID}" ) + +_TEST_MODEL_VERSION_ID = "2" +_TEST_VERSIONED_MODEL_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_ALT_ID}@{_TEST_MODEL_VERSION_ID}" + _TEST_BATCH_PREDICTION_JOB_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/batchPredictionJobs/{_TEST_ID}" _TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME = "test-batch-prediction-job" @@ -907,3 +911,39 @@ def test_batch_predict_wrong_prediction_format(self): ) assert e.match(regexp=r"accepted prediction format") + + @pytest.mark.usefixtures("get_batch_prediction_job_mock") + def test_batch_predict_job_with_versioned_model( + self, create_batch_prediction_job_mock + ): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + # Make SDK batch_predict method call + _ = jobs.BatchPredictionJob.create( + model_name=_TEST_VERSIONED_MODEL_NAME, + job_display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, + gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, + sync=True, + ) + assert ( + create_batch_prediction_job_mock.call_args_list[0][1][ + "batch_prediction_job" + ].model + == _TEST_VERSIONED_MODEL_NAME + ) + + # Make SDK batch_predict method call + _ = jobs.BatchPredictionJob.create( + model_name=f"{_TEST_ALT_ID}@{_TEST_MODEL_VERSION_ID}", + job_display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, + gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, + sync=True, + ) + assert ( + create_batch_prediction_job_mock.call_args_list[0][1][ + "batch_prediction_job" + ].model + == _TEST_VERSIONED_MODEL_NAME + ) diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index 85d6b80d8d..be4fe5996d 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -28,6 +28,17 @@ from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer +from google.cloud.aiplatform_v1 import ( + AddContextArtifactsAndExecutionsResponse, + LineageSubgraph, + Artifact as GapicArtifact, + Context as GapicContext, + Execution as GapicExecution, + MetadataServiceClient, + AddExecutionEventsResponse, + MetadataStore as GapicMetadataStore, + TensorboardServiceClient, +) from google.cloud.aiplatform.compat.types import event as gca_event from google.cloud.aiplatform.compat.types import execution as gca_execution from google.cloud.aiplatform.compat.types import ( @@ -47,16 +58,6 @@ from google.cloud.aiplatform.metadata import utils as metadata_utils from google.cloud.aiplatform import utils -from google.cloud.aiplatform_v1 import AddContextArtifactsAndExecutionsResponse -from google.cloud.aiplatform_v1 import AddExecutionEventsResponse -from google.cloud.aiplatform_v1 import Artifact as GapicArtifact -from google.cloud.aiplatform_v1 import Context as GapicContext -from google.cloud.aiplatform_v1 import Execution as GapicExecution -from google.cloud.aiplatform_v1 import LineageSubgraph -from google.cloud.aiplatform_v1 import MetadataServiceClient -from google.cloud.aiplatform_v1 import MetadataStore as GapicMetadataStore -from google.cloud.aiplatform_v1 import TensorboardServiceClient - from test_pipeline_jobs import mock_pipeline_service_get # noqa: F401 from test_pipeline_jobs import _TEST_PIPELINE_JOB_NAME # noqa: F401 diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py index f78f3a1a92..71fb7e3baa 100644 --- a/tests/unit/aiplatform/test_metadata_resources.py +++ b/tests/unit/aiplatform/test_metadata_resources.py @@ -28,15 +28,15 @@ from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import context from google.cloud.aiplatform.metadata import execution -from google.cloud.aiplatform_v1 import AddContextArtifactsAndExecutionsResponse -from google.cloud.aiplatform_v1 import Artifact as GapicArtifact -from google.cloud.aiplatform_v1 import Context as GapicContext -from google.cloud.aiplatform_v1 import Execution as GapicExecution -from google.cloud.aiplatform_v1 import LineageSubgraph from google.cloud.aiplatform_v1 import ( MetadataServiceClient, AddExecutionEventsResponse, Event, + LineageSubgraph, + Execution as GapicExecution, + Context as GapicContext, + Artifact as GapicArtifact, + AddContextArtifactsAndExecutionsResponse, ) # project diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index 20c4e6b909..fef109de18 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/tests/unit/aiplatform/test_models.py @@ -27,7 +27,7 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform -from google.cloud.aiplatform import base +from google.cloud.aiplatform import base, explain from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models from google.cloud.aiplatform import utils @@ -36,8 +36,9 @@ endpoint_service_client, model_service_client, job_service_client, + pipeline_service_client, ) -from google.cloud.aiplatform.compat.services import pipeline_service_client + from google.cloud.aiplatform.compat.types import ( batch_prediction_job as gca_batch_prediction_job, io as gca_io, @@ -54,7 +55,7 @@ encryption_spec as gca_encryption_spec, ) -from google.protobuf import field_mask_pb2 +from google.protobuf import field_mask_pb2, timestamp_pb2 from test_endpoints import create_endpoint_mock # noqa: F401 @@ -63,7 +64,14 @@ _TEST_LOCATION = "us-central1" _TEST_LOCATION_2 = "europe-west4" _TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" -_TEST_MODEL_NAME = "test-model" +_TEST_MODEL_NAME = "123" +_TEST_MODEL_NAME_ALT = "456" +_TEST_MODEL_PARENT = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_MODEL_NAME}" +) +_TEST_MODEL_PARENT_ALT = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/models/{_TEST_MODEL_NAME_ALT}" +) _TEST_ARTIFACT_URI = "gs://test/artifact/uri" _TEST_SERVING_CONTAINER_IMAGE = "gcr.io/test-serving/container:image" _TEST_SERVING_CONTAINER_PREDICTION_ROUTE = "predict" @@ -115,7 +123,8 @@ _TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) _TEST_SERVICE_ACCOUNT = "vinnys@my-project.iam.gserviceaccount.com" -_TEST_EXPLANATION_METADATA = aiplatform.explain.ExplanationMetadata( + +_TEST_EXPLANATION_METADATA = explain.ExplanationMetadata( inputs={ "features": { "input_tensor_name": "dense_input", @@ -126,7 +135,7 @@ }, outputs={"medv": {"output_tensor_name": "dense_2"}}, ) -_TEST_EXPLANATION_PARAMETERS = aiplatform.explain.ExplanationParameters( +_TEST_EXPLANATION_PARAMETERS = explain.ExplanationParameters( {"sampled_shapley_attribution": {"path_count": 10}} ) @@ -232,6 +241,63 @@ ), ] +_TEST_VERSION_ID = "2" +_TEST_VERSION_ALIAS_1 = "myalias" +_TEST_VERSION_ALIAS_2 = "youralias" +_TEST_MODEL_VERSION_DESCRIPTION = "My version description" + +_TEST_MODEL_VERSIONS_LIST = [ + gca_model.Model( + version_id="1", + create_time=timestamp_pb2.Timestamp(), + update_time=timestamp_pb2.Timestamp(), + display_name=_TEST_MODEL_NAME, + name=f"{_TEST_MODEL_PARENT}@1", + version_aliases=["default"], + version_description=_TEST_MODEL_VERSION_DESCRIPTION, + ), + gca_model.Model( + version_id="2", + create_time=timestamp_pb2.Timestamp(), + update_time=timestamp_pb2.Timestamp(), + display_name=_TEST_MODEL_NAME, + name=f"{_TEST_MODEL_PARENT}@2", + version_aliases=[_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2], + version_description=_TEST_MODEL_VERSION_DESCRIPTION, + ), + gca_model.Model( + version_id="3", + create_time=timestamp_pb2.Timestamp(), + update_time=timestamp_pb2.Timestamp(), + display_name=_TEST_MODEL_NAME, + name=f"{_TEST_MODEL_PARENT}@3", + version_aliases=[], + version_description=_TEST_MODEL_VERSION_DESCRIPTION, + ), +] + +_TEST_MODELS_LIST = _TEST_MODEL_VERSIONS_LIST + [ + gca_model.Model( + version_id="1", + create_time=timestamp_pb2.Timestamp(), + update_time=timestamp_pb2.Timestamp(), + display_name=_TEST_MODEL_NAME_ALT, + name=_TEST_MODEL_PARENT_ALT, + version_aliases=["default"], + version_description=_TEST_MODEL_VERSION_DESCRIPTION, + ), +] + +_TEST_MODEL_OBJ_WITH_VERSION = gca_model.Model( + version_id=_TEST_VERSION_ID, + create_time=timestamp_pb2.Timestamp(), + update_time=timestamp_pb2.Timestamp(), + display_name=_TEST_MODEL_NAME, + name=f"{_TEST_MODEL_PARENT}@{_TEST_VERSION_ID}", + version_aliases=[_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2], + version_description=_TEST_MODEL_VERSION_DESCRIPTION, +) + _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_ID}" @@ -242,9 +308,7 @@ def mock_model(): model._latest_future = None model._exception = None model._gca_resource = gca_model.Model( - display_name=_TEST_MODEL_NAME, - description=_TEST_DESCRIPTION, - labels=_TEST_LABEL, + display_name=_TEST_MODEL_NAME, description=_TEST_DESCRIPTION, labels=_TEST_LABEL ) yield model @@ -350,6 +414,20 @@ def get_model_with_supported_export_formats_artifact(): yield get_model_mock +@pytest.fixture +def get_model_with_supported_export_formats_artifact_and_version(): + 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, + version_id=_TEST_VERSION_ID, + ) + yield get_model_mock + + @pytest.fixture def get_model_with_both_supported_export_formats(): with mock.patch.object( @@ -376,6 +454,15 @@ def get_model_with_unsupported_export_formats(): yield get_model_mock +@pytest.fixture +def get_model_with_version(): + with mock.patch.object( + model_service_client.ModelServiceClient, "get_model" + ) as get_model_mock: + get_model_mock.return_value = _TEST_MODEL_OBJ_WITH_VERSION + yield get_model_mock + + @pytest.fixture def upload_model_mock(): with mock.patch.object( @@ -389,6 +476,19 @@ def upload_model_mock(): yield upload_model_mock +@pytest.fixture +def upload_model_with_version_mock(): + with mock.patch.object( + model_service_client.ModelServiceClient, "upload_model" + ) as upload_model_mock: + mock_lro = mock.Mock(ga_operation.Operation) + mock_lro.result.return_value = gca_model_service.UploadModelResponse( + model=_TEST_MODEL_RESOURCE_NAME, model_version_id=_TEST_VERSION_ID + ) + upload_model_mock.return_value = mock_lro + yield upload_model_mock + + @pytest.fixture def upload_model_with_custom_project_mock(): with mock.patch.object( @@ -506,6 +606,7 @@ def create_client_mock(): initializer.global_config, "create_client" ) as create_client_mock: api_client_mock = mock.Mock(spec=model_service_client.ModelServiceClient) + api_client_mock.get_model.return_value = _TEST_MODEL_OBJ_WITH_VERSION create_client_mock.return_value = api_client_mock yield create_client_mock @@ -542,6 +643,43 @@ def list_model_evaluations_mock(): yield list_model_evaluations_mock +@pytest.fixture +def list_model_versions_mock(): + with mock.patch.object( + model_service_client.ModelServiceClient, "list_model_versions" + ) as list_model_versions_mock: + list_model_versions_mock.return_value = _TEST_MODEL_VERSIONS_LIST + yield list_model_versions_mock + + +@pytest.fixture +def list_models_mock(): + with mock.patch.object( + model_service_client.ModelServiceClient, "list_models" + ) as list_models_mock: + list_models_mock.return_value = _TEST_MODELS_LIST + yield list_models_mock + + +@pytest.fixture +def delete_model_version_mock(): + with mock.patch.object( + model_service_client.ModelServiceClient, "delete_model_version" + ) as delete_model_version_mock: + mock_lro = mock.Mock(ga_operation.Operation) + delete_model_version_mock.return_value = mock_lro + yield delete_model_version_mock + + +@pytest.fixture +def merge_version_aliases_mock(): + with mock.patch.object( + model_service_client.ModelServiceClient, "merge_version_aliases" + ) as merge_version_aliases_mock: + merge_version_aliases_mock.return_value = _TEST_MODEL_OBJ_WITH_VERSION + yield merge_version_aliases_mock + + @pytest.mark.usefixtures("google_auth_mock") class TestModel: def setup_method(self): @@ -559,7 +697,7 @@ def test_constructor_creates_client(self, create_client_mock): credentials=_TEST_CREDENTIALS, ) models.Model(_TEST_ID) - create_client_mock.assert_called_once_with( + create_client_mock.assert_any_call( client_class=utils.ModelClientWithOverride, credentials=initializer.global_config.credentials, location_override=_TEST_LOCATION, @@ -572,7 +710,7 @@ def test_constructor_create_client_with_custom_location(self, create_client_mock credentials=_TEST_CREDENTIALS, ) models.Model(_TEST_ID, location=_TEST_LOCATION_2) - create_client_mock.assert_called_once_with( + create_client_mock.assert_any_call( client_class=utils.ModelClientWithOverride, credentials=initializer.global_config.credentials, location_override=_TEST_LOCATION_2, @@ -583,7 +721,7 @@ def test_constructor_creates_client_with_custom_credentials( ): creds = auth_credentials.AnonymousCredentials() models.Model(_TEST_ID, credentials=creds) - create_client_mock.assert_called_once_with( + create_client_mock.assert_any_call( client_class=utils.ModelClientWithOverride, credentials=creds, location_override=_TEST_LOCATION, @@ -639,11 +777,14 @@ def test_upload_uploads_and_gets_model( managed_model = gca_model.Model( display_name=_TEST_MODEL_NAME, container_spec=container_spec, + version_aliases=["default"], ) upload_model_mock.assert_called_once_with( - parent=initializer.global_config.common_location_path(), - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=initializer.global_config.common_location_path(), + model=managed_model, + ), timeout=None, ) @@ -670,11 +811,14 @@ def test_upload_with_timeout(self, upload_model_mock, get_model_mock, sync): managed_model = gca_model.Model( display_name=_TEST_MODEL_NAME, container_spec=container_spec, + version_aliases=["default"], ) upload_model_mock.assert_called_once_with( - parent=initializer.global_config.common_location_path(), - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=initializer.global_config.common_location_path(), + model=managed_model, + ), timeout=180.0, ) @@ -698,11 +842,14 @@ def test_upload_with_timeout_not_explicitly_set( managed_model = gca_model.Model( display_name=_TEST_MODEL_NAME, container_spec=container_spec, + version_aliases=["default"], ) upload_model_mock.assert_called_once_with( - parent=initializer.global_config.common_location_path(), - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=initializer.global_config.common_location_path(), + model=managed_model, + ), timeout=None, ) @@ -734,11 +881,14 @@ def test_upload_uploads_and_gets_model_with_labels( display_name=_TEST_MODEL_NAME, container_spec=container_spec, labels=_TEST_LABEL, + version_aliases=["default"], ) upload_model_mock.assert_called_once_with( - parent=initializer.global_config.common_location_path(), - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=initializer.global_config.common_location_path(), + model=managed_model, + ), timeout=None, ) @@ -823,11 +973,14 @@ def test_upload_uploads_and_gets_model_with_all_args( parameters=_TEST_EXPLANATION_PARAMETERS, ), labels=_TEST_LABEL, + version_aliases=["default"], ) upload_model_mock.assert_called_once_with( - parent=initializer.global_config.common_location_path(), - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=initializer.global_config.common_location_path(), + model=managed_model, + ), timeout=None, ) get_model_mock.assert_called_once_with( @@ -871,11 +1024,14 @@ def test_upload_uploads_and_gets_model_with_custom_project( display_name=_TEST_MODEL_NAME, artifact_uri=_TEST_ARTIFACT_URI, container_spec=container_spec, + version_aliases=["default"], ) upload_model_with_custom_project_mock.assert_called_once_with( - parent=f"projects/{_TEST_PROJECT_2}/locations/{_TEST_LOCATION}", - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=f"projects/{_TEST_PROJECT_2}/locations/{_TEST_LOCATION}", + model=managed_model, + ), timeout=None, ) @@ -962,11 +1118,14 @@ def test_upload_uploads_and_gets_model_with_custom_location( display_name=_TEST_MODEL_NAME, artifact_uri=_TEST_ARTIFACT_URI, container_spec=container_spec, + version_aliases=["default"], ) upload_model_with_custom_location_mock.assert_called_once_with( - parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION_2}", - model=managed_model, + request=gca_model_service.UploadModelRequest( + parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION_2}", + model=managed_model, + ), timeout=None, ) @@ -1086,6 +1245,43 @@ def test_deploy_with_timeout_not_explicitly_set(self, deploy_model_mock, sync): timeout=None, ) + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_with_version") + @pytest.mark.parametrize("sync", [True, False]) + def test_deploy_with_version(self, deploy_model_mock, sync): + + test_model = models.Model(_TEST_MODEL_NAME) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES + ) + version = _TEST_MODEL_OBJ_WITH_VERSION.version_id + + test_endpoint = models.Endpoint(_TEST_ID) + + test_endpoint = test_model.deploy( + test_endpoint, + sync=sync, + ) + + if not sync: + test_endpoint.wait() + + automatic_resources = gca_machine_resources.AutomaticResources( + min_replica_count=1, + max_replica_count=1, + ) + deployed_model = gca_endpoint.DeployedModel( + automatic_resources=automatic_resources, + model=f"{test_model.resource_name}@{version}", + display_name=None, + ) + deploy_model_mock.assert_called_once_with( + endpoint=test_endpoint.resource_name, + deployed_model=deployed_model, + traffic_split={"0": 100}, + metadata=(), + timeout=None, + ) + @pytest.mark.usefixtures( "get_endpoint_mock", "get_model_mock", "create_endpoint_mock" ) @@ -1335,13 +1531,29 @@ def test_batch_predict_gcs_source_and_dest( if not sync: batch_prediction_job.wait() + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_model_with_version", "get_batch_prediction_job_mock") + def test_batch_predict_with_version(self, sync, create_batch_prediction_job_mock): + + test_model = models.Model(_TEST_MODEL_NAME, version=_TEST_VERSION_ALIAS_1) + + # Make SDK batch_predict method call + batch_prediction_job = test_model.batch_predict( + job_display_name=_TEST_BATCH_PREDICTION_DISPLAY_NAME, + gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, + gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, + sync=sync, + create_request_timeout=None, + ) + + if not sync: + batch_prediction_job.wait() + # Construct expected request expected_gapic_batch_prediction_job = ( gca_batch_prediction_job.BatchPredictionJob( display_name=_TEST_BATCH_PREDICTION_DISPLAY_NAME, - model=model_service_client.ModelServiceClient.model_path( - _TEST_PROJECT, _TEST_LOCATION, _TEST_ID - ), + model=f"{_TEST_MODEL_PARENT}@{_TEST_VERSION_ID}", input_config=gca_batch_prediction_job.BatchPredictionJob.InputConfig( instances_format="jsonl", gcs_source=gca_io.GcsSource( @@ -1628,6 +1840,33 @@ def test_export_model_as_artifact(self, export_model_mock, sync): output_config=expected_output_config, ) + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures( + "get_model_with_supported_export_formats_artifact_and_version" + ) + def test_export_model_with_version(self, export_model_mock, sync): + test_model = models.Model(f"{_TEST_ID}@{_TEST_VERSION_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}@{_TEST_VERSION_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): @@ -1820,7 +2059,7 @@ def test_upload_xgboost_model_file_uploads_and_gets_model( upload_model_mock.assert_called_once() upload_model_call_kwargs = upload_model_mock.call_args[1] - upload_model_model = upload_model_call_kwargs["model"] + upload_model_model = upload_model_call_kwargs["request"].model # Verifying the container image selection assert ( @@ -1878,7 +2117,7 @@ def test_upload_stages_data_uploads_and_gets_model( upload_model_mock.assert_called_once() upload_model_call_kwargs = upload_model_mock.call_args[1] - upload_model_model = upload_model_call_kwargs["model"] + upload_model_model = upload_model_call_kwargs["request"].model # Verifying the staging bucket name generation assert upload_model_model.artifact_uri.startswith( @@ -1931,7 +2170,7 @@ def test_upload_scikit_learn_model_file_uploads_and_gets_model( upload_model_mock.assert_called_once() upload_model_call_kwargs = upload_model_mock.call_args[1] - upload_model_model = upload_model_call_kwargs["model"] + upload_model_model = upload_model_call_kwargs["request"].model # Verifying the container image selection assert ( @@ -1981,7 +2220,7 @@ def test_upload_tensorflow_saved_model_uploads_and_gets_model( upload_model_mock.assert_called_once() upload_model_call_kwargs = upload_model_mock.call_args[1] - upload_model_model = upload_model_call_kwargs["model"] + upload_model_model = upload_model_call_kwargs["request"].model # Verifying the container image selection assert ( @@ -2072,3 +2311,230 @@ def test_list_model_evaluations( ) assert len(eval_list) == len(_TEST_MODEL_EVAL_LIST) + + def test_init_with_version_in_resource_name(self, get_model_with_version): + model = models.Model( + model_name=models.ModelRegistry._get_versioned_name( + _TEST_MODEL_NAME, _TEST_VERSION_ALIAS_1 + ) + ) + + assert model.version_aliases == [_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2] + assert model.display_name == _TEST_MODEL_NAME + assert model.resource_name == _TEST_MODEL_PARENT + assert model.version_id == _TEST_VERSION_ID + assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION + # The Model yielded from upload should not have a version in resource name + assert "@" not in model.resource_name + # The Model yielded from upload SHOULD have a version in the versioned resource name + assert model.versioned_resource_name.endswith(f"@{_TEST_VERSION_ID}") + + def test_init_with_version_arg(self, get_model_with_version): + model = models.Model(model_name=_TEST_MODEL_NAME, version=_TEST_VERSION_ID) + + assert model.version_aliases == [_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2] + assert model.display_name == _TEST_MODEL_NAME + assert model.resource_name == _TEST_MODEL_PARENT + assert model.version_id == _TEST_VERSION_ID + assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION + # The Model yielded from upload should not have a version in resource name + assert "@" not in model.resource_name + # The Model yielded from upload SHOULD have a version in the versioned resource name + assert model.versioned_resource_name.endswith(f"@{_TEST_VERSION_ID}") + + @pytest.mark.parametrize( + "parent,location,project", + [ + (_TEST_MODEL_NAME, _TEST_LOCATION, _TEST_PROJECT), + (_TEST_MODEL_PARENT, None, None), + ], + ) + @pytest.mark.parametrize( + "aliases,default,goal", + [ + (["alias1", "alias2"], True, ["alias1", "alias2", "default"]), + (None, True, ["default"]), + (["alias1", "alias2", "default"], True, ["alias1", "alias2", "default"]), + (["alias1", "alias2", "default"], False, ["alias1", "alias2", "default"]), + (["alias1", "alias2"], False, ["alias1", "alias2"]), + (None, False, []), + ], + ) + @pytest.mark.parametrize( + "callable, model_file_path, saved_model", + [ + (models.Model.upload, None, None), + (models.Model.upload_scikit_learn_model_file, "my_model.pkl", None), + (models.Model.upload_tensorflow_saved_model, None, "saved_model.pb"), + (models.Model.upload_xgboost_model_file, "my_model.xgb", None), + ], + ) + def test_upload_new_version( + self, + upload_model_with_version_mock, + get_model_with_version, + mock_storage_blob_upload_from_filename, + parent, + location, + project, + aliases, + default, + goal, + callable, + model_file_path, + saved_model, + tmp_path: pathlib.Path, + ): + args = { + "display_name": _TEST_MODEL_NAME, + "location": location, + "project": project, + "sync": True, + "upload_request_timeout": None, + "model_id": _TEST_ID, + "parent_model": parent, + "version_description": _TEST_MODEL_VERSION_DESCRIPTION, + "version_aliases": aliases, + "is_default_version": default, + } + if model_file_path: + model_file_path = tmp_path / model_file_path + model_file_path.touch() + args["model_file_path"] = str(model_file_path) + elif saved_model: + saved_model_dir = tmp_path / "saved_model" + saved_model_dir.mkdir() + (saved_model_dir / saved_model).touch() + args["saved_model_dir"] = str(saved_model_dir) + else: + args["serving_container_image_uri"] = _TEST_SERVING_CONTAINER_IMAGE + + _ = callable(**args) + + upload_model_with_version_mock.assert_called_once() + upload_model_call_kwargs = upload_model_with_version_mock.call_args[1] + upload_model_request = upload_model_call_kwargs["request"] + + assert upload_model_request.model.display_name == _TEST_MODEL_NAME + assert upload_model_request.model.version_aliases == goal + assert ( + upload_model_request.model.version_description + == _TEST_MODEL_VERSION_DESCRIPTION + ) + assert upload_model_request.parent_model == _TEST_MODEL_PARENT + assert upload_model_request.model_id == _TEST_ID + + def test_get_model_instance_from_registry(self, get_model_with_version): + registry = models.ModelRegistry(_TEST_MODEL_PARENT) + model = registry.get_model(_TEST_VERSION_ALIAS_1) + assert model.version_aliases == [_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2] + assert model.display_name == _TEST_MODEL_NAME + assert model.resource_name == _TEST_MODEL_PARENT + assert model.version_id == _TEST_VERSION_ID + assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION + + def test_list_versions(self, list_model_versions_mock, get_model_with_version): + my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION) + versions = my_model.versioning_registry.list_versions() + + assert len(versions) == len(_TEST_MODEL_VERSIONS_LIST) + + for i in range(len(versions)): + ver = versions[i] + model = _TEST_MODEL_VERSIONS_LIST[i] + assert ver.version_id == model.version_id + assert ver.version_create_time == model.version_create_time + assert ver.version_update_time == model.version_update_time + assert ver.model_display_name == model.display_name + assert ver.version_aliases == model.version_aliases + assert ver.version_description == model.version_description + + assert model.name.startswith(ver.model_resource_name) + assert model.name.endswith(ver.version_id) + + def test_get_version_info(self, get_model_with_version): + my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION) + ver = my_model.versioning_registry.get_version_info("2") + model = _TEST_MODEL_OBJ_WITH_VERSION + + assert ver.version_id == model.version_id + assert ver.version_create_time == model.version_create_time + assert ver.version_update_time == model.version_update_time + assert ver.model_display_name == model.display_name + assert ver.version_aliases == model.version_aliases + assert ver.version_description == model.version_description + + assert model.name.startswith(ver.model_resource_name) + assert model.name.endswith(ver.version_id) + + def test_delete_version(self, delete_model_version_mock, get_model_with_version): + my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION) + my_model.versioning_registry.delete_version(_TEST_VERSION_ALIAS_1) + + delete_model_version_mock.assert_called_once_with( + name=models.ModelRegistry._get_versioned_name( + _TEST_MODEL_PARENT, _TEST_VERSION_ALIAS_1 + ) + ) + + def test_add_versions(self, merge_version_aliases_mock, get_model_with_version): + my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION) + my_model.versioning_registry.add_version_aliases( + ["new-alias", "other-new-alias"], _TEST_VERSION_ALIAS_1 + ) + + merge_version_aliases_mock.assert_called_once_with( + name=models.ModelRegistry._get_versioned_name( + _TEST_MODEL_PARENT, _TEST_VERSION_ALIAS_1 + ), + version_aliases=["new-alias", "other-new-alias"], + ) + + def test_remove_versions(self, merge_version_aliases_mock, get_model_with_version): + my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION) + my_model.versioning_registry.remove_version_aliases( + ["old-alias", "other-old-alias"], _TEST_VERSION_ALIAS_1 + ) + + merge_version_aliases_mock.assert_called_once_with( + name=models.ModelRegistry._get_versioned_name( + _TEST_MODEL_PARENT, _TEST_VERSION_ALIAS_1 + ), + version_aliases=["-old-alias", "-other-old-alias"], + ) + + @pytest.mark.parametrize( + "resource", + [ + "abc", + "abc@1", + "abc@my-alias", + pytest.param("@5", marks=pytest.mark.xfail), + pytest.param("abc@", marks=pytest.mark.xfail), + pytest.param("abc#alias", marks=pytest.mark.xfail), + ], + ) + def test_model_resource_id_validator(self, resource): + models.Model._revisioned_resource_id_validator(resource) + + def test_list(self, list_models_mock): + models_list = models.Model.list() + + assert len(models_list) == len(_TEST_MODELS_LIST) + + for i in range(len(models_list)): + listed_model = models_list[i] + ideal_model = _TEST_MODELS_LIST[i] + assert listed_model.version_id == ideal_model.version_id + assert listed_model.version_create_time == ideal_model.version_create_time + assert listed_model.version_update_time == ideal_model.version_update_time + assert listed_model.display_name == ideal_model.display_name + assert listed_model.version_aliases == ideal_model.version_aliases + assert listed_model.version_description == ideal_model.version_description + + assert ideal_model.name.startswith(listed_model.resource_name) + if "@" in ideal_model.name: + assert ideal_model.name.endswith(listed_model.version_id) + + assert listed_model.versioning_registry + assert listed_model._revisioned_resource_id_validator diff --git a/tests/unit/aiplatform/test_training_jobs.py b/tests/unit/aiplatform/test_training_jobs.py index 17d338887f..8e474e1edb 100644 --- a/tests/unit/aiplatform/test_training_jobs.py +++ b/tests/unit/aiplatform/test_training_jobs.py @@ -47,9 +47,9 @@ from google.cloud.aiplatform.utils import worker_spec_utils from google.cloud.aiplatform.compat.services import ( + job_service_client, model_service_client, pipeline_service_client, - job_service_client, ) from google.cloud.aiplatform.compat.types import ( @@ -97,7 +97,7 @@ _TEST_BASE_OUTPUT_DIR = "gs://test-base-output-dir" _TEST_SERVICE_ACCOUNT = "vinnys@my-project.iam.gserviceaccount.com" -_TEST_BIGQUERY_DESTINATION = "bq://test-project" +_TEST_BIGQUERY_DESTINATION = "bq://my-project" _TEST_RUN_ARGS = ["-v", 0.1, "--test=arg"] _TEST_REPLICA_COUNT = 1 _TEST_MACHINE_TYPE = "n1-standard-4" @@ -128,7 +128,7 @@ _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" -_TEST_ID = "12345" +_TEST_ID = "1028944691210842416" _TEST_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/{_TEST_ID}" ) @@ -138,6 +138,8 @@ _TEST_CUSTOM_JOB_RESOURCE_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/customJobs/{_TEST_ID}" ) +_TEST_MODEL_VERSION_DESCRIPTION = "My version description" +_TEST_MODEL_VERSION_ID = "2" _TEST_ALT_PROJECT = "test-project-alt" _TEST_ALT_LOCATION = "europe-west4" _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_ID}" @@ -160,10 +162,10 @@ _TEST_OUTPUT_PYTHON_PACKAGE_PATH = "gs://test/ouput/python/trainer.tar.gz" _TEST_PYTHON_MODULE_NAME = "aiplatform.task" -_TEST_MODEL_NAME = "projects/my-project/locations/us-central1/models/12345" +_TEST_MODEL_NAME = f"projects/{_TEST_PROJECT}/locations/us-central1/models/{_TEST_ID}" _TEST_PIPELINE_RESOURCE_NAME = ( - "projects/my-project/locations/us-central1/trainingPipelines/12345" + f"projects/{_TEST_PROJECT}/locations/us-central1/trainingPipelines/{_TEST_ID}" ) _TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) @@ -551,6 +553,23 @@ def mock_pipeline_service_create(): yield mock_create_training_pipeline +@pytest.fixture +def mock_pipeline_service_create_with_version(): + with mock.patch.object( + pipeline_service_client.PipelineServiceClient, "create_training_pipeline" + ) as mock_create_training_pipeline: + mock_create_training_pipeline.return_value = ( + gca_training_pipeline.TrainingPipeline( + name=_TEST_PIPELINE_RESOURCE_NAME, + state=gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED, + model_to_upload=gca_model.Model( + name=_TEST_MODEL_NAME, version_id=_TEST_MODEL_VERSION_ID + ), + ) + ) + yield mock_create_training_pipeline + + def make_training_pipeline(state, add_training_task_metadata=True): return gca_training_pipeline.TrainingPipeline( name=_TEST_PIPELINE_RESOURCE_NAME, @@ -563,6 +582,20 @@ def make_training_pipeline(state, add_training_task_metadata=True): ) +def make_training_pipeline_with_version(state, add_training_task_metadata=True): + return gca_training_pipeline.TrainingPipeline( + name=_TEST_PIPELINE_RESOURCE_NAME, + state=state, + model_to_upload=gca_model.Model( + name=_TEST_MODEL_NAME, version_id=_TEST_MODEL_VERSION_ID + ), + training_task_inputs={"tensorboard": _TEST_TENSORBOARD_RESOURCE_NAME}, + training_task_metadata={"backingCustomJob": _TEST_CUSTOM_JOB_RESOURCE_NAME} + if add_training_task_metadata + else None, + ) + + def make_training_pipeline_with_no_model_upload(state): return gca_training_pipeline.TrainingPipeline( name=_TEST_PIPELINE_RESOURCE_NAME, @@ -600,42 +633,26 @@ def make_training_pipeline_with_scheduling(state): @pytest.fixture -def mock_pipeline_service_get(): +def mock_pipeline_service_get(make_call=make_training_pipeline): with mock.patch.object( pipeline_service_client.PipelineServiceClient, "get_training_pipeline" ) as mock_get_training_pipeline: mock_get_training_pipeline.side_effect = [ - make_training_pipeline( + make_call( gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING, add_training_task_metadata=False, ), - make_training_pipeline( + make_call( gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING, ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_training_pipeline( - gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED - ), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_call(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), ] yield mock_get_training_pipeline @@ -795,6 +812,21 @@ def mock_model_service_get(): mock_get_model.return_value.supported_deployment_resources_types.append( aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES ) + mock_get_model.return_value.version_id = "1" + yield mock_get_model + + +@pytest.fixture +def mock_model_service_get_with_version(): + with mock.patch.object( + model_service_client.ModelServiceClient, "get_model" + ) as mock_get_model: + mock_get_model.return_value = gca_model.Model( + name=_TEST_MODEL_NAME, version_id=_TEST_MODEL_VERSION_ID + ) + mock_get_model.return_value.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ) yield mock_get_model @@ -987,6 +1019,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1229,6 +1262,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_and_timeout( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1390,6 +1424,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_and_timeout_not_e prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1537,6 +1572,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -1833,7 +1869,9 @@ def test_run_call_pipeline_service_create_with_no_dataset( ) true_managed_model = gca_model.Model( - display_name=_TEST_MODEL_DISPLAY_NAME, container_spec=true_container_spec + display_name=_TEST_MODEL_DISPLAY_NAME, + container_spec=true_container_spec, + version_aliases=["default"], ) true_training_pipeline = gca_training_pipeline.TrainingPipeline( @@ -2232,6 +2270,7 @@ def test_run_call_pipeline_service_create_distributed_training( parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -2648,6 +2687,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -2963,6 +3003,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -3130,6 +3171,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_and_timeout( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -3296,6 +3338,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_and_timeout_not_e prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -3443,6 +3486,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -3717,7 +3761,9 @@ def test_run_call_pipeline_service_create_with_no_dataset( ) true_managed_model = gca_model.Model( - display_name=_TEST_MODEL_DISPLAY_NAME, container_spec=true_container_spec + display_name=_TEST_MODEL_DISPLAY_NAME, + container_spec=true_container_spec, + version_aliases=["default"], ) true_training_pipeline = gca_training_pipeline.TrainingPipeline( @@ -4095,6 +4141,7 @@ def test_run_call_pipeline_service_create_distributed_training( parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -4370,6 +4417,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset( parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -4907,6 +4955,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -5080,6 +5129,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_with_timeout( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -5236,6 +5286,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_with_timeout_not_ prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -5376,6 +5427,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_without_model_dis prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -5539,6 +5591,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -5819,7 +5872,9 @@ def test_run_call_pipeline_service_create_with_no_dataset( ) true_managed_model = gca_model.Model( - display_name=_TEST_MODEL_DISPLAY_NAME, container_spec=true_container_spec + display_name=_TEST_MODEL_DISPLAY_NAME, + container_spec=true_container_spec, + version_aliases=["default"], ) true_training_pipeline = gca_training_pipeline.TrainingPipeline( @@ -6208,6 +6263,7 @@ def test_run_call_pipeline_service_create_distributed_training( parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -6486,6 +6542,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, ), + version_aliases=["default"], ) true_input_data_config = gca_training_pipeline.InputDataConfig( @@ -6578,3 +6635,131 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_raises_if_anno accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, ) + + +class TestVersionedTrainingJobs: + @pytest.mark.usefixtures("mock_pipeline_service_get") + @pytest.mark.parametrize( + "mock_pipeline_service_get", + ["make_training_pipeline_with_version"], + indirect=True, + ) + @pytest.mark.parametrize( + "parent,location,project,model_id", + [ + (_TEST_ID, _TEST_LOCATION, _TEST_PROJECT, None), + (_TEST_MODEL_NAME, None, None, None), + (None, None, None, _TEST_ID), + ], + ) + @pytest.mark.parametrize( + "aliases,default,goal", + [ + (["alias1", "alias2"], True, ["alias1", "alias2", "default"]), + (None, True, ["default"]), + (["alias1", "alias2", "default"], True, ["alias1", "alias2", "default"]), + (["alias1", "alias2", "default"], False, ["alias1", "alias2", "default"]), + (["alias1", "alias2"], False, ["alias1", "alias2"]), + (None, False, []), + ], + ) + @pytest.mark.parametrize( + "callable", + [ + training_jobs.CustomTrainingJob, + training_jobs.CustomContainerTrainingJob, + training_jobs.CustomPythonPackageTrainingJob, + ], + ) + def test_run_pipeline_for_versioned_model( + self, + mock_pipeline_service_create_with_version, + mock_python_package_to_gcs, + mock_nontabular_dataset, + mock_model_service_get_with_version, + parent, + location, + project, + model_id, + aliases, + default, + goal, + callable, + ): + aiplatform.init( + project=project, + staging_bucket=_TEST_BUCKET_NAME, + credentials=_TEST_CREDENTIALS, + location=location, + ) + job_args = { + "display_name": _TEST_DISPLAY_NAME, + "model_serving_container_image_uri": _TEST_SERVING_CONTAINER_IMAGE, + "model_serving_container_predict_route": _TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + "model_serving_container_health_route": _TEST_SERVING_CONTAINER_HEALTH_ROUTE, + "model_instance_schema_uri": _TEST_MODEL_INSTANCE_SCHEMA_URI, + "model_parameters_schema_uri": _TEST_MODEL_PARAMETERS_SCHEMA_URI, + "model_prediction_schema_uri": _TEST_MODEL_PREDICTION_SCHEMA_URI, + "model_serving_container_command": _TEST_MODEL_SERVING_CONTAINER_COMMAND, + "model_serving_container_args": _TEST_MODEL_SERVING_CONTAINER_ARGS, + "model_serving_container_environment_variables": _TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES, + "model_serving_container_ports": _TEST_MODEL_SERVING_CONTAINER_PORTS, + "model_description": _TEST_MODEL_DESCRIPTION, + "labels": _TEST_LABELS, + } + + run_args = { + "dataset": mock_nontabular_dataset, + "annotation_schema_uri": _TEST_ANNOTATION_SCHEMA_URI, + "base_output_dir": _TEST_BASE_OUTPUT_DIR, + "args": _TEST_RUN_ARGS, + "machine_type": _TEST_MACHINE_TYPE, + "accelerator_type": _TEST_ACCELERATOR_TYPE, + "accelerator_count": _TEST_ACCELERATOR_COUNT, + "training_filter_split": _TEST_TRAINING_FILTER_SPLIT, + "validation_filter_split": _TEST_VALIDATION_FILTER_SPLIT, + "test_filter_split": _TEST_TEST_FILTER_SPLIT, + "create_request_timeout": None, + "model_id": model_id, + "parent_model": parent, + "is_default_version": default, + "model_version_aliases": aliases, + "model_version_description": _TEST_MODEL_VERSION_DESCRIPTION, + } + + if issubclass(callable, (training_jobs.CustomContainerTrainingJob)): + job_args = { + "container_uri": _TEST_TRAINING_CONTAINER_IMAGE, + **job_args, + } + elif issubclass(callable, (training_jobs.CustomTrainingJob)): + job_args = { + "container_uri": _TEST_TRAINING_CONTAINER_IMAGE, + "script_path": _TEST_LOCAL_SCRIPT_FILE_NAME, + **job_args, + } + elif issubclass(callable, training_jobs.CustomPythonPackageTrainingJob): + job_args = { + "python_package_gcs_uri": _TEST_OUTPUT_PYTHON_PACKAGE_PATH, + "python_module_name": _TEST_PYTHON_MODULE_NAME, + "container_uri": _TEST_TRAINING_CONTAINER_IMAGE, + **job_args, + } + + job = callable(**job_args) + + model_from_job = job.run(**run_args) + + mock_pipeline_service_create_with_version.assert_called_once() + _, tp_kwargs = mock_pipeline_service_create_with_version.call_args_list[0] + training_pipeline = tp_kwargs["training_pipeline"] + + assert training_pipeline.model_id == (model_id if model_id else "") + assert training_pipeline.parent_model == (_TEST_MODEL_NAME if parent else "") + assert training_pipeline.model_to_upload.version_aliases == goal + assert ( + training_pipeline.model_to_upload.version_description + == _TEST_MODEL_VERSION_DESCRIPTION + ) + + assert model_from_job.version_id == _TEST_MODEL_VERSION_ID