Skip to content

Commit f34094b

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Tensorboard uploader - move remaining tensorboard uploader code to aiplatform
PiperOrigin-RevId: 627863478
1 parent b5404e7 commit f34094b

File tree

5 files changed

+29
-36
lines changed

5 files changed

+29
-36
lines changed

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

+3-3
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,14 @@
2222
import os
2323
import re
2424
from typing import DefaultDict, Dict, Generator, List, Optional, Set, Tuple
25-
from tensorboard.uploader.proto import server_info_pb2
2625

2726
from google.cloud import storage
2827
from google.cloud.aiplatform.compat.services import tensorboard_service_client
2928
from google.cloud.aiplatform.compat.types import tensorboard_data
3029
from google.cloud.aiplatform.compat.types import tensorboard_service
3130
from google.cloud.aiplatform.compat.types import tensorboard_time_series
3231
from google.cloud.aiplatform.tensorboard import upload_tracker
32+
from google.cloud.aiplatform.tensorboard import uploader_constants
3333
from google.cloud.aiplatform.tensorboard import uploader_utils
3434
import grpc
3535
import tensorflow as tf
@@ -65,7 +65,7 @@ def __init__(
6565
self,
6666
experiment_resource_name: str,
6767
api: TensorboardServiceClient,
68-
upload_limits: server_info_pb2.UploadLimits,
68+
upload_limits: uploader_constants.UploadLimits,
6969
blob_rpc_rate_limiter: uploader_utils.RateLimiter,
7070
blob_storage_bucket: storage.Bucket,
7171
blob_storage_folder: str,
@@ -81,7 +81,7 @@ def __init__(
8181
projects/{project}/locations/{location}/tensorboards/{tensorboard}/experiments/{experiment}
8282
api (TensorboardServiceClient):
8383
Required. Tensorboard service stub used to interact with experiment resource.
84-
upload_limits (server_info_pb2.UploadLimits):
84+
upload_limits (uploader_constants.UploadLimits):
8585
Required. Upload limits for for api calls.
8686
blob_rpc_rate_limiter (uploader_utils.RateLimiter):
8787
Required. A `RateLimiter` to use to limit write RPC frequency.

google/cloud/aiplatform/tensorboard/uploader.py

+4-27
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,6 @@
6262
from tensorboard.compat.proto import summary_pb2
6363
from tensorboard.compat.proto import types_pb2
6464
from tensorboard.plugins.graph import metadata as graph_metadata
65-
from tensorboard.uploader.proto import server_info_pb2
6665
from tensorboard.util import tb_logging
6766
from tensorboard.util import tensor_util
6867

@@ -96,7 +95,7 @@ def __init__(
9695
logdir: str,
9796
allowed_plugins: FrozenSet[str],
9897
experiment_display_name: Optional[str] = None,
99-
upload_limits: Optional[server_info_pb2.UploadLimits] = None,
98+
upload_limits: Optional[uploader_constants.UploadLimits] = None,
10099
logdir_poll_rate_limiter: Optional[uploader_utils.RateLimiter] = None,
101100
rpc_rate_limiter: Optional[uploader_utils.RateLimiter] = None,
102101
tensor_rpc_rate_limiter: Optional[uploader_utils.RateLimiter] = None,
@@ -121,7 +120,7 @@ def __init__(
121120
allowed_plugins: collection of string plugin names; events will only be
122121
uploaded if their time series's metadata specifies one of these plugin
123122
names
124-
upload_limits: instance of tensorboard.service.UploadLimits proto.
123+
upload_limits: dataclass of uploader_constants.UploadLimits.
125124
logdir_poll_rate_limiter: a `RateLimiter` to use to limit logdir polling
126125
frequency, to avoid thrashing disks, especially on networked file
127126
systems
@@ -160,29 +159,7 @@ def __init__(
160159

161160
self._upload_limits = upload_limits
162161
if not self._upload_limits:
163-
self._upload_limits = server_info_pb2.UploadLimits()
164-
self._upload_limits.max_scalar_request_size = (
165-
uploader_constants.DEFAULT_MAX_SCALAR_REQUEST_SIZE
166-
)
167-
self._upload_limits.min_scalar_request_interval = (
168-
uploader_constants.DEFAULT_MIN_SCALAR_REQUEST_INTERVAL
169-
)
170-
self._upload_limits.min_tensor_request_interval = (
171-
uploader_constants.DEFAULT_MIN_TENSOR_REQUEST_INTERVAL
172-
)
173-
self._upload_limits.max_tensor_request_size = (
174-
uploader_constants.DEFAULT_MAX_TENSOR_REQUEST_SIZE
175-
)
176-
self._upload_limits.max_tensor_point_size = (
177-
uploader_constants.DEFAULT_MAX_TENSOR_POINT_SIZE
178-
)
179-
self._upload_limits.min_blob_request_interval = (
180-
uploader_constants.DEFAULT_MIN_BLOB_REQUEST_INTERVAL
181-
)
182-
self._upload_limits.max_blob_request_size = (
183-
uploader_constants.DEFAULT_MAX_BLOB_REQUEST_SIZE
184-
)
185-
self._upload_limits.max_blob_size = uploader_constants.DEFAULT_MAX_BLOB_SIZE
162+
self._upload_limits = uploader_constants.UploadLimits()
186163
self._description = description
187164
self._verbosity = verbosity
188165
self._one_shot = one_shot
@@ -482,7 +459,7 @@ def __init__(
482459
experiment_resource_name: str,
483460
api: TensorboardServiceClient,
484461
allowed_plugins: Iterable[str],
485-
upload_limits: server_info_pb2.UploadLimits,
462+
upload_limits: uploader_constants.UploadLimits,
486463
rpc_rate_limiter: uploader_utils.RateLimiter,
487464
tensor_rpc_rate_limiter: uploader_utils.RateLimiter,
488465
blob_rpc_rate_limiter: uploader_utils.RateLimiter,

google/cloud/aiplatform/tensorboard/uploader_constants.py

+16
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
"""Constants shared between TensorBoard command line uploader and SDK uploader"""
22

3+
import dataclasses
4+
35
from tensorboard.plugins.distribution import (
46
metadata as distribution_metadata,
57
)
@@ -56,3 +58,17 @@
5658
DEFAULT_MAX_TENSOR_POINT_SIZE = 16 * (2**10) # 16KiB
5759

5860
DEFAULT_MAX_BLOB_SIZE = 10 * (2**30) # 10GiB
61+
62+
63+
@dataclasses.dataclass
64+
class UploadLimits:
65+
max_scalar_request_size: int = DEFAULT_MAX_SCALAR_REQUEST_SIZE
66+
max_tensor_request_size: int = DEFAULT_MAX_TENSOR_REQUEST_SIZE
67+
max_blob_request_size: int = DEFAULT_MAX_BLOB_REQUEST_SIZE
68+
69+
min_scalar_request_interval: int = DEFAULT_MIN_SCALAR_REQUEST_INTERVAL
70+
min_tensor_request_interval: int = DEFAULT_MIN_TENSOR_REQUEST_INTERVAL
71+
min_blob_request_interval: int = DEFAULT_MIN_BLOB_REQUEST_INTERVAL
72+
73+
max_blob_size: int = DEFAULT_MAX_BLOB_SIZE
74+
max_tensor_point_size: int = DEFAULT_MAX_TENSOR_POINT_SIZE

google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tensorboard_api.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,13 @@
2121
import re
2222
from typing import Tuple
2323

24-
from tensorboard.uploader.proto import server_info_pb2
2524
from google.api_core import exceptions
2625
from google.cloud import storage
2726
from google.cloud import aiplatform
2827
from google.cloud.aiplatform import training_utils
2928
from google.cloud.aiplatform.compat.types import tensorboard_experiment
3029
from google.cloud.aiplatform.tensorboard import upload_tracker
30+
from google.cloud.aiplatform.tensorboard import uploader_constants
3131
from google.cloud.aiplatform.tensorboard import uploader_utils
3232
from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
3333
from google.cloud.aiplatform.utils import TensorboardClientWithOverride
@@ -76,13 +76,13 @@ def _get_project_id() -> str:
7676
return m[1]
7777

7878

79-
def _make_upload_limits() -> server_info_pb2.UploadLimits:
79+
def _make_upload_limits() -> uploader_constants.UploadLimits:
8080
"""Creates the upload limits for tensorboard.
8181
8282
Returns:
8383
An UploadLimits object.
8484
"""
85-
upload_limits = server_info_pb2.UploadLimits()
85+
upload_limits = uploader_constants.UploadLimits()
8686
upload_limits.min_blob_request_interval = 10
8787
upload_limits.max_blob_request_size = 4 * (2**20) - 256 * (2**10)
8888
upload_limits.max_blob_size = 10 * (2**30) # 10GiB

tests/unit/aiplatform/test_uploader.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
from google.cloud.aiplatform.tensorboard import logdir_loader
4242
from google.cloud.aiplatform.tensorboard import upload_tracker
4343
from google.cloud.aiplatform.tensorboard import uploader as uploader_lib
44+
from google.cloud.aiplatform.tensorboard import uploader_constants
4445
from google.cloud.aiplatform.tensorboard import uploader_utils
4546
from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
4647
from google.cloud.aiplatform_v1.services.tensorboard_service.transports import (
@@ -61,7 +62,6 @@
6162
from tensorboard.plugins.graph import metadata as graphs_metadata
6263
from tensorboard.plugins.scalar import metadata as scalars_metadata
6364
from tensorboard.summary import v1 as summary_v1
64-
from tensorboard.uploader.proto import server_info_pb2
6565

6666
data_compat = uploader_lib.event_file_loader.data_compat
6767
dataclass_compat = uploader_lib.event_file_loader.dataclass_compat
@@ -227,7 +227,7 @@ def _create_uploader(
227227
if rpc_rate_limiter is _USE_DEFAULT:
228228
rpc_rate_limiter = uploader_utils.RateLimiter(0)
229229

230-
upload_limits = server_info_pb2.UploadLimits(
230+
upload_limits = uploader_constants.UploadLimits(
231231
max_scalar_request_size=max_scalar_request_size,
232232
max_tensor_request_size=max_tensor_request_size,
233233
max_tensor_point_size=max_tensor_point_size,
@@ -264,7 +264,7 @@ def _create_dispatcher(
264264
if allowed_plugins is _USE_DEFAULT:
265265
allowed_plugins = _SCALARS_HISTOGRAMS_AND_GRAPHS
266266

267-
upload_limits = server_info_pb2.UploadLimits(
267+
upload_limits = uploader_constants.UploadLimits(
268268
max_scalar_request_size=128000,
269269
max_tensor_request_size=128000,
270270
max_tensor_point_size=52000,

0 commit comments

Comments
 (0)