fix: Identify nested Vertex Tensorboard profile runs for uploading when standard event data is not present

PiperOrigin-RevId: 671909634

Author: vertex-sdk-bot

google/cloud/aiplatform/tensorboard/logdir_loader.py

@@ -19,6 +19,7 @@
 
 import collections
 import os
+import tensorflow as tf
 
 from tensorboard.backend.event_processing import directory_watcher
 from tensorboard.backend.event_processing import io_wrapper
@@ -28,6 +29,37 @@
 logger = tb_logging.get_logger()
 
 
+def is_plugins_subdirectory(path):
+    """Returns true if the path is a profile subdirectory."""
+    if not tf.io.gfile.isdir(path):
+        return False
+    dirs = tf.io.gfile.listdir(path)
+    return "plugins/" in dirs or "plugins" in dirs
+
+
+def get_plugins_subdirectories(path):
+    """Returns a list of plugins subdirectories within the given path."""
+    if not tf.io.gfile.exists(path):
+        # No directory to traverse.
+        logger.warning("Directory does not exist: %s", str(path))
+        return ()
+
+    current_glob_string = os.path.join(path, "*")
+    while True:
+        globs = tf.io.gfile.glob(current_glob_string)
+
+        if not globs:
+            # This subdirectory level lacks files. Terminate.
+            return
+
+        for glob in globs:
+            if is_plugins_subdirectory(glob):
+                yield glob
+
+        # Iterate to the next level of subdirectories.
+        current_glob_string = os.path.join(current_glob_string, "*")
+
+
 class LogdirLoader:
     """Loader for a root log directory, maintaining multiple DirectoryLoaders.
 
@@ -58,13 +90,11 @@ def __init__(self, logdir, directory_loader_factory):
         self._directory_loaders = {}
 
     def synchronize_runs(self):
-        """Finds new runs within `logdir` and makes `DirectoryLoaders` for
-        them.
+        """Finds new runs within `logdir` and makes `DirectoryLoaders` for them.
 
         In addition, any existing `DirectoryLoader` whose run directory
         no longer exists will be deleted.
 
-        Modify run name to work with Experiments restrictions.
         """
         logger.info("Starting logdir traversal of %s", self._logdir)
         runs_seen = set()
@@ -74,6 +104,12 @@ def synchronize_runs(self):
             if run not in self._directory_loaders:
                 logger.info("- Adding run for relative directory %s", run)
                 self._directory_loaders[run] = self._directory_loader_factory(subdir)
+        for subdir in get_plugins_subdirectories(self._logdir):
+            run = os.path.relpath(subdir, self._logdir)
+            runs_seen.add(run)
+            if run not in self._directory_loaders:
+                logger.info("- Adding run for relative directory %s", run)
+                self._directory_loaders[run] = self._directory_loader_factory(subdir)
         stale_runs = set(self._directory_loaders) - runs_seen
         if stale_runs:
             for run in stale_runs:

tests/unit/aiplatform/test_logdir_loader.py

@@ -181,6 +181,34 @@ def test_single_event_logdir(self):
         # A second load should indicate no new data for the run.
         self.assertEqual(self._extract_run_to_tags(loader.get_run_events()), {".": []})
 
+    def test_profile_logdir(self):
+        logdir = self.get_temp_dir()
+        profile_dir = os.path.join(logdir, "foo/plugins/profile")
+        os.makedirs(profile_dir, exist_ok=True)
+        tempfile.NamedTemporaryFile(
+            prefix="bar", suffix=".xplane.pb", dir=profile_dir, delete=False
+        )
+        self.assertNotEmpty(os.listdir(profile_dir))
+        loader = self._create_logdir_loader(logdir)
+        loader.synchronize_runs()
+        self.assertEqual(
+            self._extract_run_to_tags(loader.get_run_events()), {"foo": []}
+        )
+
+    def test_profile_subdirectories(self):
+        logdir = self.get_temp_dir()
+        profile_dir = os.path.join(logdir, "foo/bar/subdir/plugins/profile")
+        os.makedirs(profile_dir, exist_ok=True)
+        tempfile.NamedTemporaryFile(
+            prefix="bar", suffix=".xplane.pb", dir=profile_dir, delete=False
+        )
+        self.assertNotEmpty(os.listdir(profile_dir))
+        loader = self._create_logdir_loader(logdir)
+        loader.synchronize_runs()
+        self.assertEqual(
+            self._extract_run_to_tags(loader.get_run_events()), {"foo/bar/subdir": []}
+        )
+
     def test_multiple_writes_to_logdir(self):
         logdir = self.get_temp_dir()
         with FileWriter(os.path.join(logdir, "a")) as writer:

tests/unit/aiplatform/test_uploader.py

@@ -1375,6 +1375,45 @@ def test_profile_plugin_included_by_default(
             self.assertIn(run_name, profile_sender._run_to_file_request_sender)
             experiment_tracker_mock.set_experiment.assert_called_once()
 
+    @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_nested_profile_files_are_uploaded(
+        self, experiment_resources_mock, experiment_tracker_mock, run_resource_mock
+    ):
+        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
+        )
+        run_name = "profile_test_run"
+        run_resource_mock.return_value = _TEST_ONE_PLATFORM_RUN_NAME
+        with tempfile.TemporaryDirectory() as logdir:
+            prof_path = os.path.join(
+                logdir, run_name, profile_uploader.ProfileRequestSender.PROFILE_PATH
+            )
+            os.makedirs(prof_path)
+
+            uploader = _create_uploader(
+                self.mock_client,
+                logdir,
+                one_shot=True,
+            )
+
+            uploader.create_experiment()
+            uploader._upload_once()
+            senders = uploader._dispatcher._additional_senders
+            self.assertIn("profile", senders.keys())
+
+            profile_sender = senders["profile"]
+            self.assertIn(run_name, profile_sender._run_to_profile_loaders)
+            self.assertIn(run_name, profile_sender._run_to_file_request_sender)
+            experiment_tracker_mock.set_experiment.assert_called_once()
+
     @patch.object(metadata, "_experiment_tracker", autospec=True)
     @patch.object(experiment_resources, "Experiment", autospec=True)
     def test_active_experiment_set_experiment_not_called(