Generate the calibration scripts instead of keeping one by hand

.gitignore

@@ -21,6 +21,8 @@ data
 *.out
 tasks/tests/logs
 
-# generated by slurm/generate_slurm.py - edit the generator, not these
+# generated by slurm/generate_slurm.py and slurm/generate_calibration.py -
+# edit the generators, not these
 slurm/*_experiment.sh
 slurm/crosstask.sh
+slurm/*_calibrate.sh

README.md

@@ -82,7 +82,17 @@ sbatch slurm/crosstask.sh             # every dataset's 3 intrinsic arms, cross-
 `slurm/smoke_test.sh` has the same shape as those - serve, then run - but for one task, two memory modules, one seed and two tasks of the dataset (`--max_tasks 2 --max_trials 3`). It then prints what the run wrote and fails if the two result rows are not there — including a wrong `--model`, which otherwise fails once per experiment rather than once. It also probes whether the filesystem grants `flock`, which is what the results file's append lock needs. Worth a submission before any 24-hour job, and after any change to the cluster, the model or the environment.
 
 **Size the real jobs before submitting them**
-`slurm/calibrate.sh` sits between the smoke test and a 24-hour job: one dataset, every arm, one seed, twenty tasks at the full trial budget, in a 2-hour allocation. It prints the result table, the tokens per task and any failed tasks. That is what sizes the real jobs — every experiment in a job runs concurrently against one throughput-bound server, so the wall clock is total tokens divided by what the server sustains, and `tokens_per_task x episodes / throughput` is the estimate. Twenty tasks also takes `g-memory` past its twentieth, where `merge_insights` runs.
+`slurm/<task>_calibrate.sh` sits between the smoke test and a 24-hour job: one dataset, every arm, one seed, twenty tasks at the full trial budget, in a 2-hour allocation. `slurm/generate_calibration.py` generates them, out of the same cluster configuration and job pieces as the sweep:
+
+```bash
+uv run slurm/generate_calibration.py                  # every dataset
+uv run slurm/generate_calibration.py --task fever pddl
+sbatch slurm/fever_calibrate.sh
+```
+
+`SLURM_ACCOUNT=<account>` in front of either generator puts an `#SBATCH --account` line in whatever it writes; without it the scripts submit under your default account.
+
+Calibrations default to `$HOME/GMemory/.db-calibration`, not the sweep's directory, so calibrating several datasets fills one table of its own. Each prints that table, the tokens per task and any failed tasks. That is what sizes the real jobs — every experiment in a job runs concurrently against one throughput-bound server, so the wall clock is total tokens divided by what the server sustains, and `tokens_per_task x episodes / throughput` is the estimate. Twenty tasks also takes `g-memory` past its twentieth, where `merge_insights` runs.
 
 **Attach to an already-running vLLM/Ray cluster**
 `slurm/experiment.sh` doesn't start its own model server. It expects a vLLM/Ray serving job already running elsewhere on the cluster and resolves that job's head node from its Slurm job ID:

slurm/generate_calibration.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Generate the per-task slurm/*_calibrate.sh runs that size the sweep jobs.
+
+One dataset, every arm, one seed, twenty tasks at the full trial budget, in a
+2-hour allocation. This is the run that says whether a 24-hour job fits: the
+result rows carry the token spend, so tokens / elapsed = throughput, and
+episodes x tokens-per-task / throughput = wall clock. Twenty also takes
+g-memory past its twentieth task, where merge_insights runs.
+
+    uv run slurm/generate_calibration.py                  # every dataset
+    uv run slurm/generate_calibration.py --task fever pddl
+
+The cluster, the model and the arms are generate_slurm.py's; what a calibration
+does differently is here.
+"""
+import argparse
+
+from generate_slurm import (
+    CLEANUP,
+    SEEDS,
+    TASKS,
+    every_arm,
+    preamble,
+    run_command,
+    write_script,
+)
+
+SEEDS_CALIBRATED = SEEDS[:1]
+MAX_TASKS = 20
+TIME_LIMIT = "02:00:00"
+DB_DIR = "$HOME/GMemory/.db-calibration"
+
+# Tasks whose full budget will not calibrate inside that window, and the smaller
+# shakedown that will. Jericho runs 100 trials rather than 30 and its prompt
+# tokens grow with the square of the budget - about 1.44M per task by the curve
+# in data/data.md, so twenty tasks over ten arms is ~288M tokens, an 18-hour job
+# at the throughput the other calibrations measured. Cut to the 2-hour window it
+# would report a tenth of its arms and nothing about the rest, which reads as an
+# arm that failed rather than one that never ran. Five tasks at 20 trials is
+# ~8M tokens and answers what a calibration of Jericho can: whether it runs.
+# Sizing its real job needs a job of its own.
+OVERRIDES = {"jericho": {"max_tasks": 5, "max_trials": 20}}
+
+SUMMARY = """
+
+echo "==== calibration ===="
+column -s, -t < ${DB_DIR}/overall_results.csv
+
+python3 -c '
+import csv, sys
+rows = list(csv.DictReader(open(sys.argv[1])))
+tokens = sum(int(r["completion_tokens"]) + int(r["prompt_tokens"]) for r in rows)
+scored = sum(int(r["tasks_scored"]) for r in rows)
+print(f"{len(rows)} arms, {scored} tasks scored, {tokens:,} tokens")
+print(f"{tokens/max(scored, 1):,.0f} tokens per task")
+' ${DB_DIR}/overall_results.csv
+
+cat ${DB_DIR}/*/*/*/*/failed_tasks.csv 2>/dev/null
+"""
+
+
+def scope_flags(task: str) -> str:
+    overrides = OVERRIDES.get(task, {})
+    flags = f"\n\t--max_tasks {overrides.get('max_tasks', MAX_TASKS)} \\"
+    if "max_trials" in overrides:
+        flags += f"\n\t--max_trials {overrides['max_trials']} \\"
+    return flags
+
+
+def render(task: str) -> str:
+    return (
+        preamble(
+            f"vllm-{task}-calibrate",
+            f"out/{task}-calibrate-%x.%j.%t.out",
+            f"{task}_calibrate.sh",
+            time_limit=TIME_LIMIT,
+            db_dir=DB_DIR,
+        )
+        + "\n"
+        + run_command(
+            task,
+            every_arm(task),
+            cross_task=False,
+            seeds=SEEDS_CALIBRATED,
+            scope=scope_flags(task),
+        )
+        + SUMMARY
+        + CLEANUP
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--task",
+        nargs="+",
+        choices=TASKS,
+        default=TASKS,
+        help="the datasets to calibrate (default: all of them)",
+    )
+    for task in parser.parse_args().task:
+        write_script(f"{task}_calibrate.sh", render(task))
+
+
+if __name__ == "__main__":
+    main()

slurm/generate_slurm.py

@@ -5,7 +5,14 @@
 below and rerun to produce a new sweep:
 
     uv run slurm/generate_slurm.py
+
+The calibration runs that size these jobs are slurm/generate_calibration.py,
+which builds its scripts from the cluster configuration and the job pieces here.
+
+Set SLURM_ACCOUNT to name an account in the generated scripts; without it they
+submit under the user's default.
 """
