fix: Do not reset aiplatform.Experiment or aiplatform.ExperimentRun unnecessarily when running tensorboard uploader.

PiperOrigin-RevId: 644168539

Author: vertex-sdk-bot

google/cloud/aiplatform/tensorboard/uploader.py

@@ -30,10 +30,12 @@
 from google.cloud.aiplatform.compat.services import (
     tensorboard_service_client,
 )
+from google.cloud.aiplatform.compat.types import execution as gca_execution
 from google.cloud.aiplatform.compat.types import tensorboard_data
 from google.cloud.aiplatform.compat.types import tensorboard_service
 from google.cloud.aiplatform.compat.types import tensorboard_time_series
 from google.cloud.aiplatform.metadata import experiment_resources
+from google.cloud.aiplatform.metadata import experiment_run_resource
 from google.cloud.aiplatform.metadata import metadata
 from google.cloud.aiplatform.tensorboard import logdir_loader
 from google.cloud.aiplatform.tensorboard import upload_tracker
@@ -164,6 +166,9 @@ def __init__(
         self._one_shot = one_shot
         self._dispatcher = None
         self._additional_senders: Dict[str, uploader_utils.RequestSender] = {}
+        self._experiment_runs = []
+        self._project = None
+        self._location = None
         if logdir_poll_rate_limiter is None:
             self._logdir_poll_rate_limiter = uploader_utils.RateLimiter(
                 uploader_constants.MIN_LOGDIR_POLL_INTERVAL_SECS
@@ -221,27 +226,31 @@ def create_experiment(self):
         Vertex Experiment and associate it with a Tensorboard Experiment.
         """
         m = self._api.parse_tensorboard_path(self._tensorboard_resource_name)
+        self._project = m["project"]
+        self._location = m["location"]
 
         existing_experiment = experiment_resources.Experiment.get(
             experiment_name=self._experiment_name,
-            project=m["project"],
-            location=m["location"],
+            project=self._project,
+            location=self._location,
         )
         if not existing_experiment:
             self._is_brand_new_experiment = True
 
-        metadata._experiment_tracker.reset()
+        if metadata._experiment_tracker.experiment_name != self._experiment_name:
+            logging.info(f"Setting experiment to {self._experiment_name}")
+            metadata._experiment_tracker.reset()
+            metadata._experiment_tracker.set_experiment(
+                project=self._project,
+                location=self._location,
+                experiment=self._experiment_name,
+                description=self._description,
+                backing_tensorboard=self._tensorboard_resource_name,
+            )
         metadata._experiment_tracker.set_tensorboard(
             tensorboard=self._tensorboard_resource_name,
-            project=m["project"],
-            location=m["location"],
-        )
-        metadata._experiment_tracker.set_experiment(
-            project=m["project"],
-            location=m["location"],
-            experiment=self._experiment_name,
-            description=self._description,
-            backing_tensorboard=self._tensorboard_resource_name,
+            project=self._project,
+            location=self._location,
         )
 
         self._tensorboard_experiment_resource_name = (
@@ -309,6 +318,17 @@ def _create_additional_senders(self) -> Dict[str, uploader_utils.RequestSender]:
     def get_experiment_resource_name(self):
         return self._tensorboard_experiment_resource_name
 
+    def _end_experiment_runs(self):
+        # End all runs created by uploader
+        for run_name in self._experiment_runs:
+            if run_name:
+                logging.info("Ending run %s", run_name)
+                run = experiment_run_resource.ExperimentRun.get(
+                    project=self._project, location=self._location, run_name=run_name
+                )
+                if run:
+                    run.update_state(state=gca_execution.Execution.State.COMPLETE)
+
     def start_uploading(self):
         """Blocks forever to continuously upload data from the logdir.
 
@@ -334,6 +354,7 @@ def start_uploading(self):
             self._logdir_poll_rate_limiter.tick()
             self._upload_once()
             if self._one_shot:
+                self._end_experiment_runs()
                 break
         if self._one_shot and not self._tracker.has_data():
             logger.warning(
@@ -343,6 +364,7 @@ def start_uploading(self):
 
     def _end_uploading(self):
         self._continue_uploading = False
+        self._end_experiment_runs()
 
     def _pre_create_runs_and_time_series(self):
         """Iterates though the log dir to collect TensorboardRuns and
@@ -409,6 +431,7 @@ def _upload_once(self):
             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():

google/cloud/aiplatform/tensorboard/uploader_utils.py

@@ -234,7 +234,7 @@ def _create_or_get_run_resource(
                 run_name=run_name,
                 experiment=experiment,
                 tensorboard=tensorboard,
-                state=gca_execution.Execution.State.COMPLETE,
+                state=gca_execution.Execution.State.RUNNING,
             )
         tb_run_artifact = experiment_run._backing_tensorboard_run
         tb_run = tb_run_artifact.resource

tests/unit/aiplatform/test_uploader.py

@@ -785,7 +785,11 @@ def batch_create_time_series(parent, requests):
             with mock.patch.object(
                 uploader, "_logdir_loader_pre_create", mock_logdir_loader_pre_create
             ):
-                uploader.start_uploading()
+                with mock.patch.object(
+                    uploader, "_end_experiment_runs", return_value=None
+                ):
+                    uploader.start_uploading()
+                    uploader._end_experiment_runs.assert_called_once()
 
         self.assertEqual(existing_experiment is None, uploader._is_brand_new_experiment)
         self.assertEqual(2, mock_client.write_tensorboard_experiment_data.call_count)
@@ -797,6 +801,7 @@ def batch_create_time_series(parent, requests):
         self.assertLen(mock_tracker.scalars_tracker.call_args[0], 1)
         self.assertEqual(mock_tracker.tensors_tracker.call_count, 0)
         self.assertEqual(mock_tracker.blob_tracker.call_count, 0)
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(metadata, "_experiment_tracker", autospec=True)
     @patch.object(experiment_resources, "Experiment", autospec=True)
@@ -814,6 +819,7 @@ def test_upload_empty_logdir(
         uploader.create_experiment()
         uploader._upload_once()
         mock_client.write_tensorboard_experiment_data.assert_not_called()
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(metadata, "_experiment_tracker", autospec=True)
     @patch.object(experiment_resources, "Experiment", autospec=True)
@@ -847,6 +853,7 @@ def mock_upload_once():
         uploader.create_experiment()
         with self.assertRaises(SuccessError):
             uploader.start_uploading()
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(
         uploader_utils.OnePlatformResourceManager,
@@ -874,6 +881,7 @@ def test_upload_swallows_rpc_failure(
         mock_client.write_tensorboard_experiment_data.side_effect = error
         uploader._upload_once()
         mock_client.write_tensorboard_experiment_data.assert_called_once()
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(
         uploader_utils.OnePlatformResourceManager,
@@ -1006,10 +1014,12 @@ def test_upload_full_logdir(
         self.assertProtoEquals(expected_request3[1], request3[1])
         self.assertProtoEquals(expected_request4[0], request4[0])
         mock_client.write_tensorboard_experiment_data.reset_mock()
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
         # Empty third round
         uploader._upload_once()
         mock_client.write_tensorboard_experiment_data.assert_not_called()
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(
         uploader_utils.OnePlatformResourceManager,
@@ -1057,6 +1067,7 @@ def test_verbosity_zero_creates_upload_tracker_with_verbosity_zero(
         self.assertEqual(mock_constructor.call_count, 1)
         self.assertEqual(mock_constructor.call_args[1], {"verbosity": 0})
         self.assertEqual(mock_tracker.scalars_tracker.call_count, 1)
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(
         uploader_utils.OnePlatformResourceManager,
@@ -1160,6 +1171,7 @@ def create_time_series(tensorboard_time_series, parent=None):
         self.assertEqual(mock_tracker.scalars_tracker.call_count, 0)
         self.assertEqual(mock_tracker.tensors_tracker.call_count, 0)
         self.assertEqual(mock_tracker.blob_tracker.call_count, 12)
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
     @patch.object(
         uploader_utils.OnePlatformResourceManager,
@@ -1282,6 +1294,27 @@ def test_profile_plugin_included_by_default(
             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(
+        self, experiment_resources_mock, experiment_tracker_mock
+    ):
+        experiment_resources_mock.get.return_value = _TEST_EXPERIMENT_NAME
+        experiment_tracker_mock.set_experiment.return_value = _TEST_EXPERIMENT_NAME
+        experiment_tracker_mock.experiment_name = _TEST_EXPERIMENT_NAME
+        experiment_tracker_mock.set_tensorboard.return_value = (
+            _TEST_TENSORBOARD_RESOURCE_NAME
+        )
+        logdir = self.get_temp_dir()
+        mock_client = _create_mock_client()
+
+        uploader = _create_uploader(mock_client, logdir)
+        uploader.create_experiment()
+        uploader._upload_once()
+
+        experiment_tracker_mock.set_experiment.assert_not_called()
 
 
 # TODO(b/276368161)
@@ -1387,6 +1420,7 @@ def test_thread_continuously_uploads(
         self.assertEqual(b"12345", request2.plugin_data)
         self.assertEqual("scalars", request3.plugin_name)
         self.assertEqual("profile", request4.plugin_name)
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
         # Check write_tensorboard_experiment_data calls
         self.assertEqual(1, mock_client.write_tensorboard_experiment_data.call_count)
@@ -1425,17 +1459,22 @@ def test_thread_continuously_uploads(
         self.assertProtoEquals(expected_request1[1], request1[1])
         self.assertProtoEquals(expected_request2[0], request2[0])
 
-        uploader._end_uploading()
+        with mock.patch.object(uploader, "_end_experiment_runs", return_value=None):
+            uploader._end_uploading()
+            uploader._end_experiment_runs.assert_called_once()
         time.sleep(1)
         self.assertFalse(uploader_thread.is_alive())
         mock_client.write_tensorboard_experiment_data.reset_mock()
 
         # Empty directory
         uploader._upload_once()
         mock_client.write_tensorboard_experiment_data.assert_not_called()
-        uploader._end_uploading()
+        with mock.patch.object(uploader, "_end_experiment_runs", return_value=None):
+            uploader._end_uploading()
+            uploader._end_experiment_runs.assert_called_once()
         time.sleep(1)
         self.assertFalse(uploader_thread.is_alive())
+        experiment_tracker_mock.set_experiment.assert_called_once()
 
 
 @pytest.mark.usefixtures("google_auth_mock")