feat: Create Vertex Experiment when uploading Tensorboard logs

PiperOrigin-RevId: 626105964

google/cloud/aiplatform/metadata/metadata.py

@@ -237,7 +237,7 @@ def _get_global_tensorboard(self) -> Optional[tensorboard_resource.Tensorboard]:
         Returns:
             tensorboard_resource.Tensorboard: the global TensorBoard instance.
         """
-        if self._global_tensorboard:
+        if self._global_tensorboard and hasattr(self._global_tensorboard, "resource_name"):
             credentials, _ = google.auth.default()
             if self.experiment and self.experiment._metadata_context.credentials:
                 credentials = self.experiment._metadata_context.credentials

google/cloud/aiplatform/tensorboard/uploader.py

@@ -36,6 +36,8 @@
 from google.cloud.aiplatform.compat.types import tensorboard_experiment
 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 constants
+from google.cloud.aiplatform.metadata import metadata
 from google.cloud.aiplatform.tensorboard import logdir_loader
 from google.cloud.aiplatform.tensorboard import upload_tracker
 from google.cloud.aiplatform.tensorboard import uploader_constants
@@ -250,14 +252,16 @@ def _create_or_get_experiment(self) -> tensorboard_experiment.TensorboardExperim
         logger.info("Creating experiment")
 
         tb_experiment = tensorboard_experiment.TensorboardExperiment(
-            description=self._description, display_name=self._experiment_display_name
+            description=self._description,
+            display_name=self._experiment_display_name,
         )
 
         try:
-            experiment = self._api.create_tensorboard_experiment(
+            experiment = tensorboard_experiment.TensorboardExperiment.create(
                 parent=self._tensorboard_resource_name,
                 tensorboard_experiment=tb_experiment,
                 tensorboard_experiment_id=self._experiment_name,
+                labels=constants._VERTEX_EXPERIMENT_TB_EXPERIMENT_LABEL,
             )
             self._is_brand_new_experiment = True
         except exceptions.AlreadyExists:
@@ -271,14 +275,22 @@ def _create_or_get_experiment(self) -> tensorboard_experiment.TensorboardExperim
     def create_experiment(self):
         """Creates an Experiment for this upload session and returns the ID."""
 
+        metadata._experiment_tracker.set_tensorboard(
+            tensorboard=self._tensorboard_resource_name
+        )
+        metadata._experiment_tracker.set_experiment(
+            experiment=self._experiment_name,
+            description=self._description,
+            backing_tensorboard=self._tensorboard_resource_name,
+        )
         experiment = self._create_or_get_experiment()
-        self._experiment = experiment
+        self._experiment_resource_name = f"{self._tensorboard_resource_name}/experiments/{self._experiment_name}"
         self._one_platform_resource_manager = uploader_utils.OnePlatformResourceManager(
-            self._experiment.name, self._api
+            self._experiment_resource_name, self._api
         )
 
         self._request_sender = _BatchedRequestSender(
-            self._experiment.name,
+            self._experiment_resource_name,
             self._api,
             allowed_plugins=self._allowed_plugins,
             upload_limits=self._upload_limits,
@@ -294,7 +306,7 @@ def create_experiment(self):
         # Update partials with experiment name
         for sender in self._additional_senders.keys():
             self._additional_senders[sender] = self._additional_senders[sender](
-                experiment_resource_name=self._experiment.name,
+                experiment_resource_name=self._experiment_resource_name,
             )
 
         self._dispatcher = _Dispatcher(
@@ -333,7 +345,7 @@ def _create_additional_senders(self) -> Dict[str, uploader_utils.RequestSender]:
             )
 
     def get_experiment_resource_name(self):
-        return self._experiment.name
+        return self._experiment_resource_name
 
     def start_uploading(self):
         """Blocks forever to continuously upload data from the logdir.