Clean up _wait_for_scheduler_ready implementation (#21626)

This commit is contained in:
Liangsheng Yin
2026-03-29 01:02:33 -07:00
committed by GitHub
parent c34593f951
commit fec9961a1f
+21 -28
View File
@@ -1225,19 +1225,25 @@ def _set_gc(server_args: ServerArgs):
gc.set_threshold(*gc_threshold) gc.set_threshold(*gc_threshold)
def _scheduler_died_error(rank: int, proc) -> RuntimeError:
"""Build a descriptive error for a scheduler process that died during init."""
proc.join(timeout=10)
return RuntimeError(
f"Rank {rank} scheduler died during initialization "
f"(exit code: {proc.exitcode}). "
f"If exit code is -9 (SIGKILL), a common cause is the OS OOM killer. "
f"Run `dmesg -T | grep -i oom` to check."
)
def _wait_for_scheduler_ready( def _wait_for_scheduler_ready(
scheduler_pipe_readers: List, scheduler_pipe_readers: List,
scheduler_procs: List, scheduler_procs: List,
) -> List[Dict]: ) -> List[Dict]:
"""Wait for the model to finish loading and return scheduler infos. """Wait for the model to finish loading and return scheduler infos.
Uses polling to detect child process death quickly, rather than blocking Uses poll() with timeout instead of blocking recv(), so that child process
indefinitely on pipe recv(). This prevents the launch from hanging when death (e.g. OOM SIGKILL) is detected promptly instead of hanging forever.
a child process is killed (e.g. by OOM killer via SIGKILL) before it can
send any data through the pipe.
On each poll timeout, checks ALL processes (not just the current one) so that
a death in any rank is detected promptly regardless of iteration order.
""" """
scheduler_infos = [] scheduler_infos = []
for i in range(len(scheduler_pipe_readers)): for i in range(len(scheduler_pipe_readers)):
@@ -1246,32 +1252,19 @@ def _wait_for_scheduler_ready(
try: try:
data = scheduler_pipe_readers[i].recv() data = scheduler_pipe_readers[i].recv()
except EOFError: except EOFError:
scheduler_procs[i].join(timeout=10) raise _scheduler_died_error(i, scheduler_procs[i])
if data["status"] != "ready":
raise RuntimeError( raise RuntimeError(
f"Rank {i} scheduler died during initialization " "Initialization failed. Please see the error messages above."
f"(exit code: {scheduler_procs[i].exitcode}). "
f"If exit code is -9 (SIGKILL), a common cause is the OS OOM killer. "
f"Run `dmesg -T | grep -i oom` to check."
) )
scheduler_infos.append(data) scheduler_infos.append(data)
break break
else:
# Check ALL processes, not just the current one
for j in range(len(scheduler_procs)):
if not scheduler_procs[j].is_alive():
scheduler_procs[j].join(timeout=10)
raise RuntimeError(
f"Rank {j} scheduler died during initialization "
f"(exit code: {scheduler_procs[j].exitcode}). "
f"If exit code is -9 (SIGKILL), a common cause is the OS OOM killer. "
f"Run `dmesg -T | grep -i oom` to check."
)
for data in scheduler_infos: # Poll timed out — check all processes for early death
if data["status"] != "ready": for j in range(len(scheduler_procs)):
raise RuntimeError( if not scheduler_procs[j].is_alive():
"Initialization failed. Please see the error messages above." raise _scheduler_died_error(j, scheduler_procs[j])
)
return scheduler_infos return scheduler_infos