Skip to content

Commit b613b26

Browse files
yfang1Yicheng Fang
yfang1
and
Yicheng Fang
authored
feat: updating Tensorboard related code to use v1 (#851)
Co-authored-by: Yicheng Fang <yichengfang@google.com>
1 parent 45bd51c commit b613b26

File tree

5 files changed

+31
-52
lines changed

5 files changed

+31
-52
lines changed

‎google/cloud/aiplatform/compat/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
services.pipeline_service_client = services.pipeline_service_client_v1
9696
services.prediction_service_client = services.prediction_service_client_v1
9797
services.specialist_pool_service_client = services.specialist_pool_service_client_v1
98+
services.tensorboard_service_client = services.tensorboard_service_client_v1
9899

99100
types.accelerator_type = types.accelerator_type_v1
100101
types.annotation = types.annotation_v1

‎google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -38,20 +38,14 @@
3838
import tensorflow as tf
3939

4040
from google.cloud import storage
41-
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
42-
from google.cloud.aiplatform.compat.types import (
43-
tensorboard_data_v1beta1 as tensorboard_data,
44-
)
45-
from google.cloud.aiplatform.compat.types import (
46-
tensorboard_service_v1beta1 as tensorboard_service,
47-
)
48-
from google.cloud.aiplatform.compat.types import (
49-
tensorboard_time_series_v1beta1 as tensorboard_time_series,
50-
)
41+
from google.cloud.aiplatform.compat.services import tensorboard_service_client
42+
from google.cloud.aiplatform.compat.types import tensorboard_data
43+
from google.cloud.aiplatform.compat.types import tensorboard_service
44+
from google.cloud.aiplatform.compat.types import tensorboard_time_series
5145
from google.cloud.aiplatform.tensorboard import uploader_utils
5246
from google.protobuf import timestamp_pb2 as timestamp
5347

54-
TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
48+
TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
5549

5650
logger = tb_logging.get_logger()
5751

‎google/cloud/aiplatform/tensorboard/uploader.py‎

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -56,25 +56,17 @@
5656

5757
from google.api_core import exceptions
5858
from google.cloud import storage
59-
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
60-
from google.cloud.aiplatform.compat.types import (
61-
tensorboard_data_v1beta1 as tensorboard_data,
62-
)
63-
from google.cloud.aiplatform.compat.types import (
64-
tensorboard_experiment_v1beta1 as tensorboard_experiment,
65-
)
66-
from google.cloud.aiplatform.compat.types import (
67-
tensorboard_service_v1beta1 as tensorboard_service,
68-
)
69-
from google.cloud.aiplatform.compat.types import (
70-
tensorboard_time_series_v1beta1 as tensorboard_time_series,
71-
)
59+
from google.cloud.aiplatform.compat.services import tensorboard_service_client
60+
from google.cloud.aiplatform.compat.types import tensorboard_data
61+
from google.cloud.aiplatform.compat.types import tensorboard_experiment
62+
from google.cloud.aiplatform.compat.types import tensorboard_service
63+
from google.cloud.aiplatform.compat.types import tensorboard_time_series
7264
from google.cloud.aiplatform.tensorboard import uploader_utils
7365
from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
7466
from google.protobuf import message
7567
from google.protobuf import timestamp_pb2 as timestamp
7668

77-
TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
69+
TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
7870

7971
# Minimum length of a logdir polling cycle in seconds. Shorter cycles will
8072
# sleep to avoid spinning over the logdir, which isn't great for disks and can

‎google/cloud/aiplatform/tensorboard/uploader_utils.py‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,18 +29,12 @@
2929

3030
from google.api_core import exceptions
3131
from google.cloud import storage
32-
from google.cloud.aiplatform.compat.types import (
33-
tensorboard_run_v1beta1 as tensorboard_run,
34-
)
35-
from google.cloud.aiplatform.compat.types import (
36-
tensorboard_service_v1beta1 as tensorboard_service,
37-
)
38-
from google.cloud.aiplatform.compat.types import (
39-
tensorboard_time_series_v1beta1 as tensorboard_time_series,
40-
)
41-
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
42-
43-
TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
32+
from google.cloud.aiplatform.compat.types import tensorboard_run
33+
from google.cloud.aiplatform.compat.types import tensorboard_service
34+
from google.cloud.aiplatform.compat.types import tensorboard_time_series
35+
from google.cloud.aiplatform.compat.services import tensorboard_service_client
36+
37+
TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
4438

4539
logger = tb_logging.get_logger()
4640
logger.setLevel(logging.WARNING)
@@ -201,7 +195,7 @@ def _create_or_get_run_resource(
201195
Required. The display name of this run.
202196
203197
Returns:
204-
tb_run (google.cloud.aiplatform_v1beta1.types.TensorboardRun):
198+
tb_run (tensorboard_run.TensorboardRun):
205199
The TensorboardRun given the run_name.
206200
207201
Raises:

‎tests/unit/aiplatform/test_uploader.py‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -46,22 +46,20 @@
4646
from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
4747
import google.cloud.aiplatform.tensorboard.uploader as uploader_lib
4848
from google.cloud import storage
49-
from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
50-
from google.cloud.aiplatform_v1beta1.services.tensorboard_service.transports import (
51-
grpc as transports_grpc,
52-
)
53-
from google.cloud.aiplatform.compat.types import (
54-
tensorboard_data_v1beta1 as tensorboard_data,
55-
tensorboard_service_v1beta1 as tensorboard_service,
49+
from google.cloud.aiplatform_v1.services.tensorboard_service import (
50+
client as tensorboard_service_client,
5651
)
57-
from google.cloud.aiplatform.compat.types import (
58-
tensorboard_experiment_v1beta1 as tensorboard_experiment_type,
52+
from google.cloud.aiplatform_v1.services.tensorboard_service.transports import (
53+
grpc as transports_grpc,
5954
)
60-
from google.cloud.aiplatform.compat.types import (
61-
tensorboard_run_v1beta1 as tensorboard_run_type,
55+
from google.cloud.aiplatform_v1.types import tensorboard_data
56+
from google.cloud.aiplatform_v1.types import tensorboard_service
57+
from google.cloud.aiplatform_v1.types import (
58+
tensorboard_experiment as tensorboard_experiment_type,
6259
)
63-
from google.cloud.aiplatform.compat.types import (
64-
tensorboard_time_series_v1beta1 as tensorboard_time_series_type,
60+
from google.cloud.aiplatform_v1.types import tensorboard_run as tensorboard_run_type
61+
from google.cloud.aiplatform_v1.types import (
62+
tensorboard_time_series as tensorboard_time_series_type,
6563
)
6664
from google.protobuf import timestamp_pb2
6765
from google.protobuf import message
@@ -160,7 +158,7 @@ def create_tensorboard_time_series(
160158
service_descriptors=[], time=grpc_testing.strict_real_time()
161159
)
162160
mock_client = mock.Mock(
163-
spec=tensorboard_service_client_v1beta1.TensorboardServiceClient(
161+
spec=tensorboard_service_client.TensorboardServiceClient(
164162
transport=transports_grpc.TensorboardServiceGrpcTransport(
165163
channel=test_channel
166164
)

0 commit comments

Comments
 (0)