feat: Add display experiment run button for Ipython environments

vertex-sdk-bot · copybara-github · commit ba6582856b1d · 2024-05-31T14:10:03.000-07:00
PiperOrigin-RevId: 639147169
diff --git a/google/cloud/aiplatform/metadata/metadata.py b/google/cloud/aiplatform/metadata/metadata.py
@@ -492,6 +492,8 @@ def start_run(
                 run_name=run, experiment=self.experiment, tensorboard=tensorboard
             )
 
+        _ipython_utils.display_experiment_run_button(self._experiment_run)
+
         return self._experiment_run
 
     def end_run(
diff --git a/google/cloud/aiplatform/utils/_ipython_utils.py b/google/cloud/aiplatform/utils/_ipython_utils.py
@@ -24,6 +24,7 @@
 
 if typing.TYPE_CHECKING:
     from google.cloud.aiplatform.metadata import experiment_resources
+    from google.cloud.aiplatform.metadata import experiment_run_resource
     from google.cloud.aiplatform import model_evaluation
 
 _LOGGER = base.Logger(__name__)
@@ -168,6 +169,36 @@ def display_experiment_button(experiment: "experiment_resources.Experiment") ->
     display_link("View Experiment", uri, "science")
 
 
+def display_experiment_run_button(
+    experiment_run: "experiment_run_resource.ExperimentRun",
+) -> None:
+    """Function to generate a link bound to the Vertex experiment run"""
+    if not is_ipython_available():
+        return
+    try:
+        project = experiment_run.project
+        location = experiment_run.location
+        experiment_name = experiment_run._experiment._metadata_context.name
+        run_name = experiment_run.name
+        if (
+            run_name is None
+            or experiment_name is None
+            or project is None
+            or location is None
+        ):
+            return
+    except AttributeError:
+        _LOGGER.warning("Unable to fetch experiment run metadata")
+        return
+
+    uri = (
+        "https://console.cloud.google.com/vertex-ai/experiments/locations/"
+        + f"{location}/experiments/{experiment_name}/"
+        + f"runs/{experiment_name}-{run_name}?project={project}"
+    )
+    display_link("View Experiment Run", uri, "science")
+
+
 def display_model_evaluation_button(
     evaluation: "model_evaluation.ModelEvaluation",
 ) -> None:

Original file line number	Diff line number	Diff line change
`@@ -492,6 +492,8 @@ def start_run(`
`492`	`492`	`run_name=run, experiment=self.experiment, tensorboard=tensorboard`
`493`	`493`	`)`
`494`	`494`
	`495`	`+ _ipython_utils.display_experiment_run_button(self._experiment_run)`
	`496`	`+`
`495`	`497`	`return self._experiment_run`
`496`	`498`
`497`	`499`	`def end_run(`