From e95e6d5ffcac87eb2def855e0b51dfe22c110d9e Mon Sep 17 00:00:00 2001 From: ChristianZaccaria Date: Fri, 25 Oct 2024 13:13:53 +0100 Subject: [PATCH] Add refresh button to widgets UI --- .../common/widgets/test_widgets.py | 10 +++ src/codeflare_sdk/common/widgets/widgets.py | 70 ++++++++++++------- 2 files changed, 54 insertions(+), 26 deletions(-) diff --git a/src/codeflare_sdk/common/widgets/test_widgets.py b/src/codeflare_sdk/common/widgets/test_widgets.py index e01b91933..12c238544 100644 --- a/src/codeflare_sdk/common/widgets/test_widgets.py +++ b/src/codeflare_sdk/common/widgets/test_widgets.py @@ -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" @@ -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() @@ -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( diff --git a/src/codeflare_sdk/common/widgets/widgets.py b/src/codeflare_sdk/common/widgets/widgets.py index 8a13a4d4d..5b57a3699 100644 --- a/src/codeflare_sdk/common/widgets/widgets.py +++ b/src/codeflare_sdk/common/widgets/widgets.py @@ -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() @@ -95,6 +101,7 @@ 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_button_click(b)) def _trigger_initial_display(self): """ @@ -138,31 +145,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): """ @@ -171,15 +164,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: @@ -197,15 +187,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: @@ -214,11 +201,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,