From 0e21567d7724eff4efe3002f321327dc70b2778b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 06:53:25 +0000 Subject: [PATCH 1/2] Make the worker pool fail within a second when a worker dies A worker that exits without reporting (an illegal memory access aborts the process, or an OOM kill) was only noticed when the result timeout ran out: 600 s for a job, 120 s for the start. Until then the parent waits on a queue nobody writes to, and the other workers wait in the job, typically in a collective, for their dead peer. A start that fails, for example on a bad backend, also waited out the whole ready timeout and left the workers it had spawned behind. The pool now checks the workers' liveness once a second while it waits (_get). A dead worker fails the run, or the start, within a second, naming its rank and exit code. The pool is marked broken and the other workers are killed with their stacks printed, the way a hung job or a failed rank already ends, and get_live_pool then hands out a new pool. A start that fails kills the workers it had spawned. ready_timeout is a constructor argument (120 s by default, unchanged), so a test can reach the start timeout without waiting for it. Testing: tests/gpu/torch/utils/test_distributed_worker_pool.py runs the pool with gloo on CPU. A rank that exits with code 13 while its peer waits fails the run in about 5 s and kills the peer; a start that fails raises and leaves no worker, both when every worker dies (a bad backend) and when the wait times out while they are still starting. The three new cases take 16 s together. Each fails without its fix, see the PR description. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 39 ++++++++++++++++--- .../utils/test_distributed_worker_pool.py | 38 ++++++++++++++++++ 2 files changed, 72 insertions(+), 5 deletions(-) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 02d5f6e8869..a757fa68c38 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -131,7 +131,9 @@ def _test_fn(rank, size): ... pool.shutdown() """ - def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardown): + def __init__( + self, world_size, backend="nccl", teardown_fn=default_worker_teardown, ready_timeout=120 + ): assert world_size > 0, "World size must be greater than 0" self.world_size = world_size # Set once a job was interrupted: the workers are then killed, see ``run``. @@ -160,9 +162,36 @@ def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardo p.start() self._processes.append(p) - for _ in range(world_size): - msg = self._result_queue.get(timeout=120) - assert msg == "ready", f"Worker failed to initialize: {msg}" + try: + for _ in range(world_size): + msg = self._get(ready_timeout) + assert msg == "ready", f"Worker failed to initialize: {msg}" + except BaseException: + self._abort() # do not leave workers behind when the start fails + raise + + def _get(self, timeout): + """Return the next message from the workers; fail within a second if a worker died. + + A worker that exits (an illegal memory access aborts the process, an OOM kill) would + otherwise only be noticed when ``timeout`` runs out. + """ + deadline = time.monotonic() + timeout + while True: + try: + return self._result_queue.get( + timeout=max(0.01, min(1.0, deadline - time.monotonic())) + ) + except queue.Empty: # noqa: PERF203 + dead = [ + (rank, p.exitcode) for rank, p in enumerate(self._processes) if not p.is_alive() + ] + if dead: + raise RuntimeError( + f"Worker process(es) died (rank, exit code): {dead}" + ) from None + if time.monotonic() >= deadline: + raise @staticmethod def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teardown_fn): @@ -232,7 +261,7 @@ def run(self, fn, *args, **kwargs): q.put((fn, args, kwargs)) for _ in range(self.world_size): timeout = self.job_timeout if deadline is None else deadline - time.monotonic() - status, rank, tb = self._result_queue.get(timeout=max(timeout, 0)) + status, rank, tb = self._get(max(timeout, 0)) if status == "error": errors.append(f"--- Rank {rank} ---\n{tb}") deadline = deadline or time.monotonic() + self.failure_grace diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py index 36dce35968f..ef48d8544ae 100644 --- a/tests/gpu/torch/utils/test_distributed_worker_pool.py +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -15,6 +15,8 @@ """A hung job must not leave a pool that stalls every later job.""" +import multiprocessing +import os import queue import time @@ -41,6 +43,12 @@ def _teardown_barrier(rank, size): dist.barrier() # like megatron_worker_teardown: blocks while a peer is still in the job +def _rank_0_dies_and_rank_1_waits_for_it(rank, size): + if rank == 0: + os._exit(13) # what an illegal memory access or an OOM kill does: no exception, no result + time.sleep(600) + + def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd): pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None) try: @@ -83,3 +91,33 @@ def test_a_failed_job_retires_the_pool_even_when_every_rank_replied(): assert pool.broken finally: pool.shutdown() + + +def test_a_dead_worker_fails_the_run_at_once_and_kills_its_peers(): + pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) + try: + pool.job_timeout = 30 # a run that misses the death fails here, not at the test's own cap + start = time.monotonic() + with pytest.raises(RuntimeError, match=r"died.*13"): + pool.run(_rank_0_dies_and_rank_1_waits_for_it) + + assert time.monotonic() - start < 20 # not the 600 s job timeout + assert pool.broken + assert not any(p.is_alive() for p in pool._processes) # the peer in the job is killed too + finally: + pool.shutdown() + + +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ({"backend": "not-a-backend"}, RuntimeError), # every worker dies in init_process_group + ({"ready_timeout": 0.01}, queue.Empty), # the workers are still starting when the wait ends + ], +) +def test_a_failed_start_raises_and_leaves_no_workers(kwargs, error): + before = {p.pid for p in multiprocessing.active_children()} + with pytest.raises(error): + DistributedWorkerPool(world_size=1, **{"backend": "gloo", "teardown_fn": None, **kwargs}) + time.sleep(1) + assert {p.pid for p in multiprocessing.active_children()} <= before From 85d0fb39b09e90e19f6813e9e6c22871f3a9790c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 19:09:46 +0000 Subject: [PATCH 2/2] Keep the failed ranks' errors when another worker dies Review feedback on this PR: if rank 0 reports an exception and rank 1 then exits without reporting (an illegal memory access, an OOM kill), `_get` raises "Worker process(es) died" and `run` aborted with only that message, dropping rank 0's traceback, which is often the root cause. The death error now carries the errors the other ranks already reported. The pool is still killed at once, and only once. Testing: tests/gpu/torch/utils/test_distributed_worker_pool.py runs the pool with gloo on CPU. The dead-worker test is now parametrized: a rank that dies while its peer waits (as before) and a rank that dies two seconds after rank 0 reported an exception, which must find both the exit code and "rank 0 failed" in the error. The new case takes 8 s and fails without the change (the error then carries only the exit code); the other tests of the module pass. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/distributed/utils.py | 7 ++++++- .../utils/test_distributed_worker_pool.py | 21 ++++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index a757fa68c38..d5bbebb5404 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -261,7 +261,12 @@ def run(self, fn, *args, **kwargs): q.put((fn, args, kwargs)) for _ in range(self.world_size): timeout = self.job_timeout if deadline is None else deadline - time.monotonic() - status, rank, tb = self._get(max(timeout, 0)) + try: + status, rank, tb = self._get(max(timeout, 0)) + except RuntimeError as e: # a worker died: keep what the other ranks reported + if errors: + raise RuntimeError(f"{e}\n" + "\n".join(errors)) from None + raise if status == "error": errors.append(f"--- Rank {rank} ---\n{tb}") deadline = deadline or time.monotonic() + self.failure_grace diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py index ef48d8544ae..3d59051c74a 100644 --- a/tests/gpu/torch/utils/test_distributed_worker_pool.py +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -43,6 +43,13 @@ def _teardown_barrier(rank, size): dist.barrier() # like megatron_worker_teardown: blocks while a peer is still in the job +def _rank_0_fails_and_rank_1_dies_later(rank, size): + if rank == 0: + raise ValueError("rank 0 failed") + time.sleep(2) # rank 0 has reported by then + os._exit(13) + + def _rank_0_dies_and_rank_1_waits_for_it(rank, size): if rank == 0: os._exit(13) # what an illegal memory access or an OOM kill does: no exception, no result @@ -93,13 +100,21 @@ def test_a_failed_job_retires_the_pool_even_when_every_rank_replied(): pool.shutdown() -def test_a_dead_worker_fails_the_run_at_once_and_kills_its_peers(): +@pytest.mark.parametrize( + ("job", "message"), + [ + (_rank_0_dies_and_rank_1_waits_for_it, r"died.*13"), + # The error a rank reported before the death is kept in the message + (_rank_0_fails_and_rank_1_dies_later, r"(?s)died.*13.*rank 0 failed"), + ], +) +def test_a_dead_worker_fails_the_run_at_once_and_kills_its_peers(job, message): pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) try: pool.job_timeout = 30 # a run that misses the death fails here, not at the test's own cap start = time.monotonic() - with pytest.raises(RuntimeError, match=r"died.*13"): - pool.run(_rank_0_dies_and_rank_1_waits_for_it) + with pytest.raises(RuntimeError, match=message): + pool.run(job) assert time.monotonic() - start < 20 # not the 600 s job timeout assert pool.broken