feat: Improved metadata artifact and execution creation using python / SDK (#1430)

* rebase to master after ga merge

* add support for artifact create

* add unit tests for create from artifact parameters

* update formatting

* fix lint issues

* Add integration tests

* add support for execution types

* correct execution type in create

* add execution support and unit tests

* add support for state in artifacts

* add support for start_execution

* add support for metadata_store_id in start_execution

* lint and docs update based on review feedback

* Add e2e integraton tests and lint update

* Update google/cloud/aiplatform/metadata/artifact.py

Co-authored-by: sasha-gitg <[email protected]>

* remove the duplicate test_experiments_copy.py

* refactor based on code review feedback

* regroup tests to match module names

* fix e2e integration tests

* remove call to     _temp_prefix = tmpvrtxsdk-e2e from E2E test

* Update google/cloud/aiplatform/metadata/schema/base_execution.py

Co-authored-by: sasha-gitg <[email protected]>

* remove artifact and schema referencing the create result to self

* remove kwargs

* fix typing for container spec

* remove resouceName from system types

* metrics should default to None

* change from using resouce_name to resource_id

* fix e2e tests

* change google and system to sub folders of schema

* use create_from_base_execution_schema instead of overloading create

* update api docs

* update docstring formatting

* Update google/cloud/aiplatform/metadata/metadata.py

Co-authored-by: sasha-gitg <[email protected]>

* add return types and move args to constructor

* using forward reference for parameter and return types to resolve circular import error

* change base classes to abstract classes

* Add tests for system.artifact type

* use resouce name instead of id and populate metadata with resourceNanme accoridngly

* remove start_execution from this pr and move to a separate PR

* change all args to keyword args

* always make a copy of metadata instead of pass by reference

* auto generate uri for google types

* fix e2e tests

* switch to using Artifact.create instead of _create

* change typing for state to Optional

* change typing for state to Optional in artifact base file

* change to use the Execution.create instead of the private method

* chagne copy to deepcopy for metadata

Co-authored-by: sasha-gitg <[email protected]>
Co-authored-by: Rosie Zou <[email protected]>

Author: 3 people

google/cloud/aiplatform/metadata/artifact.py

@@ -31,6 +31,7 @@
 from google.cloud.aiplatform.metadata import metadata_store
 from google.cloud.aiplatform.metadata import resource
 from google.cloud.aiplatform.metadata import utils as metadata_utils
+from google.cloud.aiplatform.metadata.schema import base_artifact
 from google.cloud.aiplatform.utils import rest_utils
 
 
@@ -326,6 +327,56 @@ def create(
             credentials=credentials,
         )
 
+    @classmethod
+    def create_from_base_artifact_schema(
+        cls,
+        *,
+        base_artifact_schema: "base_artifact.BaseArtifactSchema",
+        metadata_store_id: Optional[str] = "default",
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> "Artifact":
+        """Creates a new Metadata Artifact from a BaseArtifactSchema class instance.
+
+        Args:
+            base_artifact_schema (BaseArtifactSchema):
+                Required. An instance of the BaseArtifactType class that can be
+                provided instead of providing artifact specific parameters.
+            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.
+        """
+
+        return cls.create(
+            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,
+            schema_version=base_artifact_schema.schema_version,
+            description=base_artifact_schema.description,
+            metadata=base_artifact_schema.metadata,
+            state=base_artifact_schema.state,
+            metadata_store_id=metadata_store_id,
+            project=project,
+            location=location,
+            credentials=credentials,
+        )
+
     @property
     def uri(self) -> Optional[str]:
         "Uri for this Artifact."

google/cloud/aiplatform/metadata/execution.py

@@ -31,6 +31,7 @@
 from google.cloud.aiplatform.metadata import artifact
 from google.cloud.aiplatform.metadata import metadata_store
 from google.cloud.aiplatform.metadata import resource
+from google.cloud.aiplatform.metadata.schema import base_execution
 
 
 class Execution(resource._Resource):
@@ -166,6 +167,57 @@ def create(
 
         return self
 
+    @classmethod
+    def create_from_base_execution_schema(
+        cls,
+        *,
+        base_execution_schema: "base_execution.BaseExecutionSchema",
+        metadata_store_id: Optional[str] = "default",
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> "Execution":
+        """
+        Creates a new Metadata Execution.
+
+        Args:
+            base_execution_schema (BaseExecutionSchema):
+                An instance of the BaseExecutionSchema class that can be
+                provided instead of providing schema specific parameters.
+            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 Execution. Overrides project set in
+                aiplatform.init.
+            location (str):
+                Optional. Location used to create this Execution. Overrides location set in
+                aiplatform.init.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials used to create this Execution. Overrides
+                credentials set in aiplatform.init.
+
+        Returns:
+            Execution: Instantiated representation of the managed Metadata Execution.
+
+        """
+        resource = Execution.create(
+            state=base_execution_schema.state,
+            schema_title=base_execution_schema.schema_title,
+            resource_id=base_execution_schema.execution_id,
+            display_name=base_execution_schema.display_name,
+            schema_version=base_execution_schema.schema_version,
+            metadata=base_execution_schema.metadata,
+            description=base_execution_schema.description,
+            metadata_store_id=metadata_store_id,
+            project=project,
+            location=location,
+            credentials=credentials,
+        )
+        return resource
+
     def __enter__(self):
         if self.state is not gca_execution.Execution.State.RUNNING:
             self.update(state=gca_execution.Execution.State.RUNNING)

google/cloud/aiplatform/metadata/schema/base_artifact.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.
+#
+
+import abc
+
+from typing import Optional, Dict
+
+from google.auth import credentials as auth_credentials
+
+from google.cloud.aiplatform.compat.types import artifact as gca_artifact
+from google.cloud.aiplatform.metadata import artifact
+from google.cloud.aiplatform.metadata import constants
+
+
+class BaseArtifactSchema(metaclass=abc.ABCMeta):
+    """Base class for Metadata Artifact types."""
+
+    @property
+    @classmethod
+    @abc.abstractmethod
+    def schema_title(cls) -> str:
+        """Identifies the Vertex Metadta schema title used by the resource."""
+        pass
+
+    def __init__(
+        self,
+        *,
+        artifact_id: 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,
+        state: Optional[gca_artifact.Artifact.State] = gca_artifact.Artifact.State.LIVE,
+    ):
+
+        """Initializes the Artifact with the given name, URI and metadata.
+
+        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:
+            resource_id (str):
+                Optional. The <resource_id> portion of the Artifact name with
+                the following format, 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.
+            state (google.cloud.gapic.types.Artifact.State):
+                Optional. The state of this Artifact. This is a
+                property of the Artifact, and does not imply or
+                capture any ongoing process. This property is
+                managed by clients (such as Vertex AI
+                Pipelines), and the system does not prescribe or
+                check the validity of state transitions.
+        """
+        self.artifact_id = artifact_id
+        self.uri = uri
+        self.display_name = display_name
+        self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
+        self.description = description
+        self.metadata = metadata
+        self.state = state
+
+    def create(
+        self,
+        *,
+        metadata_store_id: Optional[str] = "default",
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> "artifact.Artifact":
+        """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.
+        """
+        return artifact.Artifact.create_from_base_artifact_schema(
+            base_artifact_schema=self,
+            metadata_store_id=metadata_store_id,
+            project=project,
+            location=location,
+            credentials=credentials,
+        )

google/cloud/aiplatform/metadata/schema/base_execution.py

@@ -0,0 +1,114 @@
+# -*- 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.
+#
+
+import abc
+
+from typing import Optional, Dict
+
+from google.auth import credentials as auth_credentials
+
+from google.cloud.aiplatform.compat.types import execution as gca_execution
+from google.cloud.aiplatform.metadata import constants
+from google.cloud.aiplatform.metadata import execution
+
+
+class BaseExecutionSchema(metaclass=abc.ABCMeta):
+    """Base class for Metadata Execution schema."""
+
+    @property
+    @classmethod
+    @abc.abstractmethod
+    def schema_title(cls) -> str:
+        """Identifies the Vertex Metadta schema title used by the resource."""
+        pass
+
+    def __init__(
+        self,
+        *,
+        state: Optional[
+            gca_execution.Execution.State
+        ] = gca_execution.Execution.State.RUNNING,
+        execution_id: Optional[str] = None,
+        display_name: Optional[str] = None,
+        schema_version: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        description: Optional[str] = None,
+    ):
+
+        """Initializes the Execution with the given name, URI and metadata.
+
+        Args:
+            state (gca_execution.Execution.State.RUNNING):
+                Optional. State of this Execution. Defaults to RUNNING.
+            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>.
+            display_name (str):
+                Optional. The user-defined name of the Execution.
+            schema_version (str):
+                Optional. schema_version specifies the version used by the Execution.
+                If not set, defaults to use the latest version.
+            metadata (Dict):
+                Optional. Contains the metadata information that will be stored in the Execution.
+            description (str):
+                Optional. Describes the purpose of the Execution to be created.
+        """
+        self.state = state
+        self.execution_id = execution_id
+        self.display_name = display_name
+        self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
+        self.metadata = metadata
+        self.description = description
+
+    def create(
+        self,
+        *,
+        metadata_store_id: Optional[str] = "default",
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> "execution.Execution":
+        """Creates a new Metadata Execution.
+
+        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>/executions/<resource_id>
+                If not provided, the MetadataStore's ID will be set to "default".
+            project (str):
+                Optional. Project used to create this Execution. Overrides project set in
+                aiplatform.init.
+            location (str):
+                Optional. Location used to create this Execution. Overrides location set in
+                aiplatform.init.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials used to create this Execution. Overrides
+                credentials set in aiplatform.init.
+        Returns:
+            Execution: Instantiated representation of the managed Metadata Execution.
+
+        """
+        self.execution = execution.Execution.create_from_base_execution_schema(
+            base_execution_schema=self,
+            metadata_store_id=metadata_store_id,
+            project=project,
+            location=location,
+            credentials=credentials,
+        )
+        return self.execution