Skip to content

Commit 1a5c80f

Browse files
Add argument dashboard_check bool and checks
1 parent 841f3ed commit 1a5c80f

File tree

2 files changed

+16
-5
lines changed

2 files changed

+16
-5
lines changed

src/codeflare_sdk/cluster/cluster.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -258,7 +258,7 @@ def is_dashboard_ready(self) -> bool:
258258
else:
259259
return False
260260

261-
def wait_ready(self, timeout: Optional[int] = None):
261+
def wait_ready(self, timeout: Optional[int] = None, dashboard_check: bool = True):
262262
"""
263263
Waits for requested cluster to be ready, up to an optional timeout (s).
264264
Checks every five seconds.
@@ -276,19 +276,24 @@ def wait_ready(self, timeout: Optional[int] = None):
276276
)
277277
if not ready:
278278
if timeout and time >= timeout:
279-
raise TimeoutError(f"wait() timed out after waiting {timeout}s for cluster to be ready")
279+
raise TimeoutError(
280+
f"wait() timed out after waiting {timeout}s for cluster to be ready"
281+
)
280282
sleep(5)
281283
time += 5
282284
print("Requested cluster is up and running!")
283285

284-
while not dashboard_ready:
286+
while dashboard_check and not dashboard_ready:
285287
dashboard_ready = self.is_dashboard_ready()
286288
if not dashboard_ready:
287289
if timeout and time >= timeout:
288-
raise TimeoutError(f"wait() timed out after waiting {timeout}s for dashboard to be ready")
290+
raise TimeoutError(
291+
f"wait() timed out after waiting {timeout}s for dashboard to be ready"
292+
)
289293
sleep(5)
290294
time += 5
291-
print("Dashboard is ready!")
295+
if dashboard_ready:
296+
print("Dashboard is ready!")
292297

293298
def details(self, print_to_console: bool = True) -> RayCluster:
294299
cluster = _copy_to_ray(self)

tests/unit_test.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1790,6 +1790,12 @@ def test_wait_ready(mocker, capsys):
17901790
captured.out
17911791
== "Waiting for requested resources to be set up...\nRequested cluster is up and running!\nDashboard is ready!\n"
17921792
)
1793+
cf.wait_ready(dashboard_check=False)
1794+
captured = capsys.readouterr()
1795+
assert (
1796+
captured.out
1797+
== "Waiting for requested resources to be set up...\nRequested cluster is up and running!\n"
1798+
)
17931799

17941800

17951801
def test_jobdefinition_coverage():

0 commit comments

Comments
 (0)