[OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies

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

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