Fix DP launcher hangs after worker failures

Monitor process sentinels and reap sibling ranks on worker failures,
interruption, or partial startup failure. Bound graceful termination and
report global DP rank IDs. Add CPU-only subprocess regression coverage.

Co-Authored-By: Codex <codex@openai.com>

examples/_launch.py

@@ -7,10 +7,14 @@
 from __future__ import annotations
 
 import os
+import time
 from collections.abc import Callable
 from multiprocessing import get_context
+from multiprocessing.connection import wait
 from multiprocessing.process import BaseProcess
 
+_PROCESS_TERMINATION_TIMEOUT_S = 5.0
+
 
 def launch_dp(entry: Callable[[int, int], None], dp_size: int | None = None) -> None:
     """设置单机默认 RANK/WORLD_SIZE 后用 run_dp 启动子进程；dp_size 为 None 时用满本节点 GPU。"""
@@ -34,17 +38,38 @@ def run_dp(main_fn: Callable[[int, int], None], dp_size: int) -> None:
 
     if dp_size > 1:
         ctx = get_context("spawn")
-        procs: list[BaseProcess] = []
+        procs: dict[int, BaseProcess] = {}
         for dp_rank in range(dp_size):
             if dp_rank // (dp_size // num_nodes) == node_rank:
-                procs.append(ctx.Process(target=main_fn, args=(dp_size, dp_rank), daemon=False))
-        for p in procs:
-            p.start()
-        for p in procs:
-            p.join()
-
-        failed_procs = [i for i, p in enumerate(procs) if p.exitcode != 0]
-        if failed_procs:
-            raise RuntimeError(f"Some processes failed: {failed_procs}")
+                procs[dp_rank] = ctx.Process(target=main_fn, args=(dp_size, dp_rank), daemon=False)
+        try:
+            for p in procs.values():
+                p.start()
+            pending = {p.sentinel: rank for rank, p in procs.items()}
+            while pending:
+                failed_ranks = []
+                for sentinel in wait(pending):
+                    rank = pending.pop(sentinel)
+                    p = procs[rank]
+                    p.join()
+                    if p.exitcode != 0:
+                        failed_ranks.append(rank)
+                if failed_ranks:
+                    raise RuntimeError(f"Some processes failed: {sorted(failed_ranks)}")
+        finally:
+            # Also clean up after interruption or a failure partway through start().
+            started = [p for p in procs.values() if p.pid is not None]
+            for p in started:
+                if p.is_alive():
+                    p.terminate()
+            deadline = time.monotonic() + _PROCESS_TERMINATION_TIMEOUT_S
+            for p in started:
+                p.join(timeout=max(0.0, deadline - time.monotonic()))
+            # A rank may ignore SIGTERM; share one grace period across all ranks.
+            for p in started:
+                if p.is_alive():
+                    p.kill()
+            for p in started:
+                p.join()
     else:
         main_fn(dp_size, 0)

