noaa-goes-visualization/gpu_train.py

#!/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())