| """G3 audit regressions that need two real ranks (CPU, gloo): U01-3. |
| |
| Reuses the 2-rank harness pieces of test_trainer_dist (`_build_trainer`, |
| `_free_port`); only the child entry point is local, because spawn resolves the |
| body by module and `test_trainer_dist._BODIES` cannot see bodies defined here. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import sys |
| import traceback |
|
|
| import torch.multiprocessing as mp |
| from test_trainer_dist import JOIN_TIMEOUT_S, WS, _build_trainer, _free_port |
|
|
|
|
| def _body_pin_on_rank1_only(rank: int, run_dir: str) -> dict: |
| """The R5 pin flag is set on RANK 1 ONLY (a rank-local tripwire fired |
| there), then both ranks train through saves at micro 4 and 8 with keep-N=1.""" |
| trainer, _model = _build_trainer(run_dir, init_seed=0) |
| trainer.cfg.train.keep_checkpoints = 1 |
| if rank == 1: |
| trainer._pin_next_save = True |
| trainer.train() |
| root = os.path.join(run_dir, "checkpoints") |
| return { |
| "dirs": sorted(d for d in os.listdir(root) if not d.endswith(".partial")), |
| "pinned": sorted(trainer.pinned_checkpoints), |
| "flag_left": bool(trainer._pin_next_save), |
| } |
|
|
|
|
| _BODIES = {"pin_rank1": _body_pin_on_rank1_only} |
|
|
|
|
| def _child(rank: int, port: int, mode: str, run_dir: str) -> None: |
| import torch.distributed as dist |
|
|
| os.environ.update( |
| RANK=str(rank), |
| WORLD_SIZE=str(WS), |
| LOCAL_RANK=str(rank), |
| MASTER_ADDR="127.0.0.1", |
| MASTER_PORT=str(port), |
| ) |
| from align.utils.dist import init_distributed |
|
|
| init_distributed(backend="gloo") |
| code = 0 |
| try: |
| result = _BODIES[mode](rank, run_dir) |
| with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f: |
| json.dump(result, f) |
| except BaseException: |
| with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f: |
| f.write(traceback.format_exc()) |
| code = 1 |
| finally: |
| try: |
| dist.destroy_process_group() |
| except Exception: |
| pass |
| if code: |
| sys.exit(code) |
|
|
|
|
| def _run(mode: str, run_dir: str, monkeypatch) -> list[dict]: |
| monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") |
| ctx = mp.get_context("spawn") |
| port = _free_port() |
| procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)] |
| for p in procs: |
| p.start() |
| for p in procs: |
| p.join(JOIN_TIMEOUT_S) |
| hung = [p for p in procs if p.is_alive()] |
| for p in hung: |
| p.kill() |
| p.join() |
| errs = [] |
| for r in range(WS): |
| ef = os.path.join(run_dir, f"{mode}_{r}.err") |
| if os.path.isfile(ef): |
| errs.append(f"rank {r}:\n{open(ef).read()}") |
| assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs) |
| assert not errs, "\n".join(errs) |
| return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)] |
|
|
|
|
| def test_u01_3_pin_set_on_a_non_zero_rank_protects_the_checkpoint_on_rank_0( |
| tmp_workdir, monkeypatch |
| ): |
| """U01-3. `_pin_next_save` is raised by a RANK-LOCAL tripwire verdict, but |
| pruning and resume.pt are rank 0's. A pin that exists on rank 1 only left |
| rank 0's pinned set empty: keep-N deleted the very checkpoint meant as the |
| rewind point and a resume forgot the pin. The pin must be collective. |
| Mutant: drop the all_reduce in save_checkpoint -> dirs == [micro_00000008].""" |
| rows = _run("pin_rank1", tmp_workdir, monkeypatch) |
| for r in rows: |
| assert r["pinned"] == [4], rows |
| assert r["flag_left"] is False, rows |
| assert rows[0]["dirs"] == ["micro_00000004", "micro_00000008"], rows |
| import torch |
|
|
| state = torch.load( |
| os.path.join(tmp_workdir, "checkpoints", "micro_00000004", "resume.pt"), |
| weights_only=True, |
| map_location="cpu", |
| ) |
| assert state["pinned_checkpoints"] == [4] |
|
|