+import os
 from pathlib import Path
 
 SLURM_DIR = Path(__file__).parent
@@ -57,6 +64,9 @@
 PORT = 8000
 VLLM_STARTUP_SLEEP = 100
 
+# Named at generation time rather than committed: it is one user's allocation.
+ACCOUNT = os.environ.get("SLURM_ACCOUNT")
+
 # These three override the shared GPT-OSS_Hopper.yaml, which sets them to 8192,
 # 10240 and off.
 #
@@ -79,6 +89,18 @@ def intrinsic_memory_for(task: str) -> str:
     return f"intrinsicmemory-{task}"
 
 
+def intrinsic_arms(task: str) -> list[str]:
+    return [INTRINSIC_ABLATIONS[0], intrinsic_memory_for(task), INTRINSIC_ABLATIONS[1]]
+
+
+def every_arm(task: str) -> list[str]:
+    return BASELINE_MEMORIES + intrinsic_arms(task)
+
+
+def account_directive() -> str:
+    return f"#SBATCH --account={ACCOUNT}\n" if ACCOUNT else ""
+
+
 def vllm_serve_block() -> str:
     return f"""cd {VLLM_VENV_DIR}
 
@@ -122,7 +144,7 @@ def vllm_serve_block() -> str:
 
 
 def run_command(task: str, memories: list[str], cross_task: bool,
-                background: bool = False) -> str:
+                background: bool = False, seeds: list[int] = SEEDS, scope: str = "") -> str:
     """One `tasks/run.py` invocation, for one dataset and its arms.
 
     A backgrounded one records its own pid: a bare `wait` would also wait on the
