| |
| """Host guard CLI (see align/utils/hostguard.py for the design and the 2026-10-14 incident). |
| |
| hostguard.py install write + enable align.slice and align-hostguard.service (user units) |
| hostguard.py run [--cores N] [--mem-gb G] [--label L] -- CMD ... |
| run CMD inside a kernel-capped scope in align.slice |
| hostguard.py status host sample, slice usage, watchdog state, last events |
| hostguard.py stop-all kill switch: SIGKILL every process in align.slice |
| hostguard.py watch [--dry-run] the watchdog loop (run by align-hostguard.service) |
| |
| Every heavy or long command on this box (GPU jobs, Lean verification, builds, data jobs, test suites) goes |
| through `run`. Default per-job caps are deliberately below the slice caps so one job cannot take the whole |
| slice: --cores 16 (of 48), --mem-gb 60. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import dataclasses |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from align.utils import hostguard as H |
|
|
| UNIT_DIR = Path.home() / ".config/systemd/user" |
| REPO_UNITS = ROOT / "deploy/systemd" |
| PY = ROOT / ".venv/bin/python" |
|
|
|
|
| def slice_unit(p: H.Policy) -> str: |
| return f"""[Unit] |
| Description=align jobs (kernel-capped; see align/utils/hostguard.py) |
| |
| [Slice] |
| CPUAccounting=yes |
| CPUQuota={int(p.slice_cpu_cores * 100)}% |
| MemoryAccounting=yes |
| MemoryHigh={int(p.slice_mem_gb * 0.92)}G |
| MemoryMax={int(p.slice_mem_gb)}G |
| MemorySwapMax=0 |
| TasksMax=16384 |
| """ |
|
|
|
|
| def service_unit() -> str: |
| return f"""[Unit] |
| Description=align host watchdog (host CPU <= 92%, memory, D-state, GPU temp, disk) |
| After=default.target |
| |
| [Service] |
| Type=simple |
| ExecStart={PY} {ROOT / 'scripts/hostguard.py'} watch |
| Restart=always |
| RestartSec=5 |
| CPUWeight=1000 |
| MemoryMax=1G |
| Environment=PYTHONUNBUFFERED=1 |
| |
| [Install] |
| WantedBy=default.target |
| """ |
|
|
|
|
| def cmd_install(a) -> int: |
| p = H.Policy() |
| REPO_UNITS.mkdir(parents=True, exist_ok=True) |
| UNIT_DIR.mkdir(parents=True, exist_ok=True) |
| for name, text in (("align.slice", slice_unit(p)), ("align-hostguard.service", service_unit())): |
| (REPO_UNITS / name).write_text(text) |
| (UNIT_DIR / name).write_text(text) |
| for args in (("daemon-reload",), ("start", "align.slice"), ("enable", "--now", "align-hostguard.service")): |
| r = subprocess.run(["systemctl", "--user", *args], capture_output=True, text=True) |
| if r.returncode: |
| print(f"systemctl --user {' '.join(args)} failed: {r.stderr.strip()}", file=sys.stderr) |
| return 1 |
| print(json.dumps({"installed": [str(UNIT_DIR / 'align.slice'), str(UNIT_DIR / 'align-hostguard.service')], |
| "slice_cpu_cores": p.slice_cpu_cores, "slice_mem_gb": p.slice_mem_gb}, indent=1)) |
| return 0 |
|
|
|
|
| def _watchdog_active() -> bool: |
| return subprocess.run(["systemctl", "--user", "is-active", "--quiet", "align-hostguard.service"]).returncode == 0 |
|
|
|
|
| def apps_only_env() -> dict: |
| """The /apps-only path variables from env/apps_only.sh (single source of truth), as a dict.""" |
| out = subprocess.run(["bash", "-c", f". {H.APPS_ONLY_ENV} && env -0"], capture_output=True, timeout=30, |
| env={"PATH": os.environ.get("PATH", "/usr/bin:/bin")}).stdout.decode() |
| env = dict(kv.split("=", 1) for kv in out.split("\0") if "=" in kv) |
| return {k: env[k] for k in H.PATH_ENV + ("HF_HUB_DISABLE_XET",) if k in env} |
|
|
|
|
| def cap_env(cores: int) -> dict: |
| """Thread / job caps for tools that do not see the cgroup quota and would size pools by nproc, plus the |
| /apps-only cache/temp paths. Path variables the caller already set under /apps (e.g. a toolchain's |
| CARGO_HOME) are kept; anything pointing elsewhere is overridden.""" |
| n = str(max(1, cores)) |
| env = {"OMP_NUM_THREADS": n, "MKL_NUM_THREADS": n, "OPENBLAS_NUM_THREADS": n, "NUMEXPR_MAX_THREADS": n, |
| "RAYON_NUM_THREADS": n, "CARGO_BUILD_JOBS": n, "MAKEFLAGS": f"-j{n}", "OPAMJOBS": n, |
| "TORCHINDUCTOR_COMPILE_THREADS": str(min(4, cores)), "TOKENIZERS_PARALLELISM": "false"} |
| for k, v in apps_only_env().items(): |
| cur = os.environ.get(k) |
| env[k] = cur if cur and os.path.realpath(cur).startswith("/apps/") else v |
| return env |
|
|
|
|
| def cmd_run(a) -> int: |
| if not a.cmd: |
| print("usage: hostguard.py run [opts] -- CMD ...", file=sys.stderr) |
| return 2 |
| cmd = a.cmd[1:] if a.cmd[0] == "--" else a.cmd |
| p = H.Policy() |
| if not _watchdog_active() and not a.allow_no_watchdog: |
| print("REFUSED: align-hostguard.service is not active (run `scripts/hostguard.py install`)", file=sys.stderr) |
| return 3 |
| if a.cores > p.slice_cpu_cores or a.mem_gb > p.slice_mem_gb: |
| print(f"REFUSED: per-job caps exceed the slice ({p.slice_cpu_cores} cores, {p.slice_mem_gb} GB)", file=sys.stderr) |
| return 2 |
| waited = 0.0 |
| while True: |
| why = H.preflight(H.take_sample(1.0), p) |
| if not why: |
| break |
| if not a.wait or waited >= a.wait_max_s: |
| print("REFUSED (host busy): " + "; ".join(why), file=sys.stderr) |
| return 4 |
| time.sleep(30); waited += 31 |
| env = cap_env(a.cores) |
| bad = H.off_apps_paths({**os.environ, **env}) |
| if bad: |
| print("REFUSED: downloads/caches/temp must be under /apps: " + ", ".join(bad), file=sys.stderr) |
| return 5 |
| for k in ("TMPDIR", "XDG_CACHE_HOME"): |
| Path(env[k]).mkdir(parents=True, exist_ok=True) |
| unit = f"align-{a.label}-{int(time.time())}" |
| argv = ["systemd-run", "--user", "--scope", "--quiet", f"--slice={H.SLICE}", f"--unit={unit}", |
| "-p", f"CPUQuota={a.cores * 100}%", "-p", f"MemoryMax={a.mem_gb}G", "-p", f"MemoryHigh={max(1, int(a.mem_gb * 0.9))}G", |
| "-p", "MemorySwapMax=0", "-p", f"TasksMax={a.tasks}", "--", *cmd] |
| H.log_event("run", unit=unit, cores=a.cores, mem_gb=a.mem_gb, cmd=cmd[:8]) |
| os.environ.update(env) |
| os.execvp(argv[0], argv) |
|
|
|
|
| def cmd_status(a) -> int: |
| s = H.take_sample(1.0) |
| cg = H.slice_cgroup() |
| info = {"host": dataclasses.asdict(s), "preflight_refusals": H.preflight(s, H.Policy()), |
| "watchdog_active": _watchdog_active(), "slice": None} |
| if cg: |
| rd = lambda f: (cg / f).read_text().strip() if (cg / f).exists() else None |
| info["slice"] = {"pids": len(H.slice_pids()), "cpu.max": rd("cpu.max"), "memory.current_gb": |
| round(int(rd("memory.current") or 0) / 2**30, 2), "frozen": rd("cgroup.freeze"), |
| "scopes": sorted(p.name for p in cg.iterdir() if p.name.endswith(".scope"))} |
| if H.EVENT_LOG.exists(): |
| info["last_events"] = [json.loads(l) for l in H.EVENT_LOG.read_text().splitlines()[-a.events:]] |
| print(json.dumps(info, indent=1, default=str)) |
| return 0 |
|
|
|
|
| def cmd_stop_all(a) -> int: |
| r = subprocess.run(["systemctl", "--user", "kill", "--signal=SIGKILL", H.SLICE], capture_output=True, text=True) |
| H.log_event("stop_all_cli", rc=r.returncode) |
| print(f"SIGKILL sent to every process in {H.SLICE} (rc={r.returncode}); remaining pids: {H.slice_pids()}") |
| return 0 |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| sub = ap.add_subparsers(dest="cmd_name", required=True) |
| sub.add_parser("install") |
| r = sub.add_parser("run") |
| r.add_argument("--cores", type=int, default=16) |
| r.add_argument("--mem-gb", type=int, default=60) |
| r.add_argument("--tasks", type=int, default=4096) |
| r.add_argument("--label", default="job") |
| r.add_argument("--wait", action="store_true", help="wait (up to --wait-max-s) for preflight instead of refusing") |
| r.add_argument("--wait-max-s", type=float, default=6 * 3600) |
| r.add_argument("--allow-no-watchdog", action="store_true", help="tests only") |
| r.add_argument("cmd", nargs=argparse.REMAINDER) |
| st = sub.add_parser("status"); st.add_argument("--events", type=int, default=10) |
| sub.add_parser("stop-all") |
| w = sub.add_parser("watch"); w.add_argument("--dry-run", action="store_true"); w.add_argument("--interval", type=float, default=3.0) |
| a = ap.parse_args() |
| if a.cmd_name == "watch": |
| H.watch(H.Policy(), interval=a.interval, dry_run=a.dry_run) |
| return 0 |
| return {"install": cmd_install, "run": cmd_run, "status": cmd_status, "stop-all": cmd_stop_all}[a.cmd_name](a) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
