map-train-sweep — hyperparameter sweep via super map

# CLI
super map --sku gpu_1x_l4 -- python3 train.py --lr {lr} --steps 300 ::: lr=0.01,0.001

# SDK (what run_smoke.py does)
sweep = client.map("python3 train.py --lr {lr} --steps 300",
                 args={"lr": [0.01, 0.001]}, sku="gpu_1x_l4", code_dir=".")
sweep.wait()

"""Smoke driver: a 2-point LR sweep of train.py via the super map SDK.

Run locally (the smoke runner executes this on the CI box, not in a
session): submits the job, waits for the canary to gate the fan-out, waits
for completion, then asserts every task trained (loss halved, exit 0) and
reports the best learning rate.

Exit codes for tests/examples_smoke: 0 = PASS, 3 = BLOCKED (capacity /
queue timeout), anything else = FAIL. Prints a JSON summary as the last
line for the report's metrics.
"""
import json
import os
import pathlib
import sys
import time
import uuid

import supercomputer

SKU = "gpu_1x_l4"
HERE = pathlib.Path(__file__).resolve().parent
LRS = [0.01, 0.001]
# Soak knob (stress tier): fanned-out tasks keep training until this many
# seconds have elapsed. The canary (task 0) stays quick on purpose — the job
# controller only opens the fan-out once the canary SUCCEEDS, so a soaked
# canary would serialize into 2x the wall clock; the long-running sessions
# under test are the fanned-out ones. 0 = the normal quick sweep.
STRESS_DURATION_S = int(os.environ.get("SUPERCOMPUTER_STRESS_DURATION_S", "0") or "0")
# A task flips to succeeded the moment its process exits, but the host ships
# the log to the object store asynchronously; the S3 copy is only final once
# the task session's teardown drain closes it. Budget for that drain.
LOG_SETTLE_TIMEOUT_S = 180
ENDED_SESSION_STATES = ("closed", "dead", "failed", "not_found")


def durable_task_log(
    client: supercomputer.Client, task: dict
) -> supercomputer.ProcessLog:
    """Read a finished task's log once its upload has landed.

    Returns as soon as the output is complete ("sweep: OK" present) or the
    task session has ended (teardown drained the upload queue, so the copy
    is final); past the budget, returns the last read for the caller to
    judge. Never relaxes what is asserted — it only waits for durability.
    """
    deadline = time.monotonic() + LOG_SETTLE_TIMEOUT_S
    while True:
        log = client.session_process_logs(task["session_id"], task["process_id"])
        if ("sweep: OK" in log.output
                or log.session_status in ENDED_SESSION_STATES
                or time.monotonic() >= deadline):
            return log
        time.sleep(3)


