diff --git a/src/codeflare_sdk/cluster/cluster.py b/src/codeflare_sdk/cluster/cluster.py index 20aa94c90..f9bcc84f3 100644 --- a/src/codeflare_sdk/cluster/cluster.py +++ b/src/codeflare_sdk/cluster/cluster.py @@ -42,6 +42,10 @@ RayCluster, RayClusterStatus, ) +from .widgets import ( + cluster_up_down_buttons, + is_notebook, +) from kubernetes import client, config from kubernetes.utils import parse_quantity import yaml @@ -71,6 +75,8 @@ def __init__(self, config: ClusterConfiguration): self.app_wrapper_yaml = self.create_app_wrapper() self._job_submission_client = None self.app_wrapper_name = self.config.name + if is_notebook(): + cluster_up_down_buttons(self) @property def _client_headers(self): @@ -156,8 +162,12 @@ def up(self): plural="appwrappers", body=aw, ) + print(f"AppWrapper: '{self.config.name}' has successfully been created") else: self._component_resources_up(namespace, api_instance) + print( + f"Ray Cluster: '{self.config.name}' has successfully been created" + ) except Exception as e: # pragma: no cover return _kube_api_error_handling(e) @@ -198,8 +208,12 @@ def down(self): plural="appwrappers", name=self.app_wrapper_name, ) + print(f"AppWrapper: '{self.config.name}' has successfully been deleted") else: self._component_resources_down(namespace, api_instance) + print( + f"Ray Cluster: '{self.config.name}' has successfully been deleted" + ) except Exception as e: # pragma: no cover return _kube_api_error_handling(e) diff --git a/src/codeflare_sdk/cluster/widgets.py b/src/codeflare_sdk/cluster/widgets.py new file mode 100644 index 000000000..351640e04 --- /dev/null +++ b/src/codeflare_sdk/cluster/widgets.py @@ -0,0 +1,91 @@ +# Copyright 2024 IBM, Red Hat +# +# 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. + +""" +The widgets sub-module contains the ui widgets created using the ipywidgets package. +""" +import ipywidgets as widgets +from IPython.display import display +import os +import codeflare_sdk + + +def cluster_up_down_buttons(cluster: "codeflare_sdk.cluster.Cluster") -> widgets.Button: + """ + The cluster_up_down_buttons function returns two button widgets for a create and delete button. + The function uses the appwrapper bool to distinguish between resource type for the tool tip. + """ + resource = "Ray Cluster" + if cluster.config.appwrapper: + resource = "AppWrapper" + + up_button = widgets.Button( + description="Cluster Up", + tooltip=f"Create the {resource}", + icon="play", + ) + + delete_button = widgets.Button( + description="Cluster Down", + tooltip=f"Delete the {resource}", + icon="trash", + ) + + wait_ready_check = wait_ready_check_box() + output = widgets.Output() + + # Display the buttons in an HBox wrapped in a VBox which includes the wait_ready Checkbox + button_display = widgets.HBox([up_button, delete_button]) + display(widgets.VBox([button_display, wait_ready_check]), output) + + def on_up_button_clicked(b): # Handle the up button click event + with output: + output.clear_output() + cluster.up() + + # If the wait_ready Checkbox is clicked(value == True) trigger the wait_ready function + if wait_ready_check.value: + cluster.wait_ready() + + def on_down_button_clicked(b): # Handle the down button click event + with output: + output.clear_output() + cluster.down() + + up_button.on_click(on_up_button_clicked) + delete_button.on_click(on_down_button_clicked) + + +def wait_ready_check_box(): + """ + The wait_ready_check_box function will return a checkbox widget used for waiting for the resource to be in the state READY. + """ + wait_ready_check_box = widgets.Checkbox( + False, + description="Wait for Cluster?", + ) + return wait_ready_check_box + + +def is_notebook() -> bool: + """ + The is_notebook function checks if Jupyter Notebook environment variables exist in the given environment and return True/False based on that. + """ + if ( + "PYDEVD_IPYTHON_COMPATIBLE_DEBUGGING" in os.environ + or "JPY_SESSION_NAME" in os.environ + ): # If running Jupyter NBs in VsCode or RHOAI/ODH display UI buttons + return True + else: + return False diff --git a/tests/unit_test.py b/tests/unit_test.py index 8a51c6eb9..111f737c2 100644 --- a/tests/unit_test.py +++ b/tests/unit_test.py @@ -76,17 +76,21 @@ gen_names, is_openshift_cluster, ) +from codeflare_sdk.cluster.widgets import cluster_up_down_buttons import openshift from openshift.selector import Selector import ray import pytest import yaml -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from pytest_mock import MockerFixture from ray.job_submission import JobSubmissionClient from codeflare_sdk.job.ray_jobs import RayJobClient +import ipywidgets as widgets +from IPython.display import display + # For mocking openshift client results fake_res = openshift.Result("fake") @@ -2873,6 +2877,80 @@ def test_cluster_config_deprecation_conversion(mocker): assert config.worker_cpu_limits == 2 +""" + Ipywidgets tests +""" + + +@patch.dict( + "os.environ", {"JPY_SESSION_NAME": "example-test"} +) # Mock Jupyter environment variable +def test_cluster_up_down_buttons(mocker): + mocker.patch("kubernetes.client.ApisApi.get_api_versions") + mocker.patch( + "kubernetes.client.CustomObjectsApi.get_cluster_custom_object", + return_value={"spec": {"domain": "apps.cluster.awsroute.org"}}, + ) + mocker.patch( + "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object", + return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"), + ) + cluster = Cluster(createClusterConfig()) + + with patch("ipywidgets.Button") as MockButton, patch( + "ipywidgets.Checkbox" + ) as MockCheckbox, patch("ipywidgets.Output"), patch("ipywidgets.HBox"), patch( + "ipywidgets.VBox" + ), patch.object( + cluster, "up" + ) as mock_up, patch.object( + cluster, "down" + ) as mock_down, patch.object( + cluster, "wait_ready" + ) as mock_wait_ready: + # Create mock button & CheckBox instances + mock_up_button = MagicMock() + mock_down_button = MagicMock() + mock_wait_ready_check_box = MagicMock() + + # Ensure the mock Button class returns the mock button instances in sequence + MockCheckbox.side_effect = [mock_wait_ready_check_box] + MockButton.side_effect = [mock_up_button, mock_down_button] + + # Call the method under test + cluster_up_down_buttons(cluster) + + # Simulate checkbox being checked or unchecked + mock_wait_ready_check_box.value = True # Simulate checkbox being checked + + # Simulate the button clicks by calling the mock on_click handlers + mock_up_button.on_click.call_args[0][0](None) # Simulate clicking "Cluster Up" + mock_down_button.on_click.call_args[0][0]( + None + ) # Simulate clicking "Cluster Down" + + # Check if the `up` and `down` methods were called + mock_wait_ready.assert_called_once() + mock_up.assert_called_once() + mock_down.assert_called_once() + + +@patch.dict("os.environ", {}, clear=True) # Mock environment with no variables +def test_is_notebook_false(): + from codeflare_sdk.cluster.widgets import is_notebook + + assert is_notebook() is False + + +@patch.dict( + "os.environ", {"JPY_SESSION_NAME": "example-test"} +) # Mock Jupyter environment variable +def test_is_notebook_true(): + from codeflare_sdk.cluster.widgets import is_notebook + + assert is_notebook() is True + + # Make sure to always keep this function last def test_cleanup(): os.remove(f"{aw_dir}unit-test-no-kueue.yaml")