Training with checkpoints on a persistent volume (L4)

super volume create volume-train-ckpt      # once per environment
cd examples/volume-train-checkpoint
super create --sku gpu_1x_l4             # → s-xxxx
super run . --script run.sh --on s-xxxx
super download s-xxxx /workspace/out/train.json ./train.json
super stop s-xxxx

#!/usr/bin/env bash
# Train a small classifier and checkpoint it to a PERSISTENT VOLUME, INSIDE an
# supercomputer session. A volume smoke test shaped like a real training workflow: resume
# from the last checkpoint on the volume, train more, save it back. See README.md.
set -euo pipefail

# The pytorch/pytorch base ships torch in its conda env at /opt/conda, but supercomputer's
# run environment doesn't put conda on PATH (default `python` is /usr/bin, no
# torch). Prepend it so `python` resolves to the torch+CUDA interpreter.
export PATH="/opt/conda/bin:$PATH"

# Fail loudly if torch can't see the GPU (the point of running on a GPU host).
python -c "import torch; assert torch.cuda.is_available(), 'no GPU'; print('gpu:', torch.cuda.get_device_name(0))"

exec python "$(dirname "$0")/train.py"

{
    "sku": "gpu_1x_l4",
    "image": {
        "base": "pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel",
        "env": {
            "PYTHONUNBUFFERED": "1",
            "CKPT_DIR": "/workspace/ckpt",
            "EPOCHS": "5",
            "SSL_CERT_FILE": "/etc/ssl/certs/ca-certificates.crt",
            "PIP_CERT": "/etc/ssl/certs/ca-certificates.crt",
            "REQUESTS_CA_BUNDLE": "/etc/ssl/certs/ca-certificates.crt"
        }
    },
    "volumes": [
        {"name": "volume-train-ckpt", "mount": "/workspace/ckpt"}
    ]
}

"""Train a small classifier and checkpoint it to a PERSISTENT VOLUME.

This is a volume smoke test shaped like a real user's training workflow:
resume from the last checkpoint on the volume, train a few more epochs, and
save the updated checkpoint back. The dataset is synthetic and seeded so the
task is genuinely learnable — validation accuracy climbs as cumulative epochs
accumulate across runs, which is exactly the signal that proves the volume
persisted: run 1 starts from scratch (resumed=false), every later run resumes
(resumed=true) with a higher total_epochs and (until it converges) higher acc.

Volume exercised two ways every run:
  - READ  : load /workspace/ckpt/state.pt (lazy-faulted from the volume)
  - WRITE : save state.pt + append history.jsonl (committed back on stop)
"""

from __future__ import annotations

import json
import os
import time

import torch
import torch.nn as nn

CKPT_DIR = os.environ.get("CKPT_DIR", "/workspace/ckpt")
OUT_DIR = "/workspace/out"
os.makedirs(CKPT_DIR, exist_ok=True)
os.makedirs(OUT_DIR, exist_ok=True)
STATE = os.path.join(CKPT_DIR, "state.pt")
HIST = os.path.join(CKPT_DIR, "history.jsonl")

dev = "cuda" if torch.cuda.is_available() else "cpu"
assert dev == "cuda", "no GPU visible to torch"

# Fixed synthetic classification task: 8 Gaussian blobs in 64-dim. Seeded so the
# same data (and thus a meaningful, improving val accuracy) is regenerated every
# run — only the model/optimizer state carried on the volume changes.
g = torch.Generator(device=dev).manual_seed(0)
N, D, C = 8000, 64, 8
centers = torch.randn(C, D, device=dev, generator=g) * 3.0
y = torch.randint(0, C, (N,), device=dev, generator=g)
X = centers[y] + torch.randn(N, D, device=dev, generator=g)
ntr = 6000
Xtr, ytr, Xva, yva = X[:ntr], y[:ntr], X[ntr:], y[ntr:]

model = nn.Sequential(
    nn.Linear(D, 128), nn.ReLU(),
    nn.Linear(128, 128), nn.ReLU(),
    nn.Linear(128, C),
).to(dev)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

# READ the volume: resume from the last checkpoint if one is present.
resumed = False
prior_total = 0
best_acc = 0.0
if os.path.exists(STATE):
    ck = torch.load(STATE, map_location=dev)
    model.load_state_dict(ck["model"])
    opt.load_state_dict(ck["opt"])
    prior_total = int(ck.get("total_epochs", 0))
    best_acc = float(ck.get("best_acc", 0.0))
    resumed = True

EPOCHS = int(os.environ.get("EPOCHS", "5"))
# Soak knob (stress tier): keep training whole epochs until this many seconds
# have elapsed, refreshing the checkpoint on the volume every ~5 minutes — a
# long session then exercises repeated volume writes, like a real long train.
# 0 = exactly EPOCHS epochs, the daily-smoke behavior.
STRESS_DURATION_S = int(os.environ.get("SUPERCOMPUTER_STRESS_DURATION_S", "0") or "0")

lossf = nn.CrossEntropyLoss()
last_loss = float("nan")
deadline = time.monotonic() + STRESS_DURATION_S
next_save = time.monotonic() + 300
epochs_run = 0
while epochs_run < EPOCHS or time.monotonic() < deadline:
    model.train()
    perm = torch.randperm(ntr, device=dev)
    for i in range(0, ntr, 512):
        idx = perm[i:i + 512]
        opt.zero_grad()
        loss = lossf(model(Xtr[idx]), ytr[idx])
        loss.backward()
        opt.step()
    last_loss = float(loss.item())
    epochs_run += 1
    if time.monotonic() >= next_save:
        torch.save(
            {"model": model.state_dict(), "opt": opt.state_dict(),
             "total_epochs": prior_total + epochs_run, "best_acc": best_acc},
            STATE,
        )
        next_save = time.monotonic() + 300

model.eval()
with torch.no_grad():
    val_acc = (model(Xva).argmax(1) == yva).float().mean().item()
best_acc = max(best_acc, val_acc)
total = prior_total + epochs_run

# WRITE the volume: persist the checkpoint + append a run record.
torch.save(
    {"model": model.state_dict(), "opt": opt.state_dict(),
     "total_epochs": total, "best_acc": best_acc},
    STATE,
)
rec = {
    "ts": time.time(),
    "resumed": resumed,
    "prior_total_epochs": prior_total,
    "epochs_this_run": epochs_run,
    "total_epochs": total,
    "train_loss": round(last_loss, 4),
    "val_acc": round(val_acc, 4),
    "best_acc": round(best_acc, 4),
}
with open(HIST, "a") as f:
    f.write(json.dumps(rec) + "\n")
history_runs = sum(1 for _ in open(HIST))

result = {
    **rec,
    "checkpoint_bytes": os.path.getsize(STATE),
    "history_runs": history_runs,
    "device": torch.cuda.get_device_name(0),
}
with open(os.path.join(OUT_DIR, "train.json"), "w") as f:
    json.dump(result, f, indent=2)
print(json.dumps(result, indent=2))
if resumed:
    print(f"RESUMED from volume: {prior_total} -> {total} cumulative epochs "
          f"(run #{history_runs}); val_acc {val_acc:.4f}")
else:
    print(f"FIRST run (fresh volume): trained {epochs_run} epochs; val_acc {val_acc:.4f}")