feat: Add display experiment button for Ipython environments

matthew29tang · copybara-github · commit 0b33bdd873ab · 2024-03-04T11:44:48.000-08:00
PiperOrigin-RevId: 612526515
diff --git a/google/cloud/aiplatform/metadata/metadata.py b/google/cloud/aiplatform/metadata/metadata.py
@@ -38,6 +38,7 @@
 )
 from google.cloud.aiplatform.tensorboard import tensorboard_resource
 from google.cloud.aiplatform.utils import autologging_utils
+from google.cloud.aiplatform.utils import _ipython_utils
 
 from google.cloud.aiplatform_v1.types import execution as execution_v1
 
@@ -329,8 +330,32 @@ def set_experiment(
         if not current_backing_tb and backing_tb:
             experiment.assign_backing_tensorboard(tensorboard=backing_tb)
 
+        if _ipython_utils.is_ipython_available():
+            self._display_experiment_button(experiment)
+
         self._experiment = experiment
 
+    def _display_experiment_button(
+        self, experiment: experiment_resources.Experiment
+    ) -> None:
+        """Function to generate a link bound to the Vertex experiment"""
+        try:
+            project = experiment._metadata_context.project
+            location = experiment._metadata_context.location
+            experiment_name = experiment._metadata_context.name
+            if experiment_name is None or project is None or location is None:
+                return
+        except AttributeError:
+            _LOGGER.warning("Unable to fetch experiment metadata")
+            return
+
+        uri = (
+            "https://console.cloud.google.com/vertex-ai/experiments/locations/"
+            + f"{location}/experiments/{experiment_name}/"
+            + f"runs?project={project}"
+        )
+        _ipython_utils.display_link("View Experiment", uri, "science")
+
     def set_tensorboard(
         self,
         tensorboard: Union[
diff --git a/google/cloud/aiplatform/utils/_ipython_utils.py b/google/cloud/aiplatform/utils/_ipython_utils.py
@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import sys
+from uuid import uuid4
+from typing import Optional
+
+
+def _get_ipython_shell_name() -> str:
+    if "IPython" in sys.modules:
+        from IPython import get_ipython
+
+        return get_ipython().__class__.__name__
+    return ""
+
+
+def is_ipython_available() -> bool:
+    return _get_ipython_shell_name() != ""
+
+
+def _get_styles() -> None:
+    """Returns the HTML style markup to support custom buttons."""
+    return """
+    <link rel="stylesheet" href="https://fonts.googleapis.com/icon?family=Material+Icons">
+    <style>
+      .view-vertex-resource,
+      .view-vertex-resource:hover,
+      .view-vertex-resource:visited {
+        position: relative;
+        display: inline-flex;
+        flex-direction: row;
+        height: 32px;
+        padding: 0 12px;
+          margin: 4px 18px;
+        gap: 4px;
+        border-radius: 4px;
+
+        align-items: center;
+        justify-content: center;
+        background-color: rgb(255, 255, 255);
+        color: rgb(51, 103, 214);
+
+        font-family: Roboto,"Helvetica Neue",sans-serif;
+        font-size: 13px;
+        font-weight: 500;
+        text-transform: uppercase;
+        text-decoration: none !important;
+
+        transition: box-shadow 280ms cubic-bezier(0.4, 0, 0.2, 1) 0s;
+        box-shadow: 0px 3px 1px -2px rgba(0,0,0,0.2), 0px 2px 2px 0px rgba(0,0,0,0.14), 0px 1px 5px 0px rgba(0,0,0,0.12);
+      }
+      .view-vertex-resource:active {
+        box-shadow: 0px 5px 5px -3px rgba(0,0,0,0.2),0px 8px 10px 1px rgba(0,0,0,0.14),0px 3px 14px 2px rgba(0,0,0,0.12);
+      }
+      .view-vertex-resource:active .view-vertex-ripple::before {
+        position: absolute;
+        top: 0;
+        bottom: 0;
+        left: 0;
+        right: 0;
+        border-radius: 4px;
+        pointer-events: none;
+
+        content: '';
+        background-color: rgb(51, 103, 214);
+        opacity: 0.12;
+      }
+      .view-vertex-icon {
+        font-size: 18px;
+      }
+    </style>
+  """
+
+
+def display_link(text: str, url: str, icon: Optional[str] = "open_in_new") -> None:
+    """Creates and displays the link to open the Vertex resource
+
+    Args:
+        text: The text displayed on the clickable button.
+        url: The url that the button will lead to.
+          Only cloud console URIs are allowed.
+        icon: The icon name on the button (from material-icons library)
+
+    Returns:
+        Dict of custom properties with keys mapped to column names
+    """
+    CLOUD_UI_URL = "https://console.cloud.google.com"
+    if not url.startswith(CLOUD_UI_URL):
+        raise ValueError(f"Only urls starting with {CLOUD_UI_URL} are allowed.")
+
+    button_id = f"view-vertex-resource-{str(uuid4())}"
+
+    # Add the markup for the CSS and link component
+    html = f"""
+        {_get_styles()}
+        <a class="view-vertex-resource" id="{button_id}" href="#view-{button_id}">
+          <span class="material-icons view-vertex-icon">{icon}</span>
+          <span>{text}</span>
+        </a>
+        """
+
+    # Add the click handler for the link
+    html += f"""
+        <script>
+          (function () {{
+            const link = document.getElementById('{button_id}');
+            link.addEventListener('click', (e) => {{
+              if (window.google?.colab?.openUrl) {{
+                window.google.colab.openUrl('{url}');
+              }} else {{
+                window.open('{url}', '_blank');
+              }}
+              e.stopPropagation();
+              e.preventDefault();
+            }});
+          }})();
+        </script>
+    """
+
+    from IPython.core.display import display
+    from IPython.display import HTML
+
+    display(HTML(html))
diff --git a/testing/constraints-3.10.txt b/testing/constraints-3.10.txt
@@ -11,3 +11,4 @@ grpcio-testing==1.34.0
 mlflow==1.30.1 # Pinned to speed up installation
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.4.0 # Pinned until 2.9.3 is verified for Ray tests
+IPython # Added to test supernova rich html buttons
diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py
@@ -74,6 +74,7 @@
 from google.cloud.aiplatform.tensorboard import tensorboard_resource
 
 from google.cloud.aiplatform import utils
+from google.cloud.aiplatform.utils import _ipython_utils
 
 import constants as test_constants
 
@@ -187,6 +188,22 @@ def get_metadata_store_mock_raise_not_found_exception():
         yield get_metadata_store_mock
 
 
+@pytest.fixture
+def ipython_is_available_mock():
+    with patch.object(_ipython_utils, "is_ipython_available") as ipython_available_mock:
+        ipython_available_mock.return_value = True
+        yield ipython_available_mock
+
+
+@pytest.fixture
+def ipython_is_not_available_mock():
+    with patch.object(
+        _ipython_utils, "is_ipython_available"
+    ) as ipython_not_available_mock:
+        ipython_not_available_mock.return_value = False
+        yield ipython_not_available_mock
+
+
 @pytest.fixture
 def create_metadata_store_mock():
     with patch.object(
@@ -1110,6 +1127,38 @@ def setup_method(self):
     def teardown_method(self):
         initializer.global_pool.shutdown(wait=True)
 
+    @pytest.mark.usefixtures(
+        "get_metadata_store_mock",
+        "get_experiment_run_run_mock",
+    )
+    def test_init_experiment_with_ipython_environment(
+        self,
+        list_default_tensorboard_mock,
+        assign_backing_tensorboard_mock,
+        ipython_is_available_mock,
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            experiment=_TEST_EXPERIMENT,
+        )
+
+    @pytest.mark.usefixtures(
+        "get_metadata_store_mock",
+        "get_experiment_run_run_mock",
+    )
+    def test_init_experiment_with_no_ipython_environment(
+        self,
+        list_default_tensorboard_mock,
+        assign_backing_tensorboard_mock,
+        ipython_is_not_available_mock,
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            experiment=_TEST_EXPERIMENT,
+        )
+
     @pytest.mark.usefixtures("get_or_create_default_tb_none_mock")
     def test_init_experiment_with_existing_metadataStore_and_context(
         self, get_metadata_store_mock, get_experiment_run_run_mock
@@ -1496,7 +1545,6 @@ def test_start_run(
         create_experiment_run_context_mock,
         add_context_children_mock,
     ):
-
         aiplatform.init(
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
@@ -1527,7 +1575,6 @@ def test_start_run(
         "get_or_create_default_tb_none_mock",
     )
     def test_start_run_fails_when_run_name_too_long(self):
-
         aiplatform.init(
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
@@ -1846,7 +1893,6 @@ def test_start_execution_and_assign_artifact(
             display_name=_TEST_DISPLAY_NAME,
             metadata=_TEST_METADATA,
         ) as exc:
-
             exc.assign_input_artifacts([in_artifact])
             exc.assign_output_artifacts([out_artifact])
 
