| |
| """ |
| Gradio control/monitoring UI for the TinyChess research Space. |
| |
| This is NOT the training implementation. Everything runs from the CLI; this UI |
| only launches, monitors and inspects background jobs, so a browser disconnect |
| never kills an experiment. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import signal |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
|
|
| |
| |
| |
| |
| |
| |
| |
| class _NoGPU: |
| @staticmethod |
| def GPU(*a, **k): |
| def deco(fn): |
| return fn |
| return deco(a[0]) if a and callable(a[0]) else deco |
| spaces = _NoGPU() |
| _HAS_SPACES = False |
|
|
| import gradio as gr |
|
|
| ROOT = Path(__file__).parent.resolve() |
| RUNS = ROOT / "runs" |
| RUNS.mkdir(exist_ok=True) |
| JOBS = RUNS / "jobs.json" |
| PY = sys.executable |
|
|
|
|
| |
| def load_jobs() -> dict: |
| if JOBS.exists(): |
| try: |
| return json.loads(JOBS.read_text()) |
| except Exception: |
| return {} |
| return {} |
|
|
|
|
| def save_jobs(j): |
| JOBS.write_text(json.dumps(j, indent=2)) |
|
|
|
|
| def job_alive(pid) -> bool: |
| try: |
| os.kill(int(pid), 0) |
| return True |
| except Exception: |
| return False |
|
|
|
|
| def _pid_is_ours(pid: int, cmd: list[str]) -> bool: |
| """True only if `pid` is alive AND its cmdline matches the recorded job. |
| |
| After a Space rebuild the PID space restarts, so a recorded pid can be |
| re-used by an unrelated process. Checking liveness alone made launch() |
| refuse valid jobs and made job_table() report dead jobs as RUNNING. |
| """ |
| if not job_alive(pid): |
| return False |
| try: |
| with open(f"/proc/{pid}/cmdline", "rb") as fh: |
| actual = fh.read().decode(errors="ignore").split("\x00") |
| except Exception: |
| return False |
| want = os.path.basename(cmd[1]) if len(cmd) > 1 else "" |
| return any(want and want in a for a in actual) |
|
|
|
|
| def launch(name: str, cmd: list[str]) -> str: |
| jobs = load_jobs() |
| prev = jobs.get(name) |
| if prev and _pid_is_ours(prev.get("pid", -1), prev.get("cmd", [])): |
| return f"Job '{name}' is already running (pid {prev['pid']})." |
| logp = RUNS / f"{name}.log" |
| with open(logp, "a") as lf: |
| lf.write(f"\n=== launch {time.strftime('%F %T')} ===\n{' '.join(cmd)}\n") |
| p = subprocess.Popen(cmd, stdout=lf, stderr=subprocess.STDOUT, cwd=ROOT, |
| start_new_session=True) |
| jobs[name] = {"pid": p.pid, "cmd": cmd, "started": time.time(), |
| "log": str(logp)} |
| save_jobs(jobs) |
| return f"Launched '{name}' (pid {p.pid})\n$ {' '.join(cmd)}" |
|
|
|
|
| def stop_job(name: str) -> str: |
| jobs = load_jobs() |
| if name not in jobs: |
| return f"No job '{name}'." |
| pid = jobs[name]["pid"] |
| try: |
| os.killpg(os.getpgid(pid), signal.SIGTERM) |
| except Exception as e: |
| return f"Could not stop {name}: {e}" |
| return f"Sent SIGTERM to '{name}' (pid {pid})." |
|
|
|
|
| def job_table() -> str: |
| jobs = load_jobs() |
| if not jobs: |
| return "No jobs yet." |
| rows = ["| job | pid | state | runtime |", "|---|---|---|---|"] |
| for n, j in sorted(jobs.items()): |
| alive = _pid_is_ours(j.get("pid", -1), j.get("cmd", [])) |
| rt = time.time() - j["started"] |
| rows.append(f"| {n} | {j['pid']} | {'RUNNING' if alive else 'stopped'} " |
| f"| {rt/60:.1f} min |") |
| return "\n".join(rows) |
|
|
|
|
| def tail_log(name: str, n: int = 60) -> str: |
| p = RUNS / f"{name}.log" |
| if not p.exists(): |
| return f"(no log for '{name}')" |
| lines = p.read_text(errors="replace").splitlines() |
| return "\n".join(lines[-n:]) |
|
|
|
|
| |
| def list_runs(): |
| return sorted([d.name for d in RUNS.iterdir() if d.is_dir()]) |
|
|
|
|
| def read_metrics(run: str): |
| p = RUNS / run / "metrics.jsonl" |
| if not p.exists(): |
| return [], "(no metrics.jsonl)" |
| rows = [] |
| for line in p.read_text(errors="replace").splitlines(): |
| try: |
| rows.append(json.loads(line)) |
| except Exception: |
| pass |
| iters = [r for r in rows if r.get("event") == "iter"] |
| return iters, f"{len(iters)} iterations logged" |
|
|
|
|
| def metrics_plot(run: str): |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| iters, msg = read_metrics(run) |
| if not iters: |
| return None, msg |
| fig, ax = plt.subplots(2, 2, figsize=(10, 6)) |
| it = [r["iter"] for r in iters] |
| ax[0, 0].plot(it, [r.get("sp_decisive_rate", 0) for r in iters]); ax[0, 0].set_title("decisive rate") |
| ax[0, 1].plot(it, [r.get("sp_mean_entropy", 0) for r in iters]); ax[0, 1].set_title("policy entropy") |
| ax[1, 0].plot(it, [r.get("tr_value_mae", 0) for r in iters]); ax[1, 0].set_title("value MAE") |
| ax[1, 1].plot(it, [r.get("sp_mean_plies", 0) for r in iters]); ax[1, 1].set_title("mean plies") |
| for a in ax.ravel(): |
| a.grid(alpha=.3); a.set_xlabel("iteration") |
| fig.tight_layout() |
| return fig, msg |
|
|
|
|
| def summary_table(run: str) -> str: |
| iters, msg = read_metrics(run) |
| if not iters: |
| return msg |
| keys = ["iter", "opponent", "sp_decisive_rate", "sp_mean_plies", |
| "sp_mean_entropy", "tr_loss", "tr_value_mae", "t_selfplay", "t_learn"] |
| rows = ["| " + " | ".join(k.replace("sp_", "").replace("tr_", "") for k in keys) + " |", |
| "|" + "---|" * len(keys)] |
| for r in iters[-25:]: |
| rows.append("| " + " | ".join( |
| (f"{r.get(k):.3f}" if isinstance(r.get(k), float) else str(r.get(k, ""))) |
| for k in keys) + " |") |
| return "\n".join(rows) |
|
|
|
|
| |
| def list_checkpoints(): |
| out = [] |
| for d in RUNS.iterdir(): |
| if d.is_dir(): |
| out += [str(p.relative_to(ROOT)) for p in sorted(d.glob("*.pt"))] |
| return out |
|
|
|
|
| def run_trace(ckpt: str, fen: str, steps: int, adaptive: bool): |
| if not ckpt: |
| return "Select a checkpoint.", "{}" |
| cmd = [PY, "inspect_trace.py", "--checkpoint", ckpt, "--fen", fen or |
| "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", |
| "--steps", str(steps), "--format", "both"] |
| if adaptive: |
| cmd.append("--adaptive") |
| r = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True, timeout=300) |
| return (r.stdout or "") + (("\n[stderr]\n" + r.stderr) if r.stderr else ""), "" |
|
|
|
|
| |
| def gpu_smoke_test(): |
| """Tiny GPU sanity check. NOT used for training (5-min cap).""" |
| import torch |
| from tinychess.config import baseline_config |
| from tinychess.model import TinyChess |
| lines = [f"torch {torch.__version__}", f"cuda available: {torch.cuda.is_available()}"] |
| if torch.cuda.is_available(): |
| dev = torch.device("cuda") |
| lines.append(f"device: {torch.cuda.get_device_name(0)}") |
| m = TinyChess(baseline_config()).to(dev) |
| rep = m.param_report() |
| lines.append(f"model params: {rep['total']:,}") |
| sq = torch.randint(0, 13, (8, 64), device=dev) |
| ex = torch.zeros(8, 64, 4, device=dev) |
| gl = torch.zeros(8, 20, device=dev) |
| cf = torch.zeros(8, 20, 7, dtype=torch.long, device=dev) |
| cm = torch.ones(8, 20, dtype=torch.bool, device=dev) |
| import time as _t |
| t = _t.time() |
| with torch.no_grad(): |
| out = m(sq, ex, gl, cf, cm, steps=8) |
| lines.append(f"forward(8 steps, B=8): {(_t.time()-t)*1000:.1f} ms") |
| lines.append(f"logits {tuple(out['logits'].shape)} value {tuple(out['value'].shape)}") |
| lines.append("GPU OK - but training stays on CPU by design.") |
| else: |
| lines.append("no CUDA visible") |
| return "\n".join(lines) |
|
|
|
|
| ALLOWED_CMDS = {"train_selfplay.py", "inspect_trace.py", "tests/test_phase1.py", |
| "scripts/bench_threads.py", "experiments/exp_ablation.py", |
| "experiments/exp_learning.py", "experiments/exp_compute_scaling.py", |
| "experiments/exp_faithfulness.py", "experiments/exp_precision.py", |
| "experiments/analyse_grid.py", "experiments/exp_step_embedding.py", |
| "experiments/exp_halting_routing.py", |
| "experiments/exp_mate_capability.py", |
| "experiments/exp_check_feature.py", |
| "experiments/exp_mate_readout.py", |
| "experiments/exp_check_intervention.py", |
| "experiments/exp_endgame_conversion.py", |
| "scripts/bench_parallel.py", |
| "scripts/bench_learn.py", |
| "experiments/exp_greedy_health.py"} |
|
|
|
|
| def run_job(name: str, script: str, argstr: str): |
| """Launch a whitelisted research script as a detached background job.""" |
| script = script.strip() |
| if script not in ALLOWED_CMDS: |
| return f"script '{script}' not in the allowed list: {sorted(ALLOWED_CMDS)}" |
| import shlex |
| args = shlex.split(argstr or "") |
| |
| |
| if script == "train_selfplay.py" and "--push-every" not in args: |
| args += ["--push-every", "5"] |
| cmd = [PY, script] + args |
| return launch(name.strip() or "job", cmd) |
|
|
|
|
| def read_file_tail(relpath: str, n: int = 200): |
| p = (ROOT / relpath.strip()).resolve() |
| if not str(p).startswith(str(ROOT)) or not p.exists(): |
| return f"(no such file: {relpath})" |
| if p.suffix in (".pt", ".pkl"): |
| return f"{relpath}: binary, {p.stat().st_size/1e6:.2f} MB" |
| lines = p.read_text(errors="replace").splitlines() |
| return "\n".join(lines[-int(n):]) |
|
|
|
|
| def list_files(pattern: str = "runs/**/*"): |
| import glob |
| out = [] |
| for f in sorted(glob.glob(str(ROOT / pattern), recursive=True))[:400]: |
| p = Path(f) |
| if p.is_file(): |
| out.append(f"{p.relative_to(ROOT)} {p.stat().st_size/1024:.1f} KB") |
| return "\n".join(out) or "(none)" |
|
|
|
|
| def env_audit(): |
| """Report the ACTUAL hardware of this Space (the research machine).""" |
| import multiprocessing, platform, subprocess, torch |
| L = [] |
| L.append(f"python {platform.python_version()} torch {torch.__version__}") |
| L.append(f"os.cpu_count(): {os.cpu_count()}") |
| try: |
| L.append(f"sched_getaffinity: {len(os.sched_getaffinity(0))}") |
| except Exception: |
| pass |
| L.append(f"torch threads: {torch.get_num_threads()}") |
| try: |
| out = subprocess.run(["lscpu"], capture_output=True, text=True, timeout=20).stdout |
| for k in ("Model name", "CPU(s):", "Thread(s) per core", "Core(s) per socket", |
| "Socket(s)", "L3 cache"): |
| for line in out.splitlines(): |
| if line.strip().startswith(k): |
| L.append(line.strip()); break |
| except Exception as e: |
| L.append(f"lscpu failed: {e}") |
| try: |
| mt = open("/proc/meminfo").readline().split() |
| L.append(f"MemTotal: {int(mt[1])/1024/1024:.1f} GB") |
| except Exception: |
| pass |
| try: |
| import shutil |
| du = shutil.disk_usage("/data" if os.path.isdir("/data") else ".") |
| L.append(f"disk free: {du.free/1e9:.1f} GB (persistent /data: {os.path.isdir('/data')})") |
| except Exception: |
| pass |
| L.append(f"cuda available: {torch.cuda.is_available()}") |
| L.append(f"BLAS: {torch.__config__.show().split('BLAS_INFO=')[-1].split(',')[0] if 'BLAS_INFO' in torch.__config__.show() else '?'}") |
| L.append(f"cpu capability: {torch.backends.cpu.get_cpu_capability()}") |
| return "\n".join(L) |
|
|
|
|
| |
| with gr.Blocks(title="TinyChess Research Console") as demo: |
| gr.Markdown("# TinyChess — ~100K-parameter self-play research console\n" |
| "Training runs as detached CLI jobs; this UI only controls and inspects them.") |
|
|
| with gr.Tab("Train"): |
| with gr.Row(): |
| cfg_name = gr.Dropdown(["baseline", "memory", "refine", "halting", |
| "router", "full", "mlp", "transformer"], |
| value="baseline", label="config") |
| job_name = gr.Textbox("sp_baseline", label="job / run name") |
| with gr.Row(): |
| iters = gr.Number(200, label="iterations", precision=0) |
| games = gr.Number(64, label="games / iter", precision=0) |
| steps = gr.Number(4, label="recurrent steps", precision=0) |
| threads = gr.Number(4, label="torch threads", precision=0) |
| with gr.Row(): |
| updates = gr.Number(24, label="grad steps / iter", precision=0) |
| batch = gr.Number(256, label="batch", precision=0) |
| lr = gr.Number(3e-4, label="lr") |
| seed = gr.Number(0, label="seed", precision=0) |
| with gr.Row(): |
| b_start = gr.Button("Start", variant="primary") |
| b_resume = gr.Button("Resume") |
| b_stop = gr.Button("Stop", variant="stop") |
| train_out = gr.Textbox(label="result", lines=4) |
|
|
| def _start(cfg, name, it, g, st, th, up, ba, lr_, sd, resume): |
| out = f"runs/{name}" |
| cmd = [PY, "train_selfplay.py", "--config", cfg, "--out", out, |
| "--iters", str(int(it)), "--games", str(int(g)), |
| "--steps", str(int(st)), "--threads", str(int(th)), |
| "--updates", str(int(up)), "--batch", str(int(ba)), |
| "--lr", str(lr_), "--seed", str(int(sd))] |
| if resume: |
| cmd += ["--resume", f"{out}/latest.pt"] |
| return launch(name, cmd) |
|
|
| b_start.click(lambda *a: _start(*a, False), |
| [cfg_name, job_name, iters, games, steps, threads, updates, |
| batch, lr, seed], train_out) |
| b_resume.click(lambda *a: _start(*a, True), |
| [cfg_name, job_name, iters, games, steps, threads, updates, |
| batch, lr, seed], train_out) |
| b_stop.click(stop_job, job_name, train_out) |
|
|
| with gr.Tab("Jobs"): |
| jt = gr.Markdown() |
| logname = gr.Textbox("sp_baseline", label="job") |
| logbox = gr.Textbox(label="log tail", lines=25, max_lines=25) |
| with gr.Row(): |
| gr.Button("Refresh jobs").click(job_table, None, jt) |
| gr.Button("Refresh log").click(lambda n: tail_log(n, 80), logname, logbox) |
| demo.load(job_table, None, jt) |
|
|
| with gr.Tab("Metrics"): |
| run_dd = gr.Dropdown(list_runs(), label="run") |
| with gr.Row(): |
| gr.Button("Reload runs").click(lambda: gr.update(choices=list_runs()), None, run_dd) |
| gr.Button("Plot").click(metrics_plot, run_dd, [gr.Plot(), gr.Markdown()]) |
| tbl = gr.Markdown() |
| gr.Button("Table").click(summary_table, run_dd, tbl) |
|
|
| with gr.Tab("Inspect trace"): |
| ck = gr.Dropdown(list_checkpoints(), label="checkpoint") |
| fen_in = gr.Textbox("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", |
| label="FEN") |
| with gr.Row(): |
| tsteps = gr.Number(8, label="steps", precision=0) |
| tad = gr.Checkbox(False, label="adaptive halting") |
| gr.Button("Reload checkpoints").click( |
| lambda: gr.update(choices=list_checkpoints()), None, ck) |
| trace_out = gr.Textbox(label="trace", lines=30) |
| gr.Button("Run trace", variant="primary").click( |
| run_trace, [ck, fen_in, tsteps, tad], [trace_out, gr.Textbox(visible=False)]) |
|
|
| with gr.Tab("Run job"): |
| gr.Markdown("Launch a research script as a detached background job on " |
| "the Space CPUs. Jobs survive browser disconnects.") |
| rj_name = gr.Textbox("job1", label="job name") |
| rj_script = gr.Textbox("train_selfplay.py", label="script") |
| rj_args = gr.Textbox("--config baseline --iters 10 --games 32 --threads 8 --out runs/test", |
| label="arguments") |
| rj_out = gr.Textbox(label="result", lines=4) |
| gr.Button("Launch", variant="primary").click(run_job, |
| [rj_name, rj_script, rj_args], rj_out) |
| gr.Markdown("### Inspect files") |
| f_pat = gr.Textbox("runs/**/*", label="glob") |
| f_list = gr.Textbox(label="files", lines=12) |
| gr.Button("List files").click(list_files, f_pat, f_list) |
| f_path = gr.Textbox("runs/test/metrics.jsonl", label="file") |
| f_n = gr.Number(100, label="tail lines", precision=0) |
| f_body = gr.Textbox(label="contents", lines=20) |
| gr.Button("Read file").click(read_file_tail, [f_path, f_n], f_body) |
|
|
| with gr.Tab("Environment"): |
| gr.Markdown("Actual hardware of this Space - the real research machine.") |
| env_out = gr.Textbox(label="environment", lines=18) |
| gr.Button("Audit environment", variant="primary").click(env_audit, None, env_out) |
|
|
| with gr.Tab("GPU check"): |
| gr.Markdown("ZeroGPU smoke test only. Training runs on CPU as detached " |
| "jobs because the GPU allocation is capped at ~5 minutes.") |
| gpu_out = gr.Textbox(label="gpu", lines=10) |
| gr.Button("Run GPU smoke test").click(gpu_smoke_test, None, gpu_out) |
|
|
| with gr.Tab("Evaluate"): |
| ck2 = gr.Dropdown(list_checkpoints(), label="checkpoint") |
| init2 = gr.Textbox("", label="random-init checkpoint (optional)") |
| egames = gr.Number(60, label="games per match", precision=0) |
| ejob = gr.Textbox("eval1", label="job name") |
| eout = gr.Textbox(label="result", lines=4) |
| gr.Button("Reload").click(lambda: gr.update(choices=list_checkpoints()), None, ck2) |
| gr.Button("Run evaluation", variant="primary").click( |
| lambda c, i, g, n: launch(n, [PY, "experiments/exp_learning.py", |
| "--checkpoint", c, "--games", str(int(g))] |
| + (["--init", i] if i else [])), |
| [ck2, init2, egames, ejob], eout) |
|
|
| if __name__ == "__main__": |
| demo.queue().launch(server_name="0.0.0.0", |
| server_port=int(os.environ.get("PORT", 7860)), |
| ssr_mode=False, show_api=False) |
|
|
