feat: updating Tensorboard related code to use v1 (#851)

Co-authored-by: Yicheng Fang <yichengfang@google.com>
googleapis · Nov 19, 2021 · b613b26 · b613b26
1 parent 45bd51c
commit b613b26
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 52 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
             )