Upload kernels/phase4_session.py with huggingface_hub

kernels/phase4_session.py

@@ -24,11 +24,15 @@ 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). Anything here that a future session might want to change is a
-# change to the run itself, so it is asserted against the trainer's own numbers below rather than trusted.
+# 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
@@ -51,24 +55,38 @@ def sh(argv, label, timeout=None, env=None):
     e["PYTHONPATH"] = WORK
     e.update(env or {})
     t0 = time.time()
-    try:
-        p = subprocess.run(argv, cwd=WORK, capture_output=True, text=True, timeout=timeout, env=e)
-    except subprocess.TimeoutExpired as ex:
-        print(f"{label} TIMEOUT after {time.time() - t0:.0f}s", flush=True)
-        return -9, (ex.stdout or b"").decode()[-200000:], (ex.stderr or b"").decode()[-4000:]
-    out, err = p.stdout or "", p.stderr or ""
-    for line in out.splitlines():
+    # 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)
+    state = {"killed": False}
+
+    def _kill():
+        state["killed"] = True
+        try:
+            p.kill()
+        except Exception:
+            pass
+
+    timer = threading.Timer(timeout, _kill) if timeout else None
+    if timer:
+        timer.daemon = True
+        timer.start()
+    lines = []
+    for line in (p.stdout or []):
+        lines.append(line.rstrip("\n"))
         print("   |", line[:300], flush=True)
-    if p.returncode != 0:
-        lines = err.splitlines()
-        keep = []
-        for i, ln in enumerate(lines):            # first traceback, not a wrapper's tail (E-023)
-            if "Traceback" in ln or "Error" in ln:
-                keep = lines[max(0, i - 2):i + 30]
-                break
-        print("   STDERR_HEAD:\n" + "\n".join(keep or lines[:45])[-4000:], flush=True)
-    print(f"{label}_RC {p.returncode} seconds {round(time.time() - t0, 1)}", flush=True)
-    return p.returncode, out
+    rc = p.wait()
+    if timer:
+        timer.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):
@@ -88,19 +106,28 @@ def fetch(rev, want):
 def plan(start_step, free_gb):
     """Largest whole number of checkpoint intervals that fits this session's remaining hours."""
     seconds_per_ckpt = CKPT_EVERY * TOKENS_PER_STEP / PLANNING_TOK_PER_S
-    usable = min(SESSION_GPU_HOURS, max(0.0, QUOTA_LEFT_HOURS - RESERVE_HOURS))
+    # 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)
     n = int(usable * 3600.0 // seconds_per_ckpt)
     remaining_intervals = max(0, (HORIZON_STEPS - start_step + CKPT_EVERY - 1) // CKPT_EVERY)
     n = min(n, remaining_intervals)
-    stop = min(HORIZON_STEPS, start_step + n * CKPT_EVERY) if n else 0
+    target = start_step + n * CKPT_EVERY
+    # The trainer refuses a --stop-after-steps that is not a whole multiple of its own interval, and
+    # 3,814 is not one (3814 % 381 = 385). So the session that would cross the horizon is told to run
+    # to it instead -- otherwise the final leg of the run could never be launched.
+    stop = 0 if target >= HORIZON_STEPS else min((HORIZON_STEPS // CKPT_EVERY) * CKPT_EVERY, target)
     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(seconds_per_ckpt / 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": n, "stop_after_steps": stop,
-            "tokens_at_stop": stop * TOKENS_PER_STEP,
-            "pct_of_run_at_stop": round(100.0 * stop * TOKENS_PER_STEP / TOKENS, 2),
+            "runs_to_horizon_this_session": stop == 0 and n > 0,
+            "tokens_at_stop": (HORIZON_STEPS if stop == 0 else stop) * TOKENS_PER_STEP,
+            "pct_of_run_at_stop": round(100.0 * (HORIZON_STEPS if stop == 0 else stop)
+                                        * TOKENS_PER_STEP / TOKENS, 2),
             "free_gb_before_start": free_gb}