docs(samples): add Model Registry samples to Vertex AI Python SDK (#1602)

* docs(samples): add samples for vertex ai model registry python sdk

* linter test fix

* fix list model versions mock test

* nox passed

Co-authored-by: Andrew Ferlitsch <[email protected]>
Co-authored-by: nayaknishant <[email protected]>

Author: 3 people

samples/model-builder/conftest.py

@@ -114,7 +114,7 @@ def mock_create_image_dataset(mock_image_dataset):
 @pytest.fixture
 def mock_create_tabular_dataset(mock_tabular_dataset):
     with patch.object(
-        aiplatform.TabularDataset, "create"
+            aiplatform.TabularDataset, "create"
     ) as mock_create_tabular_dataset:
         mock_create_tabular_dataset.return_value = mock_tabular_dataset
         yield mock_create_tabular_dataset
@@ -123,7 +123,7 @@ def mock_create_tabular_dataset(mock_tabular_dataset):
 @pytest.fixture
 def mock_create_time_series_dataset(mock_time_series_dataset):
     with patch.object(
-        aiplatform.TimeSeriesDataset, "create"
+            aiplatform.TimeSeriesDataset, "create"
     ) as mock_create_time_series_dataset:
         mock_create_time_series_dataset.return_value = mock_time_series_dataset
         yield mock_create_time_series_dataset
@@ -462,7 +462,7 @@ def mock_get_entity_type(mock_entity_type):
 @pytest.fixture
 def mock_create_featurestore(mock_featurestore):
     with patch.object(
-        aiplatform.featurestore.Featurestore, "create"
+            aiplatform.featurestore.Featurestore, "create"
     ) as mock_create_featurestore:
         mock_create_featurestore.return_value = mock_featurestore
         yield mock_create_featurestore
@@ -471,7 +471,7 @@ def mock_create_featurestore(mock_featurestore):
 @pytest.fixture
 def mock_create_entity_type(mock_entity_type):
     with patch.object(
-        aiplatform.featurestore.EntityType, "create"
+            aiplatform.featurestore.EntityType, "create"
     ) as mock_create_entity_type:
         mock_create_entity_type.return_value = mock_entity_type
         yield mock_create_entity_type
@@ -499,7 +499,7 @@ def mock_batch_serve_to_bq(mock_featurestore):
 @pytest.fixture
 def mock_batch_create_features(mock_entity_type):
     with patch.object(
-        mock_entity_type, "batch_create_features"
+            mock_entity_type, "batch_create_features"
     ) as mock_batch_create_features:
         yield mock_batch_create_features
 
@@ -513,7 +513,7 @@ def mock_read_feature_values(mock_entity_type):
 @pytest.fixture
 def mock_import_feature_values(mock_entity_type):
     with patch.object(
-        mock_entity_type, "ingest_from_gcs"
+            mock_entity_type, "ingest_from_gcs"
     ) as mock_import_feature_values:
         yield mock_import_feature_values
 
@@ -644,7 +644,7 @@ def mock_context_list(mock_context):
 @pytest.fixture
 def mock_create_schema_base_context(mock_context):
     with patch.object(
-        aiplatform.metadata.schema.base_context.BaseContextSchema, "create"
+            aiplatform.metadata.schema.base_context.BaseContextSchema, "create"
     ) as mock_create_schema_base_context:
         mock_create_schema_base_context.return_value = mock_context
         yield mock_create_schema_base_context
@@ -702,7 +702,7 @@ def mock_create_artifact(mock_artifact):
 @pytest.fixture
 def mock_create_schema_base_artifact(mock_artifact):
     with patch.object(
-        aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create"
+            aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create"
     ) as mock_create_schema_base_artifact:
         mock_create_schema_base_artifact.return_value = mock_artifact
         yield mock_create_schema_base_artifact
@@ -711,7 +711,7 @@ def mock_create_schema_base_artifact(mock_artifact):
 @pytest.fixture
 def mock_create_schema_base_execution(mock_execution):
     with patch.object(
-        aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create"
+            aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create"
     ) as mock_create_schema_base_execution:
         mock_create_schema_base_execution.return_value = mock_execution
         yield mock_create_schema_base_execution
@@ -757,7 +757,7 @@ def mock_log_metrics():
 @pytest.fixture
 def mock_log_time_series_metrics():
     with patch.object(
-        aiplatform, "log_time_series_metrics"
+            aiplatform, "log_time_series_metrics"
     ) as mock_log_time_series_metrics:
         mock_log_time_series_metrics.return_value = None
         yield mock_log_time_series_metrics
@@ -822,7 +822,75 @@ def mock_get_params(mock_params, mock_experiment_run):
 @pytest.fixture
 def mock_get_time_series_metrics(mock_time_series_metrics, mock_experiment_run):
     with patch.object(
-        mock_experiment_run, "get_time_series_data_frame"
+            mock_experiment_run, "get_time_series_data_frame"
     ) as mock_get_time_series_metrics:
         mock_get_time_series_metrics.return_value = mock_time_series_metrics
         yield mock_get_time_series_metrics
+
+
+"""
+----------------------------------------------------------------------------
+Model Versioning Fixtures
+----------------------------------------------------------------------------
+"""
+
+
+@pytest.fixture
+def mock_model_registry():
+    mock = MagicMock(aiplatform.models.ModelRegistry)
+    yield mock
+
+
+@pytest.fixture
+def mock_version_info():
+    mock = MagicMock(aiplatform.models.VersionInfo)
+    yield mock
+
+
+@pytest.fixture
+def mock_init_model_registry(mock_model_registry):
+    with patch.object(aiplatform.models, "ModelRegistry") as mock:
+        mock.return_value = mock_model_registry
+        yield mock
+
+
+@pytest.fixture
+def mock_get_model(mock_model_registry):
+    with patch.object(mock_model_registry, "get_model") as mock_get_model:
+        mock_get_model.return_value = mock_model
+        yield mock_get_model
+
+
+@pytest.fixture
+def mock_get_model_version_info(mock_model_registry):
+    with patch.object(mock_model_registry, "get_version_info") as mock_get_model_version_info:
+        mock_get_model_version_info.return_value = mock_version_info
+        yield mock_get_model_version_info
+
+
+@pytest.fixture
+def mock_list_versions(mock_model_registry, mock_version_info):
+    with patch.object(mock_model_registry, "list_versions") as mock_list_versions:
+        mock_list_versions.return_value = [mock_version_info, mock_version_info]
+        yield mock_list_versions
+
+
+@pytest.fixture
+def mock_delete_version(mock_model_registry):
+    with patch.object(mock_model_registry, "delete_version") as mock_delete_version:
+        mock_delete_version.return_value = None
+        yield mock_delete_version
+
+
+@pytest.fixture
+def mock_add_version_aliases(mock_model_registry):
+    with patch.object(mock_model_registry, "add_version_aliases") as mock_add_version_aliases:
+        mock_add_version_aliases.return_value = None
+        yield mock_add_version_aliases
+
+
+@pytest.fixture
+def mock_remove_version_aliases(mock_model_registry):
+    with patch.object(mock_model_registry, "remove_version_aliases") as mock_remove_version_aliases:
+        mock_remove_version_aliases.return_value = None
+        yield mock_remove_version_aliases

samples/model-builder/model_registry/assign_aliases_model_version_sample.py

@@ -0,0 +1,51 @@
+# 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
+#
+#     https://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.
+
+# [START aiplatform_model_registry_assign_aliases_model_version_sample]
+
+from typing import List
+
+from google.cloud import aiplatform
+
+
+def assign_aliases_model_version_sample(
+    model_id: str,
+    version_aliases: List[str],
+    version_id: str,
+    project: str,
+    location: str,
+):
+    """
+    Assign aliases to a model version.
+    Args:
+        model_id: The ID of the model.
+        version_aliases: The version aliases to assign.
+        version_id: The version ID of the model to assign the aliases to.
+        project: The project name.
+        location: The location name.
+    Returns
+        None.
+    """
+    # Initialize the client.
+    aiplatform.init(project=project, location=location)
+
+    # Initialize the Model Registry resource with the ID 'model_id'.The parent_name of create method can be also
+    # 'projects/<your-project-id>/locations/<your-region>/models/<your-model-id>'
+    model_registry = aiplatform.models.ModelRegistry(model=model_id)
+
+    # Assign the version aliases to the model with the version 'version_id'.
+    model_registry.add_version_aliases(new_aliases=version_aliases, version=version_id)
+
+
+# [END aiplatform_model_registry_assign_aliases_model_version_sample]

samples/model-builder/model_registry/assign_aliases_model_version_sample_test.py

@@ -0,0 +1,46 @@
+# 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
+#
+#     https://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 assign_aliases_model_version_sample
+
+import test_constants as constants
+
+
+def test_assign_aliases_model_version_sample(
+    mock_sdk_init, mock_init_model_registry, mock_add_version_aliases, mock_model
+):
+
+    # Assign aliases to a model version.
+    assign_aliases_model_version_sample.assign_aliases_model_version_sample(
+        model_id=constants.MODEL_NAME,
+        version_id=constants.VERSION_ID,
+        version_aliases=constants.VERSION_ALIASES,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+    )
+
+    # Check client initialization.
+    mock_sdk_init.assert_called_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    # Check model registry initialization.
+    mock_init_model_registry.assert_called_with(model=constants.MODEL_NAME)
+
+    # Check that the model version was assigned the aliases.
+    mock_add_version_aliases.assert_called_with(
+        new_aliases=constants.VERSION_ALIASES,
+        version=constants.VERSION_ID,
+    )

samples/model-builder/model_registry/create_aliased_model_sample.py

@@ -0,0 +1,45 @@
+# 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
+#
+#     https://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.
+
+# [START aiplatform_model_registry_create_aliased_model_sample]
+
+from google.cloud import aiplatform
+
+
+def create_aliased_model_sample(
+    model_id: str, version_id: str, project: str, location: str
+):
+    """
+    Initialize a Model resource to represent an existing model version with custom alias.
+    Args:
+        model_id: The ID of the model to initialize. Parent resource name of the model is also accepted.
+        version_id: The version ID or version alias of the model to initialize.
+        project: The project.
+        location: The location.
+    Returns:
+        Model resource.
+    """
+    # Initialize the client.
+    aiplatform.init(project=project, location=location)
+
+    # Initialize the Model resource with the ID 'model_id'. The version can be also provided using @ annotation in
+    # the parent resource name:
+    # 'projects/<your-project-id>/locations/<your-region>/models/<your-model-id>@<your-version-id>'.
+
+    aliased_model = aiplatform.Model(model_name=model_id, version=version_id)
+
+    return aliased_model
+
+
+# [END aiplatform_model_registry_create_aliased_model_sample]

samples/model-builder/model_registry/create_aliased_model_sample_test.py

@@ -0,0 +1,38 @@
+# 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
+#
+#     https://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 create_aliased_model_sample
+
+import test_constants as constants
+
+
+def test_create_aliased_model_sample(mock_sdk_init, mock_init_model):
+    # Create a model with alias 'default'.
+    create_aliased_model_sample.create_aliased_model_sample(
+        model_id=constants.MODEL_NAME,
+        version_id=constants.VERSION_ID,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+    )
+
+    # Check client initialization.
+    mock_sdk_init.assert_called_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    # Check that the model was created.
+    mock_init_model.assert_called_with(
+        model_name=constants.MODEL_NAME, version=constants.VERSION_ID
+    )

samples/model-builder/model_registry/create_default_model_sample.py

@@ -0,0 +1,40 @@
+# 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
+#
+#     https://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.
+
+# [START aiplatform_model_registry_create_default_model_sample]
+
+from google.cloud import aiplatform
+
+
+def create_default_model_sample(model_id: str, project: str, location: str):
+    """
+    Initialize a Model resource to represent an existing model version with alias 'default'.
+    Args:
+        model_id: The ID of the model to initialize. Parent resource name of the model is also accepted.
+        project: The project.
+        location: The location.
+    Returns:
+        Model resource.
+    """
+    # Initialize the client.
+    aiplatform.init(project=project, location=location)
+
+    # Initialize the Model resource with the ID 'model_id'. The parent_name of create method can be also
+    # 'projects/<your-project-id>/locations/<your-region>/models/<your-model-id>'
+    default_model = aiplatform.Model(model_name=model_id)
+
+    return default_model
+
+
+# [END aiplatform_model_registry_create_default_model_sample]