fix: Tensorboard - Fix error in tensorboard batch upload of nested dirs

PiperOrigin-RevId: 675716748
googleapis · Sep 17, 2024 · ab157c8 · ab157c8
1 parent 0de2987
commit ab157c8
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 12 deletions.
diff --git a/google/cloud/aiplatform/tensorboard/uploader.py b/google/cloud/aiplatform/tensorboard/uploader.py
@@ -386,6 +386,7 @@ def _pre_create_runs_and_time_series(self):
                 if (run_name and run_name != ".")
                 else uploader_utils.DEFAULT_RUN_NAME
             )
+            run_name = uploader_utils.reformat_run_name(run_name)
             run_names.append(run_name)
             for event in events:
                 _filter_graph_defs(event)

diff --git a/tests/unit/aiplatform/test_uploader.py b/tests/unit/aiplatform/test_uploader.py
@@ -591,6 +591,11 @@ def test_start_uploading_without_create_experiment_fails(self):
         with self.assertRaisesRegex(RuntimeError, "call create_experiment()"):
             uploader.start_uploading()
 
+    @parameterized.parameters(
+        {"nested_run_dir": ""},
+        {"nested_run_dir": "nested-dir/"},
+        {"nested_run_dir": "double/nested-dir/"},
+    )
     @patch.object(
         uploader_utils.OnePlatformResourceManager,
         "get_run_resource_name",
@@ -599,7 +604,11 @@ def test_start_uploading_without_create_experiment_fails(self):
     @patch.object(metadata, "_experiment_tracker", autospec=True)
     @patch.object(experiment_resources, "Experiment", autospec=True)
     def test_start_uploading_scalars(
-        self, experiment_resources_mock, experiment_tracker_mock, run_resource_mock
+        self,
+        experiment_resources_mock,
+        experiment_tracker_mock,
+        run_resource_mock,
+        nested_run_dir,
     ):
         experiment_resources_mock.get.return_value = _TEST_EXPERIMENT_NAME
         experiment_tracker_mock.set_experiment.return_value = _TEST_EXPERIMENT_NAME
@@ -628,21 +637,21 @@ def test_start_uploading_scalars(
         mock_logdir_loader = mock.create_autospec(logdir_loader.LogdirLoader)
         mock_logdir_loader.get_run_events.side_effect = [
             {
-                "run 1": _apply_compat(
+                f"{nested_run_dir}run 1": _apply_compat(
                     [_scalar_event("1.1", 5.0), _scalar_event("1.2", 5.0)]
                 ),
-                "run 2": _apply_compat(
+                f"{nested_run_dir}run 2": _apply_compat(
                     [_scalar_event("2.1", 5.0), _scalar_event("2.2", 5.0)]
                 ),
             },
             {
-                "run 3": _apply_compat(
+                f"{nested_run_dir}run 3": _apply_compat(
                     [_scalar_event("3.1", 5.0), _scalar_event("3.2", 5.0)]
                 ),
-                "run 4": _apply_compat(
+                f"{nested_run_dir}run 4": _apply_compat(
                     [_scalar_event("4.1", 5.0), _scalar_event("4.2", 5.0)]
                 ),
-                "run 5": _apply_compat(
+                f"{nested_run_dir}run 5": _apply_compat(
                     [_scalar_event("5.1", 5.0), _scalar_event("5.2", 5.0)]
                 ),
             },
@@ -666,11 +675,20 @@ def test_start_uploading_scalars(
         self.assertEqual(mock_tracker.blob_tracker.call_count, 0)
 
     @parameterized.parameters(
-        {"existing_experiment": None, "one_platform_run_name": None},
-        {"existing_experiment": None, "one_platform_run_name": "."},
+        {
+            "existing_experiment": None,
+            "one_platform_run_name": None,
+            "nested_run_dir": "",
+        },
+        {
+            "existing_experiment": None,
+            "one_platform_run_name": ".",
+            "nested_run_dir": "nested-dir/",
+        },
         {
             "existing_experiment": _TEST_EXPERIMENT_NAME,
             "one_platform_run_name": _TEST_ONE_PLATFORM_RUN_NAME,
+            "nested_run_dir": "double/nested-dir/",
         },
     )
     @patch.object(
@@ -693,6 +711,7 @@ def test_start_uploading_scalars_one_shot(
         run_resource_mock,
         existing_experiment,
         one_platform_run_name,
+        nested_run_dir,
     ):
         """Check that one-shot uploading stops without AbortUploadError."""
 
@@ -760,10 +779,10 @@ def batch_create_time_series(parent, requests):
         mock_logdir_loader = mock.create_autospec(logdir_loader.LogdirLoader)
         mock_logdir_loader.get_run_events.side_effect = [
             {
-                "run 1": _apply_compat(
+                f"{nested_run_dir}run 1": _apply_compat(
                     [_scalar_event("tag_1.1", 5.0), _scalar_event("tag_1.2", 5.0)]
                 ),
-                "run 2": _apply_compat(
+                f"{nested_run_dir}run 2": _apply_compat(
                     [_scalar_event("tag_2.1", 5.0), _scalar_event("tag_2.2", 5.0)]
                 ),
             },
@@ -772,10 +791,10 @@ def batch_create_time_series(parent, requests):
         mock_logdir_loader_pre_create = mock.create_autospec(logdir_loader.LogdirLoader)
         mock_logdir_loader_pre_create.get_run_events.side_effect = [
             {
-                "run 1": _apply_compat(
+                f"{nested_run_dir}run 1": _apply_compat(
                     [_scalar_event("tag_1.1", 5.0), _scalar_event("tag_1.2", 5.0)]
                 ),
-                "run 2": _apply_compat(
+                f"{nested_run_dir}run 2": _apply_compat(
                     [_scalar_event("tag_2.1", 5.0), _scalar_event("tag_2.2", 5.0)]
                 ),
             },
@@ -804,6 +823,52 @@ def batch_create_time_series(parent, requests):
         self.assertEqual(mock_tracker.blob_tracker.call_count, 0)
         experiment_tracker_mock.set_experiment.assert_called_once()
 
+    @parameterized.parameters(
+        {"nested_run_dir": ""},
+        {"nested_run_dir": "nested-dir/"},
+        {"nested_run_dir": "double/nested-dir/"},
+    )
+    @patch.object(metadata, "_experiment_tracker", autospec=True)
+    @patch.object(experiment_resources, "Experiment", autospec=True)
+    def test_upload_nested_scalars_one_shot(
+        self,
+        experiment_resources_mock,
+        experiment_tracker_mock,
+        nested_run_dir,
+    ):
+        """Check that one-shot uploading stops without AbortUploadError."""
+
+        logdir = self.get_temp_dir()
+        uploader = _create_uploader(
+            logdir=logdir,
+        )
+        uploader.create_experiment()
+
+        run_1 = f"{nested_run_dir}run 1"
+        run_2 = f"{nested_run_dir}run 2"
+
+        mock_dispatcher = mock.create_autospec(uploader_lib._Dispatcher)
+        uploader._dispatcher = mock_dispatcher
+        mock_logdir_loader = mock.create_autospec(logdir_loader.LogdirLoader)
+        mock_logdir_loader.get_run_events.side_effect = [
+            {
+                run_1: _apply_compat(
+                    [_scalar_event("tag_1.1", 5.0), _scalar_event("tag_1.2", 5.0)]
+                ),
+                run_2: _apply_compat(
+                    [_scalar_event("tag_2.1", 5.0), _scalar_event("tag_2.2", 5.0)]
+                ),
+            },
+        ]
+        with mock.patch.object(uploader, "_logdir_loader", mock_logdir_loader):
+            uploader._upload_once()
+
+        self.assertEqual(1, mock_logdir_loader.get_run_events.call_count)
+        self.assertEqual(1, mock_dispatcher.dispatch_requests.call_count)
+        run_to_events = mock_dispatcher.dispatch_requests.call_args[0][0]
+        self.assertIn(run_1, run_to_events)
+        self.assertIn(run_2, run_to_events)
+
     @patch.object(metadata, "_experiment_tracker", autospec=True)
     @patch.object(experiment_resources, "Experiment", autospec=True)
     def test_upload_empty_logdir(