#!/usr/bin/env python """Start a training run on the GPU host and hand the GPU back when it ends. Detached by design: a training run is a day long, and tying it to an SSH channel means it dies with the connection -- or worse, finishes while the client stays blocked on a pipe that never closes. Both happened while this was being built. The run becomes a transient systemd unit on the host; this process only starts it, waits, and restores the LLM services afterwards. gpu_train.py --epochs 48 --steps 1500 # start and watch gpu_train.py --status # what is it doing now gpu_train.py --stop # stop it and restore the LLMs """ import argparse import os import sys import time sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from suvi import gpubox REMOTE = "/var/home/htpc/suvi" LOG = f"{REMOTE}/train.log" def command(args): return ( "bash -c 'cd /work && python -u train.py " f"--shards /work/shards --manifest /work/shards/manifest.json " f"--out /work/runs/{args.name} " f"--epochs {args.epochs} --batch {args.batch} --lr {args.lr} " f"--base {args.base} --depth {args.depth} --workers {args.workers} " f"--val-samples {args.val_samples} --eval-every {args.eval_every}" + (f" --steps-per-epoch {args.steps}" if args.steps else "") + (f" --max-steps {args.max_steps}" if args.max_steps else "") + (f" --overfit {args.overfit}" if args.overfit else "") + (f" --resume /work/runs/{args.name}/last.pt" if args.resume else "") + "'" ) def main(argv=None): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--name", default="v4") parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--steps", type=int, default=None, help="steps per epoch (default: the whole train split)") parser.add_argument("--val-samples", type=int, default=64) parser.add_argument("--eval-every", type=int, default=250) parser.add_argument("--max-steps", type=int, default=None, help="stop after this many steps (the 300-step probe gate)") parser.add_argument("--overfit", type=int, default=None, help="run the overfit gate on this many frozen samples") parser.add_argument("--batch", type=int, default=8) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--base", type=int, default=32) parser.add_argument("--depth", type=int, default=3) parser.add_argument("--workers", type=int, default=6) parser.add_argument("--resume", action="store_true") parser.add_argument("--status", action="store_true") parser.add_argument("--since", type=int, default=None, help="print history rows with step > SINCE, plus health") parser.add_argument("--stop", action="store_true") parser.add_argument("--watch", type=int, default=0, help="seconds to wait for completion (0 = start and return)") args = parser.parse_args(argv) if args.status: print(gpubox.tail(LOG, 30)) print("containers:", gpubox.run("podman ps --format '{{.Names}}'", check=False).strip()) print(f"GTT {gpubox.gtt_used() / 1024**3:.1f} GiB used") return 0 if args.since is not None: return _report_since(args) if args.stop: stopped = gpubox.stop_containers() print(f"stopped {stopped or 'nothing'}") gpubox.run("systemctl --user start llama-swap.service", check=False) return 0 state = gpubox.stop_llms() try: name, log = gpubox.launch(command(args), mounts=[(REMOTE, "/work")], log=LOG) print(f"launched {name}; log at {log}", flush=True) if not args.watch: print("Running detached. gpu_train.py --status to check, --stop to end.") print("NOTE: the LLM services stay stopped until --stop or the run finishes.") _paused_by_us(state) return 0 deadline = time.time() + args.watch while time.time() < deadline and gpubox.running(name): time.sleep(30) print(gpubox.tail(log, 40)) finally: if args.watch: gpubox.stop_containers() print("restore failures:", gpubox.start_llms(state) or "none") return 0 def _report_since(args): """New history rows past a step, plus the health a loss curve does not show. A run can go non-finite, stall with the process still alive, or lose its container entirely; each of those wastes hours if nobody looks. Returns 1 when something is wrong, so a supervising loop can react without parsing prose. """ import json raw = gpubox.run(f"cat {REMOTE}/runs/{args.name}/history.jsonl 2>/dev/null", check=False) rows = [json.loads(line) for line in raw.splitlines() if line.strip()] fresh = [row for row in rows if row.get("step", 0) > args.since] for row in fresh: if row.get("untrained"): print(f"step {row['step']:>6} val {row['val_loss']:.5f} " f"({row['val_psnr']:.2f} dB) [untrained]") else: print(f"step {row['step']:>6} train {row['train_loss']:.5f} " f"val {row['val_loss']:.5f} ({row['val_psnr']:.2f} dB) " f"grad {row.get('grad_mean', 0):.3f}/{row.get('grad_max', 0):.2f}") problems = [] for row in fresh: for key in ("train_loss", "val_loss", "val_psnr"): if key in row and not (row[key] == row[key] and abs(row[key]) < 1e30): problems.append(f"step {row['step']}: {key} is not finite") containers = gpubox.run("podman ps --format '{{.Names}}'", check=False).strip() if "suvi" not in containers: problems.append("no suvi container is running") disk = gpubox.run(f"df --output=pcent {REMOTE} | tail -1", check=False).strip() if disk.rstrip("%").strip().isdigit() and int(disk.rstrip("%").strip()) > 95: problems.append(f"host disk at {disk}") for problem in problems: print(f"PROBLEM: {problem}") return 1 if problems else 0 def _paused_by_us(state): """Record what was stopped, so a later --stop can put it back. Written to disk rather than held in memory because the run outlives this process by a day, and a restore that only works while the launcher is alive is not a restore. """ import json path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".gpu_paused.json") with open(path, "w") as handle: json.dump(state, handle) print(f" paused state recorded in {path}") if __name__ == "__main__": sys.exit(main())