diff --git a/examples/_launch.py b/examples/_launch.py index 154cfa9..4708961 100644 --- a/examples/_launch.py +++ b/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) diff --git a/tests/examples/test_launch.py b/tests/examples/test_launch.py new file mode 100644 index 0000000..66c3adc --- /dev/null +++ b/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))