feat: Add display model evaluation button for Ipython environments

PiperOrigin-RevId: 617902950

Author: matthew29tang

google/cloud/aiplatform/model_evaluation/model_evaluation_job.py

@@ -18,6 +18,7 @@
 from typing import Optional, List, Union
 
 from google.auth import credentials as auth_credentials
+import grpc
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
@@ -27,6 +28,7 @@
 )
 from google.cloud.aiplatform import model_evaluation
 from google.cloud.aiplatform import pipeline_jobs
+from google.cloud.aiplatform.utils import _ipython_utils
 
 from google.cloud.aiplatform.compat.types import (
     pipeline_state_v1 as gca_pipeline_state_v1,
@@ -380,7 +382,6 @@ def get_model_evaluation(
             return
 
         for component in self.backing_pipeline_job.task_details:
-
             # This assumes that task_details has a task with a task_name == backing_pipeline_job.name
             if not component.task_name == self.backing_pipeline_job.name:
                 continue
@@ -407,5 +408,14 @@ def get_model_evaluation(
                 evaluation_name=eval_resource_name,
                 credentials=self.credentials,
             )
-
+            _ipython_utils.display_model_evaluation_button(eval_resource)
             return eval_resource
+
+    def wait(self) -> None:
+        """Wait for the PipelineJob to complete, then get the model evaluation resource."""
+        super().wait()
+
+        try:
+            self.get_model_evaluation()
+        except grpc.RpcError as e:
+            _LOGGER.error("Get model evaluation call failed with error %s", e)

google/cloud/aiplatform/models.py

@@ -49,6 +49,7 @@
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform.utils import gcs_utils
 from google.cloud.aiplatform.utils import _explanation_utils
+from google.cloud.aiplatform.utils import _ipython_utils
 from google.cloud.aiplatform import model_evaluation
 from google.cloud.aiplatform.compat.services import endpoint_service_client
 
@@ -5136,7 +5137,8 @@ def get_model_evaluation(
                 _LOGGER.warning(
                     f"Your model has more than one model evaluation, this is returning only one evaluation resource: {evaluations[0].resource_name}"
                 )
-            return evaluations[0] if evaluations else evaluations
+            _ipython_utils.display_model_evaluation_button(evaluations[0])
+            return evaluations[0]
         else:
             resource_uri_parts = self._parse_resource_name(self.resource_name)
             evaluation_resource_name = (
@@ -5146,10 +5148,12 @@ def get_model_evaluation(
                 )
             )
 
-            return model_evaluation.ModelEvaluation(
+            evaluation = model_evaluation.ModelEvaluation(
                 evaluation_name=evaluation_resource_name,
                 credentials=self.credentials,
             )
+            _ipython_utils.display_model_evaluation_button(evaluation)
+            return evaluation
 
     def evaluate(
         self,

google/cloud/aiplatform/utils/_ipython_utils.py

@@ -16,11 +16,15 @@
 #
 
 import sys
+import typing
 from uuid import uuid4
 from typing import Optional
 
 from google.cloud.aiplatform import base
-from google.cloud.aiplatform.metadata import experiment_resources
+
+if typing.TYPE_CHECKING:
+    from google.cloud.aiplatform.metadata import experiment_resources
+    from google.cloud.aiplatform import model_evaluation
 
 _LOGGER = base.Logger(__name__)
 
@@ -142,7 +146,7 @@ def display_link(text: str, url: str, icon: Optional[str] = "open_in_new") -> No
     display(HTML(html))
 
 
-def display_experiment_button(experiment: experiment_resources.Experiment) -> None:
+def display_experiment_button(experiment: "experiment_resources.Experiment") -> None:
     """Function to generate a link bound to the Vertex experiment"""
     if not is_ipython_available():
         return
@@ -162,3 +166,34 @@ def display_experiment_button(experiment: experiment_resources.Experiment) -> No
         + f"runs?project={project}"
     )
     display_link("View Experiment", uri, "science")
+
+
+def display_model_evaluation_button(
+    evaluation: "model_evaluation.ModelEvaluation",
+) -> None:
+    """Function to generate a link bound to the Vertex model evaluation"""
+    if not is_ipython_available():
+        return
+
+    try:
+        resource_name = evaluation.resource_name
+        fields = evaluation._parse_resource_name(resource_name)
+        project = fields["project"]
+        location = fields["location"]
+        model_id = fields["model"]
+        evaluation_id = fields["evaluation"]
+    except AttributeError:
+        _LOGGER.warning("Unable to parse model evaluation metadata")
+        return
+
+    if "@" in model_id:
+        model_id, version_id = model_id.split("@")
+    else:
+        version_id = "default"
+
+    uri = (
+        "https://console.cloud.google.com/vertex-ai/models/locations/"
+        + f"{location}/models/{model_id}/versions/{version_id}/evaluations/"
+        + f"{evaluation_id}?project={project}"
+    )
+    display_link("View Model Evaluation", uri, "model_training")

tests/unit/aiplatform/test_model_evaluation.py

@@ -553,6 +553,9 @@ def mock_pipeline_service_get():
             make_pipeline_job(
                 gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
             ),
+            make_pipeline_job(
+                gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
+            ),
         ]
 
         yield mock_get_pipeline_job
@@ -797,7 +800,6 @@ def test_get_model_evaluation_metrics(self, mock_model_eval_get):
         assert eval_metrics == _TEST_MODEL_EVAL_METRICS
 
     def test_no_delete_model_evaluation_method(self, mock_model_eval_get):
-
         my_eval = aiplatform.ModelEvaluation(
             evaluation_name=_TEST_MODEL_EVAL_RESOURCE_NAME
         )
@@ -1028,6 +1030,7 @@ def test_model_evaluation_job_submit(
         mock_load_yaml_and_json,
         mock_model,
         get_model_mock,
+        mock_model_eval_get,
         mock_model_eval_job_get,
         mock_pipeline_service_get,
         mock_model_eval_job_create,
@@ -1128,6 +1131,7 @@ def test_model_evaluation_job_submit_with_experiment(
         mock_model,
         get_model_mock,
         get_experiment_mock,
+        mock_model_eval_get,
         mock_model_eval_job_get,
         mock_pipeline_service_get,
         mock_model_eval_job_create,
@@ -1308,7 +1312,6 @@ def test_model_evaluation_job_get_model_evaluation_with_failed_pipeline_run_rais
         mock_pipeline_bucket_exists,
         mock_request_urlopen,
     ):
-
         aiplatform.init(
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
@@ -1388,7 +1391,6 @@ def test_model_evaluation_job_get_model_evaluation_with_pending_pipeline_run_ret
     def test_get_template_url(
         self,
     ):
-
         template_url = model_evaluation_job._ModelEvaluationJob._get_template_url(
             model_type="automl_tabular",
             feature_attributions=False,