Skip to content

Commit 3735284

Browse files
Add argument dashboard_check bool and checks
1 parent 47c0093 commit 3735284

File tree

2 files changed

+16
-5
lines changed

2 files changed

+16
-5
lines changed

src/codeflare_sdk/cluster/cluster.py

+10-5
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,7 @@ def is_dashboard_ready(self) -> bool:
264264
else:
265265
return False
266266

267-
def wait_ready(self, timeout: Optional[int] = None):
267+
def wait_ready(self, timeout: Optional[int] = None, dashboard_check: bool = True):
268268
"""
269269
Waits for requested cluster to be ready, up to an optional timeout (s).
270270
Checks every five seconds.
@@ -282,19 +282,24 @@ def wait_ready(self, timeout: Optional[int] = None):
282282
)
283283
if not ready:
284284
if timeout and time >= timeout:
285-
raise TimeoutError(f"wait() timed out after waiting {timeout}s for cluster to be ready")
285+
raise TimeoutError(
286+
f"wait() timed out after waiting {timeout}s for cluster to be ready"
287+
)
286288
sleep(5)
287289
time += 5
288290
print("Requested cluster is up and running!")
289291

290-
while not dashboard_ready:
292+
while dashboard_check and not dashboard_ready:
291293
dashboard_ready = self.is_dashboard_ready()
292294
if not dashboard_ready:
293295
if timeout and time >= timeout:
294-
raise TimeoutError(f"wait() timed out after waiting {timeout}s for dashboard to be ready")
296+
raise TimeoutError(
297+
f"wait() timed out after waiting {timeout}s for dashboard to be ready"
298+
)
295299
sleep(5)
296300
time += 5
297-
print("Dashboard is ready!")
301+
if dashboard_ready:
302+
print("Dashboard is ready!")
298303

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

tests/unit_test.py

+6
Original file line numberDiff line numberDiff line change
@@ -1796,6 +1796,12 @@ def test_wait_ready(mocker, capsys):
17961796
captured.out
17971797
== "Waiting for requested resources to be set up...\nRequested cluster is up and running!\nDashboard is ready!\n"
17981798
)
1799+
cf.wait_ready(dashboard_check=False)
1800+
captured = capsys.readouterr()
1801+
assert (
1802+
captured.out
1803+
== "Waiting for requested resources to be set up...\nRequested cluster is up and running!\n"
1804+
)
17991805

18001806

18011807
def test_jobdefinition_coverage():

0 commit comments

Comments
 (0)