fix: EvalTask evaluation should only display experiment button once f…

…or ipython environments PiperOrigin-RevId: 711880078
googleapis · Jan 4, 2025 · 88d2656 · 88d2656
1 parent a4d98cf
commit 88d2656
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/google/cloud/aiplatform/metadata/metadata.py b/google/cloud/aiplatform/metadata/metadata.py
@@ -303,6 +303,7 @@ def set_experiment(
         ] = None,
         project: Optional[str] = None,
         location: Optional[str] = None,
+        display_button: bool = True,
     ):
         """Set the experiment. Will retrieve the Experiment if it exists or create one with the provided name.
 
@@ -326,6 +327,8 @@ def set_experiment(
             location (str):
                 Optional. Location where this experiment will be retrieved from or created. Overrides location set in
                 aiplatform.init.
+            display_button (bool):
+                Optional. If set to `True`, displays a button to the experiment in the IPython notebook.
         """
         self.reset()
 
@@ -350,7 +353,8 @@ def set_experiment(
         if not current_backing_tb and backing_tb:
             experiment.assign_backing_tensorboard(tensorboard=backing_tb)
 
-        _ipython_utils.display_experiment_button(experiment)
+        if display_button:
+            _ipython_utils.display_experiment_button(experiment)
 
         self._experiment = experiment
 

diff --git a/vertexai/evaluation/eval_task.py b/vertexai/evaluation/eval_task.py
@@ -430,7 +430,9 @@ def evaluate(
                 retry_timeout=retry_timeout,
             )
             metadata._experiment_tracker.set_experiment(
-                experiment=global_experiment_name, backing_tensorboard=False
+                experiment=global_experiment_name,
+                backing_tensorboard=False,
+                display_button=False,
             )
         elif self._experiment and not global_experiment_name:
             metadata._experiment_tracker.set_experiment(