diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 02d5f6e8869..d5bbebb5404 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,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._result_queue.get(timeout=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 36dce35968f..3d59051c74a 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,19 @@ 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 + 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 +98,41 @@ def test_a_failed_job_retires_the_pool_even_when_every_rank_replied(): assert pool.broken finally: pool.shutdown() + + +@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=message): + pool.run(job) + + 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