feat: Improved metadata artifact and execution creation using python / SDK

google/cloud/aiplatform/metadata/artifact.py

@@ -32,7 +32,7 @@
 from google.cloud.aiplatform.metadata import resource
 from google.cloud.aiplatform.metadata import utils as metadata_utils
 from google.cloud.aiplatform.utils import rest_utils
-
+from google.cloud.aiplatform.metadata.types import base as types_base
 
 _LOGGER = base.Logger(__name__)
 
@@ -176,6 +176,7 @@ def _create(
 
         """
         api_client = cls._instantiate_client(location=location, credentials=credentials)
+        api_client = cls._instantiate_client(location=location, credentials=credentials)
 
         parent = utils.full_resource_name(
             resource_name=metadata_store_id,
@@ -249,8 +250,8 @@ def _list_resources(
     @classmethod
     def create(
         cls,
-        schema_title: str,
         *,
+        schema_title: Optional[str] = None,
         resource_id: Optional[str] = None,
         uri: Optional[str] = None,
         display_name: Optional[str] = None,
@@ -262,12 +263,13 @@ def create(
         project: Optional[str] = None,
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
+        base_artifact: Optional[types_base.BaseArtifactSchema] = None,
     ) -> "Artifact":
         """Creates a new Metadata Artifact.
 
         Args:
             schema_title (str):
-                Required. schema_title identifies the schema title used by the Artifact.
+                Optional. schema_title identifies the schema title used by the Artifact.
 
                 Please reference https://cloud.google.com/vertex-ai/docs/ml-metadata/system-schemas.
             resource_id (str):
@@ -307,10 +309,28 @@ def create(
             credentials (auth_credentials.Credentials):
                 Optional. Custom credentials used to create this Artifact. Overrides
                 credentials set in aiplatform.init.
+            base_artifact (BaseArtifactType):
+                Optional. An instance of the BaseArtifactType class that can be provided instead of providing artifact specific parameters. It overrides
+                the values provided for schema_title, resource_id, uri, display_name, schema_version, description, and metadata.
 
         Returns:
             Artifact: Instantiated representation of the managed Metadata Artifact.
         """
+        if base_artifact:
+            return cls._create(
+                resource_id=base_artifact.resource_id,
+                schema_title=base_artifact.schema_title,
+                uri=base_artifact.uri,
+                display_name=base_artifact.display_name,
+                schema_version=base_artifact.schema_version,
+                description=base_artifact.description,
+                metadata=base_artifact.metadata,
+                metadata_store_id=metadata_store_id,
+                project=project,
+                location=location,
+                credentials=credentials,
+            )
+
         return cls._create(
             resource_id=resource_id,
             schema_title=schema_title,

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

@@ -0,0 +1,126 @@
+# -*- 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.
+#
+
+from google.auth import credentials as auth_credentials
+from typing import Optional, Dict
+from google.cloud.aiplatform.metadata import artifact
+from google.cloud.aiplatform import base
+
+from google.cloud.aiplatform.metadata import constants
+
+_LOGGER = base.Logger(__name__)
+
+
+class BaseArtifactSchema(object):
+    """Base class for Metadata Artifact types.
+
+    This is the base class for defining various artifact types, which can be
+    passed to google.Artifact to create a corresponding resource.
+    Artifacts carry a `metadata` field, which is a dictionary for storing
+    metadata related to this artifact. Subclasses from ArtifactType can enforce
+    various structure and field requirements for the metadata field.
+
+     Args:
+         schema_title (str):
+             Optional. The schema title used by the Artifact, defaults to "system.Artifact"
+         resource_name (str):
+             Optional. The resource name of the Artifact following the format as follows.
+             This is globally unique in a metadataStore:
+             projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>.
+         uri (str):
+             Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual
+             artifact file.
+         display_name (str):
+             Optional. The user-defined name of the Artifact.
+         schema_version (str):
+             Optional. schema_version specifies the version used by the Artifact.
+             If not set, defaults to use the latest version.
+         description (str):
+             Optional. Describes the purpose of the Artifact to be created.
+         metadata (Dict):
+             Optional. Contains the metadata information that will be stored in the Artifact.
+        **kwargs:
+            Optional. Additional Args that will be passed directly to the Artifact base method for backward compatibility.
+    """
+
+    ARTIFACT_PROPERTY_KEY_RESOURCE_NAME = "resourceName"
+    SCHEMA_TITLE = "system.Artifact"
+
+    def __init__(
+        self,
+        schema_title: Optional[str] = None,
+        resource_name: Optional[str] = None,
+        uri: Optional[str] = None,
+        display_name: Optional[str] = None,
+        schema_version: Optional[str] = None,
+        description: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        **kwargs,
+    ):
+
+        """Initializes the Artifact with the given name, URI and metadata."""
+        self.schema_title = BaseArtifactSchema.SCHEMA_TITLE
+        if schema_title:
+            self.schema_title = schema_title
+        self.resource_name = resource_name
+
+        self.resource_id = None
+        if resource_name:
+            # Temporary work around while Artifact.create takes resource_id instead of resource_name
+            self.resource_id = resource_name.split("/")[-1]
+
+        self.uri = uri
+        self.display_name = display_name
+        self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
+        self.description = description
+        self.metadata = metadata
+
+    def create(
+        self,
+        metadata_store_id: Optional[str] = "default",
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ):
+        """Creates a new Metadata Artifact.
+
+        Args:
+            metadata_store_id (str):
+                Optional. The <metadata_store_id> portion of the resource name with
+                the format:
+                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/artifacts/<resource_id>
+                If not provided, the MetadataStore's ID will be set to "default".
+            project (str):
+                Optional. Project used to create this Artifact. Overrides project set in
+                aiplatform.init.
+            location (str):
+                Optional. Location used to create this Artifact. Overrides location set in
+                aiplatform.init.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials used to create this Artifact. Overrides
+                credentials set in aiplatform.init.
+        Returns:
+            Artifact: Instantiated representation of the managed Metadata Artifact.
+        """
+        self.artifact = artifact.Artifact.create(
+            base_artifact=self,
+            metadata_store_id=metadata_store_id,
+            project=project,
+            location=location,
+            credentials=credentials,
+        )
+        return self.artifact