kernels/phase4_session.py

"""One Phase 4 session of the ounce100m main run: resume, train to a checkpoint boundary, stop.

Why this file exists as a *launcher with a planner* rather than a plain torchrun line. The run needs
35.9 GPU-hours and Kaggle gives 30 h/week on sessions whose cap I have not measured, so it will be
interrupted -- and an interruption that lands between two checkpoints throws away everything since the
last one. §2 fixes the checkpoint cadence at one per 10 % of tokens (381 steps, ~3.6 h at the measured
rate), and §3.1 forbids restarting the run once checkpoint 1 exists. The only lever left is *when a
session chooses to stop*, so this script computes the largest whole number of checkpoint intervals that
fits the session's remaining GPU hours and passes that as --stop-after-steps. Every session then ends
sitting on a pushed, Hub-verified, locally-pruned checkpoint, and no billed work is ever discarded.

Failure modes are deliberately loud: if the budget cannot cover one checkpoint interval it refuses to
start rather than burning up to 3.6 hours of quota on work that cannot be recovered, and if the Hub
pointer is unreadable for any reason other than "no checkpoint yet" it aborts (a resume that silently
starts from step 0 would train the same data twice and destroy the run's honesty).

No credentials in this file: they are fetched at run time by ounce100m_credentials (D-006).
"""

import json
import os
import signal
import subprocess
import sys
import time

WORK = "/kaggle/working"
T0 = time.time()          # session start, so every step log line can say how far into the run it is
MIX_REPO = "Cion-lab/ounce100m-mix-v1"
CKPT_REPO = "Cion-lab/ounce100m-ckpt"

