Clean up _wait_for_scheduler_ready implementation (#21626)
This commit is contained in:
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user