feat: Support for Model Versioning

google/cloud/aiplatform/jobs.py

@@ -714,7 +714,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):
@@ -742,7 +744,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

tests/system/aiplatform/test_model_version_management.py

@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 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")
+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 predictions use right version
+        endpoint = model2.deploy(machine_type="n1-standard-2")
+        shared_state["resources"].append(endpoint)
+        predict_response = endpoint.predict(instances=[[0, 0, 0]])
+
+        assert len(predict_response.predictions) == 1
+        assert predict_response.model_version_id == '2'
+        """
+
+        # Test that VersionInfo properties are correct.
+        # Currently, get_version_info and list_versions don't yield identical
+        # resource names at this time due to a Model Registry bug (b/233118690)
+        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]

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(

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(