File size: 3,405 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | #!/usr/bin/env python
"""Bounded-retry launch supervisor CLI (stdlib only; no torch import).
python scripts/supervise.py --run-dir RUN [--max-retries 2] [--backoff 30] \
-- torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml
Everything after `--` is the child command. See align/utils/supervisor.py for
the retry rules (progress + preflight + non-HALT exit code) and why each exists.
`--run-dir` must be the directory the Trainer writes history.jsonl and
checkpoints/ into; attempts are logged to RUN/supervisor.jsonl.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
# The package is not installed in the venv; without this a bare `python scripts/supervise.py` (the documented
# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py).
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from align.utils.supervisor import Supervisor
def _checkpoint_validator():
"""`checkpoint_problem` from the Trainer, imported lazily (it needs torch).
The supervisor module itself stays stdlib-only; THIS entry point is where
the cost is paid, once, so that "progress" means a checkpoint `latest_checkpoint`
would actually resume from. If torch is unavailable the supervisor still
runs, falling back to name-only progress, and says so on stderr.
"""
try:
from align.train.trainer import checkpoint_problem
except Exception as e: # noqa: BLE001
print(f"supervisor: checkpoint validation unavailable ({type(e).__name__}: {e}); "
"progress is judged by directory name only", file=sys.stderr)
return None
return checkpoint_problem
def main(argv=None) -> int:
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
ap.add_argument("--run-dir", required=True)
ap.add_argument("--max-retries", type=int, default=2)
ap.add_argument("--backoff", type=float, default=30.0, help="base seconds; doubles per retry")
ap.add_argument("--grace", type=float, default=20.0, help="SIGTERM->SIGKILL grace seconds")
ap.add_argument(
"--no-validate-checkpoints",
action="store_true",
help="count a checkpoint dir as progress by NAME only (skips importing torch for the "
"completeness check); a truncated newest dir then counts as progress",
)
ap.add_argument(
"--retry-on-stop",
action="store_true",
help="also relaunch a run that was stopped by SIGTERM/SIGINT/SIGHUP (exit 143/130/129). "
"Default: an operator's stop is a decision, not a fault, and ends the supervisor",
)
ap.add_argument("cmd", nargs=argparse.REMAINDER, help="-- child command")
a = ap.parse_args(argv)
cmd = a.cmd[1:] if a.cmd and a.cmd[0] == "--" else a.cmd
if not cmd:
ap.error("no child command given (put it after --)")
res = Supervisor(cmd, a.run_dir, max_retries=a.max_retries,
backoff_s=a.backoff, grace_s=a.grace, retry_on_stop=a.retry_on_stop,
checkpoint_validator=None if a.no_validate_checkpoints
else _checkpoint_validator()).run()
print(f"supervisor: {res.reason} after {res.attempts} attempt(s), "
f"child exits {res.child_codes}", file=sys.stderr)
return res.exit_code
if __name__ == "__main__":
raise SystemExit(main())
|