def main() -> int:
    client = supercomputer.Client()
    name = f"smoke-train-sweep-{uuid.uuid4().hex[:6]}"
    print(f"JOB_NAME={name}", flush=True)  # harness parses this for cleanup

    job = client.map(
        "python3 train.py --lr {lr} --steps 300 --min-seconds {soak}",
        args=[{"lr": lr, "soak": 0 if i == 0 else STRESS_DURATION_S}
              for i, lr in enumerate(LRS)],
        name=name,
        sku=SKU,
        code_dir=str(HERE),
        max_parallel=2,
        retries=1,
    )
    print(f"submitted {job.job_id}; waiting on canary…", flush=True)
    try:
        state = job.wait_canary(timeout=1800, poll=5)
    except TimeoutError:
        print("canary never resolved (capacity?)", flush=True)
        job.cancel()
        return 3
    if state not in ("running", "succeeded"):
        print(f"canary failed: {json.dumps(job.status().get('clusters'))}",
              flush=True)
        return 1

    # The fan-out clock starts AFTER the canary gate (which already absorbed
    # one queue wait): a fresh queue margin + the soak itself + a run margin.
    try:
        final = job.wait(timeout=3600 + STRESS_DURATION_S, poll=10)
    except TimeoutError:
        status = job.status()
        counts = status.get("counts") or {}
        print(f"job exceeded its wait budget: counts={counts}", flush=True)
        print(json.dumps(status.get("clusters")), flush=True)
        if (counts.get("failed") or 0) > 0:
            job.cancel()
            return 1
        # Tasks merely queued/running: capacity ate the margin. The job is
        # healthy server-side, so leave it to finish — canceling would record
        # a false FAIL (the 0724 run's "timed-out" jobs later succeeded).
        return 3
    if final["job"]["state"] != "succeeded":
        print(f"job ended {final['job']['state']}: counts={final['counts']}",
              flush=True)
        return 1

    # Assert each point trained and pick the winner from the durable logs.
    best = None
    for t in final["tasks"]:
        log = durable_task_log(client, t)
        assert t["exit_code"] == 0, t
        assert "sweep: OK" in log.output, log.output[-500:]
        line = next(ln for ln in log.output.splitlines()
                    if ln.startswith("sweep: first="))
        loss = float(line.rsplit("final=", 1)[1])
        print(f"task {t['index']} (lr={LRS[t['index']]}): final loss {loss}",
              flush=True)
        if best is None or loss < best[1]:
            best = (LRS[t["index"]], loss)

    cost = (final.get("cost") or {}).get("total_usd_micros")
    print(json.dumps({
        "best_lr": best[0], "best_final_loss": best[1],
        "tasks": len(final["tasks"]),
        "compute_usd_micros": cost,
    }), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""One point of a learning-rate sweep: a small-but-real PyTorch training
run (synthetic regression, minibatch Adam, checkpoint saved). `super map`
runs this N times with different --lr values, each in its own GPU session.

The script has NO supercomputer-specific code — its CLI is the interface — but it
prints the task coordinates from SUPERCOMPUTER_TASK_ID/WORLD to show the env
contract map provides.
"""
import argparse
import os
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

parser = argparse.ArgumentParser()
parser.add_argument("--lr", type=float, required=True)
parser.add_argument("--steps", type=int, default=300)
parser.add_argument("--min-seconds", type=float, default=0,
                    help="soak: keep training extra steps until this much time has passed")
args = parser.parse_args()

assert torch.cuda.is_available(), "expected a CUDA device"
dev = torch.device("cuda")
print(f"sweep: device={torch.cuda.get_device_name(0)} torch={torch.__version__} "
      f"lr={args.lr} task={os.environ.get('SUPERCOMPUTER_TASK_ID')}/"
      f"{os.environ.get('SUPERCOMPUTER_WORLD_SIZE')}", flush=True)

torch.manual_seed(0)
X = torch.randn(4096, 64, device=dev)
w_true = torch.randn(64, 1, device=dev)
y = X @ w_true + 0.1 * torch.randn(4096, 1, device=dev)

model = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 1)).to(dev)
opt = torch.optim.Adam(model.parameters(), lr=args.lr)

first = None
loss = None
deadline = time.monotonic() + args.min_seconds
# Soak steps are ms-scale, so log sparsely past --steps or the durable
# process log balloons over a multi-hour run.
log_every = 5000 if args.min_seconds else 50
step = 0
while step < args.steps or time.monotonic() < deadline:
    idx = torch.randint(0, X.shape[0], (256,), device=dev)
    loss = F.mse_loss(model(X[idx]), y[idx])
    opt.zero_grad()
    loss.backward()
    opt.step()
    if first is None:
        first = loss.item()
    if step % log_every == 0:
        print(f"step {step:4d} loss {loss.item():.4f}", flush=True)
    step += 1

final = loss.item()
torch.save(model.state_dict(), f"model-lr{args.lr}.pt")
print(f"sweep: first={first:.4f} final={final:.4f}", flush=True)
assert final < first * 0.5, f"loss did not improve: {first} -> {final}"
print("sweep: OK", flush=True)