From c40ec85e1fca2bee6813f52712d063f96264ec2c Mon Sep 17 00:00:00 2001 From: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com> Date: Thu, 18 Nov 2021 09:46:43 -0500 Subject: [PATCH] feat: Upgrade Tensorboard from v1beta1 to v1 (#849) --- google/cloud/aiplatform/compat/__init__.py | 10 +++- .../aiplatform/compat/services/__init__.py | 4 ++ .../cloud/aiplatform/compat/types/__init__.py | 18 +++++-- google/cloud/aiplatform/jobs.py | 22 +------- .../aiplatform/tensorboard/tensorboard.py | 12 ++--- google/cloud/aiplatform/utils/__init__.py | 6 ++- tests/system/aiplatform/test_tensorboard.py | 44 ++++++++++++++++ tests/unit/aiplatform/test_custom_job.py | 51 +++++++------------ .../test_hyperparameter_tuning_job.py | 46 +++++++---------- tests/unit/aiplatform/test_tensorboard.py | 4 +- 10 files changed, 117 insertions(+), 100 deletions(-) create mode 100644 tests/system/aiplatform/test_tensorboard.py diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index fd814bbc37..c8e2db1b9e 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -77,13 +77,14 @@ types.specialist_pool = types.specialist_pool_v1beta1 types.specialist_pool_service = types.specialist_pool_service_v1beta1 types.study = types.study_v1beta1 - types.training_pipeline = types.training_pipeline_v1beta1 + types.tensorboard = types.tensorboard_v1beta1 types.tensorboard_service = types.tensorboard_service_v1beta1 types.tensorboard_data = types.tensorboard_data_v1beta1 types.tensorboard_experiment = types.tensorboard_experiment_v1beta1 types.tensorboard_run = types.tensorboard_run_v1beta1 types.tensorboard_service = types.tensorboard_service_v1beta1 types.tensorboard_time_series = types.tensorboard_time_series_v1beta1 + types.training_pipeline = types.training_pipeline_v1beta1 if DEFAULT_VERSION == V1: @@ -135,6 +136,13 @@ types.specialist_pool = types.specialist_pool_v1 types.specialist_pool_service = types.specialist_pool_service_v1 types.study = types.study_v1 + types.tensorboard = types.tensorboard_v1 + types.tensorboard_service = types.tensorboard_service_v1 + types.tensorboard_data = types.tensorboard_data_v1 + types.tensorboard_experiment = types.tensorboard_experiment_v1 + types.tensorboard_run = types.tensorboard_run_v1 + types.tensorboard_service = types.tensorboard_service_v1 + types.tensorboard_time_series = types.tensorboard_time_series_v1 types.training_pipeline = types.training_pipeline_v1 __all__ = ( diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py index 098d3ff861..b0feaf6f0a 100644 --- a/google/cloud/aiplatform/compat/services/__init__.py +++ b/google/cloud/aiplatform/compat/services/__init__.py @@ -67,6 +67,9 @@ from google.cloud.aiplatform_v1.services.specialist_pool_service import ( client as specialist_pool_service_client_v1, ) +from google.cloud.aiplatform_v1.services.tensorboard_service import ( + client as tensorboard_service_client_v1, +) __all__ = ( # v1 @@ -78,6 +81,7 @@ pipeline_service_client_v1, prediction_service_client_v1, specialist_pool_service_client_v1, + tensorboard_service_client_v1, # v1beta1 dataset_service_client_v1beta1, endpoint_service_client_v1beta1, diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index 1ed76d1e25..4d64c50236 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -57,13 +57,13 @@ specialist_pool as specialist_pool_v1beta1, specialist_pool_service as specialist_pool_service_v1beta1, study as study_v1beta1, - training_pipeline as training_pipeline_v1beta1, tensorboard as tensorboard_v1beta1, tensorboard_data as tensorboard_data_v1beta1, tensorboard_experiment as tensorboard_experiment_v1beta1, tensorboard_run as tensorboard_run_v1beta1, tensorboard_service as tensorboard_service_v1beta1, tensorboard_time_series as tensorboard_time_series_v1beta1, + training_pipeline as training_pipeline_v1beta1, ) from google.cloud.aiplatform_v1.types import ( accelerator_type as accelerator_type_v1, @@ -107,6 +107,12 @@ specialist_pool as specialist_pool_v1, specialist_pool_service as specialist_pool_service_v1, study as study_v1, + tensorboard as tensorboard_v1, + tensorboard_data as tensorboard_data_v1, + tensorboard_experiment as tensorboard_experiment_v1, + tensorboard_run as tensorboard_run_v1, + tensorboard_service as tensorboard_service_v1, + tensorboard_time_series as tensorboard_time_series_v1, training_pipeline as training_pipeline_v1, ) @@ -152,6 +158,12 @@ prediction_service_v1, specialist_pool_v1, specialist_pool_service_v1, + tensorboard_v1, + tensorboard_data_v1, + tensorboard_experiment_v1, + tensorboard_run_v1, + tensorboard_service_v1, + tensorboard_time_series_v1, training_pipeline_v1, # v1beta1 accelerator_type_v1beta1, @@ -194,13 +206,11 @@ prediction_service_v1beta1, specialist_pool_v1beta1, specialist_pool_service_v1beta1, - training_pipeline_v1beta1, - metadata_service_v1beta1, tensorboard_v1beta1, - tensorboard_service_v1beta1, tensorboard_data_v1beta1, tensorboard_experiment_v1beta1, tensorboard_run_v1beta1, tensorboard_service_v1beta1, tensorboard_time_series_v1beta1, + training_pipeline_v1beta1, ) diff --git a/google/cloud/aiplatform/jobs.py b/google/cloud/aiplatform/jobs.py index c1ba9739cc..eb593c70bf 100644 --- a/google/cloud/aiplatform/jobs.py +++ b/google/cloud/aiplatform/jobs.py @@ -35,12 +35,10 @@ batch_prediction_job as gca_bp_job_compat, completion_stats as gca_completion_stats, custom_job as gca_custom_job_compat, - custom_job_v1beta1 as gca_custom_job_v1beta1, explanation as gca_explanation_compat, io as gca_io_compat, job_state as gca_job_state, hyperparameter_tuning_job as gca_hyperparameter_tuning_job_compat, - hyperparameter_tuning_job_v1beta1 as gca_hyperparameter_tuning_job_v1beta1, machine_resources as gca_machine_resources_compat, study as gca_study_compat, ) @@ -1388,17 +1386,11 @@ def run( self._gca_resource.job_spec.enable_web_access = enable_web_access if tensorboard: - v1beta1_gca_resource = gca_custom_job_v1beta1.CustomJob() - v1beta1_gca_resource._pb.MergeFromString( - self._gca_resource._pb.SerializeToString() - ) - self._gca_resource = v1beta1_gca_resource self._gca_resource.job_spec.tensorboard = tensorboard _LOGGER.log_create_with_lro(self.__class__) - version = "v1beta1" if tensorboard else "v1" - self._gca_resource = self.api_client.select_version(version).create_custom_job( + self._gca_resource = self.api_client.create_custom_job( parent=self._parent, custom_job=self._gca_resource ) @@ -1773,21 +1765,11 @@ def run( self._gca_resource.trial_job_spec.enable_web_access = enable_web_access if tensorboard: - v1beta1_gca_resource = ( - gca_hyperparameter_tuning_job_v1beta1.HyperparameterTuningJob() - ) - v1beta1_gca_resource._pb.MergeFromString( - self._gca_resource._pb.SerializeToString() - ) - self._gca_resource = v1beta1_gca_resource self._gca_resource.trial_job_spec.tensorboard = tensorboard _LOGGER.log_create_with_lro(self.__class__) - version = "v1beta1" if tensorboard else "v1" - self._gca_resource = self.api_client.select_version( - version - ).create_hyperparameter_tuning_job( + self._gca_resource = self.api_client.create_hyperparameter_tuning_job( parent=self._parent, hyperparameter_tuning_job=self._gca_resource ) diff --git a/google/cloud/aiplatform/tensorboard/tensorboard.py b/google/cloud/aiplatform/tensorboard/tensorboard.py index 3fe6507968..1e41cc9755 100644 --- a/google/cloud/aiplatform/tensorboard/tensorboard.py +++ b/google/cloud/aiplatform/tensorboard/tensorboard.py @@ -18,17 +18,13 @@ from typing import Optional, Sequence, Dict, Tuple from google.auth import credentials as auth_credentials +from google.protobuf import field_mask_pb2 from google.cloud.aiplatform import base -from google.cloud.aiplatform import compat +from google.cloud.aiplatform.compat.types import tensorboard as gca_tensorboard from google.cloud.aiplatform import initializer from google.cloud.aiplatform import utils - -from google.cloud.aiplatform.compat.types import tensorboard_v1beta1 as gca_tensorboard - -from google.protobuf import field_mask_pb2 - _LOGGER = base.Logger(__name__) @@ -156,8 +152,7 @@ def create( ) encryption_spec = initializer.global_config.get_encryption_spec( - encryption_spec_key_name=encryption_spec_key_name, - select_version=compat.V1BETA1, + encryption_spec_key_name=encryption_spec_key_name ) gapic_tensorboard = gca_tensorboard.Tensorboard( @@ -254,7 +249,6 @@ def update( if encryption_spec_key_name: encryption_spec = initializer.global_config.get_encryption_spec( encryption_spec_key_name=encryption_spec_key_name, - select_version=compat.V1BETA1, ) update_mask.append("encryption_spec") diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index 56baa42a9d..379ebfc179 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -51,6 +51,7 @@ model_service_client_v1, pipeline_service_client_v1, prediction_service_client_v1, + tensorboard_service_client_v1, ) from google.cloud.aiplatform.compat.types import ( @@ -67,6 +68,7 @@ pipeline_service_client_v1beta1.PipelineServiceClient, job_service_client_v1beta1.JobServiceClient, metadata_service_client_v1beta1.MetadataServiceClient, + tensorboard_service_client_v1beta1.TensorboardServiceClient, # v1 dataset_service_client_v1.DatasetServiceClient, endpoint_service_client_v1.EndpointServiceClient, @@ -75,6 +77,7 @@ prediction_service_client_v1.PredictionServiceClient, pipeline_service_client_v1.PipelineServiceClient, job_service_client_v1.JobServiceClient, + tensorboard_service_client_v1.TensorboardServiceClient, ) RESOURCE_NAME_PATTERN = re.compile( @@ -506,8 +509,9 @@ class MetadataClientWithOverride(ClientWithOverride): class TensorboardClientWithOverride(ClientWithOverride): _is_temporary = False - _default_version = compat.V1BETA1 + _default_version = compat.DEFAULT_VERSION _version_map = ( + (compat.V1, tensorboard_service_client_v1.TensorboardServiceClient), (compat.V1BETA1, tensorboard_service_client_v1beta1.TensorboardServiceClient), ) diff --git a/tests/system/aiplatform/test_tensorboard.py b/tests/system/aiplatform/test_tensorboard.py new file mode 100644 index 0000000000..9ec8179ca5 --- /dev/null +++ b/tests/system/aiplatform/test_tensorboard.py @@ -0,0 +1,44 @@ +# -*- 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 + + +class TestTensorboard(e2e_base.TestEndToEnd): + + _temp_prefix = "temp-vertex-sdk-e2e-test" + + def test_create_and_get_tensorboard(self, shared_state): + + aiplatform.init( + project=e2e_base._PROJECT, location=e2e_base._LOCATION, + ) + + display_name = self._make_display_name("tensorboard") + + tb = aiplatform.Tensorboard.create(display_name=display_name) + + shared_state["resources"] = [tb] + + get_tb = aiplatform.Tensorboard(tb.resource_name) + + assert tb.resource_name == get_tb.resource_name + + list_tb = aiplatform.Tensorboard.list() + + assert len(list_tb) > 0 diff --git a/tests/unit/aiplatform/test_custom_job.py b/tests/unit/aiplatform/test_custom_job.py index ea362ecda8..f865385e7d 100644 --- a/tests/unit/aiplatform/test_custom_job.py +++ b/tests/unit/aiplatform/test_custom_job.py @@ -31,18 +31,12 @@ from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform.compat.types import custom_job as gca_custom_job_compat -from google.cloud.aiplatform.compat.types import ( - custom_job_v1beta1 as gca_custom_job_v1beta1, -) from google.cloud.aiplatform.compat.types import io as gca_io_compat from google.cloud.aiplatform.compat.types import job_state as gca_job_state_compat from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec_compat, ) from google.cloud.aiplatform_v1.services.job_service import client as job_service_client -from google.cloud.aiplatform_v1beta1.services.job_service import ( - client as job_service_client_v1beta1, -) _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" @@ -114,29 +108,16 @@ ) -def _get_custom_job_proto(state=None, name=None, error=None, version="v1"): +def _get_custom_job_proto(state=None, name=None, error=None): custom_job_proto = copy.deepcopy(_TEST_BASE_CUSTOM_JOB_PROTO) custom_job_proto.name = name custom_job_proto.state = state custom_job_proto.error = error - - if version == "v1beta1": - v1beta1_custom_job_proto = gca_custom_job_v1beta1.CustomJob() - v1beta1_custom_job_proto._pb.MergeFromString( - custom_job_proto._pb.SerializeToString() - ) - custom_job_proto = v1beta1_custom_job_proto - custom_job_proto.job_spec.tensorboard = _TEST_TENSORBOARD_NAME - return custom_job_proto -def _get_custom_job_proto_with_enable_web_access( - state=None, name=None, error=None, version="v1" -): - custom_job_proto = _get_custom_job_proto( - state=state, name=name, error=error, version=version - ) +def _get_custom_job_proto_with_enable_web_access(state=None, name=None, error=None): + custom_job_proto = _get_custom_job_proto(state=state, name=name, error=error) custom_job_proto.job_spec.enable_web_access = _TEST_ENABLE_WEB_ACCESS if state == gca_job_state_compat.JobState.JOB_STATE_RUNNING: custom_job_proto.web_access_uris = _TEST_WEB_ACCESS_URIS @@ -260,24 +241,25 @@ def create_custom_job_mock_with_enable_web_access(): @pytest.fixture -def create_custom_job_mock_fail(): +def create_custom_job_mock_with_tensorboard(): with mock.patch.object( job_service_client.JobServiceClient, "create_custom_job" ) as create_custom_job_mock: - create_custom_job_mock.side_effect = RuntimeError("Mock fail") + custom_job_proto = _get_custom_job_proto( + name=_TEST_CUSTOM_JOB_NAME, + state=gca_job_state_compat.JobState.JOB_STATE_PENDING, + ) + custom_job_proto.job_spec.tensorboard = _TEST_TENSORBOARD_NAME + create_custom_job_mock.return_value = custom_job_proto yield create_custom_job_mock @pytest.fixture -def create_custom_job_v1beta1_mock(): +def create_custom_job_mock_fail(): with mock.patch.object( - job_service_client_v1beta1.JobServiceClient, "create_custom_job" + job_service_client.JobServiceClient, "create_custom_job" ) as create_custom_job_mock: - create_custom_job_mock.return_value = _get_custom_job_proto( - name=_TEST_CUSTOM_JOB_NAME, - state=gca_job_state_compat.JobState.JOB_STATE_PENDING, - version="v1beta1", - ) + create_custom_job_mock.side_effect = RuntimeError("Mock fail") yield create_custom_job_mock @@ -573,7 +555,7 @@ def test_get_web_access_uris_job_succeeded( @pytest.mark.parametrize("sync", [True, False]) def test_create_custom_job_with_tensorboard( - self, create_custom_job_v1beta1_mock, get_custom_job_mock, sync + self, create_custom_job_mock_with_tensorboard, get_custom_job_mock, sync ): aiplatform.init( @@ -601,9 +583,10 @@ def test_create_custom_job_with_tensorboard( job.wait() - expected_custom_job = _get_custom_job_proto(version="v1beta1") + expected_custom_job = _get_custom_job_proto() + expected_custom_job.job_spec.tensorboard = _TEST_TENSORBOARD_NAME - create_custom_job_v1beta1_mock.assert_called_once_with( + create_custom_job_mock_with_tensorboard.assert_called_once_with( parent=_TEST_PARENT, custom_job=expected_custom_job ) diff --git a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py index 7b4f434736..c14309ace6 100644 --- a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py +++ b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py @@ -30,14 +30,10 @@ from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec_compat, hyperparameter_tuning_job as gca_hyperparameter_tuning_job_compat, - hyperparameter_tuning_job_v1beta1 as gca_hyperparameter_tuning_job_v1beta1, job_state as gca_job_state_compat, study as gca_study_compat, ) from google.cloud.aiplatform_v1.services.job_service import client as job_service_client -from google.cloud.aiplatform_v1beta1.services.job_service import ( - client as job_service_client_v1beta1, -) import test_custom_job @@ -131,9 +127,7 @@ _TEST_BASE_TRIAL_PROTO = gca_study_compat.Trial() -def _get_hyperparameter_tuning_job_proto( - state=None, name=None, error=None, version="v1" -): +def _get_hyperparameter_tuning_job_proto(state=None, name=None, error=None): hyperparameter_tuning_job_proto = copy.deepcopy( _TEST_BASE_HYPERPARAMETER_TUNING_JOB_PROTO ) @@ -141,18 +135,6 @@ def _get_hyperparameter_tuning_job_proto( hyperparameter_tuning_job_proto.state = state hyperparameter_tuning_job_proto.error = error - if version == "v1beta1": - v1beta1_hyperparameter_tuning_job_proto = ( - gca_hyperparameter_tuning_job_v1beta1.HyperparameterTuningJob() - ) - v1beta1_hyperparameter_tuning_job_proto._pb.MergeFromString( - hyperparameter_tuning_job_proto._pb.SerializeToString() - ) - hyperparameter_tuning_job_proto = v1beta1_hyperparameter_tuning_job_proto - hyperparameter_tuning_job_proto.trial_job_spec.tensorboard = ( - test_custom_job._TEST_TENSORBOARD_NAME - ) - return hyperparameter_tuning_job_proto @@ -166,10 +148,10 @@ def _get_trial_proto(id=None, state=None): def _get_hyperparameter_tuning_job_proto_with_enable_web_access( - state=None, name=None, error=None, version="v1", trials=[] + state=None, name=None, error=None, trials=[] ): hyperparameter_tuning_job_proto = _get_hyperparameter_tuning_job_proto( - state=state, name=name, error=error, version=version + state=state, name=name, error=error, ) hyperparameter_tuning_job_proto.trial_job_spec.enable_web_access = ( test_custom_job._TEST_ENABLE_WEB_ACCESS @@ -342,14 +324,19 @@ def create_hyperparameter_tuning_job_mock_fail(): @pytest.fixture -def create_hyperparameter_tuning_job_v1beta1_mock(): +def create_hyperparameter_tuning_job_mock_with_tensorboard(): with mock.patch.object( - job_service_client_v1beta1.JobServiceClient, "create_hyperparameter_tuning_job" + job_service_client.JobServiceClient, "create_hyperparameter_tuning_job" ) as create_hyperparameter_tuning_job_mock: - create_hyperparameter_tuning_job_mock.return_value = _get_hyperparameter_tuning_job_proto( + hyperparameter_tuning_job_proto = _get_hyperparameter_tuning_job_proto( name=_TEST_HYPERPARAMETERTUNING_JOB_NAME, state=gca_job_state_compat.JobState.JOB_STATE_PENDING, - version="v1beta1", + ) + hyperparameter_tuning_job_proto.trial_job_spec.tensorboard = ( + test_custom_job._TEST_TENSORBOARD_NAME + ) + create_hyperparameter_tuning_job_mock.return_value = ( + hyperparameter_tuning_job_proto ) yield create_hyperparameter_tuning_job_mock @@ -608,7 +595,7 @@ def test_get_hyperparameter_tuning_job(self, get_hyperparameter_tuning_job_mock) @pytest.mark.parametrize("sync", [True, False]) def test_create_hyperparameter_tuning_job_with_tensorboard( self, - create_hyperparameter_tuning_job_v1beta1_mock, + create_hyperparameter_tuning_job_mock_with_tensorboard, get_hyperparameter_tuning_job_mock, sync, ): @@ -659,11 +646,12 @@ def test_create_hyperparameter_tuning_job_with_tensorboard( job.wait() - expected_hyperparameter_tuning_job = _get_hyperparameter_tuning_job_proto( - version="v1beta1" + expected_hyperparameter_tuning_job = _get_hyperparameter_tuning_job_proto() + expected_hyperparameter_tuning_job.trial_job_spec.tensorboard = ( + test_custom_job._TEST_TENSORBOARD_NAME ) - create_hyperparameter_tuning_job_v1beta1_mock.assert_called_once_with( + create_hyperparameter_tuning_job_mock_with_tensorboard.assert_called_once_with( parent=_TEST_PARENT, hyperparameter_tuning_job=expected_hyperparameter_tuning_job, ) diff --git a/tests/unit/aiplatform/test_tensorboard.py b/tests/unit/aiplatform/test_tensorboard.py index ffa4c72850..38ea935950 100644 --- a/tests/unit/aiplatform/test_tensorboard.py +++ b/tests/unit/aiplatform/test_tensorboard.py @@ -32,11 +32,11 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform import tensorboard -from google.cloud.aiplatform_v1beta1.services.tensorboard_service import ( +from google.cloud.aiplatform_v1.services.tensorboard_service import ( client as tensorboard_service_client, ) -from google.cloud.aiplatform_v1beta1.types import ( +from google.cloud.aiplatform_v1.types import ( tensorboard as gca_tensorboard, tensorboard_service as gca_tensorboard_service, encryption_spec as gca_encryption_spec,