Add isolated v5e hardware probe to NanoGPT TPU port

baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/README.md

@@ -18,6 +18,21 @@ The sibling [CUDA reference](../nanogpt_leaderboard_2026_10_06) remains byte-ver
 
 The table is **not shrunk**. Eight CPU row owners hold it in host RAM, exchange requested rows through Gloo, and upload bounded BF16 caches to their TPU chips. Dense gradients use XLA collectives. One PJRT process runs per chip (MPMD); SPMD is rejected. Dense optimizer state is replicated. The launcher requires **192 GiB available host RAM** and **160 GiB free export disk**; these guards do not prove HBM fit. `capacity.py` separately explains why the all-HBM approach fails on this machine.
 
+## Probe before preflight
+
+After installing the pinned environment with `setup_tpu.sh` and activating `.venv`, run:
+
+```bash
+export PJRT_DEVICE=TPU
+python experiment.py probe
+```
+
+The probe needs the real v5litepod-8 but no data shards. It does not allocate the 121 GiB n-gram table, build the canonical mask, or start WeightWatcher. It retains the real model's 262144-token buffers and replicated optimizer/tail state, trains twice at 128 and 49152 tokens per chip, releases optimizer/tail state, and evaluates once at 262144 tokens per chip. It uses synthetic full-softmax batches and a zero canonical mask; the loss is not benchmark NLL.
+
+Each rank writes `results/<probe-run>/probe-rank-XX.json`; rank 0 writes `PROBE_SUMMARY.json`, which the launcher finalizes after all workers exit. Reports contain hardware, package versions, host memory, HBM by stage, finite-loss checks, and step wall times plus XLA compile counts/seconds. The second step is labeled **steady**, but its first Adam update can compile an additional graph; it is not guaranteed compilation-free. Failures record the stage and exception, exit nonzero, and stop the ladder. The launcher stops the other ranks after a reported failure. A hard process kill cannot execute Python's `finally`; the summary records missing/incomplete ranks as failure.
+
+**The probe must be green on all eight ranks, including the 262144-token evaluation forward, before attempting preflight.** A successful probe is not qualification, not a 3.28 run, and not a reason to set `ready_to_train`. Preflight remains mandatory before `run`; a probe report never substitutes for its receipt. Qualification flags remain false. The probe has no 192 GiB host-RAM or 160 GiB export-disk gate, but still requires enough memory for the dense ladder and writable space for its reports. `--results-root` can select another directory.
+
 ## Run on the TPU VM
 
 Use Python 3.11 or 3.12 and `python3-venv` on the existing v5e-8 VM. No H100 Docker image, CUDA wheel or FlashAttention install is needed.
@@ -29,6 +44,7 @@ source .venv/bin/activate
 export PJRT_DEVICE=TPU
 python experiment.py verify
 python experiment.py check
