From 6c7e1fa06915ec02bbe95d8f534b2a15b3ac323c Mon Sep 17 00:00:00 2001 From: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com> Date: Thu, 21 Oct 2021 15:28:52 -0400 Subject: [PATCH] chore: upgrade GAPIC Metadata from v1beta1 to v1 and add experiment logging integration test (#778) --- google/cloud/aiplatform/compat/__init__.py | 13 +++- .../aiplatform/compat/services/__init__.py | 10 ++- .../cloud/aiplatform/compat/types/__init__.py | 25 ++++++- google/cloud/aiplatform/metadata/artifact.py | 13 ++-- google/cloud/aiplatform/metadata/context.py | 13 ++-- google/cloud/aiplatform/metadata/execution.py | 29 ++++---- .../aiplatform/metadata/metadata_store.py | 4 +- google/cloud/aiplatform/metadata/resource.py | 12 ++-- google/cloud/aiplatform/utils/__init__.py | 7 +- tests/system/aiplatform/test_metadata.py | 69 +++++++++++++++++++ tests/unit/aiplatform/test_metadata.py | 12 ++-- .../aiplatform/test_metadata_resources.py | 12 ++-- tests/unit/aiplatform/test_metadata_store.py | 8 +-- 13 files changed, 174 insertions(+), 53 deletions(-) create mode 100644 tests/system/aiplatform/test_metadata.py diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index 86dedd373e..fd814bbc37 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -40,8 +40,10 @@ types.accelerator_type = types.accelerator_type_v1beta1 types.annotation = types.annotation_v1beta1 types.annotation_spec = types.annotation_spec_v1beta1 + types.artifact = types.artifact_v1beta1 types.batch_prediction_job = types.batch_prediction_job_v1beta1 types.completion_stats = types.completion_stats_v1beta1 + types.context = types.context_v1beta1 types.custom_job = types.custom_job_v1beta1 types.data_item = types.data_item_v1beta1 types.data_labeling_job = types.data_labeling_job_v1beta1 @@ -52,6 +54,8 @@ types.endpoint = types.endpoint_v1beta1 types.endpoint_service = types.endpoint_service_v1beta1 types.env_var = types.env_var_v1beta1 + types.event = types.event_v1beta1 + types.execution = types.execution_v1beta1 types.explanation = types.explanation_v1beta1 types.explanation_metadata = types.explanation_metadata_v1beta1 types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1beta1 @@ -60,6 +64,8 @@ types.job_state = types.job_state_v1beta1 types.machine_resources = types.machine_resources_v1beta1 types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1beta1 + types.metadata_service = types.metadata_service_v1beta1 + types.metadata_store = types.metadata_store_v1beta1 types.model = types.model_v1beta1 types.model_evaluation = types.model_evaluation_v1beta1 types.model_evaluation_slice = types.model_evaluation_slice_v1beta1 @@ -72,7 +78,6 @@ types.specialist_pool_service = types.specialist_pool_service_v1beta1 types.study = types.study_v1beta1 types.training_pipeline = types.training_pipeline_v1beta1 - types.metadata_service = types.metadata_service_v1beta1 types.tensorboard_service = types.tensorboard_service_v1beta1 types.tensorboard_data = types.tensorboard_data_v1beta1 types.tensorboard_experiment = types.tensorboard_experiment_v1beta1 @@ -93,8 +98,10 @@ types.accelerator_type = types.accelerator_type_v1 types.annotation = types.annotation_v1 types.annotation_spec = types.annotation_spec_v1 + types.artifact = types.artifact_v1 types.batch_prediction_job = types.batch_prediction_job_v1 types.completion_stats = types.completion_stats_v1 + types.context = types.context_v1 types.custom_job = types.custom_job_v1 types.data_item = types.data_item_v1 types.data_labeling_job = types.data_labeling_job_v1 @@ -105,6 +112,8 @@ types.endpoint = types.endpoint_v1 types.endpoint_service = types.endpoint_service_v1 types.env_var = types.env_var_v1 + types.event = types.event_v1 + types.execution = types.execution_v1 types.explanation = types.explanation_v1 types.explanation_metadata = types.explanation_metadata_v1 types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1 @@ -113,6 +122,8 @@ types.job_state = types.job_state_v1 types.machine_resources = types.machine_resources_v1 types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1 + types.metadata_service = types.metadata_service_v1 + types.metadata_store = types.metadata_store_v1 types.model = types.model_v1 types.model_evaluation = types.model_evaluation_v1 types.model_evaluation_slice = types.model_evaluation_slice_v1 diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py index 5c104ab41f..098d3ff861 100644 --- a/google/cloud/aiplatform/compat/services/__init__.py +++ b/google/cloud/aiplatform/compat/services/__init__.py @@ -24,6 +24,9 @@ from google.cloud.aiplatform_v1beta1.services.job_service import ( client as job_service_client_v1beta1, ) +from google.cloud.aiplatform_v1beta1.services.metadata_service import ( + client as metadata_service_client_v1beta1, +) from google.cloud.aiplatform_v1beta1.services.model_service import ( client as model_service_client_v1beta1, ) @@ -36,9 +39,6 @@ from google.cloud.aiplatform_v1beta1.services.specialist_pool_service import ( client as specialist_pool_service_client_v1beta1, ) -from google.cloud.aiplatform_v1beta1.services.metadata_service import ( - client as metadata_service_client_v1beta1, -) from google.cloud.aiplatform_v1beta1.services.tensorboard_service import ( client as tensorboard_service_client_v1beta1, ) @@ -52,6 +52,9 @@ from google.cloud.aiplatform_v1.services.job_service import ( client as job_service_client_v1, ) +from google.cloud.aiplatform_v1.services.metadata_service import ( + client as metadata_service_client_v1, +) from google.cloud.aiplatform_v1.services.model_service import ( client as model_service_client_v1, ) @@ -70,6 +73,7 @@ dataset_service_client_v1, endpoint_service_client_v1, job_service_client_v1, + metadata_service_client_v1, model_service_client_v1, pipeline_service_client_v1, prediction_service_client_v1, diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index f0e4117e7e..3177e15d1c 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -19,8 +19,10 @@ accelerator_type as accelerator_type_v1beta1, annotation as annotation_v1beta1, annotation_spec as annotation_spec_v1beta1, + artifact as artifact_v1beta1, batch_prediction_job as batch_prediction_job_v1beta1, completion_stats as completion_stats_v1beta1, + context as context_v1beta1, custom_job as custom_job_v1beta1, data_item as data_item_v1beta1, data_labeling_job as data_labeling_job_v1beta1, @@ -31,6 +33,8 @@ endpoint as endpoint_v1beta1, endpoint_service as endpoint_service_v1beta1, env_var as env_var_v1beta1, + event as event_v1beta1, + execution as execution_v1beta1, explanation as explanation_v1beta1, explanation_metadata as explanation_metadata_v1beta1, hyperparameter_tuning_job as hyperparameter_tuning_job_v1beta1, @@ -39,6 +43,8 @@ job_state as job_state_v1beta1, machine_resources as machine_resources_v1beta1, manual_batch_tuning_parameters as manual_batch_tuning_parameters_v1beta1, + metadata_service as metadata_service_v1beta1, + metadata_store as metadata_store_v1beta1, model as model_v1beta1, model_evaluation as model_evaluation_v1beta1, model_evaluation_slice as model_evaluation_slice_v1beta1, @@ -52,7 +58,6 @@ specialist_pool_service as specialist_pool_service_v1beta1, study as study_v1beta1, training_pipeline as training_pipeline_v1beta1, - metadata_service as metadata_service_v1beta1, tensorboard as tensorboard_v1beta1, tensorboard_data as tensorboard_data_v1beta1, tensorboard_experiment as tensorboard_experiment_v1beta1, @@ -64,8 +69,10 @@ accelerator_type as accelerator_type_v1, annotation as annotation_v1, annotation_spec as annotation_spec_v1, + artifact as artifact_v1, batch_prediction_job as batch_prediction_job_v1, completion_stats as completion_stats_v1, + context as context_v1, custom_job as custom_job_v1, data_item as data_item_v1, data_labeling_job as data_labeling_job_v1, @@ -76,6 +83,8 @@ endpoint as endpoint_v1, endpoint_service as endpoint_service_v1, env_var as env_var_v1, + event as event_v1, + execution as execution_v1, explanation as explanation_v1, explanation_metadata as explanation_metadata_v1, hyperparameter_tuning_job as hyperparameter_tuning_job_v1, @@ -84,6 +93,8 @@ job_state as job_state_v1, machine_resources as machine_resources_v1, manual_batch_tuning_parameters as manual_batch_tuning_parameters_v1, + metadata_service as metadata_service_v1, + metadata_store as metadata_store_v1, model as model_v1, model_evaluation as model_evaluation_v1, model_evaluation_slice as model_evaluation_slice_v1, @@ -103,8 +114,10 @@ accelerator_type_v1, annotation_v1, annotation_spec_v1, + artifact_v1, batch_prediction_job_v1, completion_stats_v1, + context_v1, custom_job_v1, data_item_v1, data_labeling_job_v1, @@ -115,6 +128,8 @@ endpoint_v1, endpoint_service_v1, env_var_v1, + event_v1, + execution_v1, explanation_v1, explanation_metadata_v1, hyperparameter_tuning_job_v1, @@ -123,6 +138,8 @@ job_state_v1, machine_resources_v1, manual_batch_tuning_parameters_v1, + metadata_service_v1, + metadata_store_v1, model_v1, model_evaluation_v1, model_evaluation_slice_v1, @@ -138,8 +155,10 @@ accelerator_type_v1beta1, annotation_v1beta1, annotation_spec_v1beta1, + artifact_v1beta1, batch_prediction_job_v1beta1, completion_stats_v1beta1, + context_v1beta1, custom_job_v1beta1, data_item_v1beta1, data_labeling_job_v1beta1, @@ -150,6 +169,8 @@ endpoint_v1beta1, endpoint_service_v1beta1, env_var_v1beta1, + event_v1beta1, + execution_v1beta1, explanation_v1beta1, explanation_metadata_v1beta1, hyperparameter_tuning_job_v1beta1, @@ -158,6 +179,8 @@ job_state_v1beta1, machine_resources_v1beta1, manual_batch_tuning_parameters_v1beta1, + metadata_service_v1beta1, + metadata_store_v1beta1, model_v1beta1, model_evaluation_v1beta1, model_evaluation_slice_v1beta1, diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index b3ef6e09a2..13ac2aafbf 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -20,16 +20,17 @@ import proto from google.cloud.aiplatform import utils -from google.cloud.aiplatform.metadata.resource import _Resource -from google.cloud.aiplatform_v1beta1 import ListArtifactsRequest -from google.cloud.aiplatform_v1beta1.types import artifact as gca_artifact +from google.cloud.aiplatform.compat.types import artifact as gca_artifact +from google.cloud.aiplatform.compat.types import metadata_service +from google.cloud.aiplatform.metadata import resource -class _Artifact(_Resource): +class _Artifact(resource._Resource): """Metadata Artifact resource for Vertex AI""" _resource_noun = "artifacts" _getter_method = "get_artifact" + _delete_method = "delete_artifact" @classmethod def _create_resource( @@ -86,5 +87,7 @@ def _list_resources( filter (str): Optional. filter string to restrict the list result """ - list_request = ListArtifactsRequest(parent=parent, filter=filter,) + list_request = metadata_service.ListArtifactsRequest( + parent=parent, filter=filter, + ) return client.list_artifacts(request=list_request) diff --git a/google/cloud/aiplatform/metadata/context.py b/google/cloud/aiplatform/metadata/context.py index ddd583bbdf..670d8af844 100644 --- a/google/cloud/aiplatform/metadata/context.py +++ b/google/cloud/aiplatform/metadata/context.py @@ -20,16 +20,17 @@ import proto from google.cloud.aiplatform import utils -from google.cloud.aiplatform.metadata.resource import _Resource -from google.cloud.aiplatform_v1beta1 import ListContextsRequest -from google.cloud.aiplatform_v1beta1.types import context as gca_context +from google.cloud.aiplatform.compat.types import metadata_service +from google.cloud.aiplatform.compat.types import context as gca_context +from google.cloud.aiplatform.metadata import resource -class _Context(_Resource): +class _Context(resource._Resource): """Metadata Context resource for Vertex AI""" _resource_noun = "contexts" _getter_method = "get_context" + _delete_method = "delete_context" def add_artifacts_and_executions( self, @@ -106,5 +107,7 @@ def _list_resources( Optional. filter string to restrict the list result """ - list_request = ListContextsRequest(parent=parent, filter=filter,) + list_request = metadata_service.ListContextsRequest( + parent=parent, filter=filter, + ) return client.list_contexts(request=list_request) diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index 3605efdb4f..50655d0175 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -21,18 +21,19 @@ from google.api_core import exceptions from google.cloud.aiplatform import utils -from google.cloud.aiplatform.metadata.artifact import _Artifact -from google.cloud.aiplatform.metadata.resource import _Resource -from google.cloud.aiplatform_v1beta1 import Event -from google.cloud.aiplatform_v1beta1.types import execution as gca_execution -from google.cloud.aiplatform_v1beta1.types.metadata_service import ListExecutionsRequest +from google.cloud.aiplatform.compat.types import event as gca_event +from google.cloud.aiplatform.compat.types import execution as gca_execution +from google.cloud.aiplatform.compat.types import metadata_service +from google.cloud.aiplatform.metadata import artifact +from google.cloud.aiplatform.metadata import resource -class _Execution(_Resource): +class _Execution(resource._Resource): """Metadata Execution resource for Vertex AI""" _resource_noun = "executions" _getter_method = "get_execution" + _delete_method = "delete_execution" def add_artifact( self, artifact_resource_name: str, input: bool, @@ -46,16 +47,16 @@ def add_artifact( Required. Whether Artifact is an input event to the Execution or not. """ - event = Event( + event = gca_event.Event( artifact=artifact_resource_name, - type_=Event.Type.INPUT if input else Event.Type.OUTPUT, + type_=gca_event.Event.Type.INPUT if input else gca_event.Event.Type.OUTPUT, ) self.api_client.add_execution_events( execution=self.resource_name, events=[event], ) - def query_input_and_output_artifacts(self) -> Sequence[_Artifact]: + def query_input_and_output_artifacts(self) -> Sequence[artifact._Artifact]: """query the input and output artifacts connected to the execution. Returns: @@ -70,13 +71,13 @@ def query_input_and_output_artifacts(self) -> Sequence[_Artifact]: return [] return [ - _Artifact( - resource=artifact, + artifact._Artifact( + resource=metadata_artifact, project=self.project, location=self.location, credentials=self.credentials, ) - for artifact in artifacts + for metadata_artifact in artifacts ] @classmethod @@ -120,7 +121,9 @@ def _list_resources( Optional. filter string to restrict the list result """ - list_request = ListExecutionsRequest(parent=parent, filter=filter,) + list_request = metadata_service.ListExecutionsRequest( + parent=parent, filter=filter, + ) return client.list_executions(request=list_request) @classmethod diff --git a/google/cloud/aiplatform/metadata/metadata_store.py b/google/cloud/aiplatform/metadata/metadata_store.py index 3327f47d1f..82b6742df5 100644 --- a/google/cloud/aiplatform/metadata/metadata_store.py +++ b/google/cloud/aiplatform/metadata/metadata_store.py @@ -24,7 +24,7 @@ from google.cloud.aiplatform import base, initializer from google.cloud.aiplatform import compat from google.cloud.aiplatform import utils -from google.cloud.aiplatform_v1beta1.types import metadata_store as gca_metadata_store +from google.cloud.aiplatform.compat.types import metadata_store as gca_metadata_store class _MetadataStore(base.VertexAiResourceNounWithFutureManager): @@ -176,7 +176,7 @@ def _create( gapic_metadata_store = gca_metadata_store.MetadataStore( encryption_spec=initializer.global_config.get_encryption_spec( encryption_spec_key_name=encryption_spec_key_name, - select_version=compat.V1BETA1, + select_version=compat.DEFAULT_VERSION, ) ) diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index 37bb7327cd..b225909eee 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -27,9 +27,9 @@ from google.cloud.aiplatform import base, initializer from google.cloud.aiplatform import utils -from google.cloud.aiplatform_v1beta1 import Artifact as GapicArtifact -from google.cloud.aiplatform_v1beta1 import Context as GapicContext -from google.cloud.aiplatform_v1beta1 import Execution as GapicExecution +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 class _Resource(base.VertexAiResourceNounWithFutureManager, abc.ABC): @@ -42,7 +42,9 @@ class _Resource(base.VertexAiResourceNounWithFutureManager, abc.ABC): def __init__( self, resource_name: Optional[str] = None, - resource: Optional[Union[GapicContext, GapicArtifact, GapicExecution]] = None, + resource: Optional[ + Union[gca_context.Context, gca_artifact.Artifact, gca_execution.Execution] + ] = None, metadata_store_id: str = "default", project: Optional[str] = None, location: Optional[str] = None, @@ -56,7 +58,7 @@ def __init__( Example: "projects/123/locations/us-central1/metadataStores/default//my-resource". or "my-resource" when project and location are initialized or passed. if ``resource`` is provided, this should not be set. - resource (Union[GapicContext, GapicArtifact, GapicExecution]): + resource (Union[gca_context.Context, gca_artifact.Artifact, gca_execution.Execution]): The proto.Message that contains the full information of the resource. If both set, this field overrides ``resource_name`` field. metadata_store_id (str): diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index d239b9b441..df0d9c4c62 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -37,16 +37,17 @@ dataset_service_client_v1beta1, endpoint_service_client_v1beta1, job_service_client_v1beta1, + metadata_service_client_v1beta1, model_service_client_v1beta1, pipeline_service_client_v1beta1, prediction_service_client_v1beta1, - metadata_service_client_v1beta1, tensorboard_service_client_v1beta1, ) from google.cloud.aiplatform.compat.services import ( dataset_service_client_v1, endpoint_service_client_v1, job_service_client_v1, + metadata_service_client_v1, model_service_client_v1, pipeline_service_client_v1, prediction_service_client_v1, @@ -69,6 +70,7 @@ # v1 dataset_service_client_v1.DatasetServiceClient, endpoint_service_client_v1.EndpointServiceClient, + metadata_service_client_v1.MetadataServiceClient, model_service_client_v1.ModelServiceClient, prediction_service_client_v1.PredictionServiceClient, pipeline_service_client_v1.PipelineServiceClient, @@ -494,8 +496,9 @@ class PredictionClientWithOverride(ClientWithOverride): class MetadataClientWithOverride(ClientWithOverride): _is_temporary = True - _default_version = compat.V1BETA1 + _default_version = compat.DEFAULT_VERSION _version_map = ( + (compat.V1, metadata_service_client_v1.MetadataServiceClient), (compat.V1BETA1, metadata_service_client_v1beta1.MetadataServiceClient), ) diff --git a/tests/system/aiplatform/test_metadata.py b/tests/system/aiplatform/test_metadata.py new file mode 100644 index 0000000000..3d2116fa0a --- /dev/null +++ b/tests/system/aiplatform/test_metadata.py @@ -0,0 +1,69 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from google.cloud import aiplatform +from tests.system.aiplatform import e2e_base + + +PARAMS = {"sdk-param-test-1": 0.1, "sdk-param-test-2": 0.2} + +METRICS = {"sdk-metric-test-1": 0.8, "sdk-metric-test-2": 100} + + +class TestMetadata(e2e_base.TestEndToEnd): + + _temp_prefix = "temp-vertex-sdk-e2e-test" + + def test_experiment_logging(self, shared_state): + + # Truncating the name because of resource id constraints from the service + experiment_name = self._make_display_name("experiment")[:56] + + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + experiment=experiment_name, + ) + + shared_state["resources"] = [aiplatform.metadata.metadata_service._experiment] + + # Truncating the name because of resource id constraints from the service + run_name = self._make_display_name("run")[:56] + + aiplatform.start_run(run_name) + + shared_state["resources"].extend( + [ + aiplatform.metadata.metadata_service._run, + aiplatform.metadata.metadata_service._metrics, + ] + ) + + aiplatform.log_params(PARAMS) + + aiplatform.log_metrics(METRICS) + + df = aiplatform.get_experiment_df() + + true_df_dict = {f"metric.{key}": value for key, value in METRICS.items()} + for key, value in PARAMS.items(): + true_df_dict[f"param.{key}"] = value + + true_df_dict["experiment_name"] = experiment_name + true_df_dict["run_name"] = run_name + + assert true_df_dict == df.to_dict("records")[0] diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index 36de8938b4..70b8ea6b5c 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -28,20 +28,20 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import metadata -from google.cloud.aiplatform_v1beta1 import ( +from google.cloud.aiplatform_v1 import ( AddContextArtifactsAndExecutionsResponse, Event, LineageSubgraph, ListExecutionsRequest, ) -from google.cloud.aiplatform_v1beta1 import Artifact as GapicArtifact -from google.cloud.aiplatform_v1beta1 import Context as GapicContext -from google.cloud.aiplatform_v1beta1 import Execution as GapicExecution -from google.cloud.aiplatform_v1beta1 import ( +from google.cloud.aiplatform_v1 import Artifact as GapicArtifact +from google.cloud.aiplatform_v1 import Context as GapicContext +from google.cloud.aiplatform_v1 import Execution as GapicExecution +from google.cloud.aiplatform_v1 import ( MetadataServiceClient, AddExecutionEventsResponse, ) -from google.cloud.aiplatform_v1beta1 import MetadataStore as GapicMetadataStore +from google.cloud.aiplatform_v1 import MetadataStore as GapicMetadataStore # project diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py index 19258aef3c..7443dd695f 100644 --- a/tests/unit/aiplatform/test_metadata_resources.py +++ b/tests/unit/aiplatform/test_metadata_resources.py @@ -26,12 +26,12 @@ from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import context from google.cloud.aiplatform.metadata import execution -from google.cloud.aiplatform_v1beta1 import AddContextArtifactsAndExecutionsResponse -from google.cloud.aiplatform_v1beta1 import Artifact as GapicArtifact -from google.cloud.aiplatform_v1beta1 import Context as GapicContext -from google.cloud.aiplatform_v1beta1 import Execution as GapicExecution -from google.cloud.aiplatform_v1beta1 import LineageSubgraph -from google.cloud.aiplatform_v1beta1 import ( +from google.cloud.aiplatform_v1 import AddContextArtifactsAndExecutionsResponse +from google.cloud.aiplatform_v1 import Artifact as GapicArtifact +from google.cloud.aiplatform_v1 import Context as GapicContext +from google.cloud.aiplatform_v1 import Execution as GapicExecution +from google.cloud.aiplatform_v1 import LineageSubgraph +from google.cloud.aiplatform_v1 import ( MetadataServiceClient, AddExecutionEventsResponse, Event, diff --git a/tests/unit/aiplatform/test_metadata_store.py b/tests/unit/aiplatform/test_metadata_store.py index 516e61d849..a10a86675d 100644 --- a/tests/unit/aiplatform/test_metadata_store.py +++ b/tests/unit/aiplatform/test_metadata_store.py @@ -28,10 +28,10 @@ from google.cloud import aiplatform from google.cloud.aiplatform import initializer from google.cloud.aiplatform.metadata import metadata_store -from google.cloud.aiplatform_v1beta1 import MetadataServiceClient -from google.cloud.aiplatform_v1beta1 import MetadataStore as GapicMetadataStore -from google.cloud.aiplatform_v1beta1.types import encryption_spec as gca_encryption_spec -from google.cloud.aiplatform_v1beta1.types import metadata_service +from google.cloud.aiplatform_v1 import MetadataServiceClient +from google.cloud.aiplatform_v1 import MetadataStore as GapicMetadataStore +from google.cloud.aiplatform_v1.types import encryption_spec as gca_encryption_spec +from google.cloud.aiplatform_v1.types import metadata_service # project _TEST_PROJECT = "test-project"