use resouce name instead of id and populate metadata with resourceNa…

…nme accoridngly
googleapis · Jun 23, 2022 · 1778a34 · 1778a34
1 parent 04c9e88
commit 1778a34
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 69 deletions.
diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py
@@ -363,7 +363,7 @@ def create_from_base_artifact_schema(
         """
 
         return cls._create(
-            resource_id=base_artifact_schema.resource_id,
+            resource_id=base_artifact_schema.artifact_id,
             schema_title=base_artifact_schema.schema_title,
             uri=base_artifact_schema.uri,
             display_name=base_artifact_schema.display_name,

diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py
@@ -216,7 +216,7 @@ def create_from_base_execution_schema(
                 metadata_store=metadata_store_id,
             ),
             schema_title=base_execution_schema.schema_title,
-            resource_id=base_execution_schema.resource_id,
+            resource_id=base_execution_schema.execution_id,
             metadata=base_execution_schema.metadata,
             description=base_execution_schema.description,
             display_name=base_execution_schema.display_name,

diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -38,7 +38,7 @@ def schema_title(cls) -> str:
 
     def __init__(
         self,
-        resource_id: Optional[str] = None,
+        artifact_id: Optional[str] = None,
         uri: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
@@ -80,7 +80,7 @@ def __init__(
                 Pipelines), and the system does not prescribe or
                 check the validity of state transitions.
         """
-        self.resource_id = resource_id
+        self.artifact_id = artifact_id
         self.uri = uri
         self.display_name = display_name
         self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION

diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py
@@ -40,7 +40,7 @@ def schema_title(cls) -> str:
     def __init__(
         self,
         state: gca_execution.Execution.State = gca_execution.Execution.State.RUNNING,
-        resource_id: Optional[str] = None,
+        execution_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         metadata: Optional[Dict] = None,
@@ -52,7 +52,7 @@ def __init__(
         Args:
             state (gca_execution.Execution.State.RUNNING):
                 Optional. State of this Execution. Defaults to RUNNING.
-            resource_id (str):
+            execution_id (str):
                 Optional. The <resource_id> portion of the Execution name with
                 the following format, this is globally unique in a metadataStore.
                 projects/123/locations/us-central1/metadataStores/<metadata_store_id>/executions/<resource_id>.
@@ -67,7 +67,7 @@ def __init__(
                 Optional. Describes the purpose of the Execution to be created.
         """
         self.state = state
-        self.resource_id = resource_id
+        self.execution_id = execution_id
         self.display_name = display_name
         self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
         self.metadata = metadata

diff --git a/google/cloud/aiplatform/metadata/schema/google/artifact_schema.py b/google/cloud/aiplatform/metadata/schema/google/artifact_schema.py
@@ -20,6 +20,9 @@
 from google.cloud.aiplatform.metadata.schema import base_artifact
 from google.cloud.aiplatform.metadata.schema import utils
 
+# The artifact property key for the resource_name
+_ARTIFACT_PROPERTY_KEY_RESOURCE_NAME = "resourceName"
+
 
 class VertexDataset(base_artifact.BaseArtifactSchema):
     """An artifact representing a Vertex Dataset."""
@@ -28,23 +31,30 @@ class VertexDataset(base_artifact.BaseArtifactSchema):
 
     def __init__(
         self,
-        dataset_id: Optional[str] = None,
-        uri: Optional[str] = None,
+        dataset_name: str,
+        uri: str,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
         metadata: Optional[Dict] = None,
         state: gca_artifact.Artifact.State = gca_artifact.Artifact.State.LIVE,
     ):
         """Args:
-        dataset_id (str):
-            Optional. The <resource_id> portion of the Artifact name, in a form of
-            projects/{project}/locations/{location}/datasets/{datasets_id}. For
+        dataset_name (str):
+            The name of the Dataset resource, in a form of
+            projects/{project}/locations/{location}/datasets/{dataset}. For
             more details, see
             https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.datasets/get
         uri (str):
-            Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
-            artifact file.
+            The Vertex Dataset resource uri, in a form of
+            https://{service-endpoint}/v1/{dataset_name},
+            where {service-endpoint} is one of the supported service endpoints at
+            https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the Artifact.
         schema_version (str):
@@ -63,9 +73,11 @@ def __init__(
             check the validity of state transitions.
         """
         extended_metadata = metadata or {}
+        extended_metadata[_ARTIFACT_PROPERTY_KEY_RESOURCE_NAME] = dataset_name
+
         super(VertexDataset, self).__init__(
-            resource_id=dataset_id,
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,
@@ -81,23 +93,30 @@ class VertexModel(base_artifact.BaseArtifactSchema):
 
     def __init__(
         self,
-        vertex_model_id: Optional[str] = None,
-        uri: Optional[str] = None,
+        vertex_model_name: str,
+        uri: str,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
         metadata: Optional[Dict] = None,
         state: gca_artifact.Artifact.State = gca_artifact.Artifact.State.LIVE,
     ):
         """Args:
-        vertex_model_id (str):
-            Optional. The <resource_id> portion of the Artifact name, in a form of
-            projects/{project}/locations/{location}/models/{model_id}. For
+        vertex_model_name (str):
+            The name of the Model resource, in a form of
+            projects/{project}/locations/{location}/models/{model}. For
             more details, see
             https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.models/get
         uri (str):
-            Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
-            artifact file.
+            The Vertex Model resource uri, in a form of
+            https://{service-endpoint}/v1/{vertex_model_name},
+            where {service-endpoint} is one of the supported service endpoints at
+            https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the Artifact.
         schema_version (str):
@@ -116,10 +135,11 @@ def __init__(
             check the validity of state transitions.
         """
         extended_metadata = metadata or {}
+        extended_metadata[_ARTIFACT_PROPERTY_KEY_RESOURCE_NAME] = vertex_model_name
 
         super(VertexModel, self).__init__(
-            resource_id=vertex_model_id,
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,
@@ -135,8 +155,9 @@ class VertexEndpoint(base_artifact.BaseArtifactSchema):
 
     def __init__(
         self,
-        vertex_endpoint_id: Optional[str] = None,
-        uri: Optional[str] = None,
+        vertex_endpoint_name: str,
+        uri: str,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
@@ -145,13 +166,19 @@ def __init__(
     ):
         """Args:
         vertex_endpoint_name (str):
-            Optional. The <resource_id> portion of the Artifact name, in a form of
-            projects/{project}/locations/{location}/endpoints/{endpoint_id}. For
+            The name of the Endpoint resource, in a form of
+            projects/{project}/locations/{location}/endpoints/{endpoint}. For
             more details, see
             https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/get
         uri (str):
-            Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
-            artifact file.
+            The Vertex Endpoint resource uri, in a form of
+            https://{service-endpoint}/v1/{vertex_endpoint_name},
+            where {service-endpoint} is one of the supported service endpoints at
+            https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the Artifact.
         schema_version (str):
@@ -171,9 +198,11 @@ def __init__(
         """
         extended_metadata = metadata or {}
 
+        extended_metadata[_ARTIFACT_PROPERTY_KEY_RESOURCE_NAME] = vertex_endpoint_name
+
         super(VertexEndpoint, self).__init__(
-            resource_id=vertex_endpoint_id,
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,
@@ -191,7 +220,7 @@ def __init__(
         self,
         predict_schema_ta: utils.PredictSchemata,
         container_spec: utils.ContainerSpec,
-        unmanaged_container_model_id: Optional[str] = None,
+        artifact_id: Optional[str] = None,
         uri: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
@@ -204,7 +233,7 @@ def __init__(
             An instance of PredictSchemata which holds instance, parameter and prediction schema uris.
         container_spec (ContainerSpec):
             An instance of ContainerSpec which holds the container configuration for the model.
-        unmanaged_container_model_id (str):
+        artifact_id (str):
             Optional. The <resource_id> portion of the Artifact name with
             the format. This is globally unique in a metadataStore:
             projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
@@ -233,8 +262,8 @@ def __init__(
         extended_metadata["containerSpec"] = container_spec.to_dict()
 
         super(UnmanagedContainerModel, self).__init__(
-            resource_id=unmanaged_container_model_id,
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,

diff --git a/google/cloud/aiplatform/metadata/schema/system/artifact_schema.py b/google/cloud/aiplatform/metadata/schema/system/artifact_schema.py
@@ -29,6 +29,7 @@ class Model(base_artifact.BaseArtifactSchema):
     def __init__(
         self,
         uri: Optional[str] = None,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
@@ -39,6 +40,10 @@ def __init__(
         uri (str):
             Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
             artifact file.
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the base.
         schema_version (str):
@@ -56,13 +61,13 @@ def __init__(
             Pipelines), and the system does not prescribe or
             check the validity of state transitions.
         """
-        extended_metadata = metadata or {}
         super(Model, self).__init__(
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,
-            metadata=extended_metadata,
+            metadata=metadata,
             state=state,
         )
 
@@ -75,6 +80,7 @@ class Artifact(base_artifact.BaseArtifactSchema):
     def __init__(
         self,
         uri: Optional[str] = None,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
@@ -85,6 +91,10 @@ def __init__(
         uri (str):
             Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
             artifact file.
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the base.
         schema_version (str):
@@ -102,13 +112,13 @@ def __init__(
             Pipelines), and the system does not prescribe or
             check the validity of state transitions.
         """
-        extended_metadata = metadata or {}
         super(Artifact, self).__init__(
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,
-            metadata=extended_metadata,
+            metadata=metadata,
             state=state,
         )
 
@@ -121,6 +131,7 @@ class Dataset(base_artifact.BaseArtifactSchema):
     def __init__(
         self,
         uri: Optional[str] = None,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
@@ -131,6 +142,10 @@ def __init__(
         uri (str):
             Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
             artifact file.
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the base.
         schema_version (str):
@@ -148,13 +163,13 @@ def __init__(
             Pipelines), and the system does not prescribe or
             check the validity of state transitions.
         """
-        extended_metadata = metadata or {}
         super(Dataset, self).__init__(
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,
-            metadata=extended_metadata,
+            metadata=metadata,
             state=state,
         )
 
@@ -173,6 +188,7 @@ def __init__(
         mean_absolute_error: Optional[float] = None,
         mean_squared_error: Optional[float] = None,
         uri: Optional[str] = None,
+        artifact_id: Optional[str] = None,
         display_name: Optional[str] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
@@ -195,6 +211,10 @@ def __init__(
         uri (str):
             Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
             artifact file.
+        artifact_id (str):
+            Optional. The <resource_id> portion of the Artifact name with
+            the format. This is globally unique in a metadataStore:
+            projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
         display_name (str):
             Optional. The user-defined name of the base.
         schema_version (str):
@@ -228,6 +248,7 @@ def __init__(
 
         super(Metrics, self).__init__(
             uri=uri,
+            artifact_id=artifact_id,
             display_name=display_name,
             schema_version=schema_version,
             description=description,