feat: pre batch creating TensorboardRuns and TensorboardTimeSeries in one_shot mode to speed up uploading (#772)

* feat: pre batch creating TensorboardRuns and TensorboardTimeSeries in one_shot mode to speed up uploading Co-authored-by: Yicheng Fang <yichengfang@google.com>

Author: yfang1

google/cloud/aiplatform/tensorboard/uploader.py

@@ -29,6 +29,7 @@
  Iterable,
  Optional,
  ContextManager,
+ Tuple,
 )
 import uuid
 
@@ -195,6 +196,7 @@ def __init__(
  self._logdir = logdir
  self._allowed_plugins = frozenset(allowed_plugins)
  self._run_name_prefix = run_name_prefix
+ self._is_brand_new_experiment = False
 
  self._upload_limits = upload_limits
  if not self._upload_limits:
@@ -265,6 +267,9 @@ def active_filter(secs):
  self._logdir_loader = logdir_loader.LogdirLoader(
  self._logdir, directory_loader_factory
  )
+ self._logdir_loader_pre_create = logdir_loader.LogdirLoader(
+ self._logdir, directory_loader_factory
+ )
  self._tracker = upload_tracker.UploadTracker(verbosity=self._verbosity)
 
  self._create_additional_senders()
@@ -290,6 +295,7 @@ def _create_or_get_experiment(self) -> tensorboard_experiment.TensorboardExperim
  tensorboard_experiment=tb_experiment,
  tensorboard_experiment_id=self._experiment_name,
  )
