diff --git a/src/codeflare_sdk/__init__.py b/src/codeflare_sdk/__init__.py
index 599171a9b..a1b5535c4 100644
--- a/src/codeflare_sdk/__init__.py
+++ b/src/codeflare_sdk/__init__.py
@@ -12,7 +12,7 @@
RayJobClient,
)
-from .cluster import view_clusters
+from .common.widgets import view_clusters
from .common import (
Authentication,
diff --git a/src/codeflare_sdk/cluster/__init__.py b/src/codeflare_sdk/common/widgets/__init__.py
similarity index 100%
rename from src/codeflare_sdk/cluster/__init__.py
rename to src/codeflare_sdk/common/widgets/__init__.py
diff --git a/src/codeflare_sdk/cluster/widgets.py b/src/codeflare_sdk/common/widgets/widgets.py
similarity index 51%
rename from src/codeflare_sdk/cluster/widgets.py
rename to src/codeflare_sdk/common/widgets/widgets.py
index d827c661f..8a13a4d4d 100644
--- a/src/codeflare_sdk/cluster/widgets.py
+++ b/src/codeflare_sdk/common/widgets/widgets.py
@@ -26,15 +26,205 @@
import ipywidgets as widgets
from IPython.display import display, HTML, Javascript
import pandas as pd
-from ..ray.cluster.config import ClusterConfiguration
-from ..ray.cluster.status import RayClusterStatus
-from ..common import _kube_api_error_handling
-from ..common.kubernetes_cluster.auth import (
+from ...ray.cluster.config import ClusterConfiguration
+from ...ray.cluster.status import RayClusterStatus
+from ..kubernetes_cluster import _kube_api_error_handling
+from ..kubernetes_cluster.auth import (
config_check,
get_api_client,
)
+class RayClusterManagerWidgets:
+ """
+ The RayClusterManagerWidgets class is responsible for initialising the ToggleButtons, Button, and Output widgets.
+ It also handles the user interactions and displays the cluster details.
+ Used when calling the view_clusters function.
+ """
+
+ def __init__(self, ray_clusters_df: pd.DataFrame, namespace: str = None):
+ from ...ray.cluster.cluster import get_current_namespace
+
+ # Data
+ self.ray_clusters_df = ray_clusters_df
+ self.namespace = get_current_namespace() if not namespace else namespace
+ self.raycluster_data_output = widgets.Output()
+ self.user_output = widgets.Output()
+ self.url_output = widgets.Output()
+
+ # Widgets
+ self.classification_widget = widgets.ToggleButtons(
+ options=ray_clusters_df["Name"].tolist(),
+ value=ray_clusters_df["Name"].tolist()[0],
+ description="Select an existing cluster:",
+ )
+ self.delete_button = widgets.Button(
+ description="Delete Cluster",
+ icon="trash",
+ tooltip="Delete the selected cluster",
+ )
+ self.list_jobs_button = widgets.Button(
+ description="View Jobs",
+ icon="suitcase",
+ tooltip="Open the Ray Job Dashboard",
+ )
+ self.ray_dashboard_button = widgets.Button(
+ description="Open Ray Dashboard",
+ icon="dashboard",
+ tooltip="Open the Ray Dashboard in a new tab",
+ layout=widgets.Layout(width="auto"),
+ )
+
+ # Set up interactions
+ self._initialize_callbacks()
+ self._trigger_initial_display()
+
+ def _initialize_callbacks(self):
+ """
+ Called upon RayClusterManagerWidgets initialisation.
+ Sets up event handlers and callbacks for UI interactions.
+ """
+ # Observe cluster selection
+ self.classification_widget.observe(
+ lambda selection_change: self._on_cluster_click(selection_change),
+ names="value",
+ )
+ # Set up button clicks
+ self.delete_button.on_click(lambda b: self._on_delete_button_click(b))
+ self.list_jobs_button.on_click(lambda b: self._on_list_jobs_button_click(b))
+ self.ray_dashboard_button.on_click(
+ lambda b: self._on_ray_dashboard_button_click(b)
+ )
+
+ def _trigger_initial_display(self):
+ """
+ Called upon RayClusterManagerWidgets initialisation.
+ Triggers an initial display update with the current cluster value.
+ """
+ # Trigger display with initial cluster value
+ initial_value = self.classification_widget.value
+ self._on_cluster_click({"new": initial_value})
+
+ def _on_cluster_click(self, selection_change):
+ """
+ _on_cluster_click handles the event when a cluster is selected from the toggle buttons, updating the output with cluster details.
+ """
+ new_value = selection_change["new"]
+ self.raycluster_data_output.clear_output()
+ ray_clusters_df = _fetch_cluster_data(self.namespace)
+ self.classification_widget.options = ray_clusters_df["Name"].tolist()
+ with self.raycluster_data_output:
+ display(
+ HTML(
+ ray_clusters_df[ray_clusters_df["Name"] == new_value][
+ [
+ "Name",
+ "Namespace",
+ "Num Workers",
+ "Head GPUs",
+ "Head CPU Req~Lim",
+ "Head Memory Req~Lim",
+ "Worker GPUs",
+ "Worker CPU Req~Lim",
+ "Worker Memory Req~Lim",
+ "status",
+ ]
+ ].to_html(escape=False, index=False, border=2)
+ )
+ )
+
+ def _on_delete_button_click(self, b):
+ """
+ _on_delete_button_click handles the event when the Delete Button is clicked, deleting the selected cluster.
+ """
+ cluster_name = self.classification_widget.value
+ namespace = self.ray_clusters_df[
+ self.ray_clusters_df["Name"] == self.classification_widget.value
+ ]["Namespace"].values[0]
+
+ _delete_cluster(cluster_name, namespace)
+
+ with self.user_output:
+ self.user_output.clear_output()
+ print(
+ f"Cluster {cluster_name} in the {namespace} namespace was deleted successfully."
+ )
+
+ # Refresh the dataframe
+ new_df = _fetch_cluster_data(namespace)
+ self.ray_clusters_df = new_df
+ if new_df.empty:
+ self.classification_widget.close()
+ self.delete_button.close()
+ self.list_jobs_button.close()
+ self.ray_dashboard_button.close()
+ with self.raycluster_data_output:
+ self.raycluster_data_output.clear_output()
+ print(f"No clusters found in the {namespace} namespace.")
+ else:
+ self.classification_widget.options = new_df["Name"].tolist()
+
+ def _on_list_jobs_button_click(self, b):
+ """
+ _on_list_jobs_button_click handles the event when the View Jobs button is clicked, opening the Ray Jobs Dashboard in a new tab
+ """
+ from codeflare_sdk import Cluster
+
+ cluster_name = self.classification_widget.value
+ namespace = self.ray_clusters_df[
+ self.ray_clusters_df["Name"] == self.classification_widget.value
+ ]["Namespace"].values[0]
+
+ # Suppress from Cluster Object initialisation widgets and outputs
+ with widgets.Output(), contextlib.redirect_stdout(
+ io.StringIO()
+ ), contextlib.redirect_stderr(io.StringIO()):
+ cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
+ dashboard_url = cluster.cluster_dashboard_uri()
+
+ with self.user_output:
+ self.user_output.clear_output()
+ print(
+ f"Opening Ray Jobs Dashboard for {cluster_name} cluster:\n{dashboard_url}/#/jobs"
+ )
+ with self.url_output:
+ display(Javascript(f'window.open("{dashboard_url}/#/jobs", "_blank");'))
+
+ def _on_ray_dashboard_button_click(self, b):
+ """
+ _on_ray_dashboard_button_click handles the event when the Open Ray Dashboard button is clicked, opening the Ray Dashboard in a new tab
+ """
+ from codeflare_sdk import Cluster
+
+ cluster_name = self.classification_widget.value
+ namespace = self.ray_clusters_df[
+ self.ray_clusters_df["Name"] == self.classification_widget.value
+ ]["Namespace"].values[0]
+
+ # Suppress from Cluster Object initialisation widgets and outputs
+ with widgets.Output(), contextlib.redirect_stdout(
+ io.StringIO()
+ ), contextlib.redirect_stderr(io.StringIO()):
+ cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
+ dashboard_url = cluster.cluster_dashboard_uri()
+
+ with self.user_output:
+ self.user_output.clear_output()
+ print(f"Opening Ray Dashboard for {cluster_name} cluster:\n{dashboard_url}")
+ with self.url_output:
+ display(Javascript(f'window.open("{dashboard_url}", "_blank");'))
+
+ def display_widgets(self):
+ display(widgets.VBox([self.classification_widget, self.raycluster_data_output]))
+ display(
+ widgets.HBox(
+ [self.delete_button, self.list_jobs_button, self.ray_dashboard_button]
+ ),
+ self.url_output,
+ self.user_output,
+ )
+
+
def cluster_up_down_buttons(
cluster: "codeflare_sdk.ray.cluster.cluster.Cluster",
) -> widgets.Button:
@@ -58,7 +248,7 @@ def cluster_up_down_buttons(
icon="trash",
)
- wait_ready_check = wait_ready_check_box()
+ 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
@@ -83,7 +273,7 @@ def on_down_button_clicked(b): # Handle the down button click event
delete_button.on_click(on_down_button_clicked)
-def wait_ready_check_box():
+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.
"""
@@ -117,221 +307,23 @@ def view_clusters(namespace: str = None):
)
return # Exit function if not in Jupyter Notebook
- from ..ray.cluster.cluster import get_current_namespace
+ from ...ray.cluster.cluster import get_current_namespace
if not namespace:
namespace = get_current_namespace()
- user_output = widgets.Output()
- raycluster_data_output = widgets.Output()
- url_output = widgets.Output()
-
ray_clusters_df = _fetch_cluster_data(namespace)
if ray_clusters_df.empty:
print(f"No clusters found in the {namespace} namespace.")
return
- classification_widget = widgets.ToggleButtons(
- options=ray_clusters_df["Name"].tolist(),
- value=ray_clusters_df["Name"].tolist()[0],
- description="Select an existing cluster:",
- )
- # Setting the initial value to trigger the event handler to display the cluster details.
- initial_value = classification_widget.value
- _on_cluster_click(
- {"new": initial_value}, raycluster_data_output, namespace, classification_widget
- )
- classification_widget.observe(
- lambda selection_change: _on_cluster_click(
- selection_change, raycluster_data_output, namespace, classification_widget
- ),
- names="value",
+ # Initialize the RayClusterManagerWidgets class
+ ray_cluster_manager = RayClusterManagerWidgets(
+ ray_clusters_df=ray_clusters_df, namespace=namespace
)
- # UI table buttons
- delete_button = widgets.Button(
- description="Delete Cluster",
- icon="trash",
- tooltip="Delete the selected cluster",
- )
- delete_button.on_click(
- lambda b: _on_delete_button_click(
- b,
- classification_widget,
- ray_clusters_df,
- raycluster_data_output,
- user_output,
- delete_button,
- list_jobs_button,
- ray_dashboard_button,
- )
- )
-
- list_jobs_button = widgets.Button(
- description="View Jobs", icon="suitcase", tooltip="Open the Ray Job Dashboard"
- )
- list_jobs_button.on_click(
- lambda b: _on_list_jobs_button_click(
- b, classification_widget, ray_clusters_df, user_output, url_output
- )
- )
-
- ray_dashboard_button = widgets.Button(
- description="Open Ray Dashboard",
- icon="dashboard",
- tooltip="Open the Ray Dashboard in a new tab",
- layout=widgets.Layout(width="auto"),
- )
- ray_dashboard_button.on_click(
- lambda b: _on_ray_dashboard_button_click(
- b, classification_widget, ray_clusters_df, user_output, url_output
- )
- )
-
- display(widgets.VBox([classification_widget, raycluster_data_output]))
- display(
- widgets.HBox([delete_button, list_jobs_button, ray_dashboard_button]),
- url_output,
- user_output,
- )
-
-
-def _on_cluster_click(
- selection_change,
- raycluster_data_output: widgets.Output,
- namespace: str,
- classification_widget: widgets.ToggleButtons,
-):
- """
- _on_cluster_click handles the event when a cluster is selected from the toggle buttons, updating the output with cluster details.
- """
- new_value = selection_change["new"]
- raycluster_data_output.clear_output()
- ray_clusters_df = _fetch_cluster_data(namespace)
- classification_widget.options = ray_clusters_df["Name"].tolist()
- with raycluster_data_output:
- display(
- HTML(
- ray_clusters_df[ray_clusters_df["Name"] == new_value][
- [
- "Name",
- "Namespace",
- "Num Workers",
- "Head GPUs",
- "Head CPU Req~Lim",
- "Head Memory Req~Lim",
- "Worker GPUs",
- "Worker CPU Req~Lim",
- "Worker Memory Req~Lim",
- "status",
- ]
- ].to_html(escape=False, index=False, border=2)
- )
- )
-
-
-def _on_delete_button_click(
- b,
- classification_widget: widgets.ToggleButtons,
- ray_clusters_df: pd.DataFrame,
- raycluster_data_output: widgets.Output,
- user_output: widgets.Output,
- delete_button: widgets.Button,
- list_jobs_button: widgets.Button,
- ray_dashboard_button: widgets.Button,
-):
- """
- _on_delete_button_click handles the event when the Delete Button is clicked, deleting the selected cluster.
- """
- cluster_name = classification_widget.value
- namespace = ray_clusters_df[ray_clusters_df["Name"] == classification_widget.value][
- "Namespace"
- ].values[0]
-
- _delete_cluster(cluster_name, namespace)
-
- with user_output:
- user_output.clear_output()
- print(
- f"Cluster {cluster_name} in the {namespace} namespace was deleted successfully."
- )
-
- # Refresh the dataframe
- new_df = _fetch_cluster_data(namespace)
- if new_df.empty:
- classification_widget.close()
- delete_button.close()
- list_jobs_button.close()
- ray_dashboard_button.close()
- with raycluster_data_output:
- raycluster_data_output.clear_output()
- print(f"No clusters found in the {namespace} namespace.")
- else:
- classification_widget.options = new_df["Name"].tolist()
-
-
-def _on_ray_dashboard_button_click(
- b,
- classification_widget: widgets.ToggleButtons,
- ray_clusters_df: pd.DataFrame,
- user_output: widgets.Output,
- url_output: widgets.Output,
-):
- """
- _on_ray_dashboard_button_click handles the event when the Open Ray Dashboard button is clicked, opening the Ray Dashboard in a new tab
- """
- from codeflare_sdk.ray.cluster import Cluster
-
- cluster_name = classification_widget.value
- namespace = ray_clusters_df[ray_clusters_df["Name"] == classification_widget.value][
- "Namespace"
- ].values[0]
-
- # Suppress from Cluster Object initialisation widgets and outputs
- with widgets.Output(), contextlib.redirect_stdout(
- io.StringIO()
- ), contextlib.redirect_stderr(io.StringIO()):
- cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
- dashboard_url = cluster.cluster_dashboard_uri()
-
- with user_output:
- user_output.clear_output()
- print(f"Opening Ray Dashboard for {cluster_name} cluster:\n{dashboard_url}")
- with url_output:
- display(Javascript(f'window.open("{dashboard_url}", "_blank");'))
-
-
-def _on_list_jobs_button_click(
- b,
- classification_widget: widgets.ToggleButtons,
- ray_clusters_df: pd.DataFrame,
- user_output: widgets.Output,
- url_output: widgets.Output,
-):
- """
- _on_list_jobs_button_click handles the event when the View Jobs button is clicked, opening the Ray Jobs Dashboard in a new tab
- """
- from codeflare_sdk.ray.cluster import Cluster
-
- cluster_name = classification_widget.value
- namespace = ray_clusters_df[ray_clusters_df["Name"] == classification_widget.value][
- "Namespace"
- ].values[0]
-
- # Suppress from Cluster Object initialisation widgets and outputs
- with widgets.Output(), contextlib.redirect_stdout(
- io.StringIO()
- ), contextlib.redirect_stderr(io.StringIO()):
- cluster = Cluster(ClusterConfiguration(cluster_name, namespace))
- dashboard_url = cluster.cluster_dashboard_uri()
-
- with user_output:
- user_output.clear_output()
- print(
- f"Opening Ray Jobs Dashboard for {cluster_name} cluster:\n{dashboard_url}/#/jobs"
- )
- with url_output:
- display(Javascript(f'window.open("{dashboard_url}/#/jobs", "_blank");'))
+ # Display the UI components
+ ray_cluster_manager.display_widgets()
def _delete_cluster(
@@ -344,7 +336,7 @@ def _delete_cluster(
_delete_cluster function deletes the cluster with the given name and namespace.
It optionally waits for the cluster to be deleted.
"""
- from ..ray.cluster.cluster import _check_aw_exists
+ from ...ray.cluster.cluster import _check_aw_exists
try:
config_check()
@@ -402,7 +394,7 @@ def _fetch_cluster_data(namespace):
"""
_fetch_cluster_data function fetches all clusters and their spec in a given namespace and returns a DataFrame.
"""
- from ..ray.cluster.cluster import list_all_clusters
+ from ...ray.cluster.cluster import list_all_clusters
rayclusters = list_all_clusters(namespace, False)
if not rayclusters:
@@ -411,15 +403,19 @@ def _fetch_cluster_data(namespace):
namespaces = [item.namespace for item in rayclusters]
num_workers = [item.num_workers for item in rayclusters]
head_extended_resources = [
- f"{list(item.head_extended_resources.keys())[0]}: {list(item.head_extended_resources.values())[0]}"
- if item.head_extended_resources
- else "0"
+ (
+ f"{list(item.head_extended_resources.keys())[0]}: {list(item.head_extended_resources.values())[0]}"
+ if item.head_extended_resources
+ else "0"
+ )
for item in rayclusters
]
worker_extended_resources = [
- f"{list(item.worker_extended_resources.keys())[0]}: {list(item.worker_extended_resources.values())[0]}"
- if item.worker_extended_resources
- else "0"
+ (
+ f"{list(item.worker_extended_resources.keys())[0]}: {list(item.worker_extended_resources.values())[0]}"
+ if item.worker_extended_resources
+ else "0"
+ )
for item in rayclusters
]
head_cpu_requests = [
diff --git a/src/codeflare_sdk/ray/cluster/cluster.py b/src/codeflare_sdk/ray/cluster/cluster.py
index 4d8201f78..da87639c5 100644
--- a/src/codeflare_sdk/ray/cluster/cluster.py
+++ b/src/codeflare_sdk/ray/cluster/cluster.py
@@ -1,4 +1,4 @@
-# Copyright 2022 IBM, Red Hat
+# 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.
@@ -45,7 +45,7 @@
AppWrapper,
AppWrapperStatus,
)
-from ...cluster.widgets import (
+from ...common.widgets.widgets import (
cluster_up_down_buttons,
is_notebook,
)
diff --git a/tests/unit_test.py b/tests/unit_test.py
index 74da56b77..1f11643bd 100644
--- a/tests/unit_test.py
+++ b/tests/unit_test.py
@@ -13,85 +13,73 @@
# limitations under the License.
-from pathlib import Path
-import sys
import filecmp
import os
import re
+import sys
import uuid
-
-from codeflare_sdk.ray.cluster import cluster
+from pathlib import Path
parent = Path(__file__).resolve().parents[1]
aw_dir = os.path.expanduser("~/.codeflare/resources/")
sys.path.append(str(parent) + "/src")
-from kubernetes import client, config, dynamic
+from unittest.mock import MagicMock, patch
+
+import openshift
+import pandas as pd
+import pytest
+import ray
+import yaml
+from kubernetes import client, config
+from pytest_mock import MockerFixture
+from ray.job_submission import JobSubmissionClient
+
+import codeflare_sdk.common.widgets.widgets as cf_widgets
+from codeflare_sdk.common.kubernetes_cluster import (
+ Authentication,
+ KubeConfigFileAuthentication,
+ TokenAuthentication,
+ config_check,
+)
+from codeflare_sdk.common.utils.generate_cert import (
+ export_env,
+ generate_ca_cert,
+ generate_tls_cert,
+)
from codeflare_sdk.ray.appwrapper.awload import AWManager
+from codeflare_sdk.ray.appwrapper.status import AppWrapper, AppWrapperStatus
+from codeflare_sdk.ray.client.ray_jobs import RayJobClient
from codeflare_sdk.ray.cluster.cluster import (
Cluster,
ClusterConfiguration,
+ _app_wrapper_status,
+ _copy_to_ray,
_map_to_ray_cluster,
+ _ray_cluster_status,
+ get_cluster,
list_all_clusters,
list_all_queued,
- _copy_to_ray,
- get_cluster,
- _app_wrapper_status,
- _ray_cluster_status,
-)
-from codeflare_sdk.common.kubernetes_cluster import (
- TokenAuthentication,
- Authentication,
- KubeConfigFileAuthentication,
- config_check,
)
+from codeflare_sdk.ray.cluster.generate_yaml import gen_names, is_openshift_cluster
from codeflare_sdk.ray.cluster.pretty_print import (
- print_no_resources_found,
print_app_wrappers_status,
print_cluster_status,
print_clusters,
-)
-from codeflare_sdk.ray.appwrapper.status import (
- AppWrapper,
- AppWrapperStatus,
+ print_no_resources_found,
)
from codeflare_sdk.ray.cluster.status import (
+ CodeFlareClusterStatus,
RayCluster,
RayClusterStatus,
- CodeFlareClusterStatus,
)
-from codeflare_sdk.common.utils.generate_cert import (
- generate_ca_cert,
- generate_tls_cert,
- export_env,
-)
-
from tests.unit_test_support import (
- createClusterWithConfig,
createClusterConfig,
+ createClusterWithConfig,
createClusterWrongType,
get_package_and_version,
)
-import codeflare_sdk.common.kubernetes_cluster.kube_api_helpers
-from codeflare_sdk.ray.cluster.generate_yaml import (
- gen_names,
- is_openshift_cluster,
-)
-
-import codeflare_sdk.cluster.widgets as cf_widgets
-import pandas as pd
-
-import openshift
-from openshift.selector import Selector
-import ray
-import pytest
-import yaml
-from unittest.mock import MagicMock, patch
-from pytest_mock import MockerFixture
-from ray.job_submission import JobSubmissionClient
-from codeflare_sdk.ray.client.ray_jobs import RayJobClient
-
# For mocking openshift client results
fake_res = openshift.Result("fake")
@@ -155,7 +143,7 @@ def test_token_auth_creation():
assert token_auth.skip_tls == False
assert token_auth.ca_cert_path == f"{parent}/tests/auth-test.crt"
- except Exception as e:
+ except Exception:
assert 0 == 1
@@ -203,7 +191,7 @@ def test_config_check_no_config_file(mocker):
mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_path", None)
mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.api_client", None)
- with pytest.raises(PermissionError) as e:
+ with pytest.raises(PermissionError):
config_check()
@@ -281,7 +269,7 @@ def test_config_creation():
def test_config_creation_wrong_type():
with pytest.raises(TypeError):
- config = createClusterWrongType()
+ createClusterWrongType()
def test_cluster_creation(mocker):
@@ -889,7 +877,7 @@ def test_ray_job_wrapping(mocker):
def test_print_no_resources(capsys):
try:
print_no_resources_found()
- except:
+ except Exception:
assert 1 == 0
captured = capsys.readouterr()
assert captured.out == (
@@ -902,7 +890,7 @@ def test_print_no_resources(capsys):
def test_print_no_cluster(capsys):
try:
print_cluster_status(None)
- except:
+ except Exception:
assert 1 == 0
captured = capsys.readouterr()
assert captured.out == (
@@ -923,7 +911,7 @@ def test_print_appwrappers(capsys):
)
try:
print_app_wrappers_status([aw1, aw2])
- except:
+ except Exception:
assert 1 == 0
captured = capsys.readouterr()
assert captured.out == (
@@ -996,7 +984,7 @@ def test_ray_details(mocker, capsys):
print_clusters([ray1, ray2])
print_cluster_status(ray1)
print_cluster_status(ray2)
- except:
+ except Exception:
assert 0 == 1
captured = capsys.readouterr()
assert captured.out == (
@@ -2601,13 +2589,14 @@ def test_AWManager_submit_remove(mocker, capsys):
assert testaw.submitted == False
-from cryptography.x509 import load_pem_x509_certificate
import base64
+
from cryptography.hazmat.primitives.serialization import (
- load_pem_private_key,
Encoding,
PublicFormat,
+ load_pem_private_key,
)
+from cryptography.x509 import load_pem_x509_certificate
def test_generate_ca_cert():
@@ -2959,161 +2948,293 @@ def test_cluster_up_down_buttons(mocker):
@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
+ assert cf_widgets.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
+ assert cf_widgets.is_notebook() is True
def test_view_clusters(mocker, capsys):
- from kubernetes.client.rest import ApiException
-
- mocker.patch("codeflare_sdk.cluster.widgets.is_notebook", return_value=False)
+ # If is not a notebook environment, a warning should be raised
with pytest.warns(
UserWarning,
match="view_clusters can only be used in a Jupyter Notebook environment.",
):
- result = cf_widgets.view_clusters(namespace="default")
+ result = cf_widgets.view_clusters("default")
+
# Assert the function returns None when not in a notebook environment
assert result is None
- mocker.patch("codeflare_sdk.cluster.widgets.is_notebook", return_value=True)
+ # Prepare to run view_clusters when notebook environment is detected
+ mocker.patch("codeflare_sdk.common.widgets.widgets.is_notebook", return_value=True)
+ mock_get_current_namespace = mocker.patch(
+ "codeflare_sdk.ray.cluster.cluster.get_current_namespace",
+ return_value="default",
+ )
+ namespace = mock_get_current_namespace.return_value
+
+ # Assert the function returns None when no clusters are found
+ mock_fetch_cluster_data = mocker.patch(
+ "codeflare_sdk.common.widgets.widgets._fetch_cluster_data",
+ return_value=pd.DataFrame(),
+ )
+ result = cf_widgets.view_clusters()
+ captured = capsys.readouterr()
+ assert mock_fetch_cluster_data.return_value.empty
+ assert "No clusters found in the default namespace." in captured.out
+ assert result is None
+
+ # Prepare to run view_clusters with a test DataFrame
+ mock_fetch_cluster_data = mocker.patch(
+ "codeflare_sdk.common.widgets.widgets._fetch_cluster_data",
+ return_value=pd.DataFrame(
+ {
+ "Name": ["test-cluster"],
+ "Namespace": ["default"],
+ "Num Workers": ["1"],
+ "Head GPUs": ["0"],
+ "Worker GPUs": ["0"],
+ "Head CPU Req~Lim": ["1~1"],
+ "Head Memory Req~Lim": ["1Gi~1Gi"],
+ "Worker CPU Req~Lim": ["1~1"],
+ "Worker Memory Req~Lim": ["1Gi~1Gi"],
+ "status": ['Ready ✓'],
+ }
+ ),
+ )
+ # Create a RayClusterManagerWidgets instance
+ ray_cluster_manager_instance = cf_widgets.RayClusterManagerWidgets(
+ ray_clusters_df=mock_fetch_cluster_data.return_value, namespace=namespace
+ )
+ # Patch the constructor of RayClusterManagerWidgets to return our initialized instance
+ mock_constructor = mocker.patch(
+ "codeflare_sdk.common.widgets.widgets.RayClusterManagerWidgets",
+ return_value=ray_cluster_manager_instance,
+ )
+
+ # Use a spy to track calls to display_widgets without replacing it
+ spy_display_widgets = mocker.spy(ray_cluster_manager_instance, "display_widgets")
+
+ cf_widgets.view_clusters()
+
+ mock_constructor.assert_called_once_with(
+ ray_clusters_df=mock_fetch_cluster_data.return_value, namespace=namespace
+ )
+
+ spy_display_widgets.assert_called_once()
+
+
+def test_delete_cluster(mocker, capsys):
+ name = "test-cluster"
+ namespace = "default"
- # Mock Kubernetes API responses
+ mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+
+ mock_ray_cluster = MagicMock()
mocker.patch(
- "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
- return_value={"items": []},
+ "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
+ side_effect=[
+ mock_ray_cluster,
+ client.ApiException(status=404),
+ client.ApiException(status=404),
+ mock_ray_cluster,
+ ],
)
+
+ # In this scenario, the RayCluster exists and the AppWrapper does not.
mocker.patch(
"codeflare_sdk.ray.cluster.cluster._check_aw_exists", return_value=False
)
+ mock_delete_rc = mocker.patch(
+ "kubernetes.client.CustomObjectsApi.delete_namespaced_custom_object"
+ )
+ cf_widgets._delete_cluster(name, namespace)
- # Return empty dataframe when no clusters are found
- mocker.patch("codeflare_sdk.ray.cluster.cluster.list_all_clusters", return_value=[])
+ mock_delete_rc.assert_called_once_with(
+ group="ray.io",
+ version="v1",
+ namespace=namespace,
+ plural="rayclusters",
+ name=name,
+ )
+
+ # In this scenario, the AppWrapper exists and the RayCluster does not
mocker.patch(
- "codeflare_sdk.ray.cluster.cluster.get_current_namespace",
- return_value="default",
+ "codeflare_sdk.ray.cluster.cluster._check_aw_exists", return_value=True
)
- df = cf_widgets._fetch_cluster_data(namespace="default")
- assert df.empty
+ mock_delete_aw = mocker.patch(
+ "kubernetes.client.CustomObjectsApi.delete_namespaced_custom_object"
+ )
+ cf_widgets._delete_cluster(name, namespace)
- cf_widgets.view_clusters()
- captured = capsys.readouterr()
- assert f"No clusters found in the default namespace." in captured.out
+ mock_delete_aw.assert_called_once_with(
+ group="workload.codeflare.dev",
+ version="v1beta2",
+ namespace=namespace,
+ plural="appwrappers",
+ name=name,
+ )
- # Assert the function returns None
- assert result is None
+ # In this scenario, the deletion of the resource times out.
+ with pytest.raises(
+ TimeoutError, match=f"Timeout waiting for {name} to be deleted."
+ ):
+ cf_widgets._delete_cluster(name, namespace, 1)
- test_df = pd.DataFrame(
+
+def test_ray_cluster_manager_widgets_init(mocker, capsys):
+ namespace = "default"
+ mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
+ mocker.patch(
+ "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+ return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+ )
+ test_ray_clusters_df = pd.DataFrame(
{
- "Name": ["test-cluster"],
- "Namespace": ["default"],
- "Num Workers": ["1"],
- "Head GPUs": ["0"],
- "Worker GPUs": ["0"],
- "Head CPU Req~Lim": ["1~1"],
- "Head Memory Req~Lim": ["1Gi~1Gi"],
- "Worker CPU Req~Lim": ["1~1"],
- "Worker Memory Req~Lim": ["1Gi~1Gi"],
- "status": ['Ready ✓'],
+ "Name": ["test-cluster-1", "test-cluster-2"],
+ "Namespace": [namespace, namespace],
+ "Num Workers": ["1", "2"],
+ "Head GPUs": ["0", "0"],
+ "Worker GPUs": ["0", "0"],
+ "Head CPU Req~Lim": ["1~1", "1~1"],
+ "Head Memory Req~Lim": ["1Gi~1Gi", "1Gi~1Gi"],
+ "Worker CPU Req~Lim": ["1~1", "1~1"],
+ "Worker Memory Req~Lim": ["1Gi~1Gi", "1Gi~1Gi"],
+ "status": [
+ 'Ready ✓',
+ 'Ready ✓',
+ ],
}
)
-
- # Mock the _fetch_cluster_data function to return a test DataFrame
+ mock_fetch_cluster_data = mocker.patch(
+ "codeflare_sdk.common.widgets.widgets._fetch_cluster_data",
+ return_value=test_ray_clusters_df,
+ )
mocker.patch(
- "codeflare_sdk.cluster.widgets._fetch_cluster_data", return_value=test_df
+ "codeflare_sdk.ray.cluster.cluster.get_current_namespace",
+ return_value=namespace,
+ )
+ mock_delete_cluster = mocker.patch(
+ "codeflare_sdk.common.widgets.widgets._delete_cluster"
)
- # Mock the Cluster class and related methods
- mocker.patch("codeflare_sdk.ray.cluster.Cluster")
- mocker.patch("codeflare_sdk.ray.cluster.ClusterConfiguration")
+ # # Mock ToggleButtons
+ mock_toggle_buttons = mocker.patch("ipywidgets.ToggleButtons")
+ mock_button = mocker.patch("ipywidgets.Button")
+ mock_output = mocker.patch("ipywidgets.Output")
- with patch("ipywidgets.ToggleButtons") as MockToggleButtons, patch(
- "ipywidgets.Button"
- ) as MockButton, patch("ipywidgets.Output") as MockOutput, patch(
- "ipywidgets.HBox"
- ), patch(
- "ipywidgets.VBox"
- ), patch(
- "IPython.display.display"
- ) as mock_display, patch(
- "IPython.display.HTML"
- ), patch(
- "codeflare_sdk.cluster.widgets.Javascript"
- ) as mock_javascript:
- # Create mock widget instances
- mock_toggle = MagicMock()
- mock_delete_button = MagicMock()
- mock_list_jobs_button = MagicMock()
- mock_ray_dashboard_button = MagicMock()
- mock_output = MagicMock()
-
- # Set the return values for the mocked widgets
- MockToggleButtons.return_value = mock_toggle
- MockButton.side_effect = [
- mock_delete_button,
- mock_list_jobs_button,
- mock_ray_dashboard_button,
- ]
- MockOutput.return_value = mock_output
+ # Initialize the RayClusterManagerWidgets instance
+ ray_cluster_manager_instance = cf_widgets.RayClusterManagerWidgets(
+ ray_clusters_df=test_ray_clusters_df, namespace=namespace
+ )
- # Call the function under test
- cf_widgets.view_clusters()
+ # Assertions for DataFrame and attributes
+ assert ray_cluster_manager_instance.ray_clusters_df.equals(
+ test_ray_clusters_df
+ ), "ray_clusters_df attribute does not match the input DataFrame"
+ assert (
+ ray_cluster_manager_instance.namespace == namespace
+ ), f"Expected namespace to be '{namespace}', but got '{ray_cluster_manager_instance.namespace}'"
+ assert (
+ ray_cluster_manager_instance.classification_widget.options
+ == test_ray_clusters_df["Name"].tolist()
+ ), "classification_widget options do not match the input DataFrame"
- # Simulate selecting a cluster
- mock_toggle.value = "test-cluster"
- selection_change = {"new": "test-cluster"}
- cf_widgets._on_cluster_click(
- selection_change, mock_output, "default", mock_toggle
- )
+ # Assertions for widgets
+ mock_toggle_buttons.assert_called_once_with(
+ options=test_ray_clusters_df["Name"].tolist(),
+ value=test_ray_clusters_df["Name"].tolist()[0],
+ description="Select an existing cluster:",
+ )
+ assert (
+ ray_cluster_manager_instance.classification_widget
+ == mock_toggle_buttons.return_value
+ ), "classification_widget is not set correctly"
+ assert (
+ ray_cluster_manager_instance.delete_button == mock_button.return_value
+ ), "delete_button is not set correctly"
+ assert (
+ ray_cluster_manager_instance.list_jobs_button == mock_button.return_value
+ ), "list_jobs_button is not set correctly"
+ assert (
+ ray_cluster_manager_instance.ray_dashboard_button == mock_button.return_value
+ ), "ray_dashboard_button is not set correctly"
+ assert (
+ ray_cluster_manager_instance.raycluster_data_output == mock_output.return_value
+ ), "raycluster_data_output is not set correctly"
+ assert (
+ ray_cluster_manager_instance.user_output == mock_output.return_value
+ ), "user_output is not set correctly"
+ assert (
+ ray_cluster_manager_instance.url_output == mock_output.return_value
+ ), "url_output is not set correctly"
- # Assert that the toggle options are set correctly
- mock_toggle.observe.assert_called()
+ ### Test button click events
+ mock_delete_button = MagicMock()
+ mock_list_jobs_button = MagicMock()
+ mock_ray_dashboard_button = MagicMock()
- # Simulate clicking the list jobs button
- cf_widgets._on_list_jobs_button_click(
- None, mock_toggle, test_df, mock_output, mock_output
- )
- mock_javascript.assert_called_once()
+ mock_javascript = mocker.patch("codeflare_sdk.common.widgets.widgets.Javascript")
+ ray_cluster_manager_instance.url_output = MagicMock()
- # Simulate clicking the Ray dashboard button
- cf_widgets._on_ray_dashboard_button_click(
- None, mock_toggle, test_df, mock_output, mock_output
- )
- mock_javascript.call_count = 2
+ mock_dashboard_uri = mocker.patch(
+ "codeflare_sdk.ray.cluster.cluster.Cluster.cluster_dashboard_uri",
+ return_value="https://ray-dashboard-test-cluster-1-ns.apps.cluster.awsroute.org",
+ )
- mocker.patch(
- "kubernetes.client.CustomObjectsApi.delete_namespaced_custom_object",
- )
- mock_response = mocker.MagicMock()
- mock_response.status = 404
- mock_exception = ApiException(http_resp=mock_response)
- mocker.patch(
- "kubernetes.client.CustomObjectsApi.get_namespaced_custom_object",
- side_effect=mock_exception,
- )
+ # Simulate clicking the list jobs button
+ ray_cluster_manager_instance.classification_widget.value = "test-cluster-1"
+ ray_cluster_manager_instance._on_list_jobs_button_click(mock_list_jobs_button)
- # Simulate clicking the delete button
- cf_widgets._on_delete_button_click(
- None,
- mock_toggle,
- test_df,
- mock_output,
- mock_output,
- mock_delete_button,
- mock_list_jobs_button,
- mock_ray_dashboard_button,
- )
- MockButton.call_count = 3
+ captured = capsys.readouterr()
+ assert (
+ f"Opening Ray Jobs Dashboard for test-cluster-1 cluster:\n{mock_dashboard_uri.return_value}/#/jobs"
+ in captured.out
+ )
+ mock_javascript.assert_called_with(
+ f'window.open("{mock_dashboard_uri.return_value}/#/jobs", "_blank");'
+ )
+
+ # Simulate clicking the Ray dashboard button
+ ray_cluster_manager_instance.classification_widget.value = "test-cluster-1"
+ ray_cluster_manager_instance._on_ray_dashboard_button_click(
+ mock_ray_dashboard_button
+ )
+
+ captured = capsys.readouterr()
+ assert (
+ f"Opening Ray Dashboard for test-cluster-1 cluster:\n{mock_dashboard_uri.return_value}"
+ in captured.out
+ )
+ mock_javascript.assert_called_with(
+ f'window.open("{mock_dashboard_uri.return_value}", "_blank");'
+ )
+
+ # Simulate clicking the delete button
+ ray_cluster_manager_instance.classification_widget.value = "test-cluster-1"
+ ray_cluster_manager_instance._on_delete_button_click(mock_delete_button)
+ mock_delete_cluster.assert_called_with("test-cluster-1", namespace)
+
+ mock_fetch_cluster_data.return_value = pd.DataFrame()
+ ray_cluster_manager_instance.classification_widget.value = "test-cluster-2"
+ ray_cluster_manager_instance._on_delete_button_click(mock_delete_button)
+ mock_delete_cluster.assert_called_with("test-cluster-2", namespace)
+
+ # Assert on deletion that the dataframe is empty
+ assert (
+ ray_cluster_manager_instance.ray_clusters_df.empty
+ ), "Expected DataFrame to be empty after deletion"
+
+ captured = capsys.readouterr()
+ assert (
+ f"Cluster test-cluster-1 in the {namespace} namespace was deleted successfully."
+ in captured.out
+ )
def test_fetch_cluster_data(mocker):