#!/usr/bin/env python3 """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 # noqa: E402 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 # noqa: E731 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())