code/gpu/job1_main.py

"""HF Job 1 driver — audit + smoke (paper Wfe1iJocjF / FlashOptim).

Stages (each isolated; failures recorded, not fatal, so diagnostics persist):
  env        : GPU/torch/triton/driver report
  install    : authors' repo at pinned SHA via GitHub zip (no git in image)
  pytest     : authors' test suite (single-GPU subset; ddp/fsdp2/thirdparty excluded)
  equivalence: authors' Triton kernels vs independent flashsim (Claims 1, 6)
  mem_smoke  : tiny mem_probe runs (validates Job 2 stage A machinery)
  gpt_smoke  : 30-step tiny GPT arms ref+flash+linear on one real FineWeb shard
               (validates Job 2 stage B machinery + data path)

Writes /out/job1/results.json (+ logs). Exits 0 iff all stages ran (their
individual statuses are data).
"""

import io
import json
import os
import subprocess
import sys
import time
import urllib.request
import zipfile

SHA = "43548ceebbbdb869c4ea9db7b21bba1cf9ae7004"
OUT = os.environ.get("OUT_DIR", "/out/job1")
CODE = os.path.dirname(os.path.abspath(__file__))
RESULTS = {"stages": {}, "sha": SHA, "started": time.time()}


def stage(name):
    def deco(fn):
        def wrapper():
            t = time.time()
            try:
                out = fn()
                RESULTS["stages"][name] = {"status": "ok", "secs": round(time.time() - t, 1)}
                if out is not None:
                    RESULTS[name] = out
            except Exception as e:
                import traceback
                RESULTS["stages"][name] = {"status": "FAIL", "error": f"{type(e).__name__}: {e}",
                                           "trace": traceback.format_exc()[-2000:],
                                           "secs": round(time.time() - t, 1)}
                print(f"STAGE {name} FAILED: {e}", flush=True)
            _dump()
        return wrapper
    return deco


def _dump():
    os.makedirs(OUT, exist_ok=True)
    with open(f"{OUT}/results.json", "w") as f:
        json.dump(RESULTS, f, indent=2, default=str)


@stage("env")
def s_env():
    import torch
    import triton
    return dict(gpu=torch.cuda.get_device_name(0),
                capability=torch.cuda.get_device_capability(0),
                torch=torch.__version__, triton=triton.__version__,
                cuda=torch.version.cuda,
                driver=torch.cuda.driver_version if hasattr(torch.cuda, "driver_version") else None,
                python=sys.version.split()[0])


@stage("install")
def s_install():
    url = f"https://github.com/databricks/flashoptim/archive/{SHA}.zip"
    buf = urllib.request.urlopen(url, timeout=120).read()
    zipfile.ZipFile(io.BytesIO(buf)).extractall("/work")
    src = f"/work/flashoptim-{SHA}"
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--no-deps", src],
                   check=True)
    import flashoptim
    return dict(version=flashoptim.__version__, src=src)


@stage("pytest")
def s_pytest():
    src = f"/work/flashoptim-{SHA}"
    cmd = [sys.executable, "-m", "pytest", f"{src}/test", "-q",
           "--ignore", f"{src}/test/test_ddp.py",
           "--ignore", f"{src}/test/test_fsdp2.py",
           "-m", "not thirdparty", "-p", "no:cacheprovider",
           "--junitxml", f"{OUT}/pytest.xml"]
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=900)
    with open(f"{OUT}/pytest.log", "w") as f:
        f.write(p.stdout + "\n--- stderr ---\n" + p.stderr)
    tail = p.stdout.strip().splitlines()[-5:]
    return dict(exit=p.returncode, tail=tail)


@stage("equivalence")
def s_equiv():
    sys.path.insert(0, CODE)
    sys.path.insert(0, os.path.dirname(CODE))  # for flashsim package dir
    from equivalence import run_all
    r = run_all()
    with open(f"{OUT}/equivalence.json", "w") as f:
        json.dump(r, f, indent=2)
    return r


@stage("mem_smoke")
def s_mem():
    outs = []
    for variant in ["ref", "flash", "flash_gr", "split_only", "quant_only"]:
        p = subprocess.run([sys.executable, f"{CODE}/mem_probe.py",
                            "--variant", variant, "--model", "gpt-30m"],
                           capture_output=True, text=True, timeout=600)
        line = p.stdout.strip().splitlines()[-1] if p.stdout.strip() else ""
        outs.append(json.loads(line) if p.returncode == 0 and line.startswith("{")
                    else {"variant": variant, "exit": p.returncode,
                          "stderr": p.stderr[-800:]})
    with open(f"{OUT}/mem_smoke.json", "w") as f:
        json.dump(outs, f, indent=2)
    return outs


@stage("gpt_smoke")
def s_gpt():
    from huggingface_hub import hf_hub_download
    os.makedirs("/work/data", exist_ok=True)
    val = hf_hub_download("kjj0/fineweb10B-gpt2", "fineweb_val_000000.bin",
                          repo_type="dataset", local_dir="/work/data")
    outs = {}
    for arm in ["ref", "flash", "linear"]:
        p = subprocess.run([sys.executable, f"{CODE}/gpt2_train.py",
                            "--arm", arm, "--size", "gpt-30m", "--ctx", "256",
                            "--batch", "8", "--time-budget-s", "70",
                            "--max-steps", "40", "--compile", "0",
                            "--warmup", "10",
                            "--train-shard", val, "--val-shard", val,
                            "--out", f"{OUT}/gpt_smoke_{arm}"],
                           capture_output=True, text=True, timeout=600)
        tail = [ln for ln in p.stdout.splitlines() if ln.startswith("SUMMARY")]
        outs[arm] = (json.loads(tail[-1][8:]) if tail else
                     {"exit": p.returncode, "stderr": p.stderr[-800:]})
    return outs


if __name__ == "__main__":
    for fn in [s_env, s_install, s_pytest, s_equiv, s_mem, s_gpt]:
        fn()
    RESULTS["finished"] = time.time()
    _dump()
    print("JOB1_RESULTS " + json.dumps(RESULTS, default=str))
    bad_infra = any(v["status"] == "FAIL" for k, v in RESULTS["stages"].items()
                    if k in ("env", "install"))
    sys.exit(1 if bad_infra else 0)