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
    gpu_train.py --pool --steps 8923 ...           # train from a rotating RAM pool

With ``--pool`` a second container runs ``pool_producer.py`` beside training, reading
the archive share read-only and keeping fresh slices in host RAM (:mod:`suvi.pool`);
``--stop`` stops both and frees the pool.
"""

import argparse
import json
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"
#: Free space below which a run is at risk of failing to write a checkpoint.
DISK_FLOOR_GB = 5
#: The archive share as the host mounts it (read-only), and the RAM pool directory.
#: Both containers use the host's IPC namespace, so they share its /dev/shm.
ARCHIVE = "/var/home/htpc/NOAA_GOES"
POOL_DIR = "/dev/shm/suvi_pool"
POOL_LOG = f"{REMOTE}/pool.log"


def command(args):
    return (
        "bash -c 'cd /work && python -u train.py "
        f"--shards /work/{args.shards_dir} "
        f"--manifest /work/{args.shards_dir}/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" --lr-range {args.lr_range}" if args.lr_range else "")
        + (f" --span {args.span}" if args.span else "")
        + (f" --warmup {args.warmup}" if args.warmup is not None else "")
        + (" --relative" if args.relative else "")
        + (" --corr-flow" if args.corr_flow else "")
        + (" --attention" if args.attention 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 "")
        + (f" --pool {POOL_DIR}" if args.pool else "")
        + (f" --val-slices /work/{args.val_dir}" if args.val_dir else "")
        + (f" --evaluate /work/{args.evaluate}" if args.evaluate else "")
        + "'"
    )


def producer_command(args):
    return (f"python -u /work/pool_producer.py --manifest /work/{args.pool_manifest} "
            f"--archive /archive/Data --pool {POOL_DIR} "
            f"--budget-gb {args.pool_budget_gb} --workers {args.pool_workers}")


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--name", default="v6")
    parser.add_argument("--shards-dir", default="shards",
                        help="dataset directory under the host's suvi root")
    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("--lr-range", default=None,
                        help="'LO,HI': run the LR range test (needs --max-steps)")
    parser.add_argument("--span", type=int, default=None,
                        help="model input span in slots (rung ladder prefix)")
    parser.add_argument("--warmup", type=int, default=None)
    parser.add_argument("--relative", action="store_true")
    parser.add_argument("--corr-flow", action="store_true")
    parser.add_argument("--attention", action="store_true")
    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("--pool", action="store_true",
                        help="train from a rotating RAM pool generated from the archive "
                             "(needs --steps); validation stays on --shards-dir")
    parser.add_argument("--pool-manifest", default="pool_manifest.json.zst",
                        help="under the host's suvi root, from pool_manifest.py")
    parser.add_argument("--pool-budget-gb", type=float, default=48.0)
    parser.add_argument("--pool-workers", type=int, default=16)
    parser.add_argument("--val-dir", default=None,
                        help="validate on the fixed slices in this directory under the "
                             "host's suvi root (pool_producer.py --fixed)")
    parser.add_argument("--evaluate", default=None,
                        help="score this checkpoint (under the suvi root) and exit")
    parser.add_argument("--status", action="store_true")
    parser.add_argument("--watchdog", action="store_true",
                        help="one-line answer to 'is the box still working'")
    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.watchdog:
        return _watchdog(args)

    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()
        gpubox.run(f"rm -rf {POOL_DIR}", check=False)      # the pool is host RAM
        print(f"stopped {stopped or 'nothing'}")
        gpubox.run("systemctl --user start llama-swap.service", check=False)
        return 0

    if args.pool and not args.steps:
        raise SystemExit("--pool needs --steps (steps per epoch): a pool has no length")
    state = gpubox.stop_llms()
    try:
        if args.pool:
            producer, _ = gpubox.launch(
                producer_command(args), log=POOL_LOG,
                name=f"{gpubox.CONTAINER_PREFIX}pool-{int(time.time())}",
                mounts=[(REMOTE, "/work"), (ARCHIVE, "/archive:ro")])
            print(f"launched producer {producer}; log at {POOL_LOG}", flush=True)
        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


#: A run that has not written an eval in this long is not merely slow.
STALL_MINUTES = 30


#: Host memory below which the box is thrashing rather than training.  A run
#: that exhausts RAM keeps its process runnable at high CPU while it pages, so
#: it reads as healthy from every angle except the one that matters: it stops
#: producing output.  Naming the cause turns a baffling stall into an obvious one.
LOW_MEMORY_GB = 4.0


def progress(row):
    """Epochs and steps, separately: epochs say how far training has really got
    (passes through the corpus); steps are the fine-grained, short-term measure.

    Rows from train.py carry steps_per_epoch and total_steps; older runs only the
    epoch index, which is reported as the epoch in progress.
    """
    step, per_epoch, total = row.get("step"), row.get("steps_per_epoch"), row.get("total_steps")
    if per_epoch:
        text = f"epoch {step / per_epoch:.2f}"
        if total:
            text += f" of {total / per_epoch:.2f}"
        text += f", step {step}" + (f" of {total}" if total else "")
        return text
    if row.get("epoch") is not None:
        return f"epoch {row['epoch'] + 1} (in progress), step {step}"
    return f"step {step}"


def watchdog_verdict(containers, runs, now, stall_minutes=STALL_MINUTES,
                     busy_age=None, free_gb=None):
    """Decide what the box is doing, from one snapshot.  Pure, so it can be tested.

    `runs` maps run name to {"step", "psnr", "age"} -- plus, where the history
    has them, "epoch", "steps_per_epoch" and "total_steps" -- where age is seconds since
    its history last grew; `busy_age` is how long the training container has been
    up.  Returns (state, lines) with state one of ``training`` / ``stalled`` /
    ``idle`` / ``unknown``.

    The distinction that matters is **stalled versus finished**: both look like
    "no new rows", and telling them apart is the whole job of a watchdog.  A
    container still up with a frozen history is a wedged run burning hours; no
    container at all is simply an idle box.

    Staleness alone is not enough to call a stall, because a *freshly launched*
    run has written nothing yet, so the newest history on the box still belongs
    to the previous run and looks arbitrarily old.  Pulling the image and
    importing torch takes minutes, and the first eval is 250 steps after that.
    So a stall needs both clocks: the history frozen *and* the container up
    long enough that it should have produced a row.
    """
    if containers is None:
        return "unknown", ["host unreachable"]
    busy = [name for name in containers if name.startswith("suvi-run")]
    if not runs:
        return ("training" if busy else "idle"), ["no run history yet"]

    name = min(runs, key=lambda key: runs[key]["age"])
    row = runs[name]
    age = row["age"] / 60.0
    where = f"{name} at {progress(row)}"
    if row.get("psnr") is not None:
        where += f", val {row['psnr']:.2f} dB"

    if not busy:
        return "idle", [f"GPU idle; last was {where} ({age:.0f} min ago)"]
    starved = free_gb is not None and free_gb < LOW_MEMORY_GB
    up = (busy_age / 60.0) if busy_age is not None else age
    if age > stall_minutes and up > stall_minutes:
        lines = [f"STALLED: {where}, no eval for {age:.0f} min "
                 f"(container up {up:.0f} min)"]
        if starved:
            lines.append(f"host has {free_gb:.1f} GB free -- it is swap-thrashing, "
                         f"not computing; stop it and lower the memory ceiling")
        return "stalled", lines
    if starved:
        return "training", [f"{where} ({age:.0f} min ago)",
                            f"WARNING: only {free_gb:.1f} GB host memory free"]
    if age > up:
        return "training", [f"{busy[0]} starting up ({up:.0f} min); "
                            f"last finished run was {where}"]
    return "training", [f"{where} ({age:.0f} min ago)"]


def _watchdog(args):
    """One round trip: what is running, how far along, and is it still moving."""
    marker = "@@@"
    command = (
        "podman ps --format '{{.Names}}'; echo " + marker + "; "
        "awk '/MemAvailable/{print \"MEM\", $2}' /proc/meminfo; "
        f"for d in {REMOTE}/runs/*/; do n=$(basename $d); "
        f"h=$d/history.jsonl; [ -f $h ] || continue; "
        f"echo \"$n $(stat -c %Y $h) $(tail -1 $h)\"; done; true"
    )
    out = gpubox.run(command, check=False, timeout=60)
    if marker not in out:
        state, lines = watchdog_verdict(None, {}, time.time())
    else:
        head, _, tail = out.partition(marker)
        containers = [line.strip() for line in head.splitlines() if line.strip()]
        runs = {}
        free_gb = None
        now = time.time()
        for line in tail.splitlines():
            name, _, rest = line.strip().partition(" ")
            if name == "MEM" and rest.strip().isdigit():
                free_gb = int(rest.strip()) / 1024**2
                continue
            stamp, _, payload = rest.partition(" ")
            if not payload.strip() or not stamp.isdigit():
                continue
            try:
                row = json.loads(payload)
            except ValueError:
                continue
            runs[name] = {"step": row.get("step", -1), "psnr": row.get("val_psnr"),
                          "age": max(now - int(stamp), 0), "epoch": row.get("epoch"),
                          "steps_per_epoch": row.get("steps_per_epoch"),
                          "total_steps": row.get("total_steps")}
        # The launch timestamp is in the container's own name, which is the only
        # clock that says how long *this* run has had to produce a row.
        busy_age = None
        for name in containers:
            suffix = name.rsplit("-", 1)[-1]
            if name.startswith("suvi-run") and suffix.isdigit():
                busy_age = max(now - int(suffix), 0)
        state, lines = watchdog_verdict(containers, runs, now, busy_age=busy_age,
                                        free_gb=free_gb)

    print(f"state: {state}")
    for line in lines:
        print(f"  {line}")
    return 0 if state in ("training", "idle") else 1


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")
    # Free bytes, not percent full.  The packed corpus deliberately fills the host
    # to within ~40 GB, so "96% used" is the intended state and a percentage
    # threshold cries wolf on every poll until nobody reads the health line at all.
    # What actually matters is whether the run can still write what it must:
    # checkpoints (~30 MB each) and its log.
    disk = gpubox.run(f"df --output=avail -BG {REMOTE} | tail -1", check=False)
    free = disk.strip().rstrip("G")
    if free.isdigit() and int(free) < DISK_FLOOR_GB:
        problems.append(f"host disk has only {free} GB free")
    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())