switch to use ModelContainerSpec and PredictSchemata instead of custom_dataclasses

Author: SinaChavoshi

google/cloud/aiplatform/metadata/types/google_types.py

@@ -17,7 +17,7 @@
 from typing import Optional, Dict, NamedTuple, List
 from dataclasses import dataclass
 from google.cloud.aiplatform.metadata.types import base
-from google.cloud.aiplatform.metadata.types import utils
+from google.cloud.aiplatform.compat.types import model_v1 as model
 
 
 class VertexDataset(base.BaseArtifactSchema):
@@ -177,8 +177,8 @@ class UnmanagedContainerModel(base.BaseArtifactSchema):
 
     def __init__(
         self,
-        predict_schema_ta: utils.PredictSchemata,
-        container_spec: utils.PredictSchemata,
+        predict_schema_ta: model.PredictSchemata,
+        container_spec: model.ModelContainerSpec,
         unmanaged_container_model_name: Optional[str] = None,
         uri: Optional[str] = None,
         display_name: Optional[str] = None,
@@ -213,8 +213,35 @@ def __init__(
         """
         extended_metadata = metadata or {}
         extended_metadata["resourceName"] = unmanaged_container_model_name
-        extended_metadata["predictSchemata"] = predict_schema_ta.to_dict()
-        extended_metadata["containerSpec"] = container_spec.to_dict()
+        extended_metadata["predictSchemata"] = {}
+        extended_metadata["predictSchemata"][
+            "instanceSchemaUri"
+        ] = predict_schema_ta.instance_schema_uri
+        extended_metadata["predictSchemata"][
+            "parametersSchemaUri"
+        ] = predict_schema_ta.parameters_schema_uri
+        extended_metadata["predictSchemata"][
+            "predictionSchemaUri"
+        ] = predict_schema_ta.prediction_schema_uri
+
+        extended_metadata["containerSpec"] = {}
+        extended_metadata["containerSpec"]["imageUri"] = container_spec.image_uri
+        if container_spec.command:
+            extended_metadata["containerSpec"]["command"] = container_spec.command
+        if container_spec.args:
+            extended_metadata["containerSpec"]["args"] = container_spec.args
+        if container_spec.env:
+            extended_metadata["containerSpec"]["env"] = container_spec.env
+        if container_spec.ports:
+            extended_metadata["containerSpec"]["ports"] = container_spec.ports
+        if container_spec.predict_route:
+            extended_metadata["containerSpec"][
+                "predictRoute"
+            ] = container_spec.predict_route
+        if container_spec.health_route:
+            extended_metadata["containerSpec"][
+                "healthRoute"
+            ] = container_spec.health_route
 
         super(UnmanagedContainerModel, self).__init__(
             schema_title=self.SCHEMA_TITLE,

google/cloud/aiplatform/metadata/types/system_types.py

@@ -16,7 +16,6 @@
 #
 from typing import Optional, Dict, List
 from google.cloud.aiplatform.metadata.types import base
-from google.cloud.aiplatform.metadata.types import utils
 from itertools import zip_longest
 
 

google/cloud/aiplatform/metadata/types/utils.py

@@ -27,8 +27,8 @@
 from google.cloud.aiplatform.metadata.types import base
 from google.cloud.aiplatform.metadata.types import google_types
 from google.cloud.aiplatform.metadata.types import system_types
-from google.cloud.aiplatform.metadata.types import utils
 
+from google.cloud.aiplatform.compat.types import model_v1 as model
 from google.cloud.aiplatform_v1 import MetadataServiceClient
 from google.cloud.aiplatform_v1 import Artifact as GapicArtifact
 
@@ -212,13 +212,13 @@ def test_vertex_endpoint_constructor_parameters_are_set_correctly(self):
         assert artifact.schema_version == _TEST_SCHEMA_VERSION
 
     def test_unmanaged_container_model_title_is_set_correctly(self):
-        predict_schema_ta = utils.PredictSchemata(
+        predict_schema_ta = model.PredictSchemata(
             instance_schema_uri="instance_uri",
             prediction_schema_uri="prediction_uri",
             parameters_schema_uri="parameters_uri",
         )
 
-        container_spec = utils.ContainerSpec(
+        container_spec = model.ModelContainerSpec(
             image_uri="gcr.io/test_container_image_uri"
         )
         artifact = google_types.UnmanagedContainerModel(
@@ -228,13 +228,13 @@ def test_unmanaged_container_model_title_is_set_correctly(self):
         assert artifact.schema_title == "google.UnmanagedContainerModel"
 
     def test_unmanaged_container_model_resouce_name_is_set_in_metadata(self):
-        predict_schema_ta = utils.PredictSchemata(
+        predict_schema_ta = model.PredictSchemata(
             instance_schema_uri="instance_uri",
             prediction_schema_uri="prediction_uri",
             parameters_schema_uri="parameters_uri",
         )
 
-        container_spec = utils.ContainerSpec(
+        container_spec = model.ModelContainerSpec(
             image_uri="gcr.io/test_container_image_uri"
         )
         artifact = google_types.UnmanagedContainerModel(
@@ -245,13 +245,13 @@ def test_unmanaged_container_model_resouce_name_is_set_in_metadata(self):
         assert artifact.metadata["resourceName"] == _TEST_ARTIFACT_NAME
 
     def test_unmanaged_container_model_constructor_parameters_are_set_correctly(self):
-        predict_schema_ta = utils.PredictSchemata(
+        predict_schema_ta = model.PredictSchemata(
             instance_schema_uri="instance_uri",
             prediction_schema_uri="prediction_uri",
             parameters_schema_uri="parameters_uri",
         )
 
-        container_spec = utils.ContainerSpec(
+        container_spec = model.ModelContainerSpec(
             image_uri="gcr.io/test_container_image_uri"
         )
 
@@ -360,73 +360,3 @@ def test_system_metrics_constructor_parameters_are_set_correctly(self):
         assert artifact.metadata["f1score"] == 0.4
         assert artifact.metadata["mean_absolute_error"] == 0.5
         assert artifact.metadata["mean_squared_error"] == 0.6
-
-
-class TestMetadataUtils:
-    def setup_method(self):
-        reload(initializer)
-        reload(metadata)
-        reload(aiplatform)
-
-    def teardown_method(self):
-        initializer.global_pool.shutdown(wait=True)
-
-    def test_predict_schemata_to_dict_method_returns_correct_schema(self):
-        predict_schema_ta = utils.PredictSchemata(
-            instance_schema_uri="instance_uri",
-            prediction_schema_uri="prediction_uri",
-            parameters_schema_uri="parameters_uri",
-        )
-        expected_results = {
-            "instanceSchemaUri": "instance_uri",
-            "parametersSchemaUri": "parameters_uri",
-            "predictionSchemaUri": "prediction_uri",
-        }
-
-        assert json.dumps(predict_schema_ta.to_dict()) == json.dumps(expected_results)
-
-    def test_container_spec_to_dict_method_returns_correct_schema(self):
-        container_spec = utils.ContainerSpec(
-            image_uri="gcr.io/some_container_image_uri",
-            command=["test_command"],
-            args=["test_args"],
-            env=[{"env_var_name": "env_var_value"}],
-            ports=[1],
-            predict_route="test_prediction_rout",
-            health_route="test_health_rout",
-        )
-
-        expected_results = {
-            "imageUri": "gcr.io/some_container_image_uri",
-            "command": ["test_command"],
-            "args": ["test_args"],
-            "env": [{"env_var_name": "env_var_value"}],
-            "ports": [1],
-            "predictRoute": "test_prediction_rout",
-            "healthRoute": "test_health_rout",
-        }
-
-        assert json.dumps(container_spec.to_dict()) == json.dumps(expected_results)
-
-    def test_container_spec_to_dict_method_returns_correct_schema(self):
-        container_spec = utils.ContainerSpec(
-            image_uri="gcr.io/some_container_image_uri",
-            command=["test_command"],
-            args=["test_args"],
-            env=[{"env_var_name": "env_var_value"}],
-            ports=[1],
-            predict_route="test_prediction_rout",
-            health_route="test_health_rout",
-        )
-
-        expected_results = {
-            "imageUri": "gcr.io/some_container_image_uri",
-            "command": ["test_command"],
-            "args": ["test_args"],
-            "env": [{"env_var_name": "env_var_value"}],
-            "ports": [1],
-            "predictRoute": "test_prediction_rout",
-            "healthRoute": "test_health_rout",
-        }
-
-        assert json.dumps(container_spec.to_dict()) == json.dumps(expected_results)