feat: Add option to not use default tensorboard

PiperOrigin-RevId: 581402305

google/cloud/aiplatform/initializer.py

@@ -112,7 +112,7 @@ def init(
         experiment: Optional[str] = None,
         experiment_description: Optional[str] = None,
         experiment_tensorboard: Optional[
-            Union[str, tensorboard_resource.Tensorboard]
+            Union[str, tensorboard_resource.Tensorboard, bool]
         ] = None,
         staging_bucket: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
@@ -128,7 +128,7 @@ def init(
                 set defaults to us-central-1.
             experiment (str): Optional. The experiment name.
             experiment_description (str): Optional. The description of the experiment.
-            experiment_tensorboard (Union[str, tensorboard_resource.Tensorboard]):
+            experiment_tensorboard (Union[str, tensorboard_resource.Tensorboard, bool]):
                 Optional. The Vertex AI TensorBoard instance, Tensorboard resource name,
                 or Tensorboard resource ID to use as a backing Tensorboard for the provided
                 experiment.
@@ -141,6 +141,13 @@ def init(
                 Any subsequent calls to aiplatform.init() with `experiment` and without
                 `experiment_tensorboard` will automatically assign the global Tensorboard
                 to the `experiment`.
+
+                If `experiment_tensorboard` is ommitted or set to `True` or `None` the global
+                Tensorboard will be assigned to the `experiment`. If a global Tensorboard is
+                not set, the default Tensorboard instance will be used, and created if it deos not exist.
+
+                To disable creating and using Tensorboard with `experiment`, set `experiment_tensorboard` to False.
+                Any subsequent calls to aiplatform.init() should include this setting as well.
             staging_bucket (str): The default staging bucket to use to stage artifacts
                 when making API calls. In the form gs://...
             credentials (google.auth.credentials.Credentials): The default custom
@@ -177,7 +184,7 @@ def init(
                 "Experiment needs to be set in `init` in order to add experiment descriptions."
             )
 
-        if experiment_tensorboard:
+        if experiment_tensorboard and not isinstance(experiment_tensorboard, bool):
             metadata._experiment_tracker.set_tensorboard(
                 tensorboard=experiment_tensorboard,
                 project=project,

google/cloud/aiplatform/metadata/metadata.py

@@ -267,7 +267,7 @@ def set_experiment(
         *,
         description: Optional[str] = None,
         backing_tensorboard: Optional[
-            Union[str, tensorboard_resource.Tensorboard]
+            Union[str, tensorboard_resource.Tensorboard, bool]
         ] = None,
     ):
         """Set the experiment. Will retrieve the Experiment if it exists or create one with the provided name.
@@ -277,21 +277,30 @@ def set_experiment(
                 Required. Name of the experiment to set.
             description (str):
                 Optional. Description of an experiment.
-            backing_tensorboard Union[str, aiplatform.Tensorboard]:
+            backing_tensorboard Union[str, aiplatform.Tensorboard, bool]:
                 Optional. If provided, assigns tensorboard as backing tensorboard to support time series metrics
                 logging.
+
+                If ommitted, or set to `True` or `None`, the global tensorboard is used.
+                If no global tensorboard is set, the default tensorboard will be used, and created if it does not exist.
+
+                To disable using a backign tensorboard, set `backing_tensorboard` to `False`.
+                To maintain this behavior, set `experiment_tensorboard` to `False` in subsequent calls to aiplatform.init().
         """
         self.reset()
 
         experiment = experiment_resources.Experiment.get_or_create(
             experiment_name=experiment, description=description
         )
 
-        backing_tb = (
-            backing_tensorboard
-            or self._global_tensorboard
-            or _get_or_create_default_tensorboard()
-        )
+        if backing_tensorboard and not isinstance(backing_tensorboard, bool):
+            backing_tb = backing_tensorboard
+        elif isinstance(backing_tensorboard, bool) and not backing_tensorboard:
+            backing_tb = None
+        else:
+            backing_tb = (
+                self._global_tensorboard or _get_or_create_default_tensorboard()
+            )
 
         current_backing_tb = experiment.backing_tensorboard_resource_name
 
@@ -322,7 +331,7 @@ def set_tensorboard(
             credentials (auth_credentials.Credentials):
                 Optional. Custom credentials used to set this Tensorboard resource.
         """
-        if isinstance(tensorboard, str):
+        if tensorboard and isinstance(tensorboard, str):
             tensorboard = tensorboard_resource.Tensorboard(
                 tensorboard,
                 project=project,

tests/unit/aiplatform/test_initializer.py

@@ -114,7 +114,9 @@ def test_init_service_account_sets_service_account(self):
     def test_init_experiment_sets_experiment(self, set_experiment_mock):
         initializer.global_config.init(experiment=_TEST_EXPERIMENT)
         set_experiment_mock.assert_called_once_with(
-            experiment=_TEST_EXPERIMENT, description=None, backing_tensorboard=None
+            experiment=_TEST_EXPERIMENT,
+            description=None,
+            backing_tensorboard=None,
         )
 
     @patch.object(_experiment_tracker, "set_experiment")
@@ -162,25 +164,44 @@ def test_init_with_experiment_tensorboard_resource_sets_global_tensorboard(
             credentials=None,
         )
 
-    @patch.object(_experiment_tracker, "set_tensorboard")
     @patch.object(_experiment_tracker, "set_experiment")
     def test_init_experiment_without_tensorboard_uses_global_tensorboard(
-        self,
-        set_tensorboard_mock,
-        set_experiment_mock,
+        self, set_experiment_mock
     ):
 
-        initializer.global_config.init(experiment_tensorboard=_TEST_TENSORBOARD_NAME)
+        initializer.global_config.tensorboard = _TEST_TENSORBOARD_NAME
 
         initializer.global_config.init(
             experiment=_TEST_EXPERIMENT,
         )
 
         set_experiment_mock.assert_called_once_with(
-            tensorboard=_TEST_TENSORBOARD_NAME,
-            project=None,
-            location=None,
-            credentials=None,
+            experiment=_TEST_EXPERIMENT,
+            description=None,
+            backing_tensorboard=None,
+        )
+
+        assert initializer.global_config.tensorboard == _TEST_TENSORBOARD_NAME
+
+    @patch.object(_experiment_tracker, "set_tensorboard")
+    @patch.object(_experiment_tracker, "set_experiment")
+    def test_init_experiment_tensorboard_false_does_not_set_tensorboard(
+        self, set_experiment_mock, set_tensorboard_mock
+    ):
+
+        initializer.global_config.tensorboard = _TEST_TENSORBOARD_NAME
+
+        initializer.global_config.init(
+            experiment=_TEST_EXPERIMENT,
+            experiment_tensorboard=False,
+        )
+
+        set_tensorboard_mock.assert_not_called()
+
+        set_experiment_mock.assert_called_once_with(
+            experiment=_TEST_EXPERIMENT,
+            description=None,
+            backing_tensorboard=False,
         )
 
     def test_init_experiment_description_fail_without_experiment(self):

tests/unit/aiplatform/test_metadata.py

@@ -1140,6 +1140,57 @@ def test_init_experiment_with_default_tensorboard(
         )
         assign_backing_tensorboard_mock.assert_called_once()
 
+    @pytest.mark.usefixtures(
+        "get_metadata_store_mock",
+        "get_experiment_run_run_mock",
+    )
+    def test_init_experiment_tensorboard_false_doesNotSet_backing_tensorboard(
+        self, list_default_tensorboard_mock, assign_backing_tensorboard_mock
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            experiment=_TEST_EXPERIMENT,
+            experiment_tensorboard=False,
+        )
+
+        list_default_tensorboard_mock.assert_not_called()
+        assign_backing_tensorboard_mock.assert_not_called()
+
+    @pytest.mark.usefixtures(
+        "get_metadata_store_mock",
+        "get_experiment_run_run_mock",
+    )
+    def test_init_experiment_tensorboard_true_sets_backing_tensorboard(
+        self, list_default_tensorboard_mock, assign_backing_tensorboard_mock
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            experiment=_TEST_EXPERIMENT,
+            experiment_tensorboard=True,
+        )
+
+        list_default_tensorboard_mock.assert_called()
+        assign_backing_tensorboard_mock.assert_called()
+
+    @pytest.mark.usefixtures(
+        "get_metadata_store_mock",
+        "get_experiment_run_run_mock",
+    )
+    def test_init_experiment_tensorboard_none_sets_backing_tensorboard(
+        self, list_default_tensorboard_mock, assign_backing_tensorboard_mock
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            experiment=_TEST_EXPERIMENT,
+            experiment_tensorboard=None,
+        )
+
+        list_default_tensorboard_mock.assert_called()
+        assign_backing_tensorboard_mock.assert_called()
+
     @pytest.mark.usefixtures("get_metadata_store_mock")
     def test_create_experiment(self, create_experiment_context_mock):
         exp = aiplatform.Experiment.create(