tests/examples/test_launch.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import multiprocessing as mp
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from examples import _launch
+
+pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
+
+
+def _worker(dp_size: int, dp_rank: int) -> None:
+    directory = Path(os.environ["LAUNCH_TEST_DIR"])
+    mode = os.environ["LAUNCH_TEST_MODE"]
+    failed_rank = int(os.environ["LAUNCH_TEST_FAILED_RANK"])
+    if mode == "ignore_sigterm" and dp_rank != failed_rank:
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+    (directory / f"rank-{dp_rank}.pid").write_text(str(os.getpid()))
+    if mode == "success":
+        return
+    if mode != "interrupt" and dp_rank == failed_rank:
+        # Ensure the sibling has started (and installed its signal handler).
+        deadline = time.monotonic() + 10
+        while len(list(directory.glob("rank-*.pid"))) < dp_size // int(os.environ["WORLD_SIZE"]):
+            if time.monotonic() >= deadline:
+                raise RuntimeError("sibling failed to start")
+            time.sleep(0.01)
+        raise RuntimeError("intentional worker failure")
+    # Stand in for a rank stuck in a collective after another rank has failed.
+    time.sleep(60)
+
+
+def _run_case(
+    tmp_path: Path,
+    *,
+    mode: str,
+    failed_rank: int = 1,
+    node_rank: int = 0,
+    num_nodes: int = 1,
+    dp_size: int = 2,
+) -> subprocess.CompletedProcess[str]:
+    environment = dict(
+        os.environ,
+        PYTHONPATH=str(Path(__file__).resolve().parents[2]),
+        RANK=str(node_rank),
+        WORLD_SIZE=str(num_nodes),
+        LAUNCH_TEST_DIR=str(tmp_path),
+        LAUNCH_TEST_MODE=mode,
+        LAUNCH_TEST_FAILED_RANK=str(failed_rank),
+        LAUNCH_TEST_DP_SIZE=str(dp_size),
+    )
+    process = subprocess.Popen(
+        [sys.executable, str(Path(__file__).resolve())],
+        env=environment,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        start_new_session=True,
+    )
+    try:
+        if mode == "interrupt":
+            deadline = time.monotonic() + 10
+            while len(list(tmp_path.glob("rank-*.pid"))) < dp_size // num_nodes:
+                if process.poll() is not None or time.monotonic() >= deadline:
+                    pytest.fail("workers failed to start before interruption")
+                time.sleep(0.01)
+            process.send_signal(signal.SIGINT)
+        stdout, stderr = process.communicate(timeout=15)
+    except subprocess.TimeoutExpired:
+        pytest.fail("launcher did not exit after a worker failed")
+    finally:
+        # Also bounds the regression test when run against the broken launcher.
+        with contextlib.suppress(ProcessLookupError):
+            os.killpg(process.pid, signal.SIGKILL)
+        process.wait()
+    assert json.loads((tmp_path / "remaining_children.json").read_text()) == []
+    return subprocess.CompletedProcess(process.args, process.returncode, stdout, stderr)
+
+
+def test_successful_workers_are_joined(tmp_path: Path) -> None:
+    result = _run_case(tmp_path, mode="success")
+    assert result.returncode == 0, result.stderr
+    assert len(list(tmp_path.glob("rank-*.pid"))) == 2
+
+
+@pytest.mark.parametrize(
+    ("node_rank", "num_nodes", "dp_size", "failed_rank"),
+    [(0, 1, 2, 1), (0, 1, 2, 0), (1, 2, 4, 3)],
+)
+def test_worker_failure_reaps_siblings(tmp_path: Path, node_rank: int, num_nodes: int, dp_size: int, failed_rank: int) -> None:
+    result = _run_case(tmp_path, mode="failure", node_rank=node_rank, num_nodes=num_nodes, dp_size=dp_size, failed_rank=failed_rank)
+    assert result.returncode != 0
+    assert "intentional worker failure" in result.stderr
+    assert f"Some processes failed: [{failed_rank}]" in result.stderr
+
+
+def test_unresponsive_sibling_is_killed(tmp_path: Path) -> None:
+    result = _run_case(tmp_path, mode="ignore_sigterm")
+    assert result.returncode != 0
+    assert "Some processes failed: [1]" in result.stderr
+
+
+def test_interruption_reaps_workers(tmp_path: Path) -> None:
+    result = _run_case(tmp_path, mode="interrupt")
+    assert result.returncode != 0
+    assert "KeyboardInterrupt" in result.stderr
+
+
+def test_start_failure_reaps_started_worker(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    context = mp.get_context("spawn")
+    processes = []
+
+    def fail_start() -> None:
+        raise RuntimeError("could not start worker")
+
+    def make_process(**kwargs: object) -> mp.Process:
+        process = context.Process(target=time.sleep, args=(60,))
+        if processes:
+            monkeypatch.setattr(process, "start", fail_start)
+        processes.append(process)
+        return process
+
+    monkeypatch.setattr(_launch, "get_context", lambda _: SimpleNamespace(Process=make_process))
+    try:
+        with pytest.raises(RuntimeError, match="could not start worker"):
+            _launch.run_dp(_worker, 2)
+        assert processes[0].exitcode is not None
+        assert processes[1].pid is None
+    finally:
+        for process in processes:
+            if process.pid is not None:
+                if process.is_alive():
+                    process.kill()
+                process.join()
+
+
+def test_single_rank_runs_in_parent(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    calls = []
+    _launch.run_dp(lambda size, rank: calls.append((size, rank, os.getpid())), 1)
+    assert calls == [(1, 0, os.getpid())]
+
+
+if __name__ == "__main__":
+    _launch._PROCESS_TERMINATION_TIMEOUT_S = 0.2
+    try:
+        _launch.run_dp(_worker, int(os.environ["LAUNCH_TEST_DP_SIZE"]))
+    finally:
+        remaining = [child.pid for child in mp.active_children()]
+        (Path(os.environ["LAUNCH_TEST_DIR"]) / "remaining_children.json").write_text(json.dumps(remaining))