From a01204ebecee95750fd31e29b499d801acd78990 Mon Sep 17 00:00:00 2001 From: sina chavoshi Date: Fri, 12 Aug 2022 17:48:20 +0000 Subject: [PATCH 1/4] feat: Add support for SDK Method metrics tracking via _USER_AGENT_SDK_COMMAND --- google/cloud/aiplatform/metadata/artifact.py | 20 +++- google/cloud/aiplatform/metadata/constants.py | 5 + google/cloud/aiplatform/metadata/context.py | 19 ++- google/cloud/aiplatform/metadata/execution.py | 112 ++++++++++++++++-- .../aiplatform/metadata/metadata_store.py | 19 ++- google/cloud/aiplatform/metadata/resource.py | 33 +++++- .../metadata/schema/base_artifact.py | 9 ++ .../metadata/schema/base_context.py | 10 ++ .../metadata/schema/base_execution.py | 15 +++ tests/unit/aiplatform/test_metadata_schema.py | 41 +++++++ 10 files changed, 268 insertions(+), 15 deletions(-) diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index 5f4f94e18b..2781424db3 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -28,6 +28,7 @@ from google.cloud.aiplatform.compat.types import ( metadata_service as gca_metadata_service, ) +from google.cloud.aiplatform.metadata import constants 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 @@ -114,6 +115,7 @@ def _create_resource( artifact_id=resource_id, ) + # TODO() refactor code to move _create to _Resource class. @classmethod def _create( cls, @@ -175,7 +177,18 @@ def _create( Instantiated representation of the managed Metadata resource. """ - api_client = cls._instantiate_client(location=location, credentials=credentials) + if constants._USER_AGENT_SDK_COMMAND: + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + ) + # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + constants._USER_AGENT_SDK_COMMAND = "" + else: + api_client = cls._instantiate_client( + location=location, credentials=credentials + ) parent = utils.full_resource_name( resource_name=metadata_store_id, @@ -311,6 +324,11 @@ def create( Returns: Artifact: Instantiated representation of the managed Metadata Artifact. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.artifact.Artifact.create" + return cls._create( resource_id=resource_id, schema_title=schema_title, diff --git a/google/cloud/aiplatform/metadata/constants.py b/google/cloud/aiplatform/metadata/constants.py index 2dd0c4db86..d1070ca8a6 100644 --- a/google/cloud/aiplatform/metadata/constants.py +++ b/google/cloud/aiplatform/metadata/constants.py @@ -68,3 +68,8 @@ _TB_RUN_ARTIFACT_POST_FIX_ID = "-tb-run" _EXPERIMENT_RUN_MAX_LENGTH = 128 - len(_TB_RUN_ARTIFACT_POST_FIX_ID) + +# This field is used to pass the name of the specific SDK method that +# that is being used for usage metrics tracking purposes. +# For more details on go/oneplatform-api-analytics +_USER_AGENT_SDK_COMMAND = "" diff --git a/google/cloud/aiplatform/metadata/context.py b/google/cloud/aiplatform/metadata/context.py index d1c7dea99c..40f8e05dc8 100644 --- a/google/cloud/aiplatform/metadata/context.py +++ b/google/cloud/aiplatform/metadata/context.py @@ -23,6 +23,7 @@ from google.cloud.aiplatform import base from google.cloud.aiplatform import utils +from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import utils as metadata_utils from google.cloud.aiplatform.compat.types import context as gca_context from google.cloud.aiplatform.compat.types import ( @@ -136,6 +137,11 @@ def create( Returns: Context: Instantiated representation of the managed Metadata Context. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(cls)}" + return cls._create( resource_id=resource_id, schema_title=schema_title, @@ -202,7 +208,18 @@ def _create( Instantiated representation of the managed Metadata resource. """ - api_client = cls._instantiate_client(location=location, credentials=credentials) + if constants._USER_AGENT_SDK_COMMAND: + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + ) + # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + constants._USER_AGENT_SDK_COMMAND = "" + else: + api_client = cls._instantiate_client( + location=location, credentials=credentials + ) parent = utils.full_resource_name( resource_name=metadata_store_id, diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index 9a85bce36f..9b84224021 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -20,7 +20,6 @@ import proto from google.auth import credentials as auth_credentials -from google.cloud.aiplatform import base from google.cloud.aiplatform import models from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import event as gca_event @@ -28,6 +27,7 @@ from google.cloud.aiplatform.compat.types import ( metadata_service as gca_metadata_service, ) +from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform.metadata import resource @@ -142,18 +142,107 @@ def create( Execution: Instantiated representation of the managed Metadata Execution. """ - self = cls._empty_constructor( - project=project, location=location, credentials=credentials + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(cls)}" + + return cls._create( + resource_id=resource_id, + schema_title=schema_title, + display_name=display_name, + schema_version=schema_version, + description=description, + metadata=metadata, + state=state, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) + + # TODO() refactor code to move _create to _Resource class. + @classmethod + def _create( + cls, + schema_title: str, + *, + state: gca_execution.Execution.State = gca_execution.Execution.State.RUNNING, + resource_id: Optional[str] = None, + display_name: Optional[str] = None, + schema_version: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + description: Optional[str] = None, + metadata_store_id: str = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials=Optional[auth_credentials.Credentials], + ) -> "Execution": + """ + Creates a new Metadata Execution. + + Args: + schema_title (str): + Required. schema_title identifies the schema title used by the Execution. + state (gca_execution.Execution.State.RUNNING): + Optional. State of this Execution. Defaults to RUNNING. + resource_id (str): + Optional. The portion of the Execution name with + the format. This is globally unique in a metadataStore: + projects/123/locations/us-central1/metadataStores//executions/. + 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. + metadata_store_id (str): + Optional. The portion of the resource name with + the format: + projects/123/locations/us-central1/metadataStores//artifacts/ + 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. + + """ + if constants._USER_AGENT_SDK_COMMAND: + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + ) + # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + constants._USER_AGENT_SDK_COMMAND = "" + else: + api_client = cls._instantiate_client( + location=location, credentials=credentials + ) + + parent = utils.full_resource_name( + resource_name=metadata_store_id, + resource_noun=metadata_store._MetadataStore._resource_noun, + parse_resource_name_method=metadata_store._MetadataStore._parse_resource_name, + format_resource_name_method=metadata_store._MetadataStore._format_resource_name, + project=project, + location=location, ) - super(base.VertexAiResourceNounWithFutureManager, self).__init__() resource = Execution._create_resource( - client=self.api_client, - parent=metadata_store._MetadataStore._format_resource_name( - project=self.project, - location=self.location, - metadata_store=metadata_store_id, - ), + client=api_client, + parent=parent, schema_title=schema_title, resource_id=resource_id, metadata=metadata, @@ -162,6 +251,9 @@ def create( schema_version=schema_version, state=state, ) + self = cls._empty_constructor( + project=project, location=location, credentials=credentials + ) self._gca_resource = resource return self diff --git a/google/cloud/aiplatform/metadata/metadata_store.py b/google/cloud/aiplatform/metadata/metadata_store.py index 2f0c8e2955..ab37c11167 100644 --- a/google/cloud/aiplatform/metadata/metadata_store.py +++ b/google/cloud/aiplatform/metadata/metadata_store.py @@ -25,6 +25,7 @@ from google.cloud.aiplatform import compat from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import metadata_store as gca_metadata_store +from google.cloud.aiplatform.metadata import constants class _MetadataStore(base.VertexAiResourceNounWithFutureManager): @@ -115,6 +116,10 @@ def get_or_create( Instantiated representation of the managed metadata store resource. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(cls)}" store = cls._get( metadata_store_name=metadata_store_id, @@ -176,7 +181,19 @@ def _create( Instantiated representation of the managed metadata store resource. """ - api_client = cls._instantiate_client(location=location, credentials=credentials) + if constants._USER_AGENT_SDK_COMMAND: + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + ) + # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + constants._USER_AGENT_SDK_COMMAND = "" + else: + api_client = cls._instantiate_client( + location=location, credentials=credentials + ) + gapic_metadata_store = gca_metadata_store.MetadataStore( encryption_spec=initializer.global_config.get_encryption_spec( encryption_spec_key_name=encryption_spec_key_name, diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index 00133ba789..bd9c64cfb0 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -31,6 +31,7 @@ from google.cloud.aiplatform.compat.types import artifact as gca_artifact from google.cloud.aiplatform.compat.types import context as gca_context from google.cloud.aiplatform.compat.types import execution as gca_execution +from google.cloud.aiplatform.metadata import constants _LOGGER = base.Logger(__name__) @@ -173,6 +174,10 @@ def get_or_create( Instantiated representation of the managed Metadata resource. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.resouce._Resource.get_or_create" resource = cls._get( resource_name=resource_id, @@ -232,6 +237,11 @@ def get( Instantiated representation of the managed Metadata resource or None if no resource was found. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.resouce._Resource.get" + resource = cls._get( resource_name=resource_id, metadata_store_id=metadata_store_id, @@ -301,7 +311,15 @@ def update( if description: gca_resource.description = description - api_client = self._instantiate_client(credentials=credentials) + if constants._USER_AGENT_SDK_COMMAND: + api_client = self._instantiate_client( + credentials=credentials, + appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + ) + # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + constants._USER_AGENT_SDK_COMMAND = "" + else: + api_client = self._instantiate_client(credentials=credentials) # TODO: if etag is not valid sync and retry update_gca_resource = self._update_resource( @@ -412,7 +430,18 @@ def _create( Instantiated representation of the managed Metadata resource. """ - api_client = cls._instantiate_client(location=location, credentials=credentials) + if constants._USER_AGENT_SDK_COMMAND: + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + ) + # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + constants._USER_AGENT_SDK_COMMAND = "" + else: + api_client = cls._instantiate_client( + location=location, credentials=credentials + ) parent = ( initializer.global_config.common_location_path( diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py index 357ce9d58a..6c2f5a22f8 100644 --- a/google/cloud/aiplatform/metadata/schema/base_artifact.py +++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py @@ -114,6 +114,11 @@ def _init_with_resource_name( Artifact name with the following format, this is globally unique in a metadataStore: projects/123/locations/us-central1/metadataStores//artifacts/. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.schema.base_artifact.BaseArtifactSchema._init_with_resource_name" + super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name) def create( @@ -144,6 +149,10 @@ def create( Returns: Artifact: Instantiated representation of the managed Metadata Artifact. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.schema.base_artifact.BaseArtifactSchema.create" # Check if metadata exists to avoid proto read error metadata = None diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py index b6d7f5b4d7..9077b484c9 100644 --- a/google/cloud/aiplatform/metadata/schema/base_context.py +++ b/google/cloud/aiplatform/metadata/schema/base_context.py @@ -91,6 +91,11 @@ def _init_with_resource_name( Context name with the following format, this is globally unique in a metadataStore: projects/123/locations/us-central1/metadataStores//contexts/. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + super(BaseContextSchema, self).__init__(resource_name=context_name) def create( @@ -122,6 +127,11 @@ def create( Context: Instantiated representation of the managed Metadata Context. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + # Check if metadata exists to avoid proto read error metadata = None if self._gca_resource.metadata: diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py index 1cbf66b825..2e24aa8613 100644 --- a/google/cloud/aiplatform/metadata/schema/base_execution.py +++ b/google/cloud/aiplatform/metadata/schema/base_execution.py @@ -100,6 +100,11 @@ def _init_with_resource_name( The Execution name with the following format, this is globally unique in a metadataStore. projects/123/locations/us-central1/metadataStores//executions/. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + super(BaseExecutionSchema, self).__init__(execution_name=execution_name) def create( @@ -131,6 +136,11 @@ def create( Execution: Instantiated representation of the managed Metadata Execution. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + # Check if metadata exists to avoid proto read error metadata = None if self._gca_resource.metadata: @@ -208,6 +218,11 @@ def start_execution( Raises: ValueError: If metadata_store_id other than 'default' is provided. """ + # Add User Agent Header for metrics tracking if one is not specified + # If one is alreayd specified this call was initiated by a sub class. + if not constants._USER_AGENT_SDK_COMMAND: + constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + if metadata_store_id != "default": raise ValueError( f"metadata_store_id {metadata_store_id} is not supported. Only the default MetadataStore ID is supported." diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py index e85229c555..5aa2dfe54a 100644 --- a/tests/unit/aiplatform/test_metadata_schema.py +++ b/tests/unit/aiplatform/test_metadata_schema.py @@ -99,6 +99,22 @@ def get_artifact_mock(): yield get_artifact_mock +@pytest.fixture +def initializer_create_client_mock(): + with patch.object( + initializer.global_config, "create_client" + ) as initializer_create_client_mock: + yield initializer_create_client_mock + + +@pytest.fixture +def base_artifact_init_with_resouce_name_mock(): + with patch.object( + base_artifact.BaseArtifactSchema, "_init_with_resource_name" + ) as base_artifact_init_with_resouce_name_mock: + yield base_artifact_init_with_resouce_name_mock + + @pytest.fixture def get_execution_mock(): with patch.object(MetadataServiceClient, "get_execution") as get_execution_mock: @@ -246,6 +262,31 @@ class TestArtifact(base_artifact.BaseArtifactSchema): assert kwargs["artifact"].metadata == _TEST_UPDATED_METADATA assert kwargs["artifact"].state == _TEST_ARTIFACT_STATE + @pytest.mark.usefixtures( + "base_artifact_init_with_resouce_name_mock", + "initializer_create_client_mock", + "create_artifact_mock", + "get_artifact_mock", + ) + def test_create_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestArtifact(base_artifact.BaseArtifactSchema): + schema_title = _TEST_SCHEMA_TITLE + + artifact = TestArtifact( + uri=_TEST_URI, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + state=_TEST_ARTIFACT_STATE, + ) + artifact.create() + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs == None + @pytest.mark.usefixtures("google_auth_mock") class TestMetadataBaseExecutionSchema: From 1110e162ebec3f03401bb319f12384149b87ec75 Mon Sep 17 00:00:00 2001 From: sina chavoshi Date: Mon, 15 Aug 2022 21:50:03 +0000 Subject: [PATCH 2/4] add unit tests --- google/cloud/aiplatform/base.py | 15 +- google/cloud/aiplatform/constants/base.py | 4 + google/cloud/aiplatform/metadata/artifact.py | 18 +- google/cloud/aiplatform/metadata/constants.py | 5 - google/cloud/aiplatform/metadata/context.py | 18 +- google/cloud/aiplatform/metadata/execution.py | 18 +- .../aiplatform/metadata/metadata_store.py | 17 +- google/cloud/aiplatform/metadata/resource.py | 35 +-- .../metadata/schema/base_artifact.py | 13 +- .../metadata/schema/base_context.py | 13 +- .../metadata/schema/base_execution.py | 14 +- tests/unit/aiplatform/test_metadata_schema.py | 244 +++++++++++++++++- 12 files changed, 325 insertions(+), 89 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 19b5f7d468..868431fbdc 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -46,6 +46,7 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import encryption_spec as gca_encryption_spec +from google.cloud.aiplatform.constants import base as base_constants from google.protobuf import json_format # This is the default retry callback to be used with get methods. @@ -499,7 +500,19 @@ def __init__( self.location = location or initializer.global_config.location self.credentials = credentials or initializer.global_config.credentials - self.api_client = self._instantiate_client(self.location, self.credentials) + appended_user_agent = None + if base_constants.USER_AGENT_SDK_COMMAND: + appended_user_agent = [ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ] + # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + base_constants.USER_AGENT_SDK_COMMAND = "" + + self.api_client = self._instantiate_client( + location=self.location, + credentials=self.credentials, + appended_user_agent=appended_user_agent, + ) @classmethod def _instantiate_client( diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py index 373f66c003..328d963f2d 100644 --- a/google/cloud/aiplatform/constants/base.py +++ b/google/cloud/aiplatform/constants/base.py @@ -88,3 +88,7 @@ # Used in constructing the requests user_agent header for metrics reporting. USER_AGENT_PRODUCT = "model-builder" +# This field is used to pass the name of the specific SDK method that +# that is being used for usage metrics tracking purposes. +# For more details on go/oneplatform-api-analytics +USER_AGENT_SDK_COMMAND = "" diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index 2781424db3..a5134eb009 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -28,7 +28,7 @@ from google.cloud.aiplatform.compat.types import ( metadata_service as gca_metadata_service, ) -from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.constants import base as base_constants 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 @@ -177,14 +177,16 @@ def _create( Instantiated representation of the managed Metadata resource. """ - if constants._USER_AGENT_SDK_COMMAND: + if base_constants.USER_AGENT_SDK_COMMAND: api_client = cls._instantiate_client( location=location, credentials=credentials, - appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + appended_user_agent=[ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ], ) - # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. - constants._USER_AGENT_SDK_COMMAND = "" + # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + base_constants.USER_AGENT_SDK_COMMAND = "" else: api_client = cls._instantiate_client( location=location, credentials=credentials @@ -326,8 +328,10 @@ def create( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.artifact.Artifact.create" + if not base_constants.USER_AGENT_SDK_COMMAND: + base_constants.USER_AGENT_SDK_COMMAND = ( + "aiplatform.metadata.artifact.Artifact.create" + ) return cls._create( resource_id=resource_id, diff --git a/google/cloud/aiplatform/metadata/constants.py b/google/cloud/aiplatform/metadata/constants.py index d1070ca8a6..2dd0c4db86 100644 --- a/google/cloud/aiplatform/metadata/constants.py +++ b/google/cloud/aiplatform/metadata/constants.py @@ -68,8 +68,3 @@ _TB_RUN_ARTIFACT_POST_FIX_ID = "-tb-run" _EXPERIMENT_RUN_MAX_LENGTH = 128 - len(_TB_RUN_ARTIFACT_POST_FIX_ID) - -# This field is used to pass the name of the specific SDK method that -# that is being used for usage metrics tracking purposes. -# For more details on go/oneplatform-api-analytics -_USER_AGENT_SDK_COMMAND = "" diff --git a/google/cloud/aiplatform/metadata/context.py b/google/cloud/aiplatform/metadata/context.py index 40f8e05dc8..c50705b339 100644 --- a/google/cloud/aiplatform/metadata/context.py +++ b/google/cloud/aiplatform/metadata/context.py @@ -23,7 +23,7 @@ from google.cloud.aiplatform import base from google.cloud.aiplatform import utils -from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.constants import base as base_constants from google.cloud.aiplatform.metadata import utils as metadata_utils from google.cloud.aiplatform.compat.types import context as gca_context from google.cloud.aiplatform.compat.types import ( @@ -139,8 +139,10 @@ def create( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(cls)}" + if not base_constants.USER_AGENT_SDK_COMMAND: + base_constants.USER_AGENT_SDK_COMMAND = ( + "aiplatform.metadata.context.Context.create" + ) return cls._create( resource_id=resource_id, @@ -208,14 +210,16 @@ def _create( Instantiated representation of the managed Metadata resource. """ - if constants._USER_AGENT_SDK_COMMAND: + if base_constants.USER_AGENT_SDK_COMMAND: api_client = cls._instantiate_client( location=location, credentials=credentials, - appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + appended_user_agent=[ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ], ) - # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. - constants._USER_AGENT_SDK_COMMAND = "" + # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + base_constants.USER_AGENT_SDK_COMMAND = "" else: api_client = cls._instantiate_client( location=location, credentials=credentials diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index 9b84224021..d49778e8e9 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -27,7 +27,7 @@ from google.cloud.aiplatform.compat.types import ( metadata_service as gca_metadata_service, ) -from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.constants import base as base_constants from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform.metadata import resource @@ -144,8 +144,10 @@ def create( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(cls)}" + if not base_constants.USER_AGENT_SDK_COMMAND: + base_constants.USER_AGENT_SDK_COMMAND = ( + "aiplatform.metadata.execution.Execution.create" + ) return cls._create( resource_id=resource_id, @@ -218,14 +220,16 @@ def _create( Execution: Instantiated representation of the managed Metadata Execution. """ - if constants._USER_AGENT_SDK_COMMAND: + if base_constants.USER_AGENT_SDK_COMMAND: api_client = cls._instantiate_client( location=location, credentials=credentials, - appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + appended_user_agent=[ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ], ) - # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. - constants._USER_AGENT_SDK_COMMAND = "" + # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + base_constants.USER_AGENT_SDK_COMMAND = "" else: api_client = cls._instantiate_client( location=location, credentials=credentials diff --git a/google/cloud/aiplatform/metadata/metadata_store.py b/google/cloud/aiplatform/metadata/metadata_store.py index ab37c11167..0db41874a3 100644 --- a/google/cloud/aiplatform/metadata/metadata_store.py +++ b/google/cloud/aiplatform/metadata/metadata_store.py @@ -25,7 +25,7 @@ from google.cloud.aiplatform import compat from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import metadata_store as gca_metadata_store -from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.constants import base as base_constants class _MetadataStore(base.VertexAiResourceNounWithFutureManager): @@ -116,11 +116,6 @@ def get_or_create( Instantiated representation of the managed metadata store resource. """ - # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(cls)}" - store = cls._get( metadata_store_name=metadata_store_id, project=project, @@ -181,14 +176,16 @@ def _create( Instantiated representation of the managed metadata store resource. """ - if constants._USER_AGENT_SDK_COMMAND: + if base_constants.USER_AGENT_SDK_COMMAND: api_client = cls._instantiate_client( location=location, credentials=credentials, - appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], + appended_user_agent=[ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ], ) - # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. - constants._USER_AGENT_SDK_COMMAND = "" + # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. + base_constants.USER_AGENT_SDK_COMMAND = "" else: api_client = cls._instantiate_client( location=location, credentials=credentials diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index bd9c64cfb0..85d72bedb6 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -31,7 +31,6 @@ from google.cloud.aiplatform.compat.types import artifact as gca_artifact from google.cloud.aiplatform.compat.types import context as gca_context from google.cloud.aiplatform.compat.types import execution as gca_execution -from google.cloud.aiplatform.metadata import constants _LOGGER = base.Logger(__name__) @@ -174,10 +173,6 @@ def get_or_create( Instantiated representation of the managed Metadata resource. """ - # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.resouce._Resource.get_or_create" resource = cls._get( resource_name=resource_id, @@ -237,11 +232,6 @@ def get( Instantiated representation of the managed Metadata resource or None if no resource was found. """ - # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.resouce._Resource.get" - resource = cls._get( resource_name=resource_id, metadata_store_id=metadata_store_id, @@ -311,15 +301,7 @@ def update( if description: gca_resource.description = description - if constants._USER_AGENT_SDK_COMMAND: - api_client = self._instantiate_client( - credentials=credentials, - appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], - ) - # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. - constants._USER_AGENT_SDK_COMMAND = "" - else: - api_client = self._instantiate_client(credentials=credentials) + api_client = self._instantiate_client(credentials=credentials) # TODO: if etag is not valid sync and retry update_gca_resource = self._update_resource( @@ -430,18 +412,7 @@ def _create( Instantiated representation of the managed Metadata resource. """ - if constants._USER_AGENT_SDK_COMMAND: - api_client = cls._instantiate_client( - location=location, - credentials=credentials, - appended_user_agent=[constants._USER_AGENT_SDK_COMMAND], - ) - # Reset the value for the _USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. - constants._USER_AGENT_SDK_COMMAND = "" - else: - api_client = cls._instantiate_client( - location=location, credentials=credentials - ) + api_client = cls._instantiate_client(location=location, credentials=credentials) parent = ( initializer.global_config.common_location_path( @@ -575,4 +546,4 @@ def _extract_metadata_store_id(resource_name, resource_noun) -> str: raise ValueError( f"failed to extract metadata_store_id from resource {resource_name}" ) - return match["store"] + return match["store"] \ No newline at end of file diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py index 6c2f5a22f8..4c96e0a6d9 100644 --- a/google/cloud/aiplatform/metadata/schema/base_artifact.py +++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py @@ -23,6 +23,7 @@ from google.cloud.aiplatform.compat.types import artifact as gca_artifact from google.cloud.aiplatform.metadata import artifact +from google.cloud.aiplatform.constants import base as base_constants from google.cloud.aiplatform.metadata import constants @@ -116,8 +117,8 @@ def _init_with_resource_name( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.schema.base_artifact.BaseArtifactSchema._init_with_resource_name" + if not base_constants.USER_AGENT_SDK_COMMAND: + base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_artifact.BaseArtifactSchema._init_with_resource_name" super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name) @@ -149,10 +150,10 @@ def create( Returns: Artifact: Instantiated representation of the managed Metadata Artifact. """ - # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/aiplatform.metadata.schema.base_artifact.BaseArtifactSchema.create" + # Add User Agent Header for metrics tracking. + base_constants.USER_AGENT_SDK_COMMAND = ( + "aiplatform.metadata.schema.base_artifact.BaseArtifactSchema.create" + ) # Check if metadata exists to avoid proto read error metadata = None diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py index 9077b484c9..0ad52c2529 100644 --- a/google/cloud/aiplatform/metadata/schema/base_context.py +++ b/google/cloud/aiplatform/metadata/schema/base_context.py @@ -22,6 +22,7 @@ from google.auth import credentials as auth_credentials from google.cloud.aiplatform.compat.types import context as gca_context +from google.cloud.aiplatform.constants import base as base_constants from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import context @@ -93,8 +94,8 @@ def _init_with_resource_name( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + if not base_constants.USER_AGENT_SDK_COMMAND: + base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_context.BaseContextSchema._init_with_resource_name" super(BaseContextSchema, self).__init__(resource_name=context_name) @@ -127,10 +128,10 @@ def create( Context: Instantiated representation of the managed Metadata Context. """ - # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + # Add User Agent Header for metrics tracking. + base_constants.USER_AGENT_SDK_COMMAND = ( + "aiplatform.metadata.schema.base_context.BaseContextSchema.create" + ) # Check if metadata exists to avoid proto read error metadata = None diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py index 2e24aa8613..10a4b7d334 100644 --- a/google/cloud/aiplatform/metadata/schema/base_execution.py +++ b/google/cloud/aiplatform/metadata/schema/base_execution.py @@ -22,6 +22,7 @@ from google.auth import credentials as auth_credentials from google.cloud.aiplatform.compat.types import execution as gca_execution +from google.cloud.aiplatform.constants import base as base_constants from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import execution from google.cloud.aiplatform.metadata import metadata @@ -102,8 +103,8 @@ def _init_with_resource_name( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + if not base_constants.USER_AGENT_SDK_COMMAND: + base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_execution.BaseExecutionSchema._init_with_resource_name" super(BaseExecutionSchema, self).__init__(execution_name=execution_name) @@ -138,8 +139,9 @@ def create( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + base_constants.USER_AGENT_SDK_COMMAND = ( + "aiplatform.metadata.schema.base_execution.BaseExecutionSchema.create" + ) # Check if metadata exists to avoid proto read error metadata = None @@ -220,8 +222,8 @@ def start_execution( """ # Add User Agent Header for metrics tracking if one is not specified # If one is alreayd specified this call was initiated by a sub class. - if not constants._USER_AGENT_SDK_COMMAND: - constants._USER_AGENT_SDK_COMMAND = f"sdk_command/{type(self)}" + + base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_execution.BaseExecutionSchema.start_execution" if metadata_store_id != "default": raise ValueError( diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py index 5aa2dfe54a..38f60bc471 100644 --- a/tests/unit/aiplatform/test_metadata_schema.py +++ b/tests/unit/aiplatform/test_metadata_schema.py @@ -115,6 +115,22 @@ def base_artifact_init_with_resouce_name_mock(): yield base_artifact_init_with_resouce_name_mock +@pytest.fixture +def base_execution_init_with_resouce_name_mock(): + with patch.object( + base_execution.BaseExecutionSchema, "_init_with_resource_name" + ) as base_execution_init_with_resouce_name_mock: + yield base_execution_init_with_resouce_name_mock + + +@pytest.fixture +def base_context_init_with_resouce_name_mock(): + with patch.object( + base_context.BaseContextSchema, "_init_with_resource_name" + ) as base_context_init_with_resouce_name_mock: + yield base_context_init_with_resouce_name_mock + + @pytest.fixture def get_execution_mock(): with patch.object(MetadataServiceClient, "get_execution") as get_execution_mock: @@ -268,7 +284,7 @@ class TestArtifact(base_artifact.BaseArtifactSchema): "create_artifact_mock", "get_artifact_mock", ) - def test_create_call_sets_the_user_agent_header( + def test_artifact_create_call_sets_the_user_agent_header( self, initializer_create_client_mock ): aiplatform.init(project=_TEST_PROJECT) @@ -285,7 +301,35 @@ class TestArtifact(base_artifact.BaseArtifactSchema): ) artifact.create() _, _, kwargs = initializer_create_client_mock.mock_calls[0] - assert kwargs == None + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_artifact.BaseArtifactSchema.create" + ] + + @pytest.mark.usefixtures( + "initializer_create_client_mock", + "create_artifact_mock", + "get_artifact_mock", + ) + def test_artifact_init_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestArtifact(base_artifact.BaseArtifactSchema): + schema_title = _TEST_SCHEMA_TITLE + + artifact = TestArtifact( + uri=_TEST_URI, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + state=_TEST_ARTIFACT_STATE, + ) + artifact._init_with_resource_name(artifact_name=_TEST_ARTIFACT_NAME) + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_artifact.BaseArtifactSchema._init_with_resource_name" + ] @pytest.mark.usefixtures("google_auth_mock") @@ -357,6 +401,202 @@ class TestExecution(base_execution.BaseExecutionSchema): assert kwargs["execution"].description == _TEST_DESCRIPTION assert kwargs["execution"].metadata == _TEST_UPDATED_METADATA + @pytest.mark.usefixtures( + "base_execution_init_with_resouce_name_mock", + "initializer_create_client_mock", + "create_execution_mock", + "get_execution_mock", + ) + def test_execution_create_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestExecution(base_execution.BaseExecutionSchema): + schema_title = _TEST_SCHEMA_TITLE + + execution = TestExecution( + state=_TEST_EXECUTION_STATE, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + execution.create(metadata_store_id=_TEST_METADATA_STORE) + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_execution.BaseExecutionSchema.create" + ] + + @pytest.mark.usefixtures( + "base_execution_init_with_resouce_name_mock", + "initializer_create_client_mock", + "create_execution_mock", + "get_execution_mock", + ) + def test_execution_start_execution_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestExecution(base_execution.BaseExecutionSchema): + schema_title = _TEST_SCHEMA_TITLE + + execution = TestExecution( + state=_TEST_EXECUTION_STATE, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + execution.start_execution() + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_execution.BaseExecutionSchema.start_execution" + ] + + @pytest.mark.usefixtures( + "initializer_create_client_mock", + "create_execution_mock", + "get_execution_mock", + ) + def test_execution_init_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestExecution(base_execution.BaseExecutionSchema): + schema_title = _TEST_SCHEMA_TITLE + + execution = TestExecution( + state=_TEST_EXECUTION_STATE, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + execution._init_with_resource_name(execution_name=_TEST_EXECUTION_NAME) + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_execution.BaseExecutionSchema._init_with_resource_name" + ] + + +@pytest.mark.usefixtures("google_auth_mock") +class TestMetadataBaseContextSchema: + def setup_method(self): + reload(initializer) + reload(metadata) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_base_class_instatiated_uses_schema_title(self): + class TestContext(base_context.BaseContextSchema): + schema_title = _TEST_SCHEMA_TITLE + + context = TestContext() + assert context.schema_title == _TEST_SCHEMA_TITLE + + def test_base_class_parameters_overrides_default_values(self): + class TestContext(base_context.BaseContextSchema): + schema_title = _TEST_SCHEMA_TITLE + + context = TestContext( + schema_version=_TEST_SCHEMA_VERSION, + context_id=_TEST_CONTEXT_ID, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + assert context.schema_version == _TEST_SCHEMA_VERSION + assert context.context_id == _TEST_CONTEXT_ID + assert context.schema_title == _TEST_SCHEMA_TITLE + assert context.display_name == _TEST_DISPLAY_NAME + assert context.description == _TEST_DESCRIPTION + assert context.metadata == _TEST_UPDATED_METADATA + + def test_base_class_without_schema_title_raises_error(self): + with pytest.raises(TypeError): + base_context.BaseContextSchema() + + @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) + + class TestContext(base_context.BaseContextSchema): + schema_title = _TEST_SCHEMA_TITLE + + context = TestContext( + schema_version=_TEST_SCHEMA_VERSION, + context_id=_TEST_CONTEXT_ID, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + context.create(metadata_store_id=_TEST_METADATA_STORE) + create_context_mock.assert_called_once_with( + parent=f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}", + context=mock.ANY, + context_id=_TEST_CONTEXT_ID, + ) + _, _, kwargs = create_context_mock.mock_calls[0] + assert kwargs["context"].schema_title == _TEST_SCHEMA_TITLE + assert kwargs["context"].display_name == _TEST_DISPLAY_NAME + assert kwargs["context"].description == _TEST_DESCRIPTION + assert kwargs["context"].metadata == _TEST_UPDATED_METADATA + + @pytest.mark.usefixtures( + "base_context_init_with_resouce_name_mock", + "initializer_create_client_mock", + "create_context_mock", + "get_context_mock", + ) + def test_context_create_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestContext(base_context.BaseContextSchema): + schema_title = _TEST_SCHEMA_TITLE + + context = TestContext( + schema_version=_TEST_SCHEMA_VERSION, + context_id=_TEST_CONTEXT_ID, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + context.create() + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_context.BaseContextSchema.create" + ] + + @pytest.mark.usefixtures( + "initializer_create_client_mock", + "create_context_mock", + "get_context_mock", + ) + def test_context_init_call_sets_the_user_agent_header( + self, initializer_create_client_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + class TestContext(base_context.BaseContextSchema): + schema_title = _TEST_SCHEMA_TITLE + + context = TestContext( + schema_version=_TEST_SCHEMA_VERSION, + context_id=_TEST_CONTEXT_ID, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + metadata=_TEST_UPDATED_METADATA, + ) + context._init_with_resource_name(context_name=_TEST_CONTEXT_NAME) + _, _, kwargs = initializer_create_client_mock.mock_calls[0] + assert kwargs["appended_user_agent"] == [ + "sdk_command/aiplatform.metadata.schema.base_context.BaseContextSchema._init_with_resource_name" + ] + @pytest.mark.usefixtures("google_auth_mock") class TestMetadataGoogleArtifactSchema: From 70beec945b6c0dc3229f05bea1fc168d09d0d83a Mon Sep 17 00:00:00 2001 From: sina chavoshi Date: Mon, 15 Aug 2022 22:06:48 +0000 Subject: [PATCH 3/4] refactor for readability --- google/cloud/aiplatform/constants/base.py | 2 +- google/cloud/aiplatform/metadata/artifact.py | 21 +++++++++---------- google/cloud/aiplatform/metadata/context.py | 21 +++++++++---------- google/cloud/aiplatform/metadata/execution.py | 21 +++++++++---------- .../aiplatform/metadata/metadata_store.py | 21 +++++++++---------- google/cloud/aiplatform/metadata/resource.py | 2 +- tests/unit/aiplatform/test_metadata_schema.py | 14 +++++++------ 7 files changed, 50 insertions(+), 52 deletions(-) diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py index 328d963f2d..8c1bc1b613 100644 --- a/google/cloud/aiplatform/constants/base.py +++ b/google/cloud/aiplatform/constants/base.py @@ -88,7 +88,7 @@ # Used in constructing the requests user_agent header for metrics reporting. USER_AGENT_PRODUCT = "model-builder" -# This field is used to pass the name of the specific SDK method that +# This field is used to pass the name of the specific SDK method # that is being used for usage metrics tracking purposes. # For more details on go/oneplatform-api-analytics USER_AGENT_SDK_COMMAND = "" diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index a5134eb009..6737a1ff6b 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -177,20 +177,19 @@ def _create( Instantiated representation of the managed Metadata resource. """ + appended_user_agent = [] if base_constants.USER_AGENT_SDK_COMMAND: - api_client = cls._instantiate_client( - location=location, - credentials=credentials, - appended_user_agent=[ - f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" - ], - ) + appended_user_agent = [ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ] # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. base_constants.USER_AGENT_SDK_COMMAND = "" - else: - api_client = cls._instantiate_client( - location=location, credentials=credentials - ) + + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=appended_user_agent, + ) parent = utils.full_resource_name( resource_name=metadata_store_id, diff --git a/google/cloud/aiplatform/metadata/context.py b/google/cloud/aiplatform/metadata/context.py index c50705b339..280298645e 100644 --- a/google/cloud/aiplatform/metadata/context.py +++ b/google/cloud/aiplatform/metadata/context.py @@ -210,20 +210,19 @@ def _create( Instantiated representation of the managed Metadata resource. """ + appended_user_agent = [] if base_constants.USER_AGENT_SDK_COMMAND: - api_client = cls._instantiate_client( - location=location, - credentials=credentials, - appended_user_agent=[ - f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" - ], - ) + appended_user_agent = [ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ] # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. base_constants.USER_AGENT_SDK_COMMAND = "" - else: - api_client = cls._instantiate_client( - location=location, credentials=credentials - ) + + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=appended_user_agent, + ) parent = utils.full_resource_name( resource_name=metadata_store_id, diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index d49778e8e9..7fc24b2457 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -220,20 +220,19 @@ def _create( Execution: Instantiated representation of the managed Metadata Execution. """ + appended_user_agent = [] if base_constants.USER_AGENT_SDK_COMMAND: - api_client = cls._instantiate_client( - location=location, - credentials=credentials, - appended_user_agent=[ - f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" - ], - ) + appended_user_agent = [ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ] # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. base_constants.USER_AGENT_SDK_COMMAND = "" - else: - api_client = cls._instantiate_client( - location=location, credentials=credentials - ) + + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=appended_user_agent, + ) parent = utils.full_resource_name( resource_name=metadata_store_id, diff --git a/google/cloud/aiplatform/metadata/metadata_store.py b/google/cloud/aiplatform/metadata/metadata_store.py index 0db41874a3..ab2e7ec305 100644 --- a/google/cloud/aiplatform/metadata/metadata_store.py +++ b/google/cloud/aiplatform/metadata/metadata_store.py @@ -176,20 +176,19 @@ def _create( Instantiated representation of the managed metadata store resource. """ + appended_user_agent = [] if base_constants.USER_AGENT_SDK_COMMAND: - api_client = cls._instantiate_client( - location=location, - credentials=credentials, - appended_user_agent=[ - f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" - ], - ) + appended_user_agent = [ + f"sdk_command/{base_constants.USER_AGENT_SDK_COMMAND}" + ] # Reset the value for the USER_AGENT_SDK_COMMAND to avoid counting future unrelated api calls. base_constants.USER_AGENT_SDK_COMMAND = "" - else: - api_client = cls._instantiate_client( - location=location, credentials=credentials - ) + + api_client = cls._instantiate_client( + location=location, + credentials=credentials, + appended_user_agent=appended_user_agent, + ) gapic_metadata_store = gca_metadata_store.MetadataStore( encryption_spec=initializer.global_config.get_encryption_spec( diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index 85d72bedb6..00133ba789 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -546,4 +546,4 @@ def _extract_metadata_store_id(resource_name, resource_noun) -> str: raise ValueError( f"failed to extract metadata_store_id from resource {resource_name}" ) - return match["store"] \ No newline at end of file + return match["store"] diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py index 38f60bc471..65de1304ea 100644 --- a/tests/unit/aiplatform/test_metadata_schema.py +++ b/tests/unit/aiplatform/test_metadata_schema.py @@ -489,14 +489,14 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - def test_base_class_instatiated_uses_schema_title(self): + def test_base_context_class_instatiated_uses_schema_title(self): class TestContext(base_context.BaseContextSchema): schema_title = _TEST_SCHEMA_TITLE context = TestContext() assert context.schema_title == _TEST_SCHEMA_TITLE - def test_base_class_parameters_overrides_default_values(self): + def test_base_context_class_parameters_overrides_default_values(self): class TestContext(base_context.BaseContextSchema): schema_title = _TEST_SCHEMA_TITLE @@ -514,12 +514,14 @@ class TestContext(base_context.BaseContextSchema): assert context.description == _TEST_DESCRIPTION assert context.metadata == _TEST_UPDATED_METADATA - def test_base_class_without_schema_title_raises_error(self): + def test_base_context_class_without_schema_title_raises_error(self): with pytest.raises(TypeError): base_context.BaseContextSchema() @pytest.mark.usefixtures("create_context_mock", "get_context_mock") - def test_create_is_called_with_default_parameters(self, create_context_mock): + def test_base_context_create_is_called_with_default_parameters( + self, create_context_mock + ): aiplatform.init(project=_TEST_PROJECT) class TestContext(base_context.BaseContextSchema): @@ -550,7 +552,7 @@ class TestContext(base_context.BaseContextSchema): "create_context_mock", "get_context_mock", ) - def test_context_create_call_sets_the_user_agent_header( + def test_base_context_create_call_sets_the_user_agent_header( self, initializer_create_client_mock ): aiplatform.init(project=_TEST_PROJECT) @@ -576,7 +578,7 @@ class TestContext(base_context.BaseContextSchema): "create_context_mock", "get_context_mock", ) - def test_context_init_call_sets_the_user_agent_header( + def test_base_context_init_call_sets_the_user_agent_header( self, initializer_create_client_mock ): aiplatform.init(project=_TEST_PROJECT) From 7c5c96822ef55dbb72f8a6c9a459e51cb56855d0 Mon Sep 17 00:00:00 2001 From: sina chavoshi Date: Tue, 16 Aug 2022 17:05:09 +0000 Subject: [PATCH 4/4] correct spelling --- google/cloud/aiplatform/metadata/artifact.py | 2 +- google/cloud/aiplatform/metadata/context.py | 2 +- google/cloud/aiplatform/metadata/execution.py | 2 +- google/cloud/aiplatform/metadata/schema/base_artifact.py | 2 +- google/cloud/aiplatform/metadata/schema/base_context.py | 2 +- google/cloud/aiplatform/metadata/schema/base_execution.py | 6 +++--- 6 files changed, 8 insertions(+), 8 deletions(-) diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index 6737a1ff6b..c995ecf85b 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -326,7 +326,7 @@ def create( Artifact: Instantiated representation of the managed Metadata Artifact. """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = ( "aiplatform.metadata.artifact.Artifact.create" diff --git a/google/cloud/aiplatform/metadata/context.py b/google/cloud/aiplatform/metadata/context.py index 280298645e..c827f865d6 100644 --- a/google/cloud/aiplatform/metadata/context.py +++ b/google/cloud/aiplatform/metadata/context.py @@ -138,7 +138,7 @@ def create( Context: Instantiated representation of the managed Metadata Context. """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = ( "aiplatform.metadata.context.Context.create" diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index 7fc24b2457..5a4e19f4f8 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -143,7 +143,7 @@ def create( """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = ( "aiplatform.metadata.execution.Execution.create" diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py index 4c96e0a6d9..e4b90d8b6a 100644 --- a/google/cloud/aiplatform/metadata/schema/base_artifact.py +++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py @@ -116,7 +116,7 @@ def _init_with_resource_name( projects/123/locations/us-central1/metadataStores//artifacts/. """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_artifact.BaseArtifactSchema._init_with_resource_name" diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py index 0ad52c2529..618bda3b60 100644 --- a/google/cloud/aiplatform/metadata/schema/base_context.py +++ b/google/cloud/aiplatform/metadata/schema/base_context.py @@ -93,7 +93,7 @@ def _init_with_resource_name( projects/123/locations/us-central1/metadataStores//contexts/. """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_context.BaseContextSchema._init_with_resource_name" diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py index 10a4b7d334..2f392c856c 100644 --- a/google/cloud/aiplatform/metadata/schema/base_execution.py +++ b/google/cloud/aiplatform/metadata/schema/base_execution.py @@ -102,7 +102,7 @@ def _init_with_resource_name( projects/123/locations/us-central1/metadataStores//executions/. """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. if not base_constants.USER_AGENT_SDK_COMMAND: base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_execution.BaseExecutionSchema._init_with_resource_name" @@ -138,7 +138,7 @@ def create( """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. base_constants.USER_AGENT_SDK_COMMAND = ( "aiplatform.metadata.schema.base_execution.BaseExecutionSchema.create" ) @@ -221,7 +221,7 @@ def start_execution( ValueError: If metadata_store_id other than 'default' is provided. """ # Add User Agent Header for metrics tracking if one is not specified - # If one is alreayd specified this call was initiated by a sub class. + # If one is already specified this call was initiated by a sub class. base_constants.USER_AGENT_SDK_COMMAND = "aiplatform.metadata.schema.base_execution.BaseExecutionSchema.start_execution"