@@ -135,19 +157,20 @@ def run_command(task: str, memories: list[str], cross_task: bool,
 \t--task {task} \\
 \t--mas_type autogen \\
 \t--mas_memory {" ".join(memories)} \\
-\t--seed {" ".join(str(seed) for seed in SEEDS)} \\{flag}
+\t--seed {" ".join(str(seed) for seed in seeds)} \\{flag}{scope}
 \t--db_dir ${{DB_DIR}} \\
 \t--model ${{MODEL_NAME}} \\
 \t--max_tokens {MAX_TOKENS_OVERRIDES.get(task, DEFAULT_MAX_TOKENS)}{trailing}"""
 
 
-def preamble(job_name: str, output_pattern: str, script_name: str) -> str:
+def preamble(job_name: str, output_pattern: str, script_name: str,
+             time_limit: str = TIME_LIMIT, db_dir: str = DEFAULT_DB_DIR) -> str:
     """Everything before the run: the allocation, the server, the environment."""
     return f"""#!/bin/bash
-#SBATCH --job-name={job_name}
+{account_directive()}#SBATCH --job-name={job_name}
 #SBATCH --nodes={NODES}
 #SBATCH --gpus={GPUS}
-#SBATCH --time={TIME_LIMIT}
+#SBATCH --time={time_limit}
 #SBATCH --exclusive
 #SBATCH --output={output_pattern}
 
@@ -160,7 +183,7 @@ def preamble(job_name: str, output_pattern: str, script_name: str) -> str:
 # Every job of one experiment set must point at the same directory: they append to
 # one overall_results.csv under a lock on the file. Override at submit time with
 #   DB_DIR=/projects/<project>/results/experiment-2026-09 sbatch slurm/{script_name}
-DB_DIR=${{DB_DIR:-{DEFAULT_DB_DIR}}}
+DB_DIR=${{DB_DIR:-{db_dir}}}
 
 {vllm_serve_block()}
 
@@ -187,14 +210,10 @@ def preamble(job_name: str, output_pattern: str, script_name: str) -> str:
 
 
 def render_experiment(task: str) -> str:
-    memories = BASELINE_MEMORIES + [
-        INTRINSIC_ABLATIONS[0], intrinsic_memory_for(task), INTRINSIC_ABLATIONS[1]
-    ]
-
     return (
         preamble(f"vllm-{task}", f"out/{task}-%x.%j.%t.out", f"{task}_experiment.sh")
         + "\n"
-        + run_command(task, memories, cross_task=False)
+        + run_command(task, every_arm(task), cross_task=False)
         + CLEANUP
     )
 
@@ -209,12 +228,7 @@ def render_crosstask() -> str:
     jobs safe.
     """
     runs = "\n\n".join(
-        run_command(
-            task,
-            [INTRINSIC_ABLATIONS[0], intrinsic_memory_for(task), INTRINSIC_ABLATIONS[1]],
-            cross_task=True,
-            background=True,
-        )
+        run_command(task, intrinsic_arms(task), cross_task=True, background=True)
         for task in TASKS
     )
 
@@ -235,15 +249,19 @@ def render_crosstask() -> str:
     )
 
 
+def write_script(script_name: str, body: str) -> None:
+    path = SLURM_DIR / script_name
+    path.write_text(body)
+    path.chmod(0o755)
+    print(f"wrote {path}")
+
+
 def main() -> None:
     scripts = {f"{task}_experiment.sh": render_experiment(task) for task in TASKS}
     scripts["crosstask.sh"] = render_crosstask()
 
     for script_name, body in scripts.items():
-        path = SLURM_DIR / script_name
-        path.write_text(body)
-        path.chmod(0o755)
-        print(f"wrote {path}")
+        write_script(script_name, body)
 
 
 if __name__ == "__main__":