From b613b264524aaab2cb65e63a5487770736faa7c8 Mon Sep 17 00:00:00 2001 From: Yicheng Fang <58752348+yfang1@users.noreply.github.com> Date: Fri, 19 Nov 2021 11:02:48 -0800 Subject: [PATCH] feat: updating Tensorboard related code to use v1 (#851) Co-authored-by: Yicheng Fang --- google/cloud/aiplatform/compat/__init__.py | 1 + .../plugins/tf_profiler/profile_uploader.py | 16 ++++-------- .../cloud/aiplatform/tensorboard/uploader.py | 20 +++++--------- .../aiplatform/tensorboard/uploader_utils.py | 20 +++++--------- tests/unit/aiplatform/test_uploader.py | 26 +++++++++---------- 5 files changed, 31 insertions(+), 52 deletions(-) diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index c8e2db1b9e..057ba9344f 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -95,6 +95,7 @@ services.pipeline_service_client = services.pipeline_service_client_v1 services.prediction_service_client = services.prediction_service_client_v1 services.specialist_pool_service_client = services.specialist_pool_service_client_v1 + services.tensorboard_service_client = services.tensorboard_service_client_v1 types.accelerator_type = types.accelerator_type_v1 types.annotation = types.annotation_v1 diff --git a/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py b/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py index 41020f870e..59dedd0142 100644 --- a/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py +++ b/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py @@ -38,20 +38,14 @@ import tensorflow as tf from google.cloud import storage -from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1 -from google.cloud.aiplatform.compat.types import ( - tensorboard_data_v1beta1 as tensorboard_data, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_service_v1beta1 as tensorboard_service, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_time_series_v1beta1 as tensorboard_time_series, -) +from google.cloud.aiplatform.compat.services import tensorboard_service_client +from google.cloud.aiplatform.compat.types import tensorboard_data +from google.cloud.aiplatform.compat.types import tensorboard_service +from google.cloud.aiplatform.compat.types import tensorboard_time_series from google.cloud.aiplatform.tensorboard import uploader_utils from google.protobuf import timestamp_pb2 as timestamp -TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient +TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient logger = tb_logging.get_logger() diff --git a/google/cloud/aiplatform/tensorboard/uploader.py b/google/cloud/aiplatform/tensorboard/uploader.py index 0df67f1a06..735279b03d 100644 --- a/google/cloud/aiplatform/tensorboard/uploader.py +++ b/google/cloud/aiplatform/tensorboard/uploader.py @@ -56,25 +56,17 @@ from google.api_core import exceptions from google.cloud import storage -from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1 -from google.cloud.aiplatform.compat.types import ( - tensorboard_data_v1beta1 as tensorboard_data, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_experiment_v1beta1 as tensorboard_experiment, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_service_v1beta1 as tensorboard_service, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_time_series_v1beta1 as tensorboard_time_series, -) +from google.cloud.aiplatform.compat.services import tensorboard_service_client +from google.cloud.aiplatform.compat.types import tensorboard_data +from google.cloud.aiplatform.compat.types import tensorboard_experiment +from google.cloud.aiplatform.compat.types import tensorboard_service +from google.cloud.aiplatform.compat.types import tensorboard_time_series from google.cloud.aiplatform.tensorboard import uploader_utils from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader from google.protobuf import message from google.protobuf import timestamp_pb2 as timestamp -TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient +TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient # Minimum length of a logdir polling cycle in seconds. Shorter cycles will # sleep to avoid spinning over the logdir, which isn't great for disks and can diff --git a/google/cloud/aiplatform/tensorboard/uploader_utils.py b/google/cloud/aiplatform/tensorboard/uploader_utils.py index 679eb02ef4..86712a5542 100644 --- a/google/cloud/aiplatform/tensorboard/uploader_utils.py +++ b/google/cloud/aiplatform/tensorboard/uploader_utils.py @@ -29,18 +29,12 @@ from google.api_core import exceptions from google.cloud import storage -from google.cloud.aiplatform.compat.types import ( - tensorboard_run_v1beta1 as tensorboard_run, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_service_v1beta1 as tensorboard_service, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_time_series_v1beta1 as tensorboard_time_series, -) -from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1 - -TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient +from google.cloud.aiplatform.compat.types import tensorboard_run +from google.cloud.aiplatform.compat.types import tensorboard_service +from google.cloud.aiplatform.compat.types import tensorboard_time_series +from google.cloud.aiplatform.compat.services import tensorboard_service_client + +TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient logger = tb_logging.get_logger() logger.setLevel(logging.WARNING) @@ -201,7 +195,7 @@ def _create_or_get_run_resource( Required. The display name of this run. Returns: - tb_run (google.cloud.aiplatform_v1beta1.types.TensorboardRun): + tb_run (tensorboard_run.TensorboardRun): The TensorboardRun given the run_name. Raises: diff --git a/tests/unit/aiplatform/test_uploader.py b/tests/unit/aiplatform/test_uploader.py index fd071d0e2b..b0d41e475e 100644 --- a/tests/unit/aiplatform/test_uploader.py +++ b/tests/unit/aiplatform/test_uploader.py @@ -46,22 +46,20 @@ from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader import google.cloud.aiplatform.tensorboard.uploader as uploader_lib from google.cloud import storage -from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1 -from google.cloud.aiplatform_v1beta1.services.tensorboard_service.transports import ( - grpc as transports_grpc, -) -from google.cloud.aiplatform.compat.types import ( - tensorboard_data_v1beta1 as tensorboard_data, - tensorboard_service_v1beta1 as tensorboard_service, +from google.cloud.aiplatform_v1.services.tensorboard_service import ( + client as tensorboard_service_client, ) -from google.cloud.aiplatform.compat.types import ( - tensorboard_experiment_v1beta1 as tensorboard_experiment_type, +from google.cloud.aiplatform_v1.services.tensorboard_service.transports import ( + grpc as transports_grpc, ) -from google.cloud.aiplatform.compat.types import ( - tensorboard_run_v1beta1 as tensorboard_run_type, +from google.cloud.aiplatform_v1.types import tensorboard_data +from google.cloud.aiplatform_v1.types import tensorboard_service +from google.cloud.aiplatform_v1.types import ( + tensorboard_experiment as tensorboard_experiment_type, ) -from google.cloud.aiplatform.compat.types import ( - tensorboard_time_series_v1beta1 as tensorboard_time_series_type, +from google.cloud.aiplatform_v1.types import tensorboard_run as tensorboard_run_type +from google.cloud.aiplatform_v1.types import ( + tensorboard_time_series as tensorboard_time_series_type, ) from google.protobuf import timestamp_pb2 from google.protobuf import message @@ -160,7 +158,7 @@ def create_tensorboard_time_series( service_descriptors=[], time=grpc_testing.strict_real_time() ) mock_client = mock.Mock( - spec=tensorboard_service_client_v1beta1.TensorboardServiceClient( + spec=tensorboard_service_client.TensorboardServiceClient( transport=transports_grpc.TensorboardServiceGrpcTransport( channel=test_channel )