From cceff0c23a858b86b42a114e355a4e969f9e93cb Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 00:04:25 +0000
Subject: [PATCH 01/10] feat: Refactor code to make shema type classes a
subclass of Artifact, Execution and Context
---
google/cloud/aiplatform/metadata/artifact.py | 58 ++-----------------
google/cloud/aiplatform/metadata/resource.py | 8 +++
.../metadata/schema/base_artifact.py | 40 ++++++++++---
3 files changed, 45 insertions(+), 61 deletions(-)
diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py
index d01e9c9eb6..5f4f94e18b 100644
--- a/google/cloud/aiplatform/metadata/artifact.py
+++ b/google/cloud/aiplatform/metadata/artifact.py
@@ -31,7 +31,6 @@
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
@@ -327,60 +326,15 @@ 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."
- return self.gca_resource.uri
+ return self._gca_resource.uri
+
+ @property
+ def state(self) -> Optional[gca_artifact.Artifact.State]:
+ "The State for this Artifact."
+ return self._gca_resource.state
@classmethod
def get_with_uri(
diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py
index 89c145dcbe..69a8f94a6b 100644
--- a/google/cloud/aiplatform/metadata/resource.py
+++ b/google/cloud/aiplatform/metadata/resource.py
@@ -114,6 +114,14 @@ def schema_title(self) -> str:
def description(self) -> str:
return self._gca_resource.description
+ @property
+ def display_name(self) -> str:
+ return self._gca_resource.display_name
+
+ @property
+ def schema_version(self) -> str:
+ return self._gca_resource.schema_version
+
@classmethod
def get_or_create(
cls,
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index c89d989edd..cce0d5df31 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -16,6 +16,7 @@
#
import abc
+from copy import deepcopy
from typing import Optional, Dict
@@ -26,7 +27,7 @@
from google.cloud.aiplatform.metadata import constants
-class BaseArtifactSchema(metaclass=abc.ABCMeta):
+class BaseArtifactSchema(artifact.Artifact):
"""Base class for Metadata Artifact types."""
@property
@@ -81,13 +82,20 @@ def __init__(
Pipelines), and the system does not prescribe or
check the validity of state transitions.
"""
+ # resource_id is not stored directly in the proto. Create method uses the artifact_id
+ # to along with project_id and location to construct an artifact_name which is stored
+ # in the proto proto message.
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
+ self._gca_resource = gca_artifact.Artifact()
+ self._gca_resource.uri = uri
+ self._gca_resource.display_name = display_name
+ self._gca_resource.schema_version = (
+ schema_version or constants._DEFAULT_SCHEMA_VERSION
+ )
+ self._gca_resource.description = description
+ if metadata:
+ self._gca_resource.metadata = deepcopy(metadata)
+ self._gca_resource.state = state
def create(
self,
@@ -117,10 +125,24 @@ def create(
Returns:
Artifact: Instantiated representation of the managed Metadata Artifact.
"""
- return artifact.Artifact.create_from_base_artifact_schema(
- base_artifact_schema=self,
+
+ # Check if metadata exists to avoid proto read error
+ metadata = None
+ if self._gca_resource.metadata:
+ metadata = self.metadata
+
+ self = artifact.Artifact.create(
+ resource_id=self.artifact_id,
+ schema_title=self.schema_title,
+ uri=self.uri,
+ display_name=self.display_name,
+ schema_version=self.schema_version,
+ description=self.description,
+ metadata=metadata,
+ state=self.state,
metadata_store_id=metadata_store_id,
project=project,
location=location,
credentials=credentials,
)
+ return self
From 2ad95f262f0e05474fb8f9f2166457ac5930bb03 Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 00:52:05 +0000
Subject: [PATCH 02/10] reinstantiate schema class using resouce_name
---
.../metadata/schema/base_artifact.py | 25 ++++++++++++++++++-
1 file changed, 24 insertions(+), 1 deletion(-)
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index cce0d5df31..93b02c4579 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -97,6 +97,24 @@ def __init__(
self._gca_resource.metadata = deepcopy(metadata)
self._gca_resource.state = state
+ def initialize_with_resouce_name(
+ self,
+ *,
+ artifact_name: str,
+ ):
+
+ """Initializes the Artifact instance using an existing resource.
+
+ Args:
+ artifact_name (str):
+ 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>.
+ """
+ # resource_id is not stored directly in the proto. Create method uses the artifact_id
+ # to along with project_id and location to construct an artifact_name which is stored
+ # in the proto proto message.
+ super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name)
+
def create(
self,
*,
@@ -131,7 +149,7 @@ def create(
if self._gca_resource.metadata:
metadata = self.metadata
- self = artifact.Artifact.create(
+ new_artifact_instance = artifact.Artifact.create(
resource_id=self.artifact_id,
schema_title=self.schema_title,
uri=self.uri,
@@ -145,4 +163,9 @@ def create(
location=location,
credentials=credentials,
)
+
+ # Reinstantiate this class using the newly created resouce.
+ self.initialize_with_resouce_name(
+ resouce_name=new_artifact_instance.resouce_name
+ )
return self
From 6f6ffb83aefba3094ef840818ed500271f4a2f94 Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 00:56:31 +0000
Subject: [PATCH 03/10] fix typo in resource_name
---
google/cloud/aiplatform/metadata/schema/base_artifact.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index 93b02c4579..9a744bea95 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -166,6 +166,6 @@ def create(
# Reinstantiate this class using the newly created resouce.
self.initialize_with_resouce_name(
- resouce_name=new_artifact_instance.resouce_name
+ resouce_name=new_artifact_instance.resource_name
)
return self
From 622759cb7d8220aa68ae9f65307a31d582cf0399 Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 00:58:49 +0000
Subject: [PATCH 04/10] change resouce_name to artifact_name
---
google/cloud/aiplatform/metadata/schema/base_artifact.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index 9a744bea95..042befafc0 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -166,6 +166,6 @@ def create(
# Reinstantiate this class using the newly created resouce.
self.initialize_with_resouce_name(
- resouce_name=new_artifact_instance.resource_name
+ artifact_name=new_artifact_instance.resource_name
)
return self
From 67637aa4ef9c4963a3390d9449e3a8136766656c Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 01:45:41 +0000
Subject: [PATCH 05/10] Change constructor to internal method
---
.../cloud/aiplatform/metadata/schema/base_artifact.py | 10 +++-------
1 file changed, 3 insertions(+), 7 deletions(-)
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index 042befafc0..4a6e50c0b0 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -97,7 +97,8 @@ def __init__(
self._gca_resource.metadata = deepcopy(metadata)
self._gca_resource.state = state
- def initialize_with_resouce_name(
+ # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
+ def _init_with_resouce_name(
self,
*,
artifact_name: str,
@@ -110,9 +111,6 @@ def initialize_with_resouce_name(
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>.
"""
- # resource_id is not stored directly in the proto. Create method uses the artifact_id
- # to along with project_id and location to construct an artifact_name which is stored
- # in the proto proto message.
super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name)
def create(
@@ -165,7 +163,5 @@ def create(
)
# Reinstantiate this class using the newly created resouce.
- self.initialize_with_resouce_name(
- artifact_name=new_artifact_instance.resource_name
- )
+ self._init_with_resouce_name(artifact_name=new_artifact_instance.resource_name)
return self
From d370386a13a1308d0558530236dad48e50a25f9b Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 03:46:47 +0000
Subject: [PATCH 06/10] refactor execution.
---
google/cloud/aiplatform/metadata/execution.py | 52 ---------------
.../metadata/schema/base_artifact.py | 8 ++-
.../metadata/schema/base_execution.py | 66 ++++++++++++++++---
3 files changed, 62 insertions(+), 64 deletions(-)
diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py
index 895417fc64..9a85bce36f 100644
--- a/google/cloud/aiplatform/metadata/execution.py
+++ b/google/cloud/aiplatform/metadata/execution.py
@@ -31,7 +31,6 @@
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):
@@ -167,57 +166,6 @@ 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)
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index 4a6e50c0b0..120639817a 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -82,10 +82,12 @@ def __init__(
Pipelines), and the system does not prescribe or
check the validity of state transitions.
"""
- # resource_id is not stored directly in the proto. Create method uses the artifact_id
- # to along with project_id and location to construct an artifact_name which is stored
- # in the proto proto message.
+ # resource_id is not stored in the proto. Create method uses the
+ # resource_id along with project_id and location to construct an
+ # resource_name which is stored in the proto message.
self.artifact_id = artifact_id
+
+ # Store all other attributes using the proto structure.
self._gca_resource = gca_artifact.Artifact()
self._gca_resource.uri = uri
self._gca_resource.display_name = display_name
diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py
index c8d55a0637..d4bdf3f66e 100644
--- a/google/cloud/aiplatform/metadata/schema/base_execution.py
+++ b/google/cloud/aiplatform/metadata/schema/base_execution.py
@@ -16,6 +16,7 @@
#
import abc
+from copy import deepcopy
from typing import Optional, Dict
@@ -27,7 +28,7 @@
from google.cloud.aiplatform.metadata import metadata
-class BaseExecutionSchema(metaclass=abc.ABCMeta):
+class BaseExecutionSchema(execution.Execution):
"""Base class for Metadata Execution schema."""
@property
@@ -69,12 +70,37 @@ def __init__(
description (str):
Optional. Describes the purpose of the Execution to be created.
"""
- self.state = state
+
+ # resource_id is not stored in the proto. Create method uses the
+ # resource_id along with project_id and location to construct an
+ # resource_name which is stored in the proto message.
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
+
+ # Store all other attributes using the proto structure.
+ self._gca_resource = gca_execution.Execution()
+ self._gca_resource.state = state
+ self._gca_resource.display_name = display_name
+ self._gca_resource.schema_version = (
+ schema_version or constants._DEFAULT_SCHEMA_VERSION
+ )
+ if metadata:
+ self._gca_resource.metadata = deepcopy(metadata)
+ self._gca_resource.description = description
+
+ # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
+ def _init_with_resouce_name(
+ self,
+ *,
+ execution_name: str,
+ ):
+
+ """Initializes the Execution instance using an existing resource.
+ Args:
+ execution_name (str):
+ 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>.
+ """
+ super(BaseExecutionSchema, self).__init__(execution_name=execution_name)
def create(
self,
@@ -105,14 +131,30 @@ def create(
Execution: Instantiated representation of the managed Metadata Execution.
"""
- self.execution = execution.Execution.create_from_base_execution_schema(
+ # Check if metadata exists to avoid proto read error
+ metadata = None
+ if self._gca_resource.metadata:
+ metadata = self.metadata
+
+ new_execution_instance = execution.Execution.create(
+ resource_id=self.execution_id,
+ schema_title=self.schema_title,
+ display_name=self.display_name,
+ schema_version=self.schema_version,
+ description=self.description,
+ metadata=metadata,
+ state=self.state,
base_execution_schema=self,
metadata_store_id=metadata_store_id,
project=project,
location=location,
credentials=credentials,
)
- return self.execution
+ # Reinstantiate this class using the newly created resouce.
+ self._init_with_resouce_name(
+ execution_name=new_execution_instance.resource_name
+ )
+ return self
def start_execution(
self,
@@ -172,7 +214,7 @@ def start_execution(
f"metadata_store_id {metadata_store_id} is not supported. Only the default MetadataStore ID is supported."
)
- return metadata._ExperimentTracker().start_execution(
+ new_execution_instance = metadata._ExperimentTracker().start_execution(
schema_title=self.schema_title,
display_name=self.display_name,
resource_id=self.execution_id,
@@ -185,3 +227,9 @@ def start_execution(
location=location,
credentials=credentials,
)
+
+ # Reinstantiate this class using the newly created resouce.
+ self._init_with_resouce_name(
+ execution_name=new_execution_instance.resource_name
+ )
+ return self
From e0bd08de04837ac65cdd1648374d9d91f4350731 Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 06:25:43 +0000
Subject: [PATCH 07/10] refactor context.
---
.../metadata/schema/base_context.py | 48 ++++++++++++++++---
1 file changed, 41 insertions(+), 7 deletions(-)
diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py
index a39835da00..bea372e26b 100644
--- a/google/cloud/aiplatform/metadata/schema/base_context.py
+++ b/google/cloud/aiplatform/metadata/schema/base_context.py
@@ -16,16 +16,18 @@
#
import abc
+from copy import deepcopy
from typing import Optional, Dict
from google.auth import credentials as auth_credentials
+from google.cloud.aiplatform.compat.types import context as gca_context
from google.cloud.aiplatform.metadata import constants
from google.cloud.aiplatform.metadata import context
-class BaseContextSchema(metaclass=abc.ABCMeta):
+class BaseContextSchema(context.Context):
"""Base class for Metadata Context schema."""
@property
@@ -62,11 +64,34 @@ def __init__(
description (str):
Optional. Describes the purpose of the Context to be created.
"""
+ # resource_id is not stored in the proto. Create method uses the
+ # resource_id along with project_id and location to construct an
+ # resource_name which is stored in the proto message.
self.context_id = context_id
- self.display_name = display_name
- self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION
- self.metadata = metadata
- self.description = description
+
+ # Store all other attributes using the proto structure.
+ self._gca_resource = gca_context.Context()
+ self._gca_resource.display_name = display_name
+ self._gca_resource.schema_version = (
+ schema_version or constants._DEFAULT_SCHEMA_VERSION
+ )
+ if metadata:
+ self._gca_resource.metadata = deepcopy(metadata)
+ self._gca_resource.description = description
+
+ # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
+ def _init_with_resouce_name(
+ self,
+ *,
+ context_name: str,
+ ):
+ """Initializes the Artifact instance using an existing resource.
+ Args:
+ context_name (str):
+ Context name with the following format, this is globally unique in a metadataStore:
+ projects/123/locations/us-central1/metadataStores/<metadata_store_id>/contexts/<resource_id>.
+ """
+ super(BaseContextSchema, self).__init__(context_name=context_name)
def create(
self,
@@ -97,15 +122,24 @@ def create(
Context: Instantiated representation of the managed Metadata Context.
"""
- return context.Context.create(
+ # Check if metadata exists to avoid proto read error
+ metadata = None
+ if self._gca_resource.metadata:
+ metadata = self.metadata
+
+ new_context = context.Context.create(
resource_id=self.context_id,
schema_title=self.schema_title,
display_name=self.display_name,
schema_version=self.schema_version,
description=self.description,
- metadata=self.metadata,
+ metadata=metadata,
metadata_store_id=metadata_store_id,
project=project,
location=location,
credentials=credentials,
)
+
+ # Reinstantiate this class using the newly created resouce.
+ self._init_with_resouce_name(context_name=new_context.resource_name)
+ return self
From 66d26e697873957fc30fd71c2f0340588b67a0c4 Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 06:29:57 +0000
Subject: [PATCH 08/10] fix spelling
---
google/cloud/aiplatform/metadata/resource.py | 2 +-
.../cloud/aiplatform/metadata/schema/base_artifact.py | 6 +++---
.../cloud/aiplatform/metadata/schema/base_context.py | 6 +++---
.../cloud/aiplatform/metadata/schema/base_execution.py | 10 +++++-----
4 files changed, 12 insertions(+), 12 deletions(-)
diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py
index 7fe84ab6a3..00133ba789 100644
--- a/google/cloud/aiplatform/metadata/resource.py
+++ b/google/cloud/aiplatform/metadata/resource.py
@@ -229,7 +229,7 @@ def get(
Returns:
resource (_Resource):
- Instantiated representation of the managed Metadata resource or None if no resouce was found.
+ Instantiated representation of the managed Metadata resource or None if no resource was found.
"""
resource = cls._get(
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index 120639817a..1c93a3b7dd 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -100,7 +100,7 @@ def __init__(
self._gca_resource.state = state
# TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
- def _init_with_resouce_name(
+ def _init_with_resource_name(
self,
*,
artifact_name: str,
@@ -164,6 +164,6 @@ def create(
credentials=credentials,
)
- # Reinstantiate this class using the newly created resouce.
- self._init_with_resouce_name(artifact_name=new_artifact_instance.resource_name)
+ # Reinstantiate this class using the newly created resource.
+ self._init_with_resource_name(artifact_name=new_artifact_instance.resource_name)
return self
diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py
index bea372e26b..c17ccfaf67 100644
--- a/google/cloud/aiplatform/metadata/schema/base_context.py
+++ b/google/cloud/aiplatform/metadata/schema/base_context.py
@@ -80,7 +80,7 @@ def __init__(
self._gca_resource.description = description
# TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
- def _init_with_resouce_name(
+ def _init_with_resource_name(
self,
*,
context_name: str,
@@ -140,6 +140,6 @@ def create(
credentials=credentials,
)
- # Reinstantiate this class using the newly created resouce.
- self._init_with_resouce_name(context_name=new_context.resource_name)
+ # Reinstantiate this class using the newly created resource.
+ self._init_with_resource_name(context_name=new_context.resource_name)
return self
diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py
index d4bdf3f66e..76ebdce8d4 100644
--- a/google/cloud/aiplatform/metadata/schema/base_execution.py
+++ b/google/cloud/aiplatform/metadata/schema/base_execution.py
@@ -88,7 +88,7 @@ def __init__(
self._gca_resource.description = description
# TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
- def _init_with_resouce_name(
+ def _init_with_resource_name(
self,
*,
execution_name: str,
@@ -150,8 +150,8 @@ def create(
location=location,
credentials=credentials,
)
- # Reinstantiate this class using the newly created resouce.
- self._init_with_resouce_name(
+ # Reinstantiate this class using the newly created resource.
+ self._init_with_resource_name(
execution_name=new_execution_instance.resource_name
)
return self
@@ -228,8 +228,8 @@ def start_execution(
credentials=credentials,
)
- # Reinstantiate this class using the newly created resouce.
- self._init_with_resouce_name(
+ # Reinstantiate this class using the newly created resource.
+ self._init_with_resource_name(
execution_name=new_execution_instance.resource_name
)
return self
From 26bfbd2d4dc552be7053346b2677365b11ab1ac6 Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 07:50:31 +0000
Subject: [PATCH 09/10] fix unit tests
---
google/cloud/aiplatform/base.py | 2 +-
.../metadata/schema/base_artifact.py | 7 +-
.../metadata/schema/base_context.py | 8 +--
.../metadata/schema/base_execution.py | 7 +-
.../cloud/aiplatform/metadata/schema/utils.py | 2 +-
tests/unit/aiplatform/test_metadata_schema.py | 68 ++++++++++++++++---
6 files changed, 71 insertions(+), 23 deletions(-)
diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
index adf547e540..19b5f7d468 100644
--- a/google/cloud/aiplatform/base.py
+++ b/google/cloud/aiplatform/base.py
@@ -727,7 +727,7 @@ def __repr__(self) -> str:
def to_dict(self) -> Dict[str, Any]:
"""Returns the resource proto as a dictionary."""
- return json_format.MessageToDict(self.gca_resource._pb)
+ return json_format.MessageToDict(self._gca_resource._pb)
@classmethod
def _generate_display_name(cls, prefix: Optional[str] = None) -> str:
diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py
index 1c93a3b7dd..357ce9d58a 100644
--- a/google/cloud/aiplatform/metadata/schema/base_artifact.py
+++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py
@@ -16,7 +16,6 @@
#
import abc
-from copy import deepcopy
from typing import Optional, Dict
@@ -95,8 +94,10 @@ def __init__(
schema_version or constants._DEFAULT_SCHEMA_VERSION
)
self._gca_resource.description = description
- if metadata:
- self._gca_resource.metadata = deepcopy(metadata)
+
+ # If metadata is None covert to {}
+ metadata = metadata if metadata else {}
+ self._nested_update_metadata(self._gca_resource, metadata)
self._gca_resource.state = state
# TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py
index c17ccfaf67..b6d7f5b4d7 100644
--- a/google/cloud/aiplatform/metadata/schema/base_context.py
+++ b/google/cloud/aiplatform/metadata/schema/base_context.py
@@ -16,7 +16,6 @@
#
import abc
-from copy import deepcopy
from typing import Optional, Dict
@@ -75,8 +74,9 @@ def __init__(
self._gca_resource.schema_version = (
schema_version or constants._DEFAULT_SCHEMA_VERSION
)
- if metadata:
- self._gca_resource.metadata = deepcopy(metadata)
+ # If metadata is None covert to {}
+ metadata = metadata if metadata else {}
+ self._nested_update_metadata(self._gca_resource, metadata)
self._gca_resource.description = description
# TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
@@ -91,7 +91,7 @@ def _init_with_resource_name(
Context name with the following format, this is globally unique in a metadataStore:
projects/123/locations/us-central1/metadataStores/<metadata_store_id>/contexts/<resource_id>.
"""
- super(BaseContextSchema, self).__init__(context_name=context_name)
+ super(BaseContextSchema, self).__init__(resource_name=context_name)
def create(
self,
diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py
index 76ebdce8d4..1cbf66b825 100644
--- a/google/cloud/aiplatform/metadata/schema/base_execution.py
+++ b/google/cloud/aiplatform/metadata/schema/base_execution.py
@@ -16,7 +16,6 @@
#
import abc
-from copy import deepcopy
from typing import Optional, Dict
@@ -83,8 +82,9 @@ def __init__(
self._gca_resource.schema_version = (
schema_version or constants._DEFAULT_SCHEMA_VERSION
)
- if metadata:
- self._gca_resource.metadata = deepcopy(metadata)
+ # If metadata is None covert to {}
+ metadata = metadata if metadata else {}
+ self._nested_update_metadata(self._gca_resource, metadata)
self._gca_resource.description = description
# TODO() Switch to @singledispatchmethod constructor overload after py>=3.8
@@ -144,7 +144,6 @@ def create(
description=self.description,
metadata=metadata,
state=self.state,
- base_execution_schema=self,
metadata_store_id=metadata_store_id,
project=project,
location=location,
diff --git a/google/cloud/aiplatform/metadata/schema/utils.py b/google/cloud/aiplatform/metadata/schema/utils.py
index 72577d9324..c742824b76 100644
--- a/google/cloud/aiplatform/metadata/schema/utils.py
+++ b/google/cloud/aiplatform/metadata/schema/utils.py
@@ -159,7 +159,7 @@ def create_uri_from_resource_name(resource_name: str) -> bool:
"""
# TODO: support nested resource names such as models/123/evaluations/456
match_results = re.match(
- r"^projects\/[A-Za-z0-9-]*\/locations\/([A-Za-z0-9-]*)\/[A-Za-z0-9-]*\/[A-Za-z0-9-]*$",
+ r"^projects\/[A-Za-z0-9-]*\/locations\/([A-Za-z0-9-]*)(\/metadataStores\/[A-Za-z0-9-]*)?(\/[A-Za-z0-9-]*\/[A-Za-z0-9-]*)+$",
resource_name,
)
if not match_results:
diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py
index b133d5eebe..f550f61ab4 100644
--- a/tests/unit/aiplatform/test_metadata_schema.py
+++ b/tests/unit/aiplatform/test_metadata_schema.py
@@ -71,15 +71,61 @@
# artifact
_TEST_ARTIFACT_ID = "test-artifact-id"
-_TEST_ARTIFACT_NAME = f"{_TEST_PARENT}/artifacts/{_TEST_ARTIFACT_ID}"
+_TEST_ARTIFACT_NAME = f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}/artifacts/{_TEST_ARTIFACT_ID}"
# execution
_TEST_EXECUTION_ID = "test-execution-id"
-_TEST_EXECUTION_NAME = f"{_TEST_PARENT}/executions/{_TEST_EXECUTION_ID}"
+_TEST_EXECUTION_NAME = f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}/executions/{_TEST_EXECUTION_ID}"
# context
_TEST_CONTEXT_ID = "test-context-id"
-_TEST_CONTEXT_NAME = f"{_TEST_PARENT}/contexts/{_TEST_CONTEXT_ID}"
+_TEST_CONTEXT_NAME = (
+ f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}/contexts/{_TEST_CONTEXT_ID}"
+)
+
+
+@pytest.fixture
+def get_artifact_mock():
+ with patch.object(MetadataServiceClient, "get_artifact") as get_artifact_mock:
+ get_artifact_mock.return_value = GapicArtifact(
+ name=_TEST_ARTIFACT_NAME,
+ display_name=_TEST_DISPLAY_NAME,
+ schema_title=_TEST_SCHEMA_TITLE,
+ schema_version=_TEST_SCHEMA_VERSION,
+ description=_TEST_DESCRIPTION,
+ metadata=_TEST_METADATA,
+ state=GapicArtifact.State.STATE_UNSPECIFIED,
+ )
+ yield get_artifact_mock
+
+
+@pytest.fixture
+def get_execution_mock():
+ with patch.object(MetadataServiceClient, "get_execution") as get_execution_mock:
+ get_execution_mock.return_value = GapicExecution(
+ name=_TEST_EXECUTION_NAME,
+ display_name=_TEST_DISPLAY_NAME,
+ schema_title=_TEST_SCHEMA_TITLE,
+ schema_version=_TEST_SCHEMA_VERSION,
+ description=_TEST_DESCRIPTION,
+ metadata=_TEST_METADATA,
+ state=GapicExecution.State.RUNNING,
+ )
+ yield get_execution_mock
+
+
+@pytest.fixture
+def get_context_mock():
+ with patch.object(MetadataServiceClient, "get_context") as get_context_mock:
+ get_context_mock.return_value = GapicExecution(
+ name=_TEST_CONTEXT_NAME,
+ display_name=_TEST_DISPLAY_NAME,
+ schema_title=_TEST_SCHEMA_TITLE,
+ schema_version=_TEST_SCHEMA_VERSION,
+ description=_TEST_DESCRIPTION,
+ metadata=_TEST_METADATA,
+ )
+ yield get_context_mock
@pytest.fixture
@@ -172,7 +218,7 @@ def test_base_class_without_schema_title_raises_error(self):
with pytest.raises(TypeError):
base_artifact.BaseArtifactSchema()
- @pytest.mark.usefixtures("create_artifact_mock")
+ @pytest.mark.usefixtures("create_artifact_mock", "get_artifact_mock")
def test_create_is_called_with_default_parameters(self, create_artifact_mock):
aiplatform.init(project=_TEST_PROJECT)
@@ -242,7 +288,7 @@ def test_base_class_without_schema_title_raises_error(self):
with pytest.raises(TypeError):
base_execution.BaseExecutionSchema()
- @pytest.mark.usefixtures("create_execution_mock")
+ @pytest.mark.usefixtures("create_execution_mock", "get_execution_mock")
def test_create_method_calls_gapic_library_with_correct_parameters(
self, create_execution_mock
):
@@ -394,7 +440,7 @@ def test_unmanaged_container_model_constructor_parameters_are_set_correctly(self
metadata=_TEST_UPDATED_METADATA,
)
expected_metadata = {
- "test-param1": 2,
+ "test-param1": 2.0,
"test-param2": "test-value-1",
"test-param3": False,
"predictSchemata": {
@@ -409,7 +455,9 @@ def test_unmanaged_container_model_constructor_parameters_are_set_correctly(self
assert artifact.uri == _TEST_URI
assert artifact.display_name == _TEST_DISPLAY_NAME
assert artifact.description == _TEST_DESCRIPTION
- assert json.dumps(artifact.metadata) == json.dumps(expected_metadata)
+ assert json.dumps(artifact.metadata, sort_keys=True) == json.dumps(
+ expected_metadata, sort_keys=True
+ )
assert artifact.schema_version == _TEST_SCHEMA_VERSION
@@ -489,7 +537,7 @@ def test_system_metrics_schema_title_is_set_correctly(self):
def test_system_metrics_values_default_to_none(self):
artifact = system_artifact_schema.Metrics()
- assert artifact.metadata == {}
+ assert artifact._gca_resource.metadata is None
def test_system_metrics_constructor_parameters_are_set_correctly(self):
artifact = system_artifact_schema.Metrics(
@@ -554,7 +602,7 @@ def teardown_method(self):
initializer.global_pool.shutdown(wait=True)
# Test system.Context Schemas
- @pytest.mark.usefixtures("create_context_mock")
+ @pytest.mark.usefixtures("create_context_mock", "get_context_mock")
def test_create_is_called_with_default_parameters(self, create_context_mock):
aiplatform.init(project=_TEST_PROJECT)
@@ -650,7 +698,7 @@ def test_container_spec_to_dict_method_returns_correct_schema(self):
assert json.dumps(container_spec.to_dict()) == json.dumps(expected_results)
- @pytest.mark.usefixtures("create_execution_mock")
+ @pytest.mark.usefixtures("create_execution_mock", "get_execution_mock")
def test_start_execution_method_calls_gapic_library_with_correct_parameters(
self, create_execution_mock
):
From d8b82e29b61eef4752571978c599deddc56e6e5d Mon Sep 17 00:00:00 2001
From: sina chavoshi <sina.chavoshi@gmail.com>
Date: Fri, 22 Jul 2022 18:24:45 +0000
Subject: [PATCH 10/10] fix sample tests.
---
samples/model-builder/conftest.py | 18 +++++++++++++++++
.../create_artifact_with_sdk_sample_test.py | 20 +++----------------
.../create_execution_with_sdk_sample_test.py | 18 ++---------------
3 files changed, 23 insertions(+), 33 deletions(-)
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
index 7b6cbdb6a9..7363ec28f9 100644
--- a/samples/model-builder/conftest.py
+++ b/samples/model-builder/conftest.py
@@ -634,6 +634,24 @@ def mock_create_artifact(mock_artifact):
yield mock_create_artifact
+@pytest.fixture
+def mock_create_schema_base_artifact(mock_artifact):
+ with patch.object(
+ 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
+
+
+@pytest.fixture
+def mock_create_schema_base_execution(mock_execution):
+ with patch.object(
+ 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
+
+
@pytest.fixture
def mock_list_artifact(mock_artifact):
with patch.object(aiplatform.Artifact, "list") as mock_list_artifact:
diff --git a/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py b/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py
index 09b9249ff2..ecd013516d 100644
--- a/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py
+++ b/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py
@@ -14,12 +14,10 @@
import create_artifact_with_sdk_sample
-from google.cloud.aiplatform.compat.types import artifact as gca_artifact
-
import test_constants as constants
-def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_artifact):
+def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_schema_base_artifact):
artifact = create_artifact_with_sdk_sample.create_artifact_sample(
project=constants.PROJECT,
location=constants.LOCATION,
@@ -31,19 +29,7 @@ def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_artifact):
metadata=constants.METADATA,
)
- mock_create_artifact.assert_called_with(
- resource_id=constants.RESOURCE_ID,
- schema_title="system.Artifact",
- uri=constants.MODEL_ARTIFACT_URI,
- display_name=constants.DISPLAY_NAME,
- schema_version=constants.SCHEMA_VERSION,
- description=constants.DESCRIPTION,
- metadata=constants.METADATA,
- state=gca_artifact.Artifact.State.LIVE,
- metadata_store_id="default",
- project=constants.PROJECT,
- location=constants.LOCATION,
- credentials=None,
- )
+ mock_create_schema_base_artifact.assert_called_with(
+ project='abc', location='us-central1')
assert artifact is mock_artifact
diff --git a/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py b/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py
index 54e9563a12..4be3022862 100644
--- a/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py
+++ b/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py
@@ -14,13 +14,11 @@
import create_execution_with_sdk_sample
-from google.cloud.aiplatform.compat.types import execution as gca_execution
-
import test_constants as constants
def test_create_execution_sample(
- mock_sdk_init, mock_create_artifact, mock_create_execution, mock_execution,
+ mock_sdk_init, mock_create_artifact, mock_create_schema_base_execution, mock_execution,
):
input_art = mock_create_artifact()
@@ -42,19 +40,7 @@ def test_create_execution_sample(
project=constants.PROJECT, location=constants.LOCATION,
)
- mock_create_execution.assert_called_with(
- state=gca_execution.Execution.State.RUNNING,
- schema_title="system.ContainerExecution",
- resource_id=constants.RESOURCE_ID,
- display_name=constants.DISPLAY_NAME,
- schema_version=constants.SCHEMA_VERSION,
- metadata=constants.METADATA,
- description=constants.DESCRIPTION,
- metadata_store_id="default",
- project=None,
- location=None,
- credentials=None,
- )
+ mock_create_schema_base_execution.assert_called_with()
mock_execution.assign_input_artifacts.assert_called_with([input_art])
mock_execution.assign_output_artifacts.assert_called_with([output_art])