Merge 9a52063a8a0dc86cbbce04fae034e7e037fa3881 into 553741557a57df44a46fb3ab72b6aafeaf945337

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))