Skip to content

Commit

Permalink
Add refresh button to widgets UI
Browse files Browse the repository at this point in the history
  • Loading branch information
ChristianZaccaria committed Oct 25, 2024
1 parent 13658b8 commit 742cdeb
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 26 deletions.
10 changes: 10 additions & 0 deletions src/codeflare_sdk/common/widgets/test_widgets.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,9 @@ def test_ray_cluster_manager_widgets_init(mocker, capsys):
assert (
ray_cluster_manager_instance.ray_dashboard_button == mock_button.return_value
), "ray_dashboard_button is not set correctly"
assert (
ray_cluster_manager_instance.refresh_data_button == mock_button.return_value
), "refresh_data_button is not set correctly"
assert (
ray_cluster_manager_instance.raycluster_data_output == mock_output.return_value
), "raycluster_data_output is not set correctly"
Expand All @@ -310,6 +313,7 @@ def test_ray_cluster_manager_widgets_init(mocker, capsys):
mock_delete_button = MagicMock()
mock_list_jobs_button = MagicMock()
mock_ray_dashboard_button = MagicMock()
mock_refresh_dataframe_button = MagicMock()

mock_javascript = mocker.patch("codeflare_sdk.common.widgets.widgets.Javascript")
ray_cluster_manager_instance.url_output = MagicMock()
Expand All @@ -332,6 +336,12 @@ def test_ray_cluster_manager_widgets_init(mocker, capsys):
f'window.open("{mock_dashboard_uri.return_value}/#/jobs", "_blank");'
)

# Simulate clicking the refresh data button
ray_cluster_manager_instance._on_refresh_data_button_click(
mock_refresh_dataframe_button
)
mock_fetch_cluster_data.assert_called_with(namespace)

# 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(
Expand Down
72 changes: 46 additions & 26 deletions src/codeflare_sdk/common/widgets/widgets.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,12 @@ def __init__(self, ray_clusters_df: pd.DataFrame, namespace: str = None):
tooltip="Open the Ray Dashboard in a new tab",
layout=widgets.Layout(width="auto"),
)
self.refresh_data_button = widgets.Button(
description="Refresh Data",
icon="refresh",
tooltip="Refresh the list of Ray Clusters",
layout=widgets.Layout(width="auto", left="1em"),
)

# Set up interactions
self._initialize_callbacks()
Expand All @@ -95,6 +101,9 @@ def _initialize_callbacks(self):
self.ray_dashboard_button.on_click(
lambda b: self._on_ray_dashboard_button_click(b)
)
self.refresh_data_button.on_click(
lambda b: self._on_refresh_data_button_click(b)
)

def _trigger_initial_display(self):
"""
Expand Down Expand Up @@ -138,31 +147,17 @@ 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)
_delete_cluster(cluster_name, self.namespace)

with self.user_output:
self.user_output.clear_output()
print(
f"Cluster {cluster_name} in the {namespace} namespace was deleted successfully."
f"Cluster {cluster_name} in the {self.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()
self._refresh_dataframe()

def _on_list_jobs_button_click(self, b):
"""
Expand All @@ -171,15 +166,12 @@ def _on_list_jobs_button_click(self, b):
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))
cluster = Cluster(ClusterConfiguration(cluster_name, self.namespace))
dashboard_url = cluster.cluster_dashboard_uri()

with self.user_output:
Expand All @@ -197,15 +189,12 @@ def _on_ray_dashboard_button_click(self, b):
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))
cluster = Cluster(ClusterConfiguration(cluster_name, self.namespace))
dashboard_url = cluster.cluster_dashboard_uri()

with self.user_output:
Expand All @@ -214,11 +203,42 @@ def _on_ray_dashboard_button_click(self, b):
with self.url_output:
display(Javascript(f'window.open("{dashboard_url}", "_blank");'))

def _on_refresh_data_button_click(self, b):
"""
_on_refresh_button_click handles the event when the Refresh Data button is clicked, refreshing the list of Ray Clusters.
"""
self.refresh_data_button.disabled = True
self._refresh_dataframe()
self.refresh_data_button.disabled = False

def _refresh_dataframe(self):
"""
_refresh_data function refreshes the list of Ray Clusters.
"""
new_df = _fetch_cluster_data(self.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()
self.refresh_data_button.close()
with self.raycluster_data_output:
self.raycluster_data_output.clear_output()
print(f"No clusters found in the {self.namespace} namespace.")
else:
self.classification_widget.options = new_df["Name"].tolist()

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.delete_button,
self.list_jobs_button,
self.ray_dashboard_button,
self.refresh_data_button,
]
),
self.url_output,
self.user_output,
Expand Down

0 comments on commit 742cdeb

Please sign in to comment.