fix: Use default run_name in Tensorboard uploader for direct directory upload.

PiperOrigin-RevId: 646265028

Author: vertex-sdk-bot

google/cloud/aiplatform/tensorboard/uploader.py

@@ -67,6 +67,7 @@
 from tensorboard.util import tensor_util
 
 _LOGGER = base.Logger(__name__)
+_DEFAULT_RUN_NAME = "default"
 
 TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
 
@@ -381,6 +382,7 @@ def _pre_create_runs_and_time_series(self):
         run_names = []
         run_tag_name_to_time_series_proto = {}
         for (run_name, events) in run_to_events.items():
+            run_name = run_name if (run_name and run_name != ".") else _DEFAULT_RUN_NAME
             run_names.append(run_name)
             for event in events:
                 _filter_graph_defs(event)
@@ -427,6 +429,11 @@ def _upload_once(self):
         logger.info("Logdir sync took %.3f seconds", sync_duration_secs)
 
         run_to_events = self._logdir_loader.get_run_events()
+        run_to_events = {
+            k if (k and k != ".") else _DEFAULT_RUN_NAME: v
+            for k, v in run_to_events.items()
+            if v
+        }
         if self._run_name_prefix:
             run_to_events = {
                 self._run_name_prefix + k: v for k, v in run_to_events.items()

google/cloud/aiplatform/tensorboard/uploader_utils.py

@@ -23,7 +23,6 @@
 import re
 import time
 from typing import Callable, Dict, Generator, List, Optional, Tuple
-import uuid
 
 from absl import app
 from google.api_core import exceptions
@@ -222,8 +221,6 @@ def _create_or_get_run_resource(
         location = m[2]
         tensorboard = m[3]
         experiment = m[4]
-        if not run_name or run_name == ".":
-            run_name = str(uuid.uuid4())
         experiment_run = experiment_run_resource.ExperimentRun.get(
             project=project, location=location, run_name=run_name
         )

tests/unit/aiplatform/test_uploader.py

@@ -115,6 +115,7 @@
     _TEST_ONE_PLATFORM_RUN_NAME, _TEST_TIME_SERIES_NAME
 )
 _TEST_BLOB_STORAGE_FOLDER = "test_folder"
+_DEFAULT_RUN_NAME = "default"
 
 
 def _create_example_graph_bytes(large_attr_size):
@@ -821,6 +822,52 @@ def test_upload_empty_logdir(
         mock_client.write_tensorboard_experiment_data.assert_not_called()
         experiment_tracker_mock.set_experiment.assert_called_once()
 
+    @parameterized.parameters(
+        {"run_name_prefix": None},
+        {"run_name_prefix": "run-prefix-"},
+    )
+    @patch.object(
+        uploader_utils.OnePlatformResourceManager,
+        "get_run_resource_name",
+        autospec=True,
+    )
+    @patch.object(metadata, "_experiment_tracker", autospec=True)
+    @patch.object(experiment_resources, "Experiment", autospec=True)
+    def test_default_run_name(
+        self,
+        experiment_resources_mock,
+        experiment_tracker_mock,
+        run_resource_mock,
+        run_name_prefix,
+    ):
+        run_resource_mock.return_value = "."
+        experiment_resources_mock.get.return_value = _TEST_EXPERIMENT_NAME
+        experiment_tracker_mock.set_experiment.return_value = _TEST_EXPERIMENT_NAME
+        experiment_tracker_mock.set_tensorboard.return_value = (
+            _TEST_TENSORBOARD_RESOURCE_NAME
+        )
+        logdir = self.get_temp_dir()
+        with FileWriter(logdir) as writer:
+            writer.add_test_summary("foo")
+
+        uploader = _create_uploader(
+            logdir=logdir,
+            run_name_prefix=run_name_prefix,
+        )
+        uploader.create_experiment()
+        mock_dispatcher = mock.create_autospec(uploader_lib._Dispatcher)
+        uploader._dispatcher = mock_dispatcher
+        mock_logdir_loader = mock.create_autospec(logdir_loader.LogdirLoader)
+        mock.patch.object(uploader, "_logdir_loader", mock_logdir_loader)
+        expected_run_name = _DEFAULT_RUN_NAME
+        if run_name_prefix:
+            expected_run_name = run_name_prefix + _DEFAULT_RUN_NAME
+
+        uploader._upload_once()
+
+        run_to_events = mock_dispatcher.dispatch_requests.call_args[0][0]
+        self.assertIn(expected_run_name, run_to_events)
+
     @patch.object(metadata, "_experiment_tracker", autospec=True)
     @patch.object(experiment_resources, "Experiment", autospec=True)
     def test_upload_polls_slowly_once_done(