feat: Mark default profile run as complete when calling end_upload_tb…

…_log PiperOrigin-RevId: 646553890
googleapis · Jun 25, 2024 · 6397f30 · 6397f30
1 parent b1e9a6c
commit 6397f30
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 10 deletions.
diff --git a/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py b/google/cloud/aiplatform/tensorboard/plugins/tf_profiler/profile_uploader.py
@@ -153,7 +153,7 @@ def _profile_dir(self, run_name: str) -> str:
         Returns:
             Full path for run name.
         """
-        if run_name is None:
+        if run_name is None or run_name == uploader_utils.DEFAULT_PROFILE_RUN_NAME:
             return os.path.join(self._logdir, self.PROFILE_PATH)
         return os.path.join(self._logdir, run_name, self.PROFILE_PATH)
 
@@ -176,10 +176,7 @@ def send_request(self, run_name: str):
                 self._profile_dir(run_name)
             )
 
-        tensorboard_run_name = run_name if run_name else "profile"
-        tb_run = self._one_platform_resource_manager.get_run_resource_name(
-            tensorboard_run_name
-        )
+        tb_run = self._one_platform_resource_manager.get_run_resource_name(run_name)
 
         if run_name not in self._run_to_file_request_sender:
             self._run_to_file_request_sender[

diff --git a/google/cloud/aiplatform/tensorboard/uploader.py b/google/cloud/aiplatform/tensorboard/uploader.py
@@ -67,7 +67,6 @@
 from tensorboard.util import tensor_util
 
 _LOGGER = base.Logger(__name__)
-_DEFAULT_RUN_NAME = "default"
 
 TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
 
@@ -382,7 +381,11 @@ 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_name = (
+                run_name
+                if (run_name and run_name != ".")
+                else uploader_utils.DEFAULT_RUN_NAME
+            )
             run_names.append(run_name)
             for event in events:
                 _filter_graph_defs(event)
@@ -430,19 +433,25 @@ def _upload_once(self):
 
         run_to_events = self._logdir_loader.get_run_events()
         run_to_events = {
-            k if (k and k != ".") else _DEFAULT_RUN_NAME: v
+            k if (k and k != ".") else uploader_utils.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()
             }
-        self._experiment_runs = run_to_events.keys()
 
         # Add a profile event to trigger send_request in _additional_senders
         if self._should_profile():
-            run_to_events[self._run_name_prefix] = None
+            profile_run_name = (
+                self._run_name_prefix
+                if self._run_name_prefix
+                else uploader_utils.DEFAULT_PROFILE_RUN_NAME
+            )
+            run_to_events[profile_run_name] = None
+
+        self._experiment_runs = run_to_events.keys()
 
         with self._tracker.send_tracker():
             self._dispatcher.dispatch_requests(run_to_events)

diff --git a/google/cloud/aiplatform/tensorboard/uploader_utils.py b/google/cloud/aiplatform/tensorboard/uploader_utils.py
@@ -41,6 +41,8 @@
 from tensorboard.util import tb_logging
 
 TensorboardServiceClient = tensorboard_service_client.TensorboardServiceClient
+DEFAULT_RUN_NAME = "default"
+DEFAULT_PROFILE_RUN_NAME = "profile"
 
 logger = tb_logging.get_logger()
 logger.setLevel(logging.WARNING)