fix: Rollback change to tensorboard uploader causing increased latency

PiperOrigin-RevId: 673086243

Author: vertex-sdk-bot

google/cloud/aiplatform/tensorboard/tensorboard_resource.py

@@ -858,14 +858,6 @@ def list(
 
  return tensorboard_runs
 
- def get_tensorboard_time_series_id(self, display_name: str) -> str:
- """Returns the TensorboardTimeSeries with the given display name."""
- if display_name not in self._time_series_display_name_to_id_mapping:
- self._sync_time_series_display_name_to_id_mapping()
-
- time_series_id = self._time_series_display_name_to_id_mapping.get(display_name)
- return time_series_id
-
  def write_tensorboard_scalar_data(
  self,
  time_series_data: Dict[str, float],

google/cloud/aiplatform/tensorboard/uploader_utils.py

@@ -18,6 +18,7 @@
 """Shared utils for tensorboard log uploader."""
 import abc
 import contextlib
+import json
 import logging
 import re
 import time
@@ -120,7 +121,7 @@ def batch_create_runs(
  """
  created_runs = []
  for run_name in run_names:
- tb_run = self._get_or_create_run_resource(run_name)
+ tb_run = self._create_or_get_run_resource(run_name)
  created_runs.append(tb_run)
  if run_name not in self._run_name_to_run_resource_name:
  self._run_name_to_run_resource_name[run_name] = tb_run.resource_name
@@ -195,11 +196,11 @@ def get_run_resource_name(self, run_name: str) -> str:
  Resource name of the run.
  """
  if run_name not in self._run_name_to_run_resource_name:
- tb_run = self._get_or_create_run_resource(run_name)
+ tb_run = self._create_or_get_run_resource(run_name)
  self._run_name_to_run_resource_name[run_name] = tb_run.resource_name
  return self._run_name_to_run_resource_name[run_name]
 
- def _get_or_create_run_resource(
+ def _create_or_get_run_resource(
  self, run_name: str
  ) -> tensorboard_run.TensorboardRun:
  """Creates new experiment run and tensorboard run resources.
@@ -270,7 +271,7 @@ def get_time_series_resource_name(
  Resource name of the time series
  """
  if (run_name, tag_name) not in self._run_tag_name_to_time_series_name:
- time_series = self._get_or_create_time_series(
+ time_series = self._create_or_get_time_series(
  self.get_run_resource_name(run_name),
  tag_name,
  time_series_resource_creator,
@@ -280,7 +281,7 @@ def get_time_series_resource_name(
  ] = time_series.name
  return self._run_tag_name_to_time_series_name[(run_name, tag_name)]
 
- def _get_or_create_time_series(
+ def _create_or_get_time_series(
  self,
  run_resource_name: str,
  tag_name: str,
@@ -310,29 +311,45 @@ def _get_or_create_time_series(
  ValueError:
  More than one time series with the resource name was found.
  """
- time_series = None
- run_name = run_resource_name.split("/")[-1]
- run = self._get_or_create_run_resource(run_name)
- time_series_id = run.get_tensorboard_time_series_id(tag_name)
- if time_series_id:
- time_series = self._api.get_tensorboard_time_series(
- request=tensorboard_service.GetTensorboardTimeSeriesRequest(
- name=run_resource_name + "/timeSeries/" + time_series_id
- )
+ time_series = time_series_resource_creator()
+ time_series.display_name = tag_name
+ try:
+ time_series = self._api.create_tensorboard_time_series(
+ parent=run_resource_name, tensorboard_time_series=time_series
  )
- if not time_series:
- time_series = time_series_resource_creator()
- time_series.display_name = tag_name
- try:
- time_series = self._api.create_tensorboard_time_series(
- parent=run_resource_name, tensorboard_time_series=time_series
+ except exceptions.InvalidArgument as e:
+ # If the time series display name already exists then retrieve it
+ if "already exist" in e.message:
+ list_of_time_series = self._api.list_tensorboard_time_series(
+ request=tensorboard_service.ListTensorboardTimeSeriesRequest(
+ parent=run_resource_name,
+ filter="display_name = {}".format(json.dumps(str(tag_name))),
+ )
  )
- except exceptions.InvalidArgument as e:
- raise ValueError(
- "Could not find time series resource with display name: {}".format(
- tag_name
+ num = 0
+ time_series = None
+
+ for ts in list_of_time_series:
+ num += 1
+ if num > 1:
+ break
+ time_series = ts
+
+ if not time_series:
+ raise ExistingResourceNotFoundError(
+ "Could not find time series resource with display name: {}".format(
+ tag_name
+ )
+ )
+
+ if num != 1:
+ raise ValueError(
+ "More than one time series resource found with display_name: {}".format(
+ tag_name
+ )
  )
- ) from e
+ else:
+ raise
  return time_series
 
 
@@ -355,45 +372,6 @@ def __init__(self, run_resource_id: str, api: TensorboardServiceClient):
  str, tensorboard_time_series.TensorboardTimeSeries
  ] = {}
 
- def _get_run_resource(self) -> tensorboard_run.TensorboardRun:
- """Gets or creates new experiment run and tensorboard run resources.
-
- The experiment run will be associated with the tensorboard run resource.
- This will link all tensorboard run data to the associated experiment.
-
- Returns:
- tb_run (tensorboard_run.TensorboardRun):
- The TensorboardRun given the run_name.
-
- Raises:
- ValueError:
- run_resource_id is invalid.
- """
- m = re.match(
- "projects/(.*)/locations/(.*)/tensorboards/(.*)/experiments/(.*)/runs/(.*)",
- self._run_resource_id,
- )
- project = m[1]
- location = m[2]
- tensorboard = m[3]
- experiment = m[4]
- run_name = m[5]
- experiment_run = experiment_run_resource.ExperimentRun.get(
- project=project, location=location, run_name=run_name
- )
- if not experiment_run:
- experiment_run = experiment_run_resource.ExperimentRun.create(
- project=project,
- location=location,
- run_name=run_name,
- experiment=experiment,
- tensorboard=tensorboard,
- state=gca_execution.Execution.State.RUNNING,
- )
- tb_run_artifact = experiment_run._backing_tensorboard_run
- tb_run = tb_run_artifact.resource
- return tb_run
-
  def get_or_create(
  self,
  tag_name: str,
@@ -416,36 +394,56 @@ def get_or_create(
  A new or existing tensorboard_time_series.TensorboardTimeSeries.
 
  Raises:
- ValueError:
+ exceptions.InvalidArgument:
  The tag_name or time_series_resource_creator is an invalid argument
  to create_tensorboard_time_series api call.
+ ExistingResourceNotFoundError:
+ Could not find the resource given the tag name.
+ ValueError:
+ More than one time series with the resource name was found.
  """
  if tag_name in self._tag_to_time_series_proto:
  return self._tag_to_time_series_proto[tag_name]
 
- time_series = None
- tb_run = self._get_run_resource()
- time_series_id = tb_run.get_tensorboard_time_series_id(tag_name)
- if time_series_id:
- time_series = self._api.get_tensorboard_time_series(
- request=tensorboard_service.GetTensorboardTimeSeriesRequest(
- name=self._run_resource_id + "/timeSeries/" + time_series_id
- )
+ time_series = time_series_resource_creator()
+ time_series.display_name = tag_name
+ try:
+ time_series = self._api.create_tensorboard_time_series(
+ parent=self._run_resource_id, tensorboard_time_series=time_series
  )
- if not time_series:
- time_series = time_series_resource_creator()
- time_series.display_name = tag_name
-
- try:
- time_series = self._api.create_tensorboard_time_series(
- parent=self._run_resource_id, tensorboard_time_series=time_series
+ except exceptions.InvalidArgument as e:
+ # If the time series display name already exists then retrieve it
+ if "already exist" in e.message:
+ list_of_time_series = self._api.list_tensorboard_time_series(
+ request=tensorboard_service.ListTensorboardTimeSeriesRequest(
+ parent=self._run_resource_id,
+ filter="display_name = {}".format(json.dumps(str(tag_name))),
+ )
  )
- except exceptions.InvalidArgument as e:
- raise ValueError(
- "Could not find time series resource with display name: {}".format(
- tag_name
+ num = 0
+ time_series = None
+
+ for ts in list_of_time_series:
+ num += 1
+ if num > 1:
+ break
+ time_series = ts
+
+ if not time_series:
+ raise ExistingResourceNotFoundError(
+ "Could not find time series resource with display name: {}".format(
+ tag_name
+ )
+ )
+
+ if num != 1:
+ raise ValueError(
+ "More than one time series resource found with display_name: {}".format(
+ tag_name
+ )
  )
- ) from e
+ else:
+ raise
 
  self._tag_to_time_series_proto[tag_name] = time_series
  return time_series