diff --git a/src/codeflare_sdk/cluster/cluster.py b/src/codeflare_sdk/cluster/cluster.py index 4d8637667..cdf488911 100644 --- a/src/codeflare_sdk/cluster/cluster.py +++ b/src/codeflare_sdk/cluster/cluster.py @@ -38,6 +38,7 @@ from kubernetes import client, config import yaml import os +import requests class Cluster: @@ -250,6 +251,13 @@ def status( return status, ready + def is_dashboard_ready(self) -> bool: + response = requests.get(self.cluster_dashboard_uri(), timeout=5) + if response.status_code == 200: + return True + else: + return False + def wait_ready(self, timeout: Optional[int] = None): """ Waits for requested cluster to be ready, up to an optional timeout (s). @@ -257,20 +265,22 @@ def wait_ready(self, timeout: Optional[int] = None): """ print("Waiting for requested resources to be set up...") ready = False + dashboard_ready = False status = None time = 0 - while not ready: + while not ready or not dashboard_ready: status, ready = self.status(print_to_console=False) + dashboard_ready = self.is_dashboard_ready() if status == CodeFlareClusterStatus.UNKNOWN: print( "WARNING: Current cluster status is unknown, have you run cluster.up yet?" ) - if not ready: + if not ready or not dashboard_ready: if timeout and time >= timeout: raise TimeoutError(f"wait() timed out after waiting {timeout}s") sleep(5) time += 5 - print("Requested cluster up and running!") + print("Requested cluster and dashboard are up and running!") def details(self, print_to_console: bool = True) -> RayCluster: cluster = _copy_to_ray(self) diff --git a/tests/unit_test.py b/tests/unit_test.py index 0a2fb555b..99f5a152b 100644 --- a/tests/unit_test.py +++ b/tests/unit_test.py @@ -1753,6 +1753,21 @@ def test_wait_ready(mocker, capsys): mocker.patch("kubernetes.config.load_kube_config", return_value="ignore") mocker.patch("codeflare_sdk.cluster.cluster._app_wrapper_status", return_value=None) mocker.patch("codeflare_sdk.cluster.cluster._ray_cluster_status", return_value=None) + mocker.patch.object( + client.CustomObjectsApi, + "list_namespaced_custom_object", + return_value={ + "items": [ + { + "metadata": {"name": "ray-dashboard-test"}, + "spec": {"host": "mocked-host"}, + } + ] + }, + ) + mock_response = mocker.Mock() + mock_response.status_code = 200 + mocker.patch("requests.get", return_value=mock_response) cf = Cluster(ClusterConfiguration(name="test", namespace="ns")) try: cf.wait_ready(timeout=5) @@ -1773,7 +1788,7 @@ def test_wait_ready(mocker, capsys): captured = capsys.readouterr() assert ( captured.out - == "Waiting for requested resources to be set up...\nRequested cluster up and running!\n" + == "Waiting for requested resources to be set up...\nRequested cluster and dashboard are up and running!\n" )