fix: Tensorboard - Fixed bug in tensorboard uploader where profile logs are not uploaded from nested run directories.

PiperOrigin-RevId: 670652532

Author: vertex-sdk-bot

google/cloud/aiplatform/tensorboard/logdir_loader.py

@@ -70,7 +70,6 @@ def synchronize_runs(self):
         runs_seen = set()
         for subdir in io_wrapper.GetLogdirSubdirectories(self._logdir):
             run = os.path.relpath(subdir, self._logdir)
-            run = run.replace("/", "-").replace("_", "-")
             runs_seen.add(run)
             if run not in self._directory_loaders:
                 logger.info("- Adding run for relative directory %s", run)

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

@@ -176,7 +176,10 @@ def send_request(self, run_name: str):
                 self._profile_dir(run_name)
             )
 
-        tb_run = self._one_platform_resource_manager.get_run_resource_name(run_name)
+        experiment_run_name = uploader_utils.reformat_run_name(run_name)
+        tb_run = self._one_platform_resource_manager.get_run_resource_name(
+            experiment_run_name
+        )
 
         if run_name not in self._run_to_file_request_sender:
             self._run_to_file_request_sender[

google/cloud/aiplatform/tensorboard/uploader.py

@@ -417,7 +417,8 @@ def _pre_create_runs_and_time_series(self):
                         plugin_data=metadata.plugin_data.content,
                     )
 
-        self._one_platform_resource_manager.batch_create_runs(run_names)
+        experiment_runs = [uploader_utils.reformat_run_name(run) for run in run_names]
+        self._one_platform_resource_manager.batch_create_runs(experiment_runs)
         self._one_platform_resource_manager.batch_create_time_series(
             run_tag_name_to_time_series_proto
         )
@@ -451,7 +452,9 @@ def _upload_once(self):
             )
             run_to_events[profile_run_name] = None
 
-        self._experiment_runs = run_to_events.keys()
+        self._experiment_runs = [
+            uploader_utils.reformat_run_name(run) for run in run_to_events.keys()
+        ]
 
         with self._tracker.send_tracker():
             self._dispatcher.dispatch_requests(run_to_events)
@@ -807,6 +810,7 @@ def _add_event_internal(
         metadata: tf.compat.v1.SummaryMetadata,
     ):
         self._num_values += 1
+        run_name = uploader_utils.reformat_run_name(run_name)
         time_series_data_proto = self._run_to_tag_to_time_series_data[run_name].get(
             value.tag
         )

google/cloud/aiplatform/tensorboard/uploader_utils.py

@@ -47,6 +47,11 @@
 logger.setLevel(logging.WARNING)
 
 
+def reformat_run_name(run_name: str) -> str:
+    """Reformats the run name to be compatible with One Platform."""
+    return run_name.replace("/", "-").replace("_", "-")
+
+
 class RateLimiter:
     """Helper class for rate-limiting using a fixed minimum interval."""
 

tests/unit/aiplatform/test_logdir_loader.py

@@ -201,15 +201,15 @@ def test_multiple_writes_to_logdir(self):
             {
                 "a": ["tag_a"],
                 "b": ["tag_b"],
-                "b-x": ["tag_b_x"],
-                "b-z": ["tag_b_z"],
+                "b/x": ["tag_b_x"],
+                "b_z": ["tag_b_z"],
                 "c": ["tag_c"],
             },
         )
         # A second load should indicate no new data.
         self.assertEqual(
             self._extract_run_to_tags(loader.get_run_events()),
-            {"a": [], "b": [], "b-x": [], "b-z": [], "c": []},
+            {"a": [], "b": [], "b/x": [], "b_z": [], "c": []},
         )
         # Write some new data to both new and pre-existing event files.
         with FileWriter(os.path.join(logdir, "a"), filename_suffix=".other") as writer:
@@ -228,8 +228,8 @@ def test_multiple_writes_to_logdir(self):
             {
                 "a": ["tag_a_2", "tag_a_3", "tag_a_4"],
                 "b": [],
-                "b-x": ["tag_b_x_2"],
-                "b-z": [],
+                "b/x": ["tag_b_x_2"],
+                "b_z": [],
                 "c": ["tag_c_2"],
             },
         )