phase4_session: plan() tiles to the horizon, kill the process group, df_free_gb helper, plan before the pull, session-relative TRAIN timeout

kernels/phase4_session.py

@@ -19,6 +19,7 @@ No credentials in this file: they are fetched at run time by ounce100m_credentia
 
 import json
 import os
+import signal
 import subprocess
 import sys
 import time
@@ -63,20 +64,33 @@ def sh(argv, label, timeout=None, env=None):
     # 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)
+                         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:
-            pass
+            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"))
@@ -84,6 +98,7 @@ def sh(argv, label, timeout=None, env=None):
     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 "
@@ -106,33 +121,57 @@ 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
+    """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)
-    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)
-    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)
+    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(seconds_per_ckpt / 3600.0, 3),
+            "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": n, "stop_after_steps": stop,
-            "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),
+            "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)
@@ -143,38 +182,30 @@ def main():
         "shard_dataset.py": ("train/shard_dataset.py",
                              "adbcd96fbab9505a5e8dec4d2e83d00366a73f702786bffbde430195cf6f0471"),
         "hubckpt.py": ("train/hubckpt.py",
-                       "c8c958417dc48db92f8be5c4b505581c746be9cc10110a358230ac47c480d5bf"),
+                       "568c31b300906cb8d78a59ef890baa9731b18e064e3d69b0eaea4b45c546f23e"),
         "train_ounce100m.py": ("train/train_ounce100m.py",
-                               "487edc1be54f03b5f1a312901580cc3e1504c412322ed074ceff869febbfb8bc"),
+                               "c4c178c0ac17c68de1521d85612a0dc604e6d4748dfdc8bd3e374512bc90ce96"),
     })
 
     import ounce100m_credentials
     print("creds:", json.dumps(ounce100m_credentials.install(verify=True)), flush=True)
-    rc, out = sh(["bash", "-c", "nvidia-smi --query-gpu=name,memory.used,memory.total "
-                     "--format=csv,noheader; df -k /kaggle/working | tail -1; free -g | head -2"],
-                 "env")
-    # `df -k` not `-h`: the human-readable form prints "318G"/"500M"/"1.5T" and a float() on those raises,
-    # which would take down the session before training even starts.
-    free_gb = None
-    for line in out.splitlines():
-        if "overlay" in line or line.startswith("/dev/"):
-            parts = line.split()
-            if len(parts) >= 4:
-                try:
-                    free_gb = float(parts[3]) / 1048576.0
-                except ValueError:
-                    free_gb = None
+    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 output -- refusing to guess whether the mix and two checkpoints will fit", flush=True)
+        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:
-        print(f"VERDICT REFUSED_TO_START: {free_gb:.1f} GB free on /kaggle/working, this run needs at "
-              f"least {MIN_FREE_GB} GB (2.3 GB mix + ~1.7 GB per checkpoint + the unpacked resume dir). "
-              "A Kaggle GPU instance starts with ~20 GB, so a low number here means something is left "
-              "over from a previous session on this box.", flush=True)
+    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", flush=True)
+    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
@@ -192,20 +223,40 @@ def main():
         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')\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:,} "
@@ -219,21 +270,13 @@ def main():
               "implicitly by launching anyway.", flush=True)
         raise SystemExit(9)
 
-    pl = plan(start_step, free_gb)
-    print("PLAN_JSON", json.dumps(pl, indent=1), flush=True)
-    if not pl["intervals_this_session"]:
-        print(f"VERDICT REFUSED_TO_START: {pl['usable_hours']} usable hours will not reach one "
-              f"{pl['hours_per_checkpoint_interval']} h checkpoint interval. 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)
     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).
         print("VERDICT PREP_ONLY_OK stop_after_steps", pl["stop_after_steps"],
-              "intervals", pl["intervals_this_session"], flush=True)
+              "planned_steps", pl["planned_steps"], flush=True)
         return
 
     argv = [sys.executable, "-u", "train_ounce100m.py",
@@ -245,15 +288,20 @@ def main():
             "--log-every", "20", "--val-tokens", "2000000",
             "--resume", "auto", "--stop-after-steps", str(pl["stop_after_steps"])]
     rc, out = sh(["torchrun", "--nproc_per_node=2"] + argv, "TRAIN",
-                 timeout=int((pl["intervals_this_session"] * pl["hours_per_checkpoint_interval"] + 1.5)
-                             * 3600))
+                 # 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)
-    print("VERDICT PHASE4_SESSION rc", rc, "steps", start_step, "->", pl["stop_after_steps"],
-          "tokens_at_stop", pl["tokens_at_stop"], flush=True)
-    sh(["bash", "-c", "df -h /kaggle/working | tail -1"], "disk_after")
+    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)