+ self._is_brand_new_experiment = True
  except exceptions.AlreadyExists:
  logger.info("Creating experiment failed. Retrieving experiment.")
  experiment_name = os.path.join(
@@ -303,7 +309,11 @@ def create_experiment(self):
 
  experiment = self._create_or_get_experiment()
  self._experiment = experiment
- request_sender = _BatchedRequestSender(
+ self._one_platform_resource_manager = uploader_utils.OnePlatformResourceManager(
+ self._experiment.name, self._api
+ )
+
+ self._request_sender = _BatchedRequestSender(
  self._experiment.name,
  self._api,
  allowed_plugins=self._allowed_plugins,
@@ -313,6 +323,7 @@ def create_experiment(self):
  blob_rpc_rate_limiter=self._blob_rpc_rate_limiter,
  blob_storage_bucket=self._blob_storage_bucket,
  blob_storage_folder=self._blob_storage_folder,
+ one_platform_resource_manager=self._one_platform_resource_manager,
  tracker=self._tracker,
  )
 
@@ -323,7 +334,8 @@ def create_experiment(self):
  )
 
  self._dispatcher = _Dispatcher(
- request_sender=request_sender, additional_senders=self._additional_senders,
+ request_sender=self._request_sender,
+ additional_senders=self._additional_senders,
  )
 
  def _create_additional_senders(self) -> Dict[str, uploader_utils.RequestSender]:
@@ -366,6 +378,17 @@ def start_uploading(self):
  """
  if self._dispatcher is None:
  raise RuntimeError("Must call create_experiment() before start_uploading()")
+
+ if self._one_shot:
+ if self._is_brand_new_experiment:
+ self._pre_create_runs_and_time_series()
+ else:
+ logger.warning(
+ "Please consider uploading to a new experiment instead of "
+ "an existing one, as the former allows for better upload "
+ "performance."
+ )
+
  while True:
  self._logdir_poll_rate_limiter.tick()
  self._upload_once()
@@ -377,6 +400,58 @@ def start_uploading(self):
  "without any uploadable data" % self._logdir
  )
 
+ def _pre_create_runs_and_time_series(self):
+ """
+ Iterates though the log dir to collect TensorboardRuns and
+ TensorboardTimeSeries that need to be created, and creates them in batch
+ to speed up uploading later on.
+ """
+ self._logdir_loader_pre_create.synchronize_runs()
+ run_to_events = self._logdir_loader_pre_create.get_run_events()
+ if self._run_name_prefix:
+ run_to_events = {
+ self._run_name_prefix + k: v for k, v in run_to_events.items()
+ }
+
+ run_names = []
+ run_tag_name_to_time_series_proto = {}
+ for (run_name, events) in run_to_events.items():
+ run_names.append(run_name)
+ for event in events:
+ _filter_graph_defs(event)
+ for value in event.summary.value:
+ metadata, is_valid = self._request_sender.get_metadata_and_validate(
+ run_name, value
+ )
+ if not is_valid:
+ continue
+ if metadata.data_class == summary_pb2.DATA_CLASS_SCALAR:
+ value_type = (
+ tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR
+ )
+ elif metadata.data_class == summary_pb2.DATA_CLASS_TENSOR:
+ value_type = (
+ tensorboard_time_series.TensorboardTimeSeries.ValueType.TENSOR
+ )
+ elif metadata.data_class == summary_pb2.DATA_CLASS_BLOB_SEQUENCE:
+ value_type = (
+ tensorboard_time_series.TensorboardTimeSeries.ValueType.BLOB_SEQUENCE
+ )
+
+ run_tag_name_to_time_series_proto[
+ (run_name, value.tag)
+ ] = tensorboard_time_series.TensorboardTimeSeries(
+ display_name=value.tag,
+ value_type=value_type,
+ plugin_name=metadata.plugin_data.plugin_name,
+ plugin_data=metadata.plugin_data.content,
+ )
+
+ self._one_platform_resource_manager.batch_create_runs(run_names)
+ self._one_platform_resource_manager.batch_create_time_series(
+ run_tag_name_to_time_series_proto
+ )
+
  def _upload_once(self):
  """Runs one upload cycle, sending zero or more RPCs."""
  logger.info("Starting an upload cycle")
@@ -439,6 +514,7 @@ def __init__(
  blob_rpc_rate_limiter: util.RateLimiter,
  blob_storage_bucket: storage.Bucket,
  blob_storage_folder: str,
+ one_platform_resource_manager: uploader_utils.OnePlatformResourceManager,
  tracker: upload_tracker.UploadTracker,
  ):
  """Constructs _BatchedRequestSender for the given experiment resource.
@@ -456,16 +532,16 @@ def __init__(
  Note the chunk stream is internally rate-limited by backpressure from
  the server, so it is not a concern that we do not explicitly rate-limit
  within the stream here.
+ one_platform_resource_manager: An instance of the One Platform
+ resource management class.
  tracker: Upload tracker to track information about uploads.
  """
  self._experiment_resource_name = experiment_resource_name
  self._api = api
  self._tag_metadata = {}
  self._allowed_plugins = frozenset(allowed_plugins)
  self._tracker = tracker
- self._one_platform_resource_manager = uploader_utils.OnePlatformResourceManager(
- self._experiment_resource_name, self._api
- )
+ self._one_platform_resource_manager = one_platform_resource_manager
  self._scalar_request_sender = _ScalarBatchedRequestSender(
  experiment_resource_id=experiment_resource_name,
  api=api,
@@ -516,6 +592,37 @@ def send_request(
  RuntimeError: If no progress can be made because even a single
  point is too large (say, due to a gigabyte-long tag name).
  """
+ metadata, is_valid = self.get_metadata_and_validate(run_name, value)
+ if not is_valid:
+ return
+ plugin_name = metadata.plugin_data.plugin_name
+ self._tracker.add_plugin_name(plugin_name)
+
+ if metadata.data_class == summary_pb2.DATA_CLASS_SCALAR:
+ self._scalar_request_sender.add_event(run_name, event, value, metadata)
+ elif metadata.data_class == summary_pb2.DATA_CLASS_TENSOR:
+ self._tensor_request_sender.add_event(run_name, event, value, metadata)
+ elif metadata.data_class == summary_pb2.DATA_CLASS_BLOB_SEQUENCE:
+ self._blob_request_sender.add_event(run_name, event, value, metadata)
+
+ def flush(self):
+ """Flushes any events that have been stored."""
+ self._scalar_request_sender.flush()
+ self._tensor_request_sender.flush()
+ self._blob_request_sender.flush()
+
+ def get_metadata_and_validate(
+ self, run_name: str, value: tf.compat.v1.Summary.Value
+ ) -> Tuple[tf.compat.v1.SummaryMetadata, bool]:
+ """
+
+ :param run_name: Name of the run retrieved by
+ `LogdirLoader.get_run_events`
+ :param value: A single `tf.compat.v1.Summary.Value` from the event,
+ where there can be multiple values per event.
+ :return: (metadata, is_valid): a metadata derived from the value, and
+ whether the value itself is valid.
+ """
 
  time_series_key = (run_name, value.tag)
 
@@ -539,29 +646,16 @@ def send_request(
  metadata.plugin_data.plugin_name,
  value.metadata.plugin_data.plugin_name,
  )
- return
+ return metadata, False
  if plugin_name not in self._allowed_plugins:
  if first_in_time_series:
  logger.info(
  "Skipping time series %r with unsupported plugin name %r",
  time_series_key,
  plugin_name,
  )
- return
- self._tracker.add_plugin_name(plugin_name)
-
- if metadata.data_class == summary_pb2.DATA_CLASS_SCALAR:
- self._scalar_request_sender.add_event(run_name, event, value, metadata)
- elif metadata.data_class == summary_pb2.DATA_CLASS_TENSOR:
- self._tensor_request_sender.add_event(run_name, event, value, metadata)
- elif metadata.data_class == summary_pb2.DATA_CLASS_BLOB_SEQUENCE:
- self._blob_request_sender.add_event(run_name, event, value, metadata)
-
- def flush(self):
- """Flushes any events that have been stored."""
- self._scalar_request_sender.flush()
- self._tensor_request_sender.flush()
- self._blob_request_sender.flush()
+ return metadata, False
+ return metadata, True
 
 
 class _Dispatcher(object):

google/cloud/aiplatform/tensorboard/uploader_utils.py

@@ -22,7 +22,7 @@
 import logging
 import re
 import time
-from typing import Callable, Dict, Generator, Optional
+from typing import Callable, Dict, Generator, Optional, List, Tuple
 import uuid
 
 from tensorboard.util import tb_logging
@@ -39,7 +39,6 @@
  tensorboard_time_series_v1beta1 as tensorboard_time_series,
 )
 from google.cloud.aiplatform.compat.services import tensorboard_service_client_v1beta1
-from google.cloud.aiplatform_v1beta1.types import TensorboardRun
 
 TensorboardServiceClient = tensorboard_service_client_v1beta1.TensorboardServiceClient
 
@@ -66,6 +65,9 @@ def send_requests(run_name: str):
 class OnePlatformResourceManager(object):
  """Helper class managing One Platform resources."""
 
+ CREATE_RUN_BATCH_SIZE = 1000
+ CREATE_TIME_SERIES_BATCH_SIZE = 1000
+
  def __init__(self, experiment_resource_name: str, api: TensorboardServiceClient):
  """Constructor for OnePlatformResourceManager.
 
@@ -81,6 +83,96 @@ def __init__(self, experiment_resource_name: str, api: TensorboardServiceClient)
  self._run_name_to_run_resource_name: Dict[str, str] = {}
  self._run_tag_name_to_time_series_name: Dict[(str, str), str] = {}
 
+ def batch_create_runs(
+ self, run_names: List[str]
+ ) -> List[tensorboard_run.TensorboardRun]:
+ """Batch creates TensorboardRuns.
+
+ Args:
+ run_names: a list of run_names for creating the TensorboardRuns.
+ Returns:
+ the created TensorboardRuns
+ """
+ batch_size = OnePlatformResourceManager.CREATE_RUN_BATCH_SIZE
+ created_runs = []
+ for i in range(0, len(run_names), batch_size):
+ one_batch_run_names = run_names[i : i + batch_size]
+ tb_run_requests = [
+ tensorboard_service.CreateTensorboardRunRequest(
+ parent=self._experiment_resource_name,
+ tensorboard_run=tensorboard_run.TensorboardRun(
+ display_name=run_name
+ ),
+ tensorboard_run_id=str(uuid.uuid4()),
+ )
+ for run_name in one_batch_run_names
+ ]
+
+ tb_runs = self._api.batch_create_tensorboard_runs(
+ parent=self._experiment_resource_name, requests=tb_run_requests,
+ ).tensorboard_runs
+
+ self._run_name_to_run_resource_name.update(
+ {run.display_name: run.name for run in tb_runs}
+ )
+
+ created_runs.extend(tb_runs)
+
+ return created_runs
+
+ def batch_create_time_series(
+ self,
+ run_tag_name_to_time_series: Dict[
+ Tuple[str, str], tensorboard_time_series.TensorboardTimeSeries
+ ],
+ ) -> List[tensorboard_time_series.TensorboardTimeSeries]:
+ """Batch creates TensorboardTimeSeries.
+
+ Args:
+ run_tag_name_to_time_series: a dictionary of
+ (run_name, tag_name) to TensorboardTimeSeries proto, containing
+ the TensorboardTimeSeries to create.
+ Returns:
+ the created TensorboardTimeSeries
+ """
+ batch_size = OnePlatformResourceManager.CREATE_TIME_SERIES_BATCH_SIZE
+ run_tag_name_to_time_series_entries = list(run_tag_name_to_time_series.items())
+ run_resource_name_to_run_name = {
+ v: k for k, v in self._run_name_to_run_resource_name.items()
+ }
+ created_time_series = []
+ for i in range(0, len(run_tag_name_to_time_series_entries), batch_size):
+ requests = [
+ tensorboard_service.CreateTensorboardTimeSeriesRequest(
+ parent=self._run_name_to_run_resource_name[run_name],
+ tensorboard_time_series=time_series,
+ )
+ for (
+ (run_name, tag_name),
+ time_series,
+ ) in run_tag_name_to_time_series_entries[i : i + batch_size]
+ ]
+
+ time_series = self._api.batch_create_tensorboard_time_series(
+ parent=self._experiment_resource_name, requests=requests,
+ ).tensorboard_time_series
+
+ self._run_tag_name_to_time_series_name.update(
+ {
+ (
+ run_resource_name_to_run_name[
+ ts.name[: ts.name.index("/timeSeries")]
+ ],
+ ts.display_name,
+ ): ts.name
+ for ts in time_series
+ }
+ )
+
+ created_time_series.extend(time_series)
+
+ return created_time_series
+
  def get_run_resource_name(self, run_name: str) -> str:
  """
  Get the resource name of the run if it exists, otherwise creates the run
@@ -99,7 +191,9 @@ def get_run_resource_name(self, run_name: str) -> str:
  self._run_name_to_run_resource_name[run_name] = tb_run.name
  return self._run_name_to_run_resource_name[run_name]
 
- def _create_or_get_run_resource(self, run_name: str) -> TensorboardRun:
+ def _create_or_get_run_resource(
+ self, run_name: str
+ ) -> tensorboard_run.TensorboardRun:
  """Creates a new run resource in current tensorboard experiment resource.
 
  Args: