fix: Tensorboard uploader - move remaining tensorboard uploader code …

…to aiplatform PiperOrigin-RevId: 627863478
googleapis · Apr 24, 2024 · f34094b · f34094b
1 parent b5404e7
commit f34094b
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 36 deletions.
diff --git a/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py b/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py
@@ -22,14 +22,14 @@
 import os
 import re
 from typing import DefaultDict, Dict, Generator, List, Optional, Set, Tuple
-from tensorboard.uploader.proto import server_info_pb2
 
 from google.cloud import storage
 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 upload_tracker
+from google.cloud.aiplatform.tensorboard import uploader_constants
 from google.cloud.aiplatform.tensorboard import uploader_utils
 import grpc
 import tensorflow as tf
@@ -65,7 +65,7 @@ def __init__(
         self,
         experiment_resource_name: str,
         api: TensorboardServiceClient,
-        upload_limits: server_info_pb2.UploadLimits,
+        upload_limits: uploader_constants.UploadLimits,
         blob_rpc_rate_limiter: uploader_utils.RateLimiter,
         blob_storage_bucket: storage.Bucket,
         blob_storage_folder: str,
@@ -81,7 +81,7 @@ def __init__(
                     projects/{project}/locations/{location}/tensorboards/{tensorboard}/experiments/{experiment}
             api (TensorboardServiceClient):
                 Required. Tensorboard service stub used to interact with experiment resource.
-            upload_limits (server_info_pb2.UploadLimits):
+            upload_limits (uploader_constants.UploadLimits):
                 Required. Upload limits for for api calls.
             blob_rpc_rate_limiter (uploader_utils.RateLimiter):
                 Required. A `RateLimiter` to use to limit write RPC frequency.

diff --git a/google/cloud/aiplatform/tensorboard/uploader.py b/google/cloud/aiplatform/tensorboard/uploader.py
@@ -62,7 +62,6 @@
 from tensorboard.compat.proto import summary_pb2
 from tensorboard.compat.proto import types_pb2
 from tensorboard.plugins.graph import metadata as graph_metadata
-from tensorboard.uploader.proto import server_info_pb2
 from tensorboard.util import tb_logging
 from tensorboard.util import tensor_util
 
@@ -96,7 +95,7 @@ def __init__(
         logdir: str,
         allowed_plugins: FrozenSet[str],
         experiment_display_name: Optional[str] = None,
-        upload_limits: Optional[server_info_pb2.UploadLimits] = None,
+        upload_limits: Optional[uploader_constants.UploadLimits] = None,
         logdir_poll_rate_limiter: Optional[uploader_utils.RateLimiter] = None,
         rpc_rate_limiter: Optional[uploader_utils.RateLimiter] = None,
         tensor_rpc_rate_limiter: Optional[uploader_utils.RateLimiter] = None,
@@ -121,7 +120,7 @@ def __init__(
           allowed_plugins: collection of string plugin names; events will only be
             uploaded if their time series's metadata specifies one of these plugin
             names
-          upload_limits: instance of tensorboard.service.UploadLimits proto.
+          upload_limits: dataclass of uploader_constants.UploadLimits.
           logdir_poll_rate_limiter: a `RateLimiter` to use to limit logdir polling
             frequency, to avoid thrashing disks, especially on networked file
             systems
@@ -160,29 +159,7 @@ def __init__(
 
         self._upload_limits = upload_limits
         if not self._upload_limits:
-            self._upload_limits = server_info_pb2.UploadLimits()
-            self._upload_limits.max_scalar_request_size = (
-                uploader_constants.DEFAULT_MAX_SCALAR_REQUEST_SIZE
-            )
-            self._upload_limits.min_scalar_request_interval = (
-                uploader_constants.DEFAULT_MIN_SCALAR_REQUEST_INTERVAL
-            )
-            self._upload_limits.min_tensor_request_interval = (
-                uploader_constants.DEFAULT_MIN_TENSOR_REQUEST_INTERVAL
-            )
-            self._upload_limits.max_tensor_request_size = (
-                uploader_constants.DEFAULT_MAX_TENSOR_REQUEST_SIZE
-            )
-            self._upload_limits.max_tensor_point_size = (
-                uploader_constants.DEFAULT_MAX_TENSOR_POINT_SIZE
-            )
-            self._upload_limits.min_blob_request_interval = (
-                uploader_constants.DEFAULT_MIN_BLOB_REQUEST_INTERVAL
-            )
-            self._upload_limits.max_blob_request_size = (
-                uploader_constants.DEFAULT_MAX_BLOB_REQUEST_SIZE
-            )
-            self._upload_limits.max_blob_size = uploader_constants.DEFAULT_MAX_BLOB_SIZE
+            self._upload_limits = uploader_constants.UploadLimits()
         self._description = description
         self._verbosity = verbosity
         self._one_shot = one_shot
@@ -482,7 +459,7 @@ def __init__(
         experiment_resource_name: str,
         api: TensorboardServiceClient,
         allowed_plugins: Iterable[str],
-        upload_limits: server_info_pb2.UploadLimits,
+        upload_limits: uploader_constants.UploadLimits,
         rpc_rate_limiter: uploader_utils.RateLimiter,
         tensor_rpc_rate_limiter: uploader_utils.RateLimiter,
         blob_rpc_rate_limiter: uploader_utils.RateLimiter,

diff --git a/google/cloud/aiplatform/tensorboard/uploader_constants.py b/google/cloud/aiplatform/tensorboard/uploader_constants.py
@@ -1,5 +1,7 @@
 """Constants shared between TensorBoard command line uploader and SDK uploader"""
 
+import dataclasses
+
 from tensorboard.plugins.distribution import (
     metadata as distribution_metadata,
 )
@@ -56,3 +58,17 @@
 DEFAULT_MAX_TENSOR_POINT_SIZE = 16 * (2**10)  # 16KiB
 
 DEFAULT_MAX_BLOB_SIZE = 10 * (2**30)  # 10GiB
+
+
+@dataclasses.dataclass
+class UploadLimits:
+    max_scalar_request_size: int = DEFAULT_MAX_SCALAR_REQUEST_SIZE
+    max_tensor_request_size: int = DEFAULT_MAX_TENSOR_REQUEST_SIZE
+    max_blob_request_size: int = DEFAULT_MAX_BLOB_REQUEST_SIZE
+
+    min_scalar_request_interval: int = DEFAULT_MIN_SCALAR_REQUEST_INTERVAL
+    min_tensor_request_interval: int = DEFAULT_MIN_TENSOR_REQUEST_INTERVAL
+    min_blob_request_interval: int = DEFAULT_MIN_BLOB_REQUEST_INTERVAL
+
+    max_blob_size: int = DEFAULT_MAX_BLOB_SIZE
+    max_tensor_point_size: int = DEFAULT_MAX_TENSOR_POINT_SIZE
diff --git a/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tensorboard_api.py b/google/cloud/aiplatform/training_utils/cloud_profiler/plugins/tensorflow/tensorboard_api.py
@@ -21,13 +21,13 @@
 import re
 from typing import Tuple
 
-from tensorboard.uploader.proto import server_info_pb2
 from google.api_core import exceptions
 from google.cloud import storage
 from google.cloud import aiplatform
 from google.cloud.aiplatform import training_utils
 from google.cloud.aiplatform.compat.types import tensorboard_experiment
 from google.cloud.aiplatform.tensorboard import upload_tracker
+from google.cloud.aiplatform.tensorboard import uploader_constants
 from google.cloud.aiplatform.tensorboard import uploader_utils
 from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
 from google.cloud.aiplatform.utils import TensorboardClientWithOverride
@@ -76,13 +76,13 @@ def _get_project_id() -> str:
     return m[1]
 
 
-def _make_upload_limits() -> server_info_pb2.UploadLimits:
+def _make_upload_limits() -> uploader_constants.UploadLimits:
     """Creates the upload limits for tensorboard.
 
     Returns:
         An UploadLimits object.
     """
-    upload_limits = server_info_pb2.UploadLimits()
+    upload_limits = uploader_constants.UploadLimits()
     upload_limits.min_blob_request_interval = 10
     upload_limits.max_blob_request_size = 4 * (2**20) - 256 * (2**10)
     upload_limits.max_blob_size = 10 * (2**30)  # 10GiB

diff --git a/tests/unit/aiplatform/test_uploader.py b/tests/unit/aiplatform/test_uploader.py
@@ -41,6 +41,7 @@
 from google.cloud.aiplatform.tensorboard import logdir_loader
 from google.cloud.aiplatform.tensorboard import upload_tracker
 from google.cloud.aiplatform.tensorboard import uploader as uploader_lib
+from google.cloud.aiplatform.tensorboard import uploader_constants
 from google.cloud.aiplatform.tensorboard import uploader_utils
 from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader
 from google.cloud.aiplatform_v1.services.tensorboard_service.transports import (
@@ -61,7 +62,6 @@
 from tensorboard.plugins.graph import metadata as graphs_metadata
 from tensorboard.plugins.scalar import metadata as scalars_metadata
 from tensorboard.summary import v1 as summary_v1
-from tensorboard.uploader.proto import server_info_pb2
 
 data_compat = uploader_lib.event_file_loader.data_compat
 dataclass_compat = uploader_lib.event_file_loader.dataclass_compat
@@ -227,7 +227,7 @@ def _create_uploader(
     if rpc_rate_limiter is _USE_DEFAULT:
         rpc_rate_limiter = uploader_utils.RateLimiter(0)
 
-    upload_limits = server_info_pb2.UploadLimits(
+    upload_limits = uploader_constants.UploadLimits(
         max_scalar_request_size=max_scalar_request_size,
         max_tensor_request_size=max_tensor_request_size,
         max_tensor_point_size=max_tensor_point_size,
@@ -264,7 +264,7 @@ def _create_dispatcher(
     if allowed_plugins is _USE_DEFAULT:
         allowed_plugins = _SCALARS_HISTOGRAMS_AND_GRAPHS
 
-    upload_limits = server_info_pb2.UploadLimits(
+    upload_limits = uploader_constants.UploadLimits(
         max_scalar_request_size=128000,
         max_tensor_request_size=128000,
         max_tensor_point_size=52000,