+python experiment.py probe  # require a green PROBE_SUMMARY.json before proceeding
 python experiment.py prepare --data-root /data/nanogpt-leaderboard
 python experiment.py preflight --data-root /data/nanogpt-leaderboard
 ```

baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/experiment.py

@@ -3,11 +3,13 @@
 import datetime
 import hashlib
 import importlib.util
+import importlib.metadata
 import json
 import os
 from pathlib import Path
 import secrets
 import shutil
+import signal
 import subprocess
 import sys
 import time
@@ -70,11 +72,12 @@ def runtime_environment():
         raise RuntimeError('Production/preflight requires PJRT_DEVICE=TPU')
     os.environ['PJRT_DEVICE']='TPU'
     os.environ.setdefault('OMP_NUM_THREADS','4')
-    import torch, torch_xla
-    versions=[torch.__version__.split('+')[0],torch_xla.__version__.split('+')[0]]
+    torch_version=importlib.metadata.version('torch')
+    xla_version=importlib.metadata.version('torch_xla')
+    versions=[torch_version.split('+')[0],xla_version.split('+')[0]]
     if versions!=['2.9.0','2.9.0']:
         raise RuntimeError('Use matching torch==2.9.0 and torch_xla==2.9.0: '+repr(versions))
-    return {'torch':torch.__version__,'torch_xla':torch_xla.__version__,
+    return {'torch':torch_version,'torch_xla':xla_version,
             'python':sys.version,'host_available_bytes':available_host_memory()}
 
 
@@ -89,16 +92,40 @@ def tracking_environment():
     return env
 
 
+def probe_subprocess(cmd,root):
+    # Poll failure reports even when compilation emits no stdout. On an OOM,
+    # stop the entire launch process group, including ranks blocked in a collective.
+    with (root/'console.log').open('w') as log:
+        process=subprocess.Popen(cmd,stdout=log,stderr=subprocess.STDOUT,start_new_session=True)
+        try:
+            while process.poll() is None:
+                for path in root.glob('probe-rank-*.json'):
+                    report=json.loads(path.read_text())
+                    if report.get('status')=='failed':
+                        raise RuntimeError(f"Probe rank {report['rank']} failed at {report['failed_stage']}: {report['exception']}")
+                try: process.wait(timeout=1)
+                except subprocess.TimeoutExpired: pass
+            return process.returncode
+        finally:
+            # Kill remaining descendants too if a rank or the launch parent failed.
+            if process.poll() is None or process.returncode:
+                try: os.killpg(process.pid,signal.SIGTERM)
+                except ProcessLookupError: pass
+                try: process.wait(timeout=10)
+                except subprocess.TimeoutExpired:
+                    os.killpg(process.pid,signal.SIGKILL); process.wait()
+
+
 def launch(action,data_root,results_root,receipt=None):
     verify_reference()
     results_root=Path(results_root).resolve(); results_root.mkdir(parents=True,exist_ok=True)
     needed=(160 if action=='run' else 2 if action=='preflight' else 0)*GIB
     if shutil.disk_usage(results_root).free<needed:
         raise RuntimeError(f'Need at least {needed//GIB} GiB free results disk; no TPU allocation made')
     environment=runtime_environment()
-    if action!='check' and environment['host_available_bytes']<192*GIB:
+    if action not in ('check','probe') and environment['host_available_bytes']<192*GIB:
         raise RuntimeError('Need at least 192 GiB available host RAM for full table, data and CPU buffers; no allocation made')
-    data=None if action=='check' else data_manifest(data_root)
+    data=None if action in ('check','probe') else data_manifest(data_root)
     fingerprint=source_fingerprint()
     if action=='run':
         if receipt is None:
@@ -127,25 +154,32 @@ def launch(action,data_root,results_root,receipt=None):
     (root/'run_manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
     tracking_process=None; tracking_log=None
     try:
-        if action!='check':
+        if action not in ('check','probe'):
             # CPU-only subprocess: analysis cannot mutate training tensors or consume training RNG.
             subprocess.run(tracking_command('check',root),cwd=HERE,env=tracking_environment(),check=True)
             tracking_log=(root/'weightwatcher.log').open('w')
             tracking_process=subprocess.Popen(tracking_command('watch',root),cwd=HERE,
                 env=tracking_environment(),stdout=tracking_log,stderr=subprocess.STDOUT)
         # Execute in a fresh interpreter: libtpu must not be initialized in the spawn parent.
         cmd=[sys.executable,str(HERE/'experiment.py'),'_worker-launch','--options',json.dumps(options)]
-        with (root/'console.log').open('w') as log:
-            process=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,text=True,bufsize=1)
-            try:
-                for line in process.stdout:
-                    print(line,end='',flush=True); log.write(line); log.flush()
-                    if tracking_process is not None and tracking_process.poll() is not None:
-                        raise RuntimeError('WeightWatcher worker exited early; inspect '+str(root/'weightwatcher.log'))
-                code=process.wait()
-            except BaseException:
-                process.terminate(); process.wait(); raise
+        if action=='probe':
+            code=probe_subprocess(cmd,root)
+        else:
+            with (root/'console.log').open('w') as log:
+                process=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,text=True,bufsize=1)
+                try:
+                    for line in process.stdout:
+                        print(line,end='',flush=True); log.write(line); log.flush()
+                        if tracking_process is not None and tracking_process.poll() is not None:
+                            raise RuntimeError('WeightWatcher worker exited early; inspect '+str(root/'weightwatcher.log'))
+                    code=process.wait()
+                except BaseException:
+                    process.terminate(); process.wait(); raise
         if code: raise RuntimeError('TPU subprocess failed; inspect '+str(root/'console.log'))
+        if action=='probe':
+            from tpu_port.probe_report import summarize
+            if summarize(root,final=True)['status']!='complete':
+                raise RuntimeError('Probe did not complete all eight ranks, including full-size evaluation')
         if tracking_process is not None:
             (root/'tracking/TRAINING_DONE').touch()
             while tracking_process.poll() is None:
@@ -172,14 +206,17 @@ def launch(action,data_root,results_root,receipt=None):
         (root/'run_manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
         raise
     finally:
+        if action=='probe':
+            from tpu_port.probe_report import summarize
+            summarize(root,final=True,exception=manifest.get('error'))
         if tracking_process is not None and tracking_process.poll() is None:
             tracking_process.terminate(); tracking_process.wait()
         if tracking_log is not None: tracking_log.close()
 
 
 def main():
     parser=argparse.ArgumentParser(description=__doc__)
-    parser.add_argument('action',choices=['plan','verify','prepare','check','preflight','run','cpu-smoke','_worker-launch'])
+    parser.add_argument('action',choices=['plan','verify','prepare','check','probe','preflight','run','cpu-smoke','_worker-launch'])
     parser.add_argument('--data-root',type=Path,default=HERE/'cache')
     parser.add_argument('--results-root',type=Path,default=HERE/'results')
     parser.add_argument('--preflight-receipt',type=Path)

baseline/experiments/nanogpt_leaderboard_2026_10_06_tpu/tests/test_probe.py

@@ -0,0 +1,139 @@
+"""CPU-only hardware-probe routing, synthetic batches and failure reporting. No TPU launch."""
+import builtins
+import json
+from pathlib import Path
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+HERE=Path(__file__).resolve().parents[1]
+sys.path.insert(0,str(HERE))
+import experiment
+from tpu_port.probe import synthetic_batch
+from tpu_port.probe_report import FLAGS, STEP_STAGES, ProbeReport, summarize
+
+
+@pytest.mark.parametrize('tokens',[16,49152])
+def test_synthetic_batch_shapes_and_bounds(tokens):
+    batch=synthetic_batch(tokens)
+    assert batch.inputs.shape==batch.targets.shape==(tokens,)
+    assert batch.inputs.dtype==torch.int32 and batch.targets.dtype==torch.int64
+    assert batch.cache.shape==batch.sink.shape==(2*tokens,768)
+    assert batch.cache.dtype==batch.sink.dtype==torch.bfloat16
+    assert batch.sink.is_leaf and batch.sink.requires_grad
+    assert batch.slots.shape==(2*tokens,)
+    assert int(batch.slots.min())>=0 and int(batch.slots.max())<batch.cache.shape[0]
+    assert batch.seqlens[0]==0 and batch.seqlens[-1]==tokens
+    assert ((batch.seqlens>=0)&(batch.seqlens<=tokens)).all()
+    assert ((batch.seqlens.diff()>0)&(batch.seqlens.diff()<=2560)).all()
+    assert batch.inputs.min()>=0 and batch.inputs.max()<50257
+    assert batch.targets.min()>=0 and batch.targets.max()<50257
+    expected=batch.inputs.clone(); del batch
+    # Different process/model RNG does not change the synthetic batch.
+    torch.manual_seed(42)
+    evaluation=synthetic_batch(tokens,training=False)
+    assert evaluation.sink is None
+    assert torch.equal(evaluation.inputs,expected)
+
+
+def test_environment_versions_without_xla_import(monkeypatch):
+    for key in ('XLA_USE_SPMD','XLA_AUTO_SPMD','XLA_USE_BF16','XLA_DOWNCAST_BF16',
+                'NUM_SCHEDULED_ITERATIONS','TRAIN_SEED'):
+        monkeypatch.delenv(key,raising=False)
+    monkeypatch.setenv('PJRT_DEVICE','TPU')
+    monkeypatch.setattr(experiment.importlib.metadata,'version',lambda name:'2.9.0')
+    original=builtins.__import__
+    def guarded(name,*args,**kwargs):
+        if name=='torch_xla' or name.startswith('torch_xla.'):
+            pytest.fail('Launch grandparent imported torch_xla')
+        return original(name,*args,**kwargs)
+    monkeypatch.setattr(builtins,'__import__',guarded)
+    assert experiment.runtime_environment()['torch_xla']=='2.9.0'
+
+
+def complete_reports(root):
+    for rank in range(8):
+        value={'rank':rank,'status':'complete','finite_loss':True,
+               'stages':[{'stage':stage,'status':'complete','finite_loss':True} for stage in STEP_STAGES],**FLAGS}
+        (root/f'probe-rank-{rank:02d}.json').write_text(json.dumps(value))
+
+
+@pytest.mark.parametrize('action',['check','probe'])
+def test_probe_and_check_skip_data_tracking_and_capacity_gates(monkeypatch,tmp_path,action):
+    def forbidden(*args,**kwargs): pytest.fail('Data, tracking or extra subprocess requested')
+    monkeypatch.setattr(experiment,'verify_reference',lambda:None)
+    monkeypatch.setattr(experiment,'runtime_environment',lambda:{'torch':'2.9.0','torch_xla':'2.9.0','host_available_bytes':0})
+    monkeypatch.setattr(experiment,'source_fingerprint',lambda:{})
+    monkeypatch.setattr(experiment.shutil,'disk_usage',lambda path:SimpleNamespace(free=0))
+    monkeypatch.setattr(experiment,'data_manifest',forbidden)
+    monkeypatch.setattr(experiment,'tracking_command',forbidden)
+    monkeypatch.setattr(experiment.subprocess,'run',forbidden)
+    seen=[]
+    def fake_probe(cmd,root):
+        seen.append(json.loads(cmd[-1])['action']); complete_reports(root); return 0
+    monkeypatch.setattr(experiment,'probe_subprocess',fake_probe)
+    def fake_popen(cmd,**kwargs):
+        assert action=='check'
+        seen.append(json.loads(cmd[-1])['action'])
+        return SimpleNamespace(stdout=[],wait=lambda:0)
+    monkeypatch.setattr(experiment.subprocess,'Popen',fake_popen)
+    root=experiment.launch(action,tmp_path/'nonexistent-data',tmp_path/'results')
+    manifest=json.loads((root/'run_manifest.json').read_text())
+    assert manifest['data'] is None and manifest['status']=='complete'
+    assert seen==[action]
+    assert not (root/'PREFLIGHT_COMPLETE.json').exists()
+    assert 'weightwatcher_complete' not in manifest
+
+
+def test_probe_cli_routes_action(monkeypatch):
+    seen=[]
+    monkeypatch.setattr(sys,'argv',['experiment.py','probe'])
+    monkeypatch.setattr(experiment,'launch',lambda *args:seen.append(args[0]))
+    experiment.main()
+    assert seen==['probe']
+
+
+def test_worker_oom_records_failure_and_does_not_continue(monkeypatch,tmp_path):
+    from tpu_port import runtime
+    def fail(index,options,report):
+        report.begin('train_49152_compile')
+        raise torch.OutOfMemoryError('simulated CPU-only OOM')
+    monkeypatch.setattr(runtime,'_worker',fail)
+    with pytest.raises(torch.OutOfMemoryError):
+        runtime.worker(0,{'action':'probe','output':str(tmp_path)})
+    value=json.loads((tmp_path/'probe-rank-00.json').read_text())
+    assert value['status']=='failed' and value['failed_stage']=='train_49152_compile'
+    assert 'simulated CPU-only OOM' in value['exception']
+    assert all(value[key] is expected for key,expected in FLAGS.items())
+    assert summarize(tmp_path,final=True)['status']=='failed'
+
+
+def test_incomplete_workers_cannot_be_green(tmp_path):
+    report=ProbeReport(tmp_path,0)
+    report.begin('eval_262144_compile')
+    assert summarize(tmp_path,final=True)['status']=='failed'
+
+
+def test_full_eval_is_required_for_green_summary(tmp_path):
+    complete_reports(tmp_path)
+    assert summarize(tmp_path,final=True)['status']=='complete'
+    path=tmp_path/'probe-rank-07.json'
+    report=json.loads(path.read_text())
+    report['stages']=[s for s in report['stages'] if s['stage']!='eval_262144_compile']
+    path.write_text(json.dumps(report))
+    assert summarize(tmp_path,final=True)['status']=='failed'
+
+
+def test_failed_launch_finalizes_json_without_success(monkeypatch,tmp_path,capsys):
+    monkeypatch.setattr(experiment,'verify_reference',lambda:None)
+    monkeypatch.setattr(experiment,'runtime_environment',lambda:{'host_available_bytes':0})
+    monkeypatch.setattr(experiment,'source_fingerprint',lambda:{})
+    monkeypatch.setattr(experiment,'probe_subprocess',lambda *args:1)
+    with pytest.raises(RuntimeError,match='TPU subprocess failed'):
+        experiment.launch('probe',tmp_path/'no-data',tmp_path/'results')
+    root=next((tmp_path/'results').iterdir())
+    assert json.loads((root/'PROBE_SUMMARY.json').read_text())['status']=='failed'
+    assert json.loads((root/'run_manifest.json').read_text())['status']=='failed'
+    assert 'Results:' not in capsys.readouterr().out