# The frozen recipe (D-011, D-012, D-015). TOKENS_PER_STEP and HORIZON_STEPS below are re-derived here
# with the same integer arithmetic the trainer uses (`seqs_per_step = micro*world*accum`,
# `steps_planned = int(tokens/tokens_per_step)` in train_ounce100m.py), so the two agree by construction
# rather than by copy -- if that formula ever changes in the trainer, change it here too and the mix-size
# check below will disagree.
SEQ_LEN = 1024
MICRO_BATCH = 4
ACCUM = 32
WORLD = 2
TOKENS = 1_000_000_000
TOKENS_PER_STEP = MICRO_BATCH * SEQ_LEN * ACCUM * WORLD      # 262,144
HORIZON_STEPS = int(TOKENS / TOKENS_PER_STEP)                 # 3,814 -- the trainer's own int() division
CKPT_EVERY = max(1, HORIZON_STEPS // 10)                      # 381 -- §2's one-per-10 % cadence

# Planning rate: the defaults below are P3's measured frozen-geometry figures, so a session launched with
# no environment set still plans at the rate the gate measured rather than the pre-Gate 3 estimate.
# D-018's planning rate: the *sustained* measurement from the two probes (12,221 and 12,312 tok/s over
# 180 and 120 steps respectively), not the 20-step progress-bar figure of 12,977. Planning on the optimistic
# number is how a session ends up killed by its own ceiling with an interval still to push.
PLANNING_TOK_PER_S = float(os.environ.get("PLANNING_TOK_PER_S", "12300"))
SESSION_GPU_HOURS = float(os.environ.get("SESSION_GPU_HOURS", "6.9"))
QUOTA_LEFT_HOURS = float(os.environ.get("QUOTA_LEFT_HOURS", "30.0"))
RESERVE_HOURS = float(os.environ.get("RESERVE_HOURS", "0.6"))  # startup, val eval, the last push
MIN_FREE_GB = float(os.environ.get("MIN_FREE_GB", "8"))       # mix 2.3 GB + two checkpoints + slack


def sh(argv, label, timeout=None, env=None):
    print(f"=== {label}", flush=True)
    e = dict(os.environ)
    e["PYTHONPATH"] = WORK
    e["PYTHONUNBUFFERED"] = "1"   # the -u torchrun cannot carry anymore
    e.update(env or {})
    t0 = time.time()
    # Streamed, not buffered: capture_output=True would hold an eleven-hour session's loss curve until
    # the child exited, so nobody could see the run going wrong while it was still running (a review
    # finding; §5 requires monitoring loss on every wake). stderr is merged into stdout so a torchrun
    # child's traceback arrives in the same stream instead of a separate tail that hides it (E-023).
    import threading
    p = subprocess.Popen(argv, cwd=WORK, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                         text=True, bufsize=1, env=e, start_new_session=True)
    state = {"killed": False}

    def _kill():
        # torchrun is not the process tree. SIGKILLing only the child leaves both ranks alive on the two
        # T4s, still stepping and still rolling latest.json, while the next session is launched into the
        # same container -- so the whole group goes, politely first (review finding).
        state["killed"] = True
        try:
            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
        except Exception:
            p.kill()

    def _hardkill():
        try:
            os.killpg(os.getpgid(p.pid), signal.SIGKILL)
        except Exception:
            p.kill()
        p.kill()

    timer = threading.Timer(timeout, _kill) if timeout else None
    if timer:
        timer.daemon = True
        timer.start()
        hard = threading.Timer(timeout + 120, _hardkill)
        hard.daemon = True
        hard.start()
    lines = []
    for line in (p.stdout or []):
        lines.append(line.rstrip("\n"))
        print("   |", line[:300], flush=True)
    rc = p.wait()
    if timer:
        timer.cancel()
        hard.cancel()
    if state["killed"]:
        rc = -9
    print(f"{label}_RC {rc}{' TIMEOUT' if state['killed'] else ''} seconds "
          f"{round(time.time() - t0, 1)} elapsed {round(time.time() - T0)}", flush=True)
    return rc, "\n".join(lines)[-400000:]


def fetch(rev, want):
    """Download the run's code at a pinned revision and assert every sha256 before executing any of it."""
    base = f"https://huggingface.co/Cion-lab/ounce100m-code/resolve/{rev}"
    import hashlib
    for local, (repo_path, sha) in sorted(want.items()):
        rc = subprocess.run(["curl", "-sfL", f"{base}/{repo_path}", "-o", local],
                            capture_output=True, text=True).returncode
        assert rc == 0, f"fetch failed: {repo_path} at {rev}"
        got = hashlib.sha256(open(os.path.join(WORK, local), "rb").read()).hexdigest()
        assert got == sha, f"SHA MISMATCH {repo_path}: hub {got[:16]} != expected {sha[:16]}"
        print("OK", local, got[:12], flush=True)
    print("REV_OK", rev, flush=True)


def plan(start_step, free_gb):
    """How far this session can train and still end on a checkpoint the Hub has verified.

    The segment has to stop where a checkpoint is pushed: the 381-step grid, or the horizon. Running to the
    horizon is preferred whenever it fits, because 3,814 is not a multiple of 381 -- stopping at the last
    grid point leaves four steps orphaned and demands one more session whose only budget test is "can you
    afford a whole 2.97 h interval", which the run could fail forever while its last seven minutes of work
    stayed unstarted (review finding).
    """
    sec_per_step = TOKENS_PER_STEP / PLANNING_TOK_PER_S
    # The reserve comes off the total, not off one term: `min(session, quota - reserve)` let an 11 h
    # session plan 11 h of training and leave nothing for the code fetch, the 2.3 GB mix download, the
    # push/verify cycles and the end-of-run evaluation (review finding).
    usable = max(0.0, min(SESSION_GPU_HOURS, QUOTA_LEFT_HOURS) - RESERVE_HOURS)
    steps_that_fit = int(usable * 3600.0 // sec_per_step)
    remaining = HORIZON_STEPS - start_step
    if steps_that_fit >= remaining:
        stop, planned = 0, remaining                      # the trainer runs to its own horizon
    else:
        n = max(0, steps_that_fit // CKPT_EVERY)
        stop, planned = start_step + n * CKPT_EVERY, n * CKPT_EVERY
    end = HORIZON_STEPS if stop == 0 else stop
    return {"start_step": start_step, "horizon_steps": HORIZON_STEPS, "ckpt_every": CKPT_EVERY,
            "tokens_per_step": TOKENS_PER_STEP, "planning_tok_per_s": PLANNING_TOK_PER_S,
            "hours_per_checkpoint_interval": round(CKPT_EVERY * sec_per_step / 3600.0, 3),
            "session_gpu_hours": SESSION_GPU_HOURS, "quota_left_hours": QUOTA_LEFT_HOURS,
            "reserve_hours": RESERVE_HOURS, "usable_hours": round(usable, 3),
            "intervals_this_session": max(1, planned // CKPT_EVERY) if planned else 0,
            "planned_steps": planned, "stop_after_steps": stop,
            "runs_to_horizon_this_session": stop == 0 and planned > 0,
            "tokens_at_stop": end * TOKENS_PER_STEP,
            "pct_of_run_at_stop": round(100.0 * end * TOKENS_PER_STEP / TOKENS, 2),
            "free_gb_before_start": free_gb}


def df_free_gb(path=WORK):
    """Free space on the working volume, in GB, or None if `df` could not be read.

    `df -k`, not `-h`: the human-readable form prints "318G"/"500M"/"1.5T" and float() on those raises,
    which would take the session down before training started.
    """
    rc, out = sh(["df", "-k", path], "df", timeout=60)
    for line in out.splitlines():
        parts = line.split()
        if len(parts) >= 4 and ("overlay" in line or line.startswith("/dev/")):
            try:
                return float(parts[3]) / 1048576.0
            except ValueError:
                return None
    return None


def main():
    os.chdir(WORK)
    sys.path.insert(0, WORK)
    rev = os.environ.get("OUNCE100M_REV") or sys.argv[1]
    fetch(rev, {
        "ounce100m_credentials.py": ("ounce100m_credentials.py",
                                     "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
        "shard_dataset.py": ("train/shard_dataset.py",
                             "f35653bf4c8f2cfe7bb0c2c7835e308505fb5a84b4eb002d0f82f2cada768ca6"),
        "hubckpt.py": ("train/hubckpt.py",
                       "d4b50ed0928c678c94ce6764612a4655f2959c0a7b91d6fd588e8efddff298b2"),
        "train_ounce100m.py": ("train/train_ounce100m.py",
                               "dcb0ac199c0616575423ddea0d6b76d2257c06089d3f73aae0e0359de958ff03"),
    })

    import ounce100m_credentials
    print("creds:", json.dumps(ounce100m_credentials.install(verify=True)), flush=True)
    sh(["bash", "-c", "nvidia-smi --query-gpu=name,memory.used,memory.total "
          "--format=csv,noheader; free -g | head -2"], "env")
    # The mix is 2.3 GB and it lands during this session, so the pre-download floor is the run's floor plus
    # the pull. Measuring again after it, rather than trusting the arithmetic, is the point.
    free_gb = df_free_gb()
    if free_gb is None:
        print("VERDICT REFUSED_TO_START: could not read the free space on /kaggle/working from `df` -- "
              "refusing to guess whether the mix and two checkpoints will fit", flush=True)
        raise SystemExit(10)
    if free_gb < MIN_FREE_GB + 2.3:
        print(f"VERDICT REFUSED_TO_START: {free_gb:.1f} GB free on /kaggle/working before the pull; this "
              f"run needs {MIN_FREE_GB} GB free once the 2.3 GB mix has landed (~1.7 GB per checkpoint and "
              "the two unpacked resume directories, on top of it). A Kaggle instance starts near 19.5 GB, "
              "so a low number here means something is left over from a previous session on this box.",
              flush=True)
        raise SystemExit(10)
    print(f"disk: {free_gb:.1f} GB free before the mix lands", flush=True)

    # The Hub pointer is the only truth about where the run is (never local state -- §3.13).
    import hubckpt
    tok = os.environ["HF_TOKEN"]
    ptr = hubckpt.latest_pointer(CKPT_REPO, repo_type="dataset", token=tok)
    if ptr.get("error"):
        print(f"FATAL: cannot read latest.json from {CKPT_REPO}: {ptr['error'][:300]}\n"
              "Refusing to guess a start step: the alternative is training data that was already read.",
              flush=True)
        raise SystemExit(6)
    start_step = int(ptr.get("step", 0)) if not ptr.get("not_found") else 0
    print(f"hub says: {'no checkpoint yet -- starting from step 0' if start_step == 0 else ptr}",
          flush=True)
    if start_step >= HORIZON_STEPS:
        print("VERDICT RUN_ALREADY_COMPLETE at step", start_step, flush=True)
        return

    # Plan before downloading anything: an unusable budget should refuse in seconds, not after pulling
    # 2.3 GB onto a billed GPU session.
    pl = plan(start_step, free_gb)
    print("PLAN_JSON", json.dumps(pl, indent=1), flush=True)
    if not pl["planned_steps"]:
        print(f"VERDICT REFUSED_TO_START: {pl['usable_hours']} usable hours will not reach one "
              f"{pl['hours_per_checkpoint_interval']} h checkpoint interval ({pl['start_step']} → "
              f"{HORIZON_STEPS} left). Launching anyway would burn quota on work that cannot be resumed "
              "from. Raise SESSION_GPU_HOURS or wait for quota.", flush=True)
        raise SystemExit(8)

    # Data: the published mix, fetched by anonymous read so this works on any instance. The checks run
    # here rather than inside the downloader so a quoting mistake cannot turn into a silent pass.
    rc, _ = sh([sys.executable, "-c",
                "import sys; sys.path.insert(0, '/kaggle/working');\n"
                "from huggingface_hub import snapshot_download\n"
                f"p = snapshot_download(repo_id='{MIX_REPO}', repo_type='dataset',\n"
                "                        local_dir='/kaggle/working/mixroot', max_workers=4)\n"
                "print('downloaded to', p)\n"],
               "fetch_mix", 3600)
    if rc != 0:
        raise SystemExit(7)
    free_gb = df_free_gb()
    if free_gb is not None and free_gb < MIN_FREE_GB:
        print(f"VERDICT REFUSED_TO_START: only {free_gb:.1f} GB free after the mix landed, below the "
              f"{MIN_FREE_GB} GB the checkpoints and the two resume directories need", flush=True)
        raise SystemExit(10)
    man = json.load(open(os.path.join(WORK, "mixroot", "manifest.json")))
    windows = int(man["total_tokens"]) // SEQ_LEN
    needed_windows = HORIZON_STEPS * MICRO_BATCH * ACCUM * WORLD
    if len(man.get("shards") or []) != int(man.get("n_shards") or -1):
        print(f"VERDICT REFUSED_TO_START: the published manifest lists {len(man.get('shards') or [])} "
              f"shard records but claims n_shards={man.get('n_shards')}", flush=True)
        raise SystemExit(9)
    print(f"mix: {man['total_tokens']:,} train tokens in {man['n_shards']} shards, "
          f">= {man['distinct_sources_per_shard_min']} sources per shard, "
          f"{windows:,} windows of {SEQ_LEN}; the horizon needs {needed_windows:,} "
          f"({100.0 * windows / needed_windows - 100:.1f} % margin)", flush=True)
    if int(man["total_tokens"]) < TOKENS or windows < needed_windows:
        # Training past the end of the dataset would end the epoch early and report a token count that
        # was never reached. The token target is a claim in the final report, so it is checked here.
        print(f"VERDICT REFUSED_TO_START: mix holds {man['total_tokens']:,} tokens / {windows:,} "
              f"windows but the run needs {TOKENS:,} / {needed_windows:,}. Either the target drops to "
              "0.9 B (D-012's fallback) or more sources are staged -- not a decision to make "
              "implicitly by launching anyway.", flush=True)
        raise SystemExit(9)

    argv = ["train_ounce100m.py",
            "--root", f"{WORK}/mixroot", "--out", f"{WORK}/run",
            "--hub-repo", CKPT_REPO, "--prune",
            "--seq-len", str(SEQ_LEN), "--attn", "eager",
            # D-018, from probe 2 (soak): 180 steps without checkpointing at the frozen geometry held at
            # 12.84 GiB reserved cross-rank with no drift between step 10 and step 120, and the two probes'
            # sustained rates (12,221-12,312 tok/s) are +27 % over the checkpointed 9,358. Gradient
            # checkpointing computes identical gradients, so this changes no frozen hyperparameter -- and
            # the checkpointed rate does not fit the week's quota at all (D-017 addendum, plan sweep).
            "--no-grad-ckpt",
            "--micro-batch", str(MICRO_BATCH), "--accum", str(ACCUM),
            "--tokens", str(TOKENS), "--lr", "6e-4", "--warmup-frac", "0.02", "--decay-frac", "0.80",
            # Every 127 steps: three pushes per 762-step session, and 127 divides 762 so the plan's tiling
            # and the trainer's whole-multiple guard stay satisfied. D-017-reopened pre-registered the
            # finer cadence and E-041 killed it at 797 s a push; the bounded read-back made a push cycle
            # 13-40 s, which buys a worst-case loss of 45 minutes instead of 2.3 hours.
            "--push-every-steps", "127",
            "--log-every", "20", "--val-tokens", "2000000",
            "--resume", "auto", "--stop-after-steps", str(pl["stop_after_steps"])]

    if os.environ.get("PHASE4_PREP_ONLY") == "1":
        # Everything a session can get wrong before it bills GPU time has now been checked: the pinned
        # code hashes, the credentials resolve, the checkpoint pointer is readable, the published mix is
        # downloadable and large enough, and the plan lands on a checkpoint boundary. A CPU instance can
        # run all of that for free (E-031's arithmetic is exactly the kind of thing to find there).
        #
        # `TRAIN_ARGV` is handed over for the same reason: the argument list below is the one thing on the
        # session's critical path that PREP_ONLY otherwise never touches (E-037 was a torchrun line that
        # had never run, and the rehearsal could not see it because the line lived after the return). The
        # rehearsal parses these exact flags against the pinned trainer instead of trusting they match.
        print("TRAIN_ARGV " + json.dumps(argv), flush=True)
        print("VERDICT PREP_ONLY_OK stop_after_steps", pl["stop_after_steps"],
              "planned_steps", pl["planned_steps"], flush=True)
        return

    rc, out = sh(["torchrun", "--nproc_per_node=2"] + argv, "TRAIN",
                 # Counted from the session's own start, not the child's: the ceiling is on the container,
                 # and by now the code fetch and the 2.3 GB pull have already spent part of the reserve.
                 # Killing at ceiling-minus-15-min takes the whole process group with it (see sh) instead
                 # of being caught mid-checkpoint by a platform kill that leaves both ranks running.
                 timeout=max(600, int(SESSION_GPU_HOURS * 3600 - (time.time() - T0) - 900)))
    for line in out.splitlines():
        if line.startswith(("CKPT ", "RUN_JSON ", "resume from", "auto-resume", "precision:",
                            "params:", "validation loss", "segment boundary")):
            print("KEY>", line[:400], flush=True)
    end_step = HORIZON_STEPS if pl["stop_after_steps"] == 0 else pl["stop_after_steps"]
    print("VERDICT PHASE4_SESSION rc", rc, "steps", start_step, "->", end_step,
          "tokens_at_stop", pl["tokens_at_stop"], "session_seconds",
          round(time.time() - T0, 1), flush=True)
    print(f"disk at exit: {df_free_gb()} GB free", flush=True)
    # The session's own wall clock is the measurement that sharpens every later plan.
    raise SystemExit(rc)


if __name__ == "__main__":
    main()