diff --git a/.claude/settings.json b/.claude/settings.json index 467457c..4ccb2ad 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -14,8 +14,8 @@ "Bash(.venv/bin/python -m pruning_example.final_eval *)", "Edit(./src/**)", "Edit(./slurm/**)", "Edit(./configs/**)", "Edit(./prompts/**)", "Edit(./data/**)", "Edit(./tests/**)", "Edit(./.claude/**)", "Edit(./.github/**)", - "Edit(./state/**)", "Edit(./results/**)", "Edit(./local/**)", "Edit(./logs/**)", - "Edit(./AGENTS.md)", "Edit(./CLAUDE.md)", "Edit(./experiment.yaml)", + "Edit(./state/**)", "Edit(./state_multi/**)", "Edit(./results/**)", "Edit(./local/**)", "Edit(./logs/**)", + "Edit(./AGENTS.md)", "Edit(./CLAUDE.md)", "Edit(./experiment.yaml)", "Edit(./experiment_multi.yaml)", "Edit(./pyproject.toml)", "Edit(./uv.lock)", "Edit(./README.md)", "Edit(./LICENSE)", "Edit(./CITATION.cff)", "Edit(./.gitignore)", "Edit(./.python-version)", "Edit(./cluster.env)", "Edit(./cluster.env.example)", diff --git a/.claude/skills/check-round/SKILL.md b/.claude/skills/check-round/SKILL.md index ec5f15a..c7a9da0 100644 --- a/.claude/skills/check-round/SKILL.md +++ b/.claude/skills/check-round/SKILL.md @@ -19,6 +19,12 @@ The command compares: - The nonzero weights with the count that the configuration gives. - The pass decision with the accuracy threshold in `experiment.yaml`. +For a round of the multi-agent loop, set the state folder and the experiment file first: + +```bash +STATE_DIR=state_multi EXPERIMENT_FILE=experiment_multi.yaml .venv/bin/python -m pruning_example.check_round $ARGUMENTS +``` + Report each line that starts with `DIFF`. If all lines start with `OK`, say that the records agree. Do not change any file. If a check fails, do not try to repair the records. Tell the user which check failed. diff --git a/.gitignore b/.gitignore index f3d2a24..9b5d6d4 100644 --- a/.gitignore +++ b/.gitignore @@ -7,9 +7,10 @@ __pycache__/ .pytest_cache/ *.egg-info/ -# Outputs of the jobs and the loop script +# Outputs of the jobs and the loop scripts logs/ state/ +state_multi/ agent/ results/ diff --git a/AGENTS.md b/AGENTS.md index 60a130b..419b990 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,7 +10,7 @@ This project prunes ResNet-18 on CIFAR-10 on the Kempner AI cluster. Read this f - **Prunable weights.** The 20 Conv2d layers and the Linear layer: 11,164,352 weights. `conv1` has 1,728 weights, and `fc` has 5,120. Excluded layers stay dense, and they still count. - **Weights left.** Nonzero prunable weights divided by 11,164,352. With sparsity s and no excluded layers, a round keeps 11,164,352 - ceil(11,164,352 x s) weights. - **A round.** The round job prunes in `steps` equal-ratio steps with global magnitude pruning, and it fine-tunes after each step. The learning rate restarts at `lr` in each fine-tuning block and decays with a cosine curve. -- **Files to read.** `experiment.yaml`, `configs/round.schema.json`, `state/baseline/metrics.json`, `state/history.tsv`, and `state/notes.md`. +- **Files to read.** The experiment file (`experiment.yaml`, or the file that your prompt names), `configs/round.schema.json`, `state/baseline/metrics.json`, and `history.tsv` and `notes.md` in the state folder (`state/`, or the folder that your prompt names). - **Data.** Use only validation results. The test set is for the final steps only. - **Files.** Do not change tracked files. The round job stops if a tracked file changed. @@ -28,3 +28,10 @@ This project prunes ResNet-18 on CIFAR-10 on the Kempner AI cluster. Read this f - Return the next configuration only in the required JSON format, with a short reason. - Do not try to run commands or write files. The loop script runs the round. - Do not repeat a configuration that already ran. Learn from the rows and notes of the earlier rounds. + +## Roles in the multi-agent loop + +- Each role step has a role prompt from `prompts/multi/roles/`: the planner, the critic, the writer, or the auditor. Follow it. +- The multi-agent loop uses `experiment_multi.yaml` and the state folder `state_multi/`. Read only the state folder that your prompt names, and `state/baseline/metrics.json`. +- A wave has up to 3 rounds that run at the same time. The loop script stops after `target.plateau_rounds` complete rounds in a row, in round order, without a smaller passing model. +- Return only the required JSON or text. Do not try to run commands or write files. The loop script runs the rounds. diff --git a/README.md b/README.md index eb1efc0..b33e37c 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,8 @@ A Claude Code agent prunes a ResNet-18 model for CIFAR-10 on the Kempner AI clus In the unattended phase, the agent has read-only tools. It cannot run commands, write files, or submit jobs. The loop script, which is committed code, does these tasks. +A second version of the loop uses four agent roles and runs 3 rounds at a time; see [Multi-agent version](#multi-agent-version). + ## Before you start - A Kempner AI cluster account, and a SLURM account that can use `kempner_requeue` and `shared`. @@ -72,7 +74,7 @@ source cluster.env sbatch slurm/loop.sbatch ``` -Watch the run with `tail -f logs/prune-loop_.out` and `state/history.tsv`. To stop the run, run `scancel ` on the loop job. The loop script cancels its own round job and writes `state/summary.json`. To run the final steps later, run `sbatch slurm/loop.sbatch --finish`. +Watch the run with `tail -f logs/prune-loop_.out` and `state/history.tsv`. To stop the run, run `scancel ` on the loop job. The loop script cancels its own round job and writes `state/summary.json`. The next loop job counts the GPU minutes of the canceled job. To run the final steps later, run `sbatch slurm/loop.sbatch --finish`. After the final steps, use a new state folder to search again. If round jobs end at the time limit with no epoch line in their logs, check the storage of `DATA_DIR`. ## Limits @@ -80,7 +82,7 @@ Watch the run with `tail -f logs/prune-loop_.out` and `state/history.tsv |---|---| | Accuracy threshold | Validation accuracy of at least 0.905. The loop looks for the fewest weights that keep it, and stops after 3 rounds in a row without a smaller passing model. | | Tokens | 1,000,000 for the unattended phase, with 10% kept for the report step. All four token types count, cache reads too. | -| Agent step | `--max-turns 8`, `--max-budget-usd 1.00`, read-only tools, `dontAsk` | +| Agent step | `claude-opus-5-5` at the `medium` effort level, `--max-turns 8`, `--max-budget-usd 1.00`, read-only tools, `dontAsk`. The loop script stops if a step used another model. | | Rounds | 15 round attempts, round 1 included | | GPU | 60 GPU minutes for the unattended phase. One round job at a time, 1 GPU, 10 minutes at most. | | Queue | A job can wait at most 45 minutes in the queue | @@ -91,6 +93,7 @@ Watch the run with `tail -f logs/prune-loop_.out` and `state/history.tsv | File | Content | |---|---| | `state/summary.json` | Stop reason, totals, the selected round, both seeds, validation and test accuracy with 95% intervals | +| `state/selection.json` | The selected round, saved before the confirmation round, so that `--finish` keeps it | | `state/final_report.md` | The draft report of the agent. Check each number against the other files. | | `state/history.tsv` | One row for each round attempt | | `state/ledger.tsv` | Tokens of each agent step, and GPU minutes of each job | @@ -108,7 +111,7 @@ The round checkpoints are in `RUN_ROOT` on scratch storage. FASRC deletes scratc | Writes | The agent asks first. Deny rules block edits of tracked files and output folders. | None | | Shell | Sandbox on, SLURM commands outside it, `scancel` and `rm -r` blocked by a hook | No shell | | Job submission | `sbatch` asks | The loop script submits | -| Reads | File tools: only the project folder. Sandboxed commands: not the home folder. | Only the project folder | +| Reads | File tools: only the project folder. Sandboxed commands: not the home folder. | Only the project folder, and not this README, which holds example results | | Integrity | The round job and the loop script stop if a tracked file changed | The same | ## Skills @@ -122,7 +125,6 @@ The project has three skills in `.claude/skills/`. Claude Code loads a skill whe | `testing-and-continuous-integration` | The Coding Agent Skills page of the handbook | Practices for tests and CI. Useful when you change the code or the tests. | - In this example, only the interactive session of round 1 can use the skills. The agent step in the loop has only the Read, Grep, and Glob tools, so it cannot load a skill. -- In our run, round 1 used `check-round`. The rounds did not use the two handbook skills. - Read each skill before you trust the project folder, because Claude Code follows its instructions. To update the two handbook skills, download them again, and commit the change. The round job stops if a tracked file changed. @@ -135,6 +137,45 @@ done git diff --stat .claude/skills ``` +## Multi-agent version + +Four agent roles run the same task. The loop script still checks each output, submits the jobs, and writes all records. The chapter [Multi-Agent Workflow Example](https://handbook.eng.kempnerinstitute.harvard.edu/s5b_agentic_ai_workflows/multi_agent_example.html) describes it step by step. + +| Role | Model | Effort | Task | +|---|---|---|---| +| Planner | `claude-opus-5-5` | `xhigh` | Plans 3 configurations for each wave | +| Critic | `claude-sonnet-5-5` | `high` | Accepts or rejects each configuration before any GPU job runs | +| Writer | `claude-sonnet-5-5` | `medium` | Writes the draft report | +| Auditor | `claude-opus-5-5` | `max` | Checks each claim of the report against the records | + +`experiment_multi.yaml` sets the roles, the waves, and the limits. With `roles.critic.enabled: false`, one agent plans the waves. + +After the setup, the baseline, and round 1 above: + +```bash +source cluster.env +sbatch --time=01:30:00 slurm/multi_loop.sbatch --plan-only # 1. plan wave 1, then stop +cat state_multi/waves/wave_01/plan.md # 2. read the plan +sbatch slurm/multi_loop.sbatch # 3. run the waves and the final steps +``` + +- **Plan again.** Remove `state_multi/waves/wave_01/`, and run step 1 again. +- **Watch.** Run `tail -f logs/prune-multi-loop_.out`. +- **Stop.** Run `scancel ` on the loop job. Check with `squeue --me` that no `prune-round` task is left. The next loop job counts the GPU minutes of the canceled tasks. +- **Finish later.** Run `sbatch slurm/multi_loop.sbatch --finish`. After the final steps, use a new state folder to search again. + +| Part | Single-agent loop | Multi-agent loop | +|---|---|---| +| Rounds | One round job at a time | Waves of up to 3 round jobs, as one array job (`slurm/round_array.sbatch`) | +| GPUs | Full GPUs | A100 MIG slices (3g.20gb), 1 to 3 minutes for each round | +| Checks | The schema | The schema and a repeat check in code, then the critic. One revision of a rejected configuration. | +| Limits | 1,000,000 tokens, 60 GPU minutes, 15 round attempts | 1,250,000 tokens (450,000 kept for the writer and the auditor; plan for 10% more), 150 GPU minutes, 45 round attempts | +| Plateau | 3 rounds | 6 rounds | +| Final steps | One confirmation seed, a draft report | Two confirmation seeds and 3-seed means, a draft report, the auditor, `check_report`, and `check_round` | +| Records | `state/` | `state_multi/`, with `waves/`, `hypotheses.tsv`, and `audit.json` | + +To audit the report of the single-agent loop, run `STATE_DIR=state sbatch slurm/multi_loop.sbatch --audit-only`. To compare designs, run each one in its own clone with only the baseline and round 1; see [Comparison of the designs](#comparison-of-the-designs). + ## Test the loop with small limits Put test copies in `local/`, which Git ignores. `EXPERIMENT_FILE` and `STATE_DIR` point the jobs to them: @@ -160,57 +201,63 @@ The tests use generated images, so they need no GPU and no data. CI runs them wi ```text AGENTS.md, CLAUDE.md agent instructions experiment.yaml accuracy threshold, recipe, and limits +experiment_multi.yaml roles, models, effort levels, waves, and limits of the multi-agent loop cluster.env.example your account and paths .claude/ settings, hook, and skills -configs/ round configuration schema and an example +configs/ schemas of a round configuration, a wave plan, a review, and an audit prompts/ prompts of the agent step and the report step +prompts/multi/ role prompts and task prompts of the multi-agent loop data/split.json fixed validation split -slurm/ baseline, round, and loop jobs -src/pruning_example/ code: data, model, evaluator, pruning, records, loop, final steps +slurm/ baseline, round, round array, loop, and multi-agent loop jobs +src/pruning_example/ code: data, model, evaluator, pruning, records, loops, final steps, checks tests/ pytest tests ``` -## Our run - -One full run on October 17, 2026. Round 1 was supervised. The loop job ran rounds 2 to 15 and stopped at the round limit. +## Example results -| Round | Sparsity | Steps | Epochs | Learning rate | Excluded | Validation accuracy | Weights left | -|---|---|---|---|---|---|---|---| -| 1 | 0.5 | 1 | 2 | 0.01 | none | 0.9320 | 50% | -| 2 | 0.85 | 2 | 4 | 0.01 | none | 0.9356 | 15% | -| 3 | 0.95 | 3 | 6 | 0.01 | conv1 | 0.9180 | 5.0% | -| 5 | 0.98 | 5 | 10 | 0.01 | conv1 | 0.9074 | 2.0% | -| 7 | 0.988 | 5 | 20 | 0.01 | conv1 | 0.9052 | 1.2% | -| 8 | 0.989 | 4 | 20 | 0.01 | conv1 | 0.8976 (below) | 1.1% | -| 10 | 0.9895 | 5 | 20 | 0.03 | conv1 | 0.9154 | 1.07% | -| 15 | 0.9927 | 5 | 20 | 0.08 | conv1 | 0.9076 | 0.75% | - -From round 9, the agent raised the learning rate by 0.01 in each round. This kept the threshold at higher sparsities. +One run of the single-agent loop with the default settings. Your results can differ. | Result | Value | |---|---| -| Selected model | Round 15: 83,215 nonzero weights (0.75%), sparse size 1.56 MB (dense: 44.7 MB) | -| Test accuracy | 0.9045 (95% interval 0.899 to 0.910) | +| Rounds | Round 1 supervised, rounds 2 to 15 in the loop job. The loop stopped at the round limit. | +| Selected model | Round 15: sparsity 0.9927 in 5 steps, 20 epochs, learning rate 0.08, `conv1` kept dense. 83,215 nonzero weights (0.75%), sparse size 1.56 MB (dense: 44.7 MB). | +| Accuracy | Validation 0.9076, test 0.9045 (95% interval 0.899 to 0.910) | | Confirmation round (seed 1) | Validation 0.9050, test 0.9032 | | Tokens | 323,072, estimated cost $1.11 | | GPU minutes | 16.1 for the rounds, 1.0 for the confirmation round | | Time of the loop job | 66 minutes | -The confirmation round passed exactly at the threshold, and both test values are close to 0.900. Read the result as a lead to check, not as a finding. +## Comparison of the designs + +Each design ran one time, from the same baseline and round 1, with the same search limits. A result holds if its 3-seed means are at least 0.905 for validation and 0.900 for test. The chapter [Multi-Agent Workflow Example](https://handbook.eng.kempnerinstitute.harvard.edu/s5b_agentic_ai_workflows/multi_agent_example.html) gives the details. + +| Result | One agent, one round at a time | One agent, waves of 3 | Planner and critic, waves of 3 | +|---|---|---|---| +| Selected round | Round 9: 0.700% weights left | Round 16: 0.485% | Round 21: 0.661% | +| Means of 3 seeds: validation, test | 0.9043, 0.9001 | 0.9038, 0.8978 | 0.9049, 0.9017 | +| Result holds | No | No | No | +| Rounds after round 1, time of the search | 14, 62 minutes | 21, 45 minutes | 21, 52 minutes | +| Tokens, estimated cost, audit included | 783,766, $2.79 | 879,101, $4.22 | 1,348,571, $5.99 | + +- **Result.** No result holds, so the comparison shows no gain from the extra agents. Each selected round passed with seed 0, but its 3-seed mean of the validation accuracy was below 0.905. +- **Critic.** It accepted all 21 configurations. +- **Auditor.** In the reports of both wave designs, it found wrong claims that `check_report` did not find. ## Tested with -Claude Code 2.1.302, uv 0.8.19, Python 3.12.11, PyTorch 2.14.1 with CUDA 13.0, torchvision 0.29.1, and NVIDIA driver 610.57.04, on A100, H100, H200, and RTX PRO 6000 GPUs, in October 2026. +Claude Code 2.1.302 to 2.1.309 (the multi-agent loop: 2.1.307 to 2.1.309), uv 0.8.19, Python 3.12.11, PyTorch 2.14.1 with CUDA 13.0, torchvision 0.29.1, and NVIDIA driver 610.57.04. GPUs: A100 (full GPUs and 3g.20gb MIG slices), H100, H200, and RTX PRO 6000. ## Adapt the example -- **Change** the task: the data, the model, the metric, the round configuration schema, the accuracy threshold, and the limits. -- **Keep** the fixed evaluator, the read-only agent step, the loop script that submits the jobs, the limits outside the model, the run records, the confirmation round, and the one-time test evaluation. +- **Task.** Change the data, the model, the metric, the round configuration schema, the accuracy threshold, and the limits. +- **GPUs.** To run the round jobs on other GPUs, set `round.constraint` in the experiment file. For example, set `a100-mig` in `experiment.yaml` for MIG slices, or `a100|h100|h200|rtx6000pro` in `experiment_multi.yaml` for full GPUs. The loop scripts pass it to `sbatch`, and it replaces the `--constraint` line of the job scripts. A MIG slice is 3 to 4 times slower than a full GPU, so set `gpu_minutes` to match. +- **Changes.** Commit each change to a tracked file. The jobs refuse to start while a tracked file has changes. +- **Keep.** Keep the fixed evaluator, the read-only agent step, and the loop script that submits the jobs. Also keep the limits outside the model, the run records, the confirmation round, and the one-time test evaluation. ## Sources -The skills in `.claude/skills/reproducible-research/` and `.claude/skills/testing-and-continuous-integration/`, and the hook `.claude/hooks/guard.sh`, come from the Kempner Institute Computing Handbook. +The skills in `.claude/skills/reproducible-research/` and `.claude/skills/testing-and-continuous-integration/` come from [Coding Agent Skills](https://handbook.eng.kempnerinstitute.harvard.edu/s2_swe_for_research/coding_agent_skills.html) in the Kempner Institute Computing Handbook. The hook `.claude/hooks/guard.sh` comes from [Hooks](https://handbook.eng.kempnerinstitute.harvard.edu/s5b_agentic_ai_workflows/configuring_agents.html#hooks) in the same handbook. ## License -MIT. Copyright (c) 2026 Harvard University. See `LICENSE`. To cite this example, see `CITATION.cff`. +MIT. Copyright (c) 2026 Harvard University. See [LICENSE](LICENSE). To cite this example, see [CITATION.cff](CITATION.cff). diff --git a/configs/audit.schema.json b/configs/audit.schema.json new file mode 100644 index 0000000..57458cc --- /dev/null +++ b/configs/audit.schema.json @@ -0,0 +1,60 @@ +{ + "title": "Audit of a draft report", + "type": "object", + "properties": { + "claims": { + "type": "array", + "description": "Each claim of the report that has a number or a conclusion.", + "items": { + "type": "object", + "properties": { + "claim": { + "type": "string", + "maxLength": 400, + "description": "The claim, as the report states it." + }, + "verdict": { + "enum": [ + "supported", + "wrong", + "unsupported" + ], + "description": "supported: the records show it. wrong: the records show something else. unsupported: the records do not show it." + }, + "evidence": { + "type": "string", + "maxLength": 400, + "description": "The file and the value that decide the verdict." + }, + "correction": { + "type": "string", + "maxLength": 400, + "description": "The correct value for a wrong claim; empty otherwise." + } + }, + "required": [ + "claim", + "verdict", + "evidence", + "correction" + ], + "additionalProperties": false + } + }, + "limits_stated": { + "type": "boolean", + "description": "True if the report states the seeds, the intervals, the one test evaluation, and the number of runs." + }, + "summary": { + "type": "string", + "maxLength": 800, + "description": "Your audit in a few sentences." + } + }, + "required": [ + "claims", + "limits_stated", + "summary" + ], + "additionalProperties": false +} diff --git a/configs/review.schema.json b/configs/review.schema.json new file mode 100644 index 0000000..3cf453f --- /dev/null +++ b/configs/review.schema.json @@ -0,0 +1,46 @@ +{ + "title": "Critic review", + "type": "object", + "properties": { + "verdicts": { + "type": "array", + "minItems": 1, + "description": "One verdict for each configuration, by its item number.", + "items": { + "type": "object", + "properties": { + "item": { + "type": "integer", + "minimum": 1, + "description": "The item number of the configuration." + }, + "accept": { + "type": "boolean", + "description": "True to run the configuration, false to reject it." + }, + "reason": { + "type": "string", + "maxLength": 400, + "description": "Why you accept or reject it." + } + }, + "required": [ + "item", + "accept", + "reason" + ], + "additionalProperties": false + } + }, + "summary": { + "type": "string", + "maxLength": 600, + "description": "Your review of the wave in a few sentences." + } + }, + "required": [ + "verdicts", + "summary" + ], + "additionalProperties": false +} diff --git a/configs/wave.schema.json b/configs/wave.schema.json new file mode 100644 index 0000000..e4f4d54 --- /dev/null +++ b/configs/wave.schema.json @@ -0,0 +1,103 @@ +{ + "title": "Wave plan", + "type": "object", + "properties": { + "configurations": { + "type": "array", + "minItems": 1, + "maxItems": 3, + "description": "One configuration for each round of the wave. The loop script sets minItems and maxItems to the wave size.", + "items": { + "type": "object", + "properties": { + "sparsity": { + "type": "number", + "minimum": 0.5, + "maximum": 0.999, + "description": "Fraction of the weights in the pruned layers to remove." + }, + "steps": { + "type": "integer", + "minimum": 1, + "maximum": 10, + "description": "Number of pruning steps. 1 is one-shot pruning." + }, + "finetune_epochs": { + "type": "integer", + "minimum": 1, + "maximum": 20, + "description": "Total fine-tuning epochs, divided into one block for each step. Must be at least steps." + }, + "lr": { + "type": "number", + "minimum": 0.0001, + "maximum": 0.1, + "description": "Learning rate at the start of each fine-tuning block. It decays with a cosine curve in the block." + }, + "weight_decay": { + "type": "number", + "minimum": 0, + "maximum": 0.001, + "description": "Weight decay." + }, + "exclude_layers": { + "type": "array", + "items": { + "enum": [ + "conv1", + "fc" + ] + }, + "uniqueItems": true, + "description": "Layers that the round job does not prune. These layers still count in the weights left." + }, + "reason": { + "type": "string", + "maxLength": 600, + "description": "Why you chose this configuration." + }, + "hypothesis": { + "type": "string", + "maxLength": 300, + "description": "The idea that this configuration tests, in one or two sentences." + }, + "predicted_val_accuracy": { + "type": "number", + "minimum": 0, + "maximum": 1, + "description": "The validation accuracy that you expect." + }, + "pass_probability": { + "type": "number", + "minimum": 0, + "maximum": 1, + "description": "The probability that the round keeps the accuracy threshold." + } + }, + "required": [ + "sparsity", + "steps", + "finetune_epochs", + "lr", + "weight_decay", + "exclude_layers", + "reason", + "hypothesis", + "predicted_val_accuracy", + "pass_probability" + ], + "additionalProperties": false + } + }, + "summary": { + "type": "string", + "maxLength": 600, + "description": "How the configurations of this wave differ, and why." + } + }, + "required": [ + "configurations", + "summary" + ], + "additionalProperties": false +} diff --git a/experiment.yaml b/experiment.yaml index 922cadf..d5be90b 100644 --- a/experiment.yaml +++ b/experiment.yaml @@ -17,6 +17,7 @@ round: batch_size: 256 momentum: 0.9 time_minutes: 10 # must agree with --time in slurm/round.sbatch + # constraint: a100-mig # optional: replaces the --constraint line of the round job scripts budget: tokens: 1000000 # all four token types of all models, cache reads included @@ -29,11 +30,15 @@ budget: final_reserve_minutes: 65 # time kept for the final steps agent_step: + model: claude-opus-5-5 # a full model name; an alias such as opus moves to newer models + effort: medium # low, medium, high, xhigh, or max; medium is the default of Opus 5.5 max_turns: 8 max_budget_usd: 1.00 timeout_minutes: 15 report_step: + model: claude-opus-5-5 + effort: medium max_turns: 6 max_budget_usd: 1.00 timeout_minutes: 15 diff --git a/experiment_multi.yaml b/experiment_multi.yaml new file mode 100644 index 0000000..ea58ac2 --- /dev/null +++ b/experiment_multi.yaml @@ -0,0 +1,68 @@ +# Settings and limits of the multi-agent design. +# Read this file before you start the multi-agent loop job: it sets the roles, the models, the effort levels, +# and the budget. The sections target (val_accuracy), baseline, and round (except constraint) must agree with +# experiment.yaml; a test checks this. + +target: + val_accuracy: 0.905 # accuracy threshold: 0.900 plus a margin for noise in 5,000 validation images + test_accuracy: 0.900 # the goal on the test set; the 3-seed means in the summary use it + plateau_rounds: 6 # stop after 6 complete rounds in a row, in round order, without a smaller passing model + +baseline: + epochs: 40 + lr: 0.1 + warmup_epochs: 2 + momentum: 0.9 + weight_decay: 5.0e-4 + batch_size: 256 + +round: + batch_size: 256 + momentum: 0.9 + time_minutes: 10 # must agree with --time in slurm/round.sbatch and slurm/round_array.sbatch + # constraint: a100|h100|h200|rtx6000pro # optional: full GPUs in place of the MIG slices of slurm/round_array.sbatch + +wave: + size: 3 # configurations in each wave; fewer when a limit allows fewer + max_parallel: 3 # round jobs of one wave that can run at the same time + +# One model and one effort level for each role. Use full model names: an alias such as opus moves to newer models. +# Effort levels: low, medium, high, xhigh, max. A checker uses a different model from the role that it checks. +roles: + planner: # plans each wave, and runs the revision of rejected configurations + model: claude-opus-5-5 + effort: xhigh + max_turns: 10 + max_budget_usd: 3.00 + timeout_minutes: 20 + critic: # checks each plan before any GPU job runs; enabled: false gives one agent with waves + enabled: true + model: claude-sonnet-5-5 + effort: high + max_turns: 8 + max_budget_usd: 1.00 + timeout_minutes: 15 + writer: # writes the draft report + model: claude-sonnet-5-5 + effort: medium + max_turns: 6 + max_budget_usd: 1.00 + timeout_minutes: 15 + auditor: # checks each claim of the draft report against the records + model: claude-opus-5-5 + effort: max + max_turns: 12 + max_budget_usd: 5.00 + timeout_minutes: 30 + +budget: + tokens: 1250000 # all four token types of all models, cache reads included, supervised wave included + final_reserve_fraction: 0.36 # 450,000 tokens kept for the writer and the auditor; the search can use 800,000 + wave_token_estimate: 150000 # used before the first measured wave; then the largest wave so far + gpu_minutes: 150 # rounds after round 1, on A100 MIG slices; failed and requeued runs count + max_round_attempts: 45 # round 1 included; all agents of a run share them + max_failed_waves: 2 # waves in a row with no complete round + queue_limit_minutes: 45 # total queue time of one array job, requeues included + final_reserve_minutes: 110 # time kept for the final steps + +confirmation_seeds: [1, 2] # the selected configuration runs again with each seed, as one array job diff --git a/prompts/multi/auditor.md b/prompts/multi/auditor.md new file mode 100644 index 0000000..68b3a82 --- /dev/null +++ b/prompts/multi/auditor.md @@ -0,0 +1 @@ +Check the draft report {REPORT} against the records in {STATE}: summary.json, history.tsv, hypotheses.tsv if it exists, the metrics.json files in rounds/ and confirm/, and the files in test_eval/. Give a verdict for each claim. Return the result in the required JSON format. diff --git a/prompts/multi/critic.md b/prompts/multi/critic.md new file mode 100644 index 0000000..f894942 --- /dev/null +++ b/prompts/multi/critic.md @@ -0,0 +1,3 @@ +Check the configurations that the planner proposes for wave {W} of a pruning study. They are in {PROPOSED}, with their item numbers. + +Read {EXPERIMENT}, configs/round.schema.json, {STATE}/history.tsv, {STATE}/hypotheses.tsv, and {STATE}/notes.md. Give one verdict for each configuration, by its item number. Return the verdicts in the required JSON format. diff --git a/prompts/multi/planner.md b/prompts/multi/planner.md new file mode 100644 index 0000000..ccc7e0c --- /dev/null +++ b/prompts/multi/planner.md @@ -0,0 +1,4 @@ +You plan wave {W} of a pruning study. Plan {K} configurations. Round attempts left: {R}. GPU minutes left: {G}. +Best round so far that keeps the accuracy threshold: {BEST}. Rounds in a row without a smaller model: {STALE} of {PLATEAU}. + +Read {EXPERIMENT}, configs/round.schema.json, state/baseline/metrics.json, {STATE}/history.tsv, {STATE}/hypotheses.tsv, and {STATE}/notes.md. Plan {K} configurations that can keep the accuracy threshold with fewer weights than the best round so far. Return them in the required JSON format. Give each one a hypothesis, a predicted validation accuracy, a pass probability, and a short reason. diff --git a/prompts/multi/revision.md b/prompts/multi/revision.md new file mode 100644 index 0000000..aa059b1 --- /dev/null +++ b/prompts/multi/revision.md @@ -0,0 +1,5 @@ +You plan wave {W} of a pruning study. The checks rejected {K} of the configurations that you proposed in {PROPOSED}: + +{REJECTED} + +Plan {K} new configurations to replace them. Read {EXPERIMENT}, configs/round.schema.json, state/baseline/metrics.json, {STATE}/history.tsv, {STATE}/hypotheses.tsv, and {STATE}/notes.md. Do not repeat a configuration that already ran or that this wave already has. Return them in the required JSON format. Give each one a hypothesis, a predicted validation accuracy, a pass probability, and a short reason. diff --git a/prompts/multi/roles/auditor.md b/prompts/multi/roles/auditor.md new file mode 100644 index 0000000..bb06ac9 --- /dev/null +++ b/prompts/multi/roles/auditor.md @@ -0,0 +1,10 @@ +You are the auditor of a pruning study. You check a draft report that another agent wrote. You did not write it. + +- Read the report and the records that your task prompt names. +- Find each claim of the report that has a number or a conclusion. Give each claim a verdict: + - supported: the records show it. + - wrong: the records show something else. Give the correct value. + - unsupported: the records do not show it. +- For each verdict, give the evidence: the file and the value. +- Check that the report states the limits of the result: the number of seeds, the intervals, that the test set was used one time for each model, and the number of runs. +- Return only the required JSON. Do not try to run commands or write files. diff --git a/prompts/multi/roles/critic.md b/prompts/multi/roles/critic.md new file mode 100644 index 0000000..2022ace --- /dev/null +++ b/prompts/multi/roles/critic.md @@ -0,0 +1,12 @@ +You are the critic of a pruning study. You check the configurations that the planner proposes, before any GPU job runs. You did not plan them. + +- Read the files that your task prompt names. +- Give one verdict for each configuration: accept or reject, with a short reason. +- Reject a configuration in these cases: + - It breaks the keys or the ranges in configs/round.schema.json. + - It repeats an earlier configuration. + - The history shows that it will most likely not keep the accuracy threshold. + - It changes many settings at once, and its reason does not say why. + - It does not test its own hypothesis. +- Accept each other configuration. Do not reject a configuration only because it is bold. +- You cannot change a configuration. Return only the required JSON. Do not try to run commands or write files. diff --git a/prompts/multi/roles/planner.md b/prompts/multi/roles/planner.md new file mode 100644 index 0000000..07c4342 --- /dev/null +++ b/prompts/multi/roles/planner.md @@ -0,0 +1,8 @@ +You are the planner of a pruning study. You plan the next wave of experiments. + +- Read the files that your task prompt names. +- Plan one configuration for each round of the wave. Use the keys and ranges in configs/round.schema.json. +- Give each configuration its own hypothesis, so that one wave tests different ideas. +- Predict the validation accuracy of each configuration. Give the probability that it keeps the accuracy threshold. +- Do not repeat a configuration that already ran. +- Return only the required JSON. Do not try to run commands or write files. diff --git a/prompts/multi/roles/writer.md b/prompts/multi/roles/writer.md new file mode 100644 index 0000000..398356d --- /dev/null +++ b/prompts/multi/roles/writer.md @@ -0,0 +1,6 @@ +You are the writer of a pruning study. You write a short draft report for a researcher. + +- Read the files that your task prompt names. +- Use only numbers that are in these files. Copy each number exactly. +- State the limits of the result: the number of seeds, the intervals, that the test set was used one time for each model, and that the study is one run. +- Do not try to run commands or write files. diff --git a/prompts/multi/writer.md b/prompts/multi/writer.md new file mode 100644 index 0000000..8347f11 --- /dev/null +++ b/prompts/multi/writer.md @@ -0,0 +1 @@ +Read {STATE}/summary.json, {STATE}/history.tsv, {STATE}/hypotheses.tsv, and {STATE}/notes.md. Write a short report in Markdown: the result, the path of the waves, what the critic changed, and the limits used. Use only numbers from these files. Do not try to run commands or write files. diff --git a/pyproject.toml b/pyproject.toml index af1f9bf..0e0977b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pruning-example" -version = "0.1.0" +version = "0.2.0" description = "End-to-end agentic workflow example: a Claude Code agent prunes ResNet-18 on CIFAR-10 on the Kempner AI cluster." readme = "README.md" license = "MIT" diff --git a/slurm/multi_loop.sbatch b/slurm/multi_loop.sbatch new file mode 100644 index 0000000..f622f10 --- /dev/null +++ b/slurm/multi_loop.sbatch @@ -0,0 +1,29 @@ +#!/bin/bash +#SBATCH --job-name=prune-multi-loop +#SBATCH --partition=shared +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=2 +#SBATCH --mem=8G +#SBATCH --time=06:00:00 +#SBATCH --signal=B:SIGTERM@120 +#SBATCH --output=logs/%x_%j.out + +# Run the multi-agent loop, then the final steps. Submit from the project folder, on a login node: +# sbatch --time=01:30:00 slurm/multi_loop.sbatch --plan-only # plan wave 1, then stop so that you can read it +# sbatch slurm/multi_loop.sbatch # wave 1 and the next waves, then the final steps +# sbatch slurm/multi_loop.sbatch --finish # only the final steps, after an early stop +# sbatch slurm/multi_loop.sbatch --audit-only # the auditor and check_report, for STATE_DIR +# Read experiment_multi.yaml first: it sets the roles, the models, the effort levels, and the budget of this job. + +set -euo pipefail +cd "$SLURM_SUBMIT_DIR" +unset ANTHROPIC_API_KEY CLAUDE_CODE_EFFORT_LEVEL # print mode uses your Harvard login; the experiment file sets the effort +source cluster.env +export STATE_DIR="${STATE_DIR:-state_multi}" +export EXPERIMENT_FILE="${EXPERIMENT_FILE:-experiment_multi.yaml}" +export PATH="$HOME/.local/bin:$PATH" +CLAUDE_VERSION="$(claude --version)" +export CLAUDE_VERSION +echo "start $(date -Is) on $(hostname -s), ${CLAUDE_VERSION}, state folder ${STATE_DIR}, experiment file ${EXPERIMENT_FILE}" +exec .venv/bin/python -m pruning_example.multi_loop "$@" diff --git a/slurm/round_array.sbatch b/slurm/round_array.sbatch new file mode 100644 index 0000000..2b75d8f --- /dev/null +++ b/slurm/round_array.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=prune-round +#SBATCH --partition=kempner_requeue +#SBATCH --constraint=a100-mig +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=32G +#SBATCH --time=00:10:00 +#SBATCH --requeue +#SBATCH --open-mode=append +#SBATCH --output=logs/%x_%j.out + +# Run one round for each task of an array job. The loop scripts submit it: +# sbatch --array=0-2%3 slurm/round_array.sbatch TASKS_FILE +# Each line of TASKS_FILE, with no header, has 4 fields separated by tabs: +# configuration path, round number, kind (round or confirm), seed +# Task i runs line i+1. %j gives each task its own job ID, so each task has its own log. +# kempner_requeue can preempt a task. The task saves a checkpoint after each epoch and continues after a requeue. +# Keep --time equal to round.time_minutes in the experiment file. +# The tasks use A100 MIG slices (3g.20gb). A round is small, and the slices are often free when the full GPUs +# are busy. A round takes 1 to 3 minutes on a slice, 3 to 4 times longer than on a full GPU, so +# experiment_multi.yaml has a larger GPU budget. The single-agent loop uses this script only for several +# confirmation seeds; its rounds keep the full GPUs of slurm/round.sbatch. + +set -euo pipefail +cd "$SLURM_SUBMIT_DIR" +source cluster.env +IFS=$'\t' read -r config round kind seed < <(sed -n "$((SLURM_ARRAY_TASK_ID + 1))p" "$1") +echo "start $(date -Is) on $(hostname -s), task ${SLURM_ARRAY_TASK_ID}, restart ${SLURM_RESTART_COUNT:-0}" +exec .venv/bin/python -m pruning_example.run_round "$config" --round "$round" --kind "$kind" --seed "$seed" --no-history diff --git a/src/pruning_example/agent.py b/src/pruning_example/agent.py index 9949313..ff78583 100644 --- a/src/pruning_example/agent.py +++ b/src/pruning_example/agent.py @@ -1,13 +1,20 @@ -"""Run claude -p for an agent step or for the report step, and count its tokens.""" +"""Run claude -p for an unattended agent step, and count its tokens. + +The single-agent loop runs the agent step and the report step. The multi-agent loop runs the role steps. +""" from __future__ import annotations import json import os +import re import subprocess TOKEN_FIELDS = ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens") READ_ONLY_TOOLS = "Read,Grep,Glob" +EFFORT_LEVELS = ("low", "medium", "high", "xhigh", "max") +# The README holds example results. An unattended step must find its own way, so it cannot read them. +UNATTENDED_SETTINGS = {"permissions": {"deny": ["Read(./README.md)"]}} class AgentError(RuntimeError): @@ -18,6 +25,10 @@ class AgentFormatError(AgentError): """The JSON result does not have a required field. The loop script stops.""" +class ModelCheckError(AgentError): + """The JSON result does not show the configured model. The loop script stops.""" + + def tokens_used(result: dict) -> int: """Add the four token types of all models in modelUsage. Cache reads also count.""" usage = result.get("modelUsage") @@ -29,8 +40,25 @@ def tokens_used(result: dict) -> int: raise AgentFormatError(f"modelUsage has no token count: {error}") from None +def check_model(result: dict, model: str | None) -> None: + """Raise ModelCheckError if the step did not use the configured model. + + Claude Code replaces a model that the organization does not allow with the default model, + and it gives only a warning. An alias such as opus names no fixed model, so it is not checked. + """ + if not model or not model.startswith("claude-"): + return + used = sorted((result.get("modelUsage") or {}).keys()) + # The name itself, with a date such as -20261001, or with a suffix such as [1m]. claude-opus-5 does not + # match claude-opus-5-5, which is another model. + same = re.compile(re.escape(model) + r"(-\d{8})?(\[[^\]]*\])?") + if not any(same.fullmatch(name) for name in used): + raise ModelCheckError(f"the step used {', '.join(used) or 'no model'}, not the configured model {model}") + + def command(prompt: str, caps: dict, schema_text: str | None = None) -> list[str]: - """Return the claude -p command: read-only tools, dontAsk, no MCP tools, and the caps.""" + """Return the claude -p command: read-only tools, dontAsk, no MCP tools, the caps, and the optional + model, effort level, and role prompt.""" args = [ "claude", "-p", prompt, "--permission-mode", "dontAsk", @@ -39,12 +67,31 @@ def command(prompt: str, caps: dict, schema_text: str | None = None) -> list[str "--max-turns", str(caps["max_turns"]), "--max-budget-usd", f"{caps['max_budget_usd']:.2f}", "--output-format", "json", + "--settings", json.dumps(UNATTENDED_SETTINGS), ] + if caps.get("model"): + args += ["--model", str(caps["model"])] + if caps.get("effort"): + args += ["--effort", str(caps["effort"])] + if caps.get("role_prompt"): + args += ["--append-system-prompt-file", str(caps["role_prompt"])] if schema_text is not None: args += ["--json-schema", schema_text] return args +def step_env(caps: dict) -> dict: + """Return the environment of a step. + + Without ANTHROPIC_API_KEY, print mode uses your Harvard login. When the step sets an effort level, + CLAUDE_CODE_EFFORT_LEVEL is removed too, because it would override --effort. + """ + drop = {"ANTHROPIC_API_KEY"} + if caps.get("effort"): + drop.add("CLAUDE_CODE_EFFORT_LEVEL") + return {name: value for name, value in os.environ.items() if name not in drop} + + class Agent: """Runs claude -p in the project folder. stop() ends a running call.""" @@ -53,9 +100,8 @@ def __init__(self, project): self.process = None def run(self, prompt: str, caps: dict, schema_text: str | None = None) -> dict: - env = {name: value for name, value in os.environ.items() if name != "ANTHROPIC_API_KEY"} process = subprocess.Popen( - command(prompt, caps, schema_text), cwd=self.project, env=env, + command(prompt, caps, schema_text), cwd=self.project, env=step_env(caps), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) self.process = process diff --git a/src/pruning_example/check_report.py b/src/pruning_example/check_report.py new file mode 100644 index 0000000..687dffb --- /dev/null +++ b/src/pruning_example/check_report.py @@ -0,0 +1,114 @@ +"""Check the numbers of a draft report against the records of a state folder. + +Usage: python -m pruning_example.check_report [REPORT] + +The command lists each number of the report that it cannot find in the records: the files of the state +folder and the baseline metrics, which a report often cites. It checks each number +with a decimal point, each percentage, and each whole number of 4 or more digits. It skips small whole +numbers, such as round numbers and seeds, and years. A record value matches if the report number is +a rounding of it, and a percentage matches if it is a rounding of 100 times a record value. The sign +counts. The list is a hint for you, not a verdict: +a correct number can be missing from the records, for example a difference of two values. +""" + +from __future__ import annotations + +import csv +import json +import re +import sys +from pathlib import Path + +from . import config + +# A number in prose: 83,215 or 0.9076 or -0.0123 or 0.75%. Version numbers such as 2.1.302 do not match. +NUMBER = re.compile(r"(? None: + """Add each number in a JSON value to found.""" + if isinstance(value, bool): + return + if isinstance(value, (int, float)): + found.append(float(value)) + elif isinstance(value, str): + try: + found.append(float(value)) + except ValueError: + pass + elif isinstance(value, dict): + for item in value.values(): + numbers_in(item, found) + elif isinstance(value, list): + for item in value: + numbers_in(item, found) + + +def record_values(state: Path, baseline: Path | None = None) -> list[float]: + """Return the numbers in the records of a state folder, and in the baseline metrics if given.""" + found: list[float] = [] + paths = [path for pattern in JSON_GLOBS for path in sorted(state.glob(pattern))] + for path in paths + ([baseline] if baseline is not None else []): + try: + numbers_in(json.loads(path.read_text()), found) + except (OSError, json.JSONDecodeError): + continue + for name in TABLES: + path = state / name + if path.exists(): + with open(path, newline="") as f: + for row in csv.DictReader(f, delimiter="\t"): + numbers_in(list(row.values()), found) + return found + + +def matches(sign: str, whole: str, fraction: str | None, percent: str, values: list[float]) -> bool: + """Return True if the number is a rounding of a record value. A percentage must be a rounding of + 100 times a record value, so a fraction written as a percentage does not match.""" + number = float(sign + whole.replace(",", "") + ("." + fraction if fraction else "")) + half_step = 0.5 * 10 ** -(len(fraction) if fraction else 0) + 1e-12 + factor = 100 if percent else 1 + return any(abs(factor * value - number) <= half_step for value in values) + + +def checked(whole: str, fraction: str | None, percent: str) -> bool: + """Return True for the numbers that the check covers.""" + if percent and whole == "95" and not fraction: + return False # the level of the intervals, as in "95% interval" + if fraction or percent: + return True + digits = whole.replace(",", "") + is_year = "," not in whole and len(digits) == 4 and 1900 <= int(digits) <= 2100 + return len(digits) >= 4 and not is_year + + +def unmatched_numbers(text: str, state: Path, baseline: Path | None = None) -> list[dict]: + """Return the numbers of the report that the records do not show, with their line numbers.""" + values = record_values(state, baseline) + missing = [] + for line_no, line in enumerate(text.splitlines(), 1): + for match in NUMBER.finditer(line): + sign, whole, fraction, percent = match.groups() + if checked(whole, fraction, percent) and not matches(sign, whole, fraction, percent, values): + missing.append({"number": match.group(0), "line": line_no}) + return missing + + +def main(argv=None) -> int: + argv = sys.argv[1:] if argv is None else argv + paths = config.paths_from_env() + report = Path(argv[0]) if argv else paths.state / "final_report.md" + if not report.exists(): + print(f"There is no report: {report}") + return 2 + missing = unmatched_numbers(report.read_text(), paths.state, paths.project / config.BASELINE_METRICS) + for item in missing: + print(f"line {item['line']}: {item['number']} is not in the records") + print(f"{len(missing)} number(s) not found in {paths.state.name}/" if missing else "Each checked number is in the records.") + return 1 if missing else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/pruning_example/config.py b/src/pruning_example/config.py index e61c264..6764faa 100644 --- a/src/pruning_example/config.py +++ b/src/pruning_example/config.py @@ -14,6 +14,7 @@ SCHEMA_PATH = Path("configs/round.schema.json") SPLIT_PATH = Path("data/split.json") BASELINE_MODEL = Path("results/baseline/model.pt") +BASELINE_METRICS = Path("state/baseline/metrics.json") class ConfigError(ValueError): diff --git a/src/pruning_example/final_eval.py b/src/pruning_example/final_eval.py index 8503a59..9dfd121 100644 --- a/src/pruning_example/final_eval.py +++ b/src/pruning_example/final_eval.py @@ -1,9 +1,10 @@ """One-time test evaluation of one model. -Usage: python -m pruning_example.final_eval --round N [--kind round|confirm] +Usage: python -m pruning_example.final_eval --round N [--kind round|confirm] [--seed S] The command refuses a second evaluation of the same model, and it logs each call. The loop script runs this command in the final steps. The agent cannot run it. +--seed selects a confirmation round with seed 2 or more; seed 1 keeps the paths of a single confirmation round. """ from __future__ import annotations @@ -24,19 +25,20 @@ def main(argv=None) -> int: parser = argparse.ArgumentParser(description="Evaluate one model on the test set, one time.") parser.add_argument("--round", type=int, required=True) parser.add_argument("--kind", choices=["round", "confirm"], default="round") + parser.add_argument("--seed", type=int, default=None, help="seed of a confirmation round (default: 1)") args = parser.parse_args(argv) paths = config.paths_from_env() paths.logs.mkdir(parents=True, exist_ok=True) with open(paths.logs / "final_eval.log", "a") as log: log.write(f"{time.strftime('%Y-%m-%dT%H:%M:%S')} job={records.job_id()} " - f"kind={args.kind} round={args.round}\n") - out = paths.state / "test_eval" / f"{args.kind}_{args.round:02d}.json" + f"kind={args.kind} round={args.round}{'' if args.seed is None else f' seed={args.seed}'}\n") + out = records.test_eval_path(paths.state, args.kind, args.round, args.seed) if out.exists(): print(f"Refused: {out} exists. The test set evaluates each model one time only.") return 2 - record_dir = paths.state / ("rounds" if args.kind == "round" else "confirm") / f"round_{args.round:02d}" + record_dir = records.record_dir(paths.state, args.kind, args.round, args.seed) meta = records.read_json(record_dir / "run_meta.json") torch.set_num_threads(len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else 1) model = build_model() @@ -51,6 +53,8 @@ def main(argv=None) -> int: "test_images": len(test), "time": time.strftime("%Y-%m-%dT%H:%M:%S"), } + if args.seed is not None: + result["seed"] = args.seed records.write_json(out, result) print(f"test evaluation, {args.kind} {args.round}: test_accuracy={test_accuracy:.4f}") return 0 diff --git a/src/pruning_example/finish.py b/src/pruning_example/finish.py index 523d998..58ffd06 100644 --- a/src/pruning_example/finish.py +++ b/src/pruning_example/finish.py @@ -1,4 +1,4 @@ -"""The final steps: select a round, run the confirmation round, evaluate on the test set, +"""The final steps: select a round, run the confirmation rounds, evaluate on the test set, write the summary, run the report step, and copy the model.""" from __future__ import annotations @@ -11,10 +11,13 @@ from pathlib import Path from . import config, final_eval, metrics, records -from .agent import AgentError, tokens_used -from .slurm import SlurmError +from .agent import AgentError, ModelCheckError, check_model, tokens_used +from .slurm import Accounting, SlurmError RESULT_KEYS = ("val_accuracy", "val_interval", "nonzero", "weights_left", "sparse_size_mb", "passes") +SELECTION_FILE = "selection.json" +PENDING_FILE = "pending_jobs.tsv" # jobs whose GPU minutes are not in the ledger yet +DEFAULT_TEST_TARGET = 0.900 # the goal of the example: a test accuracy of at least 90% def row_passes(row: dict, target: dict) -> bool: @@ -51,39 +54,200 @@ def select_round(history: list[dict], target: dict) -> dict | None: return max(ok, key=lambda row: float(row["val_accuracy"]), default=None) +def confirmation_seeds(experiment: dict) -> list[int]: + """Return the seeds of the confirmation rounds. The default is one confirmation round with seed 1.""" + return [int(seed) for seed in experiment.get("confirmation_seeds", [1])] + + +def saved_selection(ctx, history: list[dict]) -> int | None: + """Return the selected round. Select it the first time, and save it before the first confirmation round + or test evaluation, so that a later --finish uses the same round.""" + path = ctx.paths.state / SELECTION_FILE + if path.exists(): + return records.read_json(path)["round"] + row = select_round(history, ctx.experiment["target"]) + round_no = int(row["round"]) if row is not None else None + records.write_json(path, {"round": round_no, "time": datetime.now().isoformat(timespec="seconds")}) + return round_no + + +def job_accounting(ctx, job: str) -> Accounting: + """Return the accounting of a job. If sacct has no record of it, count its time limit, + so that the GPU budget stays safe, and continue.""" + try: + return ctx.slurm.accounting(job) + except (SlurmError, subprocess.TimeoutExpired) as error: + ctx.log(f"job {job}: no accounting ({error}); counting its time limit") + return Accounting("UNKNOWN", float(ctx.experiment["round"]["time_minutes"]), 0) + + +def refuse_after_final_steps(ctx) -> None: + """Refuse a new search in a state folder whose final steps are done. Its test set is used, and its + selection is saved, so new rounds would not change the result.""" + path = ctx.paths.state / "summary.json" + if path.exists() and records.read_json(path).get("final_steps_done"): + raise SystemExit(f"The final steps of {ctx.paths.state.name}/ are done. To search again, use a new state " + "folder (STATE_DIR). To run the final steps again, use --finish.") + + +def round_job_options(ctx) -> tuple: + """Return the sbatch options of each round job. An optional round.constraint in the experiment file + replaces the --constraint line of the job scripts, so that all designs of a comparison use the same GPUs.""" + constraint = ctx.experiment["round"].get("constraint") + return (f"--constraint={constraint}",) if constraint else () + + +def note_pending_job(ctx, job: str, kind: str, round_no: int, tasks: int = 0) -> None: + """Note a submitted job until its GPU minutes are in the ledger. tasks is the number of array tasks, + or 0 for one job. If the loop job stops first, for example after scancel, the next loop job counts it.""" + with open(ctx.paths.state / PENDING_FILE, "a") as pending: + pending.write(f"{job}\t{kind}\t{round_no}\t{tasks}\n") + + +def mark_counted(ctx, job: str) -> None: + """Remove a job from the pending jobs after its ledger rows are written.""" + path = ctx.paths.state / PENDING_FILE + if not path.exists(): + return + lines = [line for line in path.read_text().splitlines() if line.split("\t")[0] != job] + if lines: + path.write_text("\n".join(lines) + "\n") + else: + path.unlink() + + +def count_pending_jobs(ctx) -> int: + """Write the ledger rows of the jobs that an earlier loop job submitted but did not count, with the + status interrupted. Cancel each job first, in case it still runs. Return the number of jobs.""" + path = ctx.paths.state / PENDING_FILE + if not path.exists(): + return 0 + jobs = [line.split("\t") for line in path.read_text().splitlines() if line.count("\t") == 3] + for job, kind, round_no, tasks in jobs: + ctx.slurm.cancel(job, wait_seconds=15) + if int(tasks): + usage = ctx.slurm.array_accounting(job, int(tasks)) + for index in range(int(tasks)): + task = usage.get(index) + ctx.ledger.add(round=int(round_no), kind=kind, status="interrupted", + job_id=task.job_id if task else f"{job}_{index}", + # without accounting, count the time limit, so that the GPU total stays safe + gpu_minutes=task.gpu_minutes if task else ctx.experiment["round"]["time_minutes"], + slurm_state=task.state if task else "UNKNOWN") + else: + usage = job_accounting(ctx, job) + ctx.ledger.add(round=int(round_no), kind=kind, status="interrupted", job_id=job, + gpu_minutes=usage.gpu_minutes, slurm_state=usage.state) + mark_counted(ctx, job) # one job at a time, so that an error cannot count a job two times + path.unlink(missing_ok=True) # lines that are not jobs + if jobs: + ctx.log(f"counted the GPU minutes of {len(jobs)} job(s) that an earlier loop job did not count") + return len(jobs) + + def confirmation_round(ctx, round_no: int) -> dict | None: """Run the selected configuration again with seed 1. Return its metrics, or None.""" state = ctx.paths.state - record = state / "confirm" / f"round_{round_no:02d}" / "metrics.json" + record = records.record_dir(state, "confirm", round_no, 1) / "metrics.json" if record.exists(): return records.read_json(record) config_path = state.relative_to(ctx.paths.project) / "rounds" / f"round_{round_no:02d}" / "config.json" args = [str(config_path), "--round", str(round_no), "--kind", "confirm", "--seed", "1"] try: - job = ctx.slurm.submit("slurm/round.sbatch", args, ctx.paths.project) + job = ctx.slurm.submit("slurm/round.sbatch", args, ctx.paths.project, options=round_job_options(ctx)) except (SlurmError, subprocess.TimeoutExpired) as error: ctx.log(f"confirmation round: sbatch failed: {error}") return None (state / "current_job").write_text(job + "\n") + note_pending_job(ctx, job, "confirm", round_no) ctx.log(f"confirmation round: submitted job {job}") outcome = ctx.slurm.wait(job, 60 * ctx.budget["queue_limit_minutes"]) if outcome == "queue_limit": ctx.slurm.cancel(job) - usage = ctx.slurm.accounting(job) + usage = job_accounting(ctx, job) (state / "current_job").unlink(missing_ok=True) ctx.ledger.add(round=round_no, kind="confirm", status=outcome, job_id=job, gpu_minutes=usage.gpu_minutes, slurm_state=usage.state) + mark_counted(ctx, job) return records.read_json(record) if record.exists() else None -def evaluate_on_test_set(ctx, kind: str, round_no: int) -> dict | None: +def confirmation_rounds(ctx, round_no: int, seeds: list[int]) -> dict[int, dict | None]: + """Run the confirmation rounds of the selected configuration. Return the metrics of each seed, or None. + + One seed runs as one round job, as in chapter 28. Several seeds run as one array job. A seed that + already has its metrics does not run again. + """ + if seeds == [1]: + return {1: confirmation_round(ctx, round_no)} + state, project = ctx.paths.state, ctx.paths.project + missing = [seed for seed in seeds + if not (records.record_dir(state, "confirm", round_no, seed) / "metrics.json").exists()] + if missing: + config_path = (records.record_dir(state, "round", round_no) / "config.json").relative_to(project) + tasks = state / "confirm" / f"round_{round_no:02d}_tasks.tsv" + tasks.parent.mkdir(parents=True, exist_ok=True) + tasks.write_text("".join(f"{config_path}\t{round_no}\tconfirm\t{seed}\n" for seed in missing)) + max_parallel = ctx.experiment.get("wave", {}).get("max_parallel", len(missing)) + try: + job = ctx.slurm.submit_array("slurm/round_array.sbatch", str(tasks.relative_to(project)), + len(missing), max_parallel, project, options=round_job_options(ctx)) + except (SlurmError, subprocess.TimeoutExpired) as error: + ctx.log(f"confirmation rounds: sbatch failed: {error}") + else: + (state / "current_job").write_text(job + "\n") + note_pending_job(ctx, job, "confirm", round_no, len(missing)) + ctx.log(f"confirmation rounds: submitted array job {job} for seeds {', '.join(map(str, missing))}") + outcome = ctx.slurm.wait_array(job, 60 * ctx.budget["queue_limit_minutes"]) + usage = ctx.slurm.array_accounting(job, len(missing)) + (state / "current_job").unlink(missing_ok=True) + for index, seed in enumerate(missing): + task = usage.get(index) + ctx.ledger.add(round=round_no, kind="confirm", status=outcome, + job_id=task.job_id if task else f"{job}_{index}", + # without accounting, count the time limit, so that the GPU total stays safe + gpu_minutes=task.gpu_minutes if task else ctx.experiment["round"]["time_minutes"], + slurm_state=task.state if task else "UNKNOWN") + mark_counted(ctx, job) + results = {} + for seed in seeds: + record = records.record_dir(state, "confirm", round_no, seed) / "metrics.json" + results[seed] = records.read_json(record) if record.exists() else None + return results + + +def evaluate_on_test_set(ctx, kind: str, round_no: int, seed: int | None = None) -> dict | None: """Return the one-time test evaluation of a model. Run it if it does not exist yet.""" - out = ctx.paths.state / "test_eval" / f"{kind}_{round_no:02d}.json" + out = records.test_eval_path(ctx.paths.state, kind, round_no, seed) if not out.exists(): - final_eval.main(["--round", str(round_no), "--kind", kind]) + final_eval.main(["--round", str(round_no), "--kind", kind] + ([] if seed is None else ["--seed", str(seed)])) return records.read_json(out) if out.exists() else None +def seed_means(selected: dict, confirmations: dict[int, dict | None], tests: dict, target: dict) -> dict: + """Return the means of the selected round and its confirmation rounds on the validation and the test set. + + The result holds if both means keep their thresholds. Only the test set is independent of the selection. + """ + complete = all(confirmations.values()) + vals = [selected["val_accuracy"]] + [result["val_accuracy"] for result in confirmations.values() if result] + test_keys = ["selected"] + [("confirmation" if seed == 1 else f"confirmation_seed{seed}") for seed in confirmations] + test_values = [tests[key]["test_accuracy"] for key in test_keys if tests.get(key)] + val_mean = sum(vals) / len(vals) if complete else None + test_mean = sum(test_values) / len(test_values) if complete and len(test_values) == len(test_keys) else None + test_target = target.get("test_accuracy", DEFAULT_TEST_TARGET) + holds = None if val_mean is None or test_mean is None else ( + val_mean >= target["val_accuracy"] and test_mean >= test_target) + return { + "seeds": 1 + len(confirmations), + "complete": complete and test_mean is not None, + "val_accuracy": None if val_mean is None else round(val_mean, 5), + "test_accuracy": None if test_mean is None else round(test_mean, 5), + "test_target": test_target, + "holds": holds, + } + + def previous_stop_reason(ctx) -> str | None: """Return the stop reason of an earlier loop job, or None.""" path = ctx.paths.state / "summary.json" @@ -99,7 +263,9 @@ def agent_models(ctx) -> list[str]: def write_summary(ctx, stop_reason: str, final: bool, selected=None, confirmation=None, tests=None, - terminated: bool = False) -> dict: + terminated: bool = False, confirmations=None, means=None, extra=None) -> dict: + """Write state/summary.json. confirmations and means appear only with more than one confirmation seed, + and extra adds the fields of the multi-agent loop.""" budget = ctx.budget history = records.read_history(ctx.paths.state) experiment_file = config.experiment_path(ctx.paths.project) @@ -130,6 +296,11 @@ def write_summary(ctx, stop_reason: str, final: bool, selected=None, confirmatio "confirmation": confirmation, "test": tests or None, } + if confirmations is not None: + summary["confirmations"] = confirmations + if means is not None: + summary["means"] = means + summary.update(extra or {}) records.write_json(ctx.paths.state / "summary.json", summary) return summary @@ -150,31 +321,54 @@ def run_report_step(ctx) -> None: records.write_json(ctx.paths.state / "agent_steps" / "report.json", result) ctx.ledger.add(round="", kind="report", status=result.get("subtype", ""), tokens=tokens_used(result), cost_usd=result.get("total_cost_usd", 0)) + try: + check_model(result, ctx.experiment["report_step"].get("model")) + except ModelCheckError as error: + ctx.log(f"report step: {error}") (ctx.paths.state / "final_report.md").write_text(str(result.get("result", "")).strip() + "\n") -def run_final_steps(ctx, stop_reason: str) -> dict: +def run_final_steps(ctx, stop_reason: str, report_step=None, extra=None) -> dict: + """Select a round, confirm it, evaluate it on the test set, copy the model, and write the summary. + + report_step replaces the report step of chapter 28, and extra(ctx) returns more summary fields; + the multi-agent loop uses both. + """ state = ctx.paths.state - row = select_round(records.read_history(state), ctx.experiment["target"]) - selected = confirmation = None + seeds = confirmation_seeds(ctx.experiment) + round_no = saved_selection(ctx, records.read_history(state)) + selected = confirmation = confirmations = means = None tests = {} - if row is not None: - round_no = int(row["round"]) - folder = state / "rounds" / f"round_{round_no:02d}" + if round_no is not None: + folder = records.record_dir(state, "round", round_no) result = records.read_json(folder / "metrics.json") selected = {"round": round_no, "seed": result.get("seed", 0), "config": records.read_json(folder / "config.json"), **{key: result[key] for key in RESULT_KEYS}} - confirm = confirmation_round(ctx, round_no) - if confirm: - confirmation = {"seed": confirm["seed"], **{key: confirm[key] for key in RESULT_KEYS}} + confirms = confirmation_rounds(ctx, round_no, seeds) + if confirms.get(1): + confirmation = {"seed": confirms[1]["seed"], **{key: confirms[1][key] for key in RESULT_KEYS}} tests["selected"] = evaluate_on_test_set(ctx, "round", round_no) - if confirm: - tests["confirmation"] = evaluate_on_test_set(ctx, "confirm", round_no) + for seed, confirm in confirms.items(): + if not confirm: + continue + if seed == 1: + tests["confirmation"] = evaluate_on_test_set(ctx, "confirm", round_no) + else: + tests[f"confirmation_seed{seed}"] = evaluate_on_test_set(ctx, "confirm", round_no, seed) + if len(seeds) > 1: + confirmations = [{"seed": seed, **({key: confirm[key] for key in RESULT_KEYS} if confirm else {"missing": True})} + for seed, confirm in confirms.items()] + means = seed_means(result, confirms, tests, ctx.experiment["target"]) model = Path(records.read_json(folder / "run_meta.json")["run_folder"]) / "model.pt" if model.exists(): # scratch files are temporary, so keep a copy on lab storage target = ctx.paths.project / "results" / state.name target.mkdir(parents=True, exist_ok=True) shutil.copyfile(model, target / f"round_{round_no:02d}_model.pt") - write_summary(ctx, stop_reason, True, selected, confirmation, tests) - run_report_step(ctx) - return write_summary(ctx, stop_reason, True, selected, confirmation, tests) + + def summary() -> dict: + return write_summary(ctx, stop_reason, True, selected, confirmation, tests, confirmations=confirmations, + means=means, extra=extra(ctx) if extra else None) + + summary() + (report_step or run_report_step)(ctx) + return summary() diff --git a/src/pruning_example/loop.py b/src/pruning_example/loop.py index 7ae01aa..8042438 100644 --- a/src/pruning_example/loop.py +++ b/src/pruning_example/loop.py @@ -21,12 +21,14 @@ from pathlib import Path from . import config, finish, records -from .agent import Agent, AgentError, AgentFormatError, tokens_used +from .agent import Agent, AgentError, AgentFormatError, ModelCheckError, check_model, tokens_used from .config import ConfigError from .records import IntegrityError from .slurm import Slurm, SlurmError LEDGER_FIELDS = ["time", "round", "kind", "status", "tokens", "cost_usd", "job_id", "gpu_minutes", "slurm_state"] +TOKEN_KINDS = frozenset({"agent_step", "report"}) # the rows of the ledger that hold tokens +STEP_KINDS = frozenset({"agent_step"}) # the rows that decide the next round class Terminated(Exception): @@ -48,10 +50,16 @@ def log(message: str) -> None: class Ledger: - """The budget record in state/ledger.tsv: tokens of each claude -p call and GPU minutes of each job.""" + """The budget record in state/ledger.tsv: tokens of each claude -p call and GPU minutes of each job. - def __init__(self, path: Path): + The multi-agent loop gives its own columns and the kinds of its role steps. + """ + + def __init__(self, path: Path, fields: list[str] = LEDGER_FIELDS, token_kinds=TOKEN_KINDS, step_kinds=STEP_KINDS): self.path = path + self.fields = list(fields) + self.token_kinds = frozenset(token_kinds) + self.step_kinds = frozenset(step_kinds) self.rows = [] if path.exists(): with open(path, newline="") as f: @@ -62,31 +70,31 @@ def add(self, **row) -> None: self.path.parent.mkdir(parents=True, exist_ok=True) new = not self.path.exists() with open(self.path, "a", newline="") as f: - writer = csv.DictWriter(f, fieldnames=LEDGER_FIELDS, delimiter="\t", extrasaction="ignore") + writer = csv.DictWriter(f, fieldnames=self.fields, delimiter="\t", extrasaction="ignore") if new: writer.writeheader() - writer.writerow({name: row.get(name, "") for name in LEDGER_FIELDS}) - self.rows.append({name: str(row.get(name, "")) for name in LEDGER_FIELDS}) + writer.writerow({name: row.get(name, "") for name in self.fields}) + self.rows.append({name: str(row.get(name, "")) for name in self.fields}) - def total(self, column: str, kinds: set[str]) -> float: + def total(self, column: str, kinds) -> float: return sum(float(row[column] or 0) for row in self.rows if row["kind"] in kinds) @property def tokens(self) -> int: - return int(self.total("tokens", {"agent_step", "report"})) + return int(self.total("tokens", self.token_kinds)) @property def largest_step(self) -> int: - return max((int(float(row["tokens"] or 0)) for row in self.rows if row["kind"] == "agent_step"), default=0) + return max((int(float(row["tokens"] or 0)) for row in self.rows if row["kind"] in self.step_kinds), default=0) @property def gpu_minutes(self) -> float: - """GPU minutes of the unattended rounds. The confirmation round is separate.""" + """GPU minutes of the unattended rounds. The confirmation rounds are separate.""" return self.total("gpu_minutes", {"round"}) @property def cost_usd(self) -> float: - return self.total("cost_usd", {"agent_step", "report"}) + return self.total("cost_usd", self.token_kinds) @dataclass @@ -132,7 +140,10 @@ def stop_reason(ctx: Context) -> str | None: next_step = ctx.ledger.largest_step or budget["step_token_estimate"] if ctx.ledger.tokens + next_step > step_limit: return "token_budget" - if ctx.ledger.gpu_minutes >= budget["gpu_minutes"]: + if budget.get("gpu_reservation"): # start a round only if its full time limit fits in the budget + if ctx.ledger.gpu_minutes + ctx.experiment["round"]["time_minutes"] > budget["gpu_minutes"]: + return "gpu_budget" + elif ctx.ledger.gpu_minutes >= budget["gpu_minutes"]: return "gpu_budget" if failed_in_a_row(history) >= budget["max_failed_in_a_row"]: return "failed_attempts" @@ -150,6 +161,14 @@ def render(path: Path, values: dict) -> str: return text +def describe_best(best: dict | None) -> str: + """Describe the best round so far for a prompt.""" + if not best: + return "none yet" + return (f"round {best['round']}, {int(best['nonzero']):,} nonzero weights " + f"({100 * float(best['weights_left']):.2f}% left), validation accuracy {best['val_accuracy']}") + + def fail(ctx: Context, round_no: int, status: str, note: str, round_config: dict | None = None) -> None: """Write the history row of a failed round attempt.""" records.append_history(ctx.paths.state, records.history_row(round_no, status, round_config, None, note[:200])) @@ -168,9 +187,7 @@ def run_attempt(ctx: Context, round_no: int) -> None: "N": round_no, "R": ctx.budget["max_round_attempts"] - len(history), "G": round(max(0.0, ctx.budget["gpu_minutes"] - ctx.ledger.gpu_minutes)), - "BEST": (f"round {best['round']}, {int(best['nonzero']):,} nonzero weights " - f"({100 * float(best['weights_left']):.2f}% left), validation accuracy {best['val_accuracy']}" - if best else "none yet"), + "BEST": describe_best(best), "STALE": finish.rounds_without_improvement(history, target), "PLATEAU": target["plateau_rounds"], }) @@ -186,6 +203,7 @@ def run_attempt(ctx: Context, round_no: int) -> None: tokens = tokens_used(result) ctx.ledger.add(round=round_no, kind="agent_step", status=result.get("subtype", ""), tokens=tokens, cost_usd=result.get("total_cost_usd", 0)) + check_model(result, ctx.experiment["agent_step"].get("model")) round_config = result.get("structured_output") try: @@ -201,19 +219,22 @@ def run_attempt(ctx: Context, round_no: int) -> None: notes.write(f"## Round {round_no}\n\n{round_config['reason']}\n\n") try: - job = ctx.slurm.submit("slurm/round.sbatch", [str(config_path), "--round", str(round_no)], ctx.paths.project) + job = ctx.slurm.submit("slurm/round.sbatch", [str(config_path), "--round", str(round_no)], ctx.paths.project, + options=finish.round_job_options(ctx)) except (SlurmError, subprocess.TimeoutExpired) as error: fail(ctx, round_no, "sbatch_failed", str(error), round_config) return (state / "current_job").write_text(job + "\n") + finish.note_pending_job(ctx, job, "round", round_no) ctx.log(f"round {round_no}: submitted job {job}") outcome = ctx.slurm.wait(job, 60 * ctx.budget["queue_limit_minutes"]) if outcome == "queue_limit": ctx.slurm.cancel(job) - usage = ctx.slurm.accounting(job) + usage = finish.job_accounting(ctx, job) (state / "current_job").unlink(missing_ok=True) ctx.ledger.add(round=round_no, kind="round", status=outcome, job_id=job, gpu_minutes=usage.gpu_minutes, slurm_state=usage.state) + finish.mark_counted(ctx, job) rows = [row for row in records.read_history(state) if int(row["round"]) == round_no] if outcome == "queue_limit": @@ -245,20 +266,33 @@ def run_rounds(ctx: Context) -> str: except AgentFormatError as error: ctx.log(f"{error}. Check the Claude Code version.") return "agent_format" + except ModelCheckError as error: + ctx.log(f"{error}. Check which models your organization allows.") + return "model_check" + + +def check_time_lines(ctx: Context, scripts: tuple[str, ...]) -> None: + """Refuse to start if the --time line of a round job script does not agree with round.time_minutes.""" + experiment_file = config.experiment_path(ctx.paths.project).relative_to(ctx.paths.project) + for script in scripts: + text = (ctx.paths.project / script).read_text() + match = re.search(r"^#SBATCH --time=(\d+):(\d+):(\d+)", text, re.M) + minutes = int(match.group(1)) * 60 + int(match.group(2)) if match else None + if minutes != ctx.experiment["round"]["time_minutes"]: + raise SystemExit(f"{script} has --time of {minutes} minutes, but {experiment_file} " + f"has round.time_minutes {ctx.experiment['round']['time_minutes']}.") def check_start(ctx: Context) -> None: - """Refuse to start without the baseline and round 1, or with settings that do not agree.""" + """Refuse to start without the baseline and round 1, with settings that do not agree, or after the final steps.""" project = ctx.paths.project + finish.refuse_after_final_steps(ctx) if not (project / "state" / "baseline" / "metrics.json").exists(): raise SystemExit("There is no baseline. Run 'sbatch slurm/baseline.sbatch' first.") if not (ctx.paths.state / "rounds" / "round_01" / "metrics.json").exists(): raise SystemExit("There is no round 1. Run the supervised round first.") - match = re.search(r"^#SBATCH --time=(\d+):(\d+):(\d+)", (project / "slurm" / "round.sbatch").read_text(), re.M) - minutes = int(match.group(1)) * 60 + int(match.group(2)) if match else None - if minutes != ctx.experiment["round"]["time_minutes"]: - raise SystemExit(f"slurm/round.sbatch has --time of {minutes} minutes, but experiment.yaml " - f"has round.time_minutes {ctx.experiment['round']['time_minutes']}.") + several_seeds = len(finish.confirmation_seeds(ctx.experiment)) > 1 # their confirmation rounds run as an array job + check_time_lines(ctx, ("slurm/round.sbatch",) + (("slurm/round_array.sbatch",) if several_seeds else ())) records.check_integrity(project) notes = ctx.paths.state / "notes.md" if not notes.exists(): @@ -288,9 +322,11 @@ def run(ctx: Context, finish_only: bool = False) -> int: f"state folder: {ctx.paths.state.relative_to(ctx.paths.project)}") try: if finish_only: + finish.count_pending_jobs(ctx) reason = finish.previous_stop_reason(ctx) or "finish_only" else: check_start(ctx) + finish.count_pending_jobs(ctx) reason = run_rounds(ctx) ctx.log(f"stop: {reason}") finish.run_final_steps(ctx, reason) diff --git a/src/pruning_example/multi_loop.py b/src/pruning_example/multi_loop.py new file mode 100644 index 0000000..4c2832a --- /dev/null +++ b/src/pruning_example/multi_loop.py @@ -0,0 +1,1060 @@ +"""The multi-agent loop script: plan, check, and run waves of rounds, then the final steps. + +Usage: python -m pruning_example.multi_loop [--plan-only | --finish | --audit-only] + +slurm/multi_loop.sbatch starts this script with exec, so the script gets the SLURM signals. +--plan-only plans wave 1 and stops, so that you can read the plan before any GPU job runs. +--finish runs only the final steps, for example after an early stop. +--audit-only runs the auditor and check_report on the report of STATE_DIR. + +Four roles do the agent work, each in its own claude -p call with read-only tools: the planner, the critic, +the writer, and the auditor. This script runs the role steps, checks their output in code, submits one +array job for each wave, writes all records, and applies the stop rules. experiment_multi.yaml sets the +model, the effort level, and the caps of each role. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +import signal +import subprocess +from dataclasses import dataclass +from pathlib import Path + +import jsonschema + +from . import check_report, config, finish, records +from .agent import Agent, AgentError, AgentFormatError, EFFORT_LEVELS, ModelCheckError, check_model, tokens_used +from .config import ConfigError +from .loop import ( + LEDGER_FIELDS, Context, Ledger, StopLoop, Terminated, check_time_lines, describe_best, job_end_time, log, + on_sigterm, render, +) +from .records import IntegrityError +from .slurm import Slurm, SlurmError + +ROLES = ("planner", "critic", "writer", "auditor") +# The kinds of role steps in the ledger, and the role of each kind +ROLE_OF = {"planner": "planner", "critic": "critic", "revision": "planner", "critic_2": "critic", + "writer": "writer", "auditor": "auditor"} +WAVE_KINDS = ("planner", "critic", "revision", "critic_2") +MULTI_LEDGER_FIELDS = LEDGER_FIELDS + ["wave", "model", "effort"] +EXTRA_KEYS = ("hypothesis", "predicted_val_accuracy", "pass_probability") +HYPOTHESIS_FIELDS = ["wave", "item", "round", "status", "verdict", "hypothesis", "predicted_val_accuracy", + "pass_probability", "val_accuracy", "passes", "error"] +ROUND_SCRIPTS = ("slurm/round.sbatch", "slurm/round_array.sbatch") +SCHEMA_FILES = {"wave": "configs/wave.schema.json", "review": "configs/review.schema.json", + "audit": "configs/audit.schema.json"} + + +@dataclass +class MultiContext(Context): + """The context of the loop of chapter 28, plus the schemas of the role outputs.""" + + schemas: dict | None = None + plan_only: bool = False + + +# Settings and limits + +def roles(ctx) -> dict: + return ctx.experiment["roles"] + + +def critic_on(ctx) -> bool: + return bool(roles(ctx)["critic"].get("enabled", True)) + + +def check_roles(experiment: dict) -> None: + """Refuse to start without a model, a valid effort level, and the caps of each role.""" + for role in ROLES: + settings = (experiment.get("roles") or {}).get(role) or {} + if not settings.get("model"): + raise SystemExit(f"Set roles.{role}.model in the experiment file.") + if settings.get("effort") not in EFFORT_LEVELS: + raise SystemExit(f"roles.{role}.effort must be one of {', '.join(EFFORT_LEVELS)}.") + for key in ("max_turns", "max_budget_usd", "timeout_minutes"): + if key not in settings: + raise SystemExit(f"Set roles.{role}.{key} in the experiment file.") + + +def role_caps(ctx, role: str) -> dict: + caps = {key: value for key, value in roles(ctx)[role].items() if key != "enabled"} + caps["role_prompt"] = f"prompts/multi/roles/{role}.md" + return caps + + +def rel(ctx, path: Path) -> Path: + return path.relative_to(ctx.paths.project) + + +def wave_dir(ctx, wave: int) -> Path: + return ctx.paths.state / "waves" / f"wave_{wave:02d}" + + +def wave_numbers(ctx) -> list[int]: + folders = (ctx.paths.state / "waves").glob("wave_*") + return sorted(int(folder.name.split("_")[1]) for folder in folders if folder.is_dir()) + + +def search_limit(ctx) -> float: + return ctx.budget["tokens"] * (1 - ctx.budget["final_reserve_fraction"]) + + +def wave_tokens(ctx, wave: int) -> int: + return int(sum(float(row["tokens"] or 0) for row in ctx.ledger.rows + if row["kind"] in WAVE_KINDS and row.get("wave") == str(wave))) + + +def wave_estimate(ctx) -> int: + """The tokens of the next wave plan: the largest wave so far, or the estimate before the first wave.""" + done = [wave for wave in wave_numbers(ctx) if (wave_dir(ctx, wave) / "outcome.json").exists()] + return max((wave_tokens(ctx, wave) for wave in done), default=0) or ctx.budget["wave_token_estimate"] + + +def gpu_room(ctx) -> int: + """The rounds whose full time limit still fits in the GPU budget.""" + left = ctx.budget["gpu_minutes"] - ctx.ledger.gpu_minutes + return max(0, math.floor(left / ctx.experiment["round"]["time_minutes"] + 1e-9)) + + +def wave_size(ctx, history: list[dict]) -> int: + """The smallest of the wave size, the GPU room, the round attempts left, and the rounds left before the + plateau rule stops the run. So no limit is passed by a part of a wave.""" + target = ctx.experiment["target"] + return max(0, min(ctx.experiment["wave"]["size"], gpu_room(ctx), + ctx.budget["max_round_attempts"] - len(history), + target["plateau_rounds"] - finish.rounds_without_improvement(history, target))) + + +def time_for(ctx, minutes: float) -> bool: + return ctx.job_end is None or ctx.job_end - ctx.clock() >= 60 * minutes + + +def plan_minutes(ctx) -> float: + """The longest time of the role steps of one wave: planner and revision, and two critic steps.""" + minutes = 2 * roles(ctx)["planner"]["timeout_minutes"] + if critic_on(ctx): + minutes += 2 * roles(ctx)["critic"]["timeout_minutes"] + return minutes + + +def array_minutes(ctx) -> float: + """The time that the loop job needs to run one array job and the final steps.""" + return (ctx.experiment["round"]["time_minutes"] + ctx.budget["queue_limit_minutes"] + + ctx.budget["final_reserve_minutes"]) + + +def failed_waves(ctx) -> int: + """Count the waves in a row, from the last one, that ended with no complete round.""" + count = 0 + for wave in reversed(wave_numbers(ctx)): + outcome = wave_dir(ctx, wave) / "outcome.json" + if not outcome.exists(): + continue + if records.read_json(outcome)["complete_rounds"] > 0: + break + count += 1 + return count + + +def stop_reason(ctx, planned: bool = False) -> str | None: + """Return the reason to stop before the next wave, or None. A wave that is already planned needs no + tokens and no time for its plan.""" + budget = ctx.budget + target = ctx.experiment["target"] + history = records.read_history(ctx.paths.state) + if len(history) >= budget["max_round_attempts"]: + return "round_limit" + if finish.rounds_without_improvement(history, target) >= target["plateau_rounds"]: + return "plateau" + if not planned and ctx.ledger.tokens + wave_estimate(ctx) > search_limit(ctx): + return "token_budget" + if gpu_room(ctx) < 1: + return "gpu_budget" + if failed_waves(ctx) >= budget["max_failed_waves"]: + return "failed_attempts" + if not planned and not ctx.plan_only and not time_for(ctx, plan_minutes(ctx) + budget["final_reserve_minutes"]): + return "time_left" + return None + + +# Role steps + +def run_role(ctx, kind: str, prompt: str, schema_text: str | None = None, wave: int | None = None) -> dict | None: + """Run one role step. Return the JSON result, or None if the step failed. + + A result without token counts stops the loop (AgentFormatError), and so does a result that does not + show the configured model (ModelCheckError). + """ + caps = role_caps(ctx, ROLE_OF[kind]) + common = {"wave": wave or "", "model": caps["model"], "effort": caps["effort"]} + if ctx.plan_only and not time_for(ctx, caps["timeout_minutes"]): + raise StopLoop("time_left") + try: + result = ctx.agent.run(prompt, caps, schema_text) + except AgentFormatError: + raise + except AgentError as error: + ctx.ledger.add(round="", kind=kind, status="failed", tokens=0, **common) + ctx.log(f"{f'wave {wave}: ' if wave else ''}{kind} step failed: {error}") + return None + name = f"wave_{wave:02d}_{kind}.json" if wave else f"{kind}.json" + records.write_json(ctx.paths.state / "agent_steps" / name, result) + ctx.ledger.add(round="", kind=kind, status=result.get("subtype", ""), tokens=tokens_used(result), + cost_usd=result.get("total_cost_usd", 0), **common) + check_model(result, caps["model"]) + return result + + +def structured(result: dict | None, schema: dict) -> dict | None: + """Return the structured output of a role step if it matches the schema. The documentation does not + list the schema keywords that Claude Code supports, so the code checks the output again.""" + output = (result or {}).get("structured_output") + if not isinstance(output, dict): + return None + try: + jsonschema.validate(output, schema) + except jsonschema.ValidationError: + return None + return output + + +def wave_schema(ctx, count: int) -> dict: + schema = json.loads(ctx.schemas["wave"]) + schema["properties"]["configurations"]["minItems"] = count + schema["properties"]["configurations"]["maxItems"] = count + return schema + + +# The wave plan + +def split_item(item: dict) -> tuple[dict, dict]: + """Split a planned item into the round configuration and the extra keys, which round.schema.json does not allow.""" + round_config = {key: value for key, value in item.items() if key not in EXTRA_KEYS} + return round_config, {key: item.get(key) for key in EXTRA_KEYS} + + +def config_key(round_config: dict) -> tuple: + """The six settings of a configuration, with the numbers rounded and the excluded layers as a set.""" + return (round(float(round_config["sparsity"]), 6), int(round_config["steps"]), + int(round_config["finetune_epochs"]), round(float(round_config["lr"]), 6), + round(float(round_config["weight_decay"]), 6), frozenset(round_config["exclude_layers"])) + + +def earlier_keys(ctx) -> set: + """The settings of each configuration that already ran or was submitted.""" + state = ctx.paths.state + keys = set() + for path in list((state / "configs").glob("round_*.json")) + list((state / "rounds").glob("round_*/config.json")): + try: + keys.add(config_key(records.read_json(path))) + except (KeyError, TypeError, ValueError): + continue + return keys + + +def code_check(ctx, entries: list[dict], earlier: set, in_wave: set) -> None: + """Check each entry in code: the schema, the extra rule, and no repeat. Reject the entries that fail.""" + for entry in entries: + try: + config.validate_round_config(dict(entry["config"]), ctx.schema) + except ConfigError as error: + entry.update(status="rejected", reasons=[f"code: {error}"]) + continue + key = config_key(entry["config"]) + if key in earlier: + entry.update(status="rejected", reasons=["code: it repeats an earlier configuration"]) + elif key in in_wave: + entry.update(status="rejected", reasons=["code: it repeats another configuration of this wave"]) + else: + in_wave.add(key) + + +def new_entries(items: list[dict], first_item: int, stage: str) -> list[dict]: + entries = [] + for offset, item in enumerate(items): + round_config, extras = split_item(item) + entries.append({"item": first_item + offset, "stage": stage, "config": round_config, "extras": extras, + "status": "proposed", "reasons": []}) + return entries + + +def describe_config(round_config: dict) -> str: + excluded = ", ".join(round_config["exclude_layers"]) or "none" + return (f"sparsity {round_config['sparsity']}, steps {round_config['steps']}, " + f"epochs {round_config['finetune_epochs']}, lr {round_config['lr']}, " + f"weight decay {round_config['weight_decay']}, excluded {excluded}") + + +def write_proposed(path: Path, entries: list[dict]) -> None: + records.write_json(path, [{"item": entry["item"], **entry["config"], **entry["extras"]} for entry in entries]) + + +def prompt_values(ctx, wave: int, count: int) -> dict: + state = ctx.paths.state + history = records.read_history(state) + target = ctx.experiment["target"] + return { + "W": wave, + "K": count, + "R": ctx.budget["max_round_attempts"] - len(history), + "G": round(max(0.0, ctx.budget["gpu_minutes"] - ctx.ledger.gpu_minutes)), + "BEST": describe_best(finish.best_round(history, target)), + "STALE": finish.rounds_without_improvement(history, target), + "PLATEAU": target["plateau_rounds"], + "STATE": rel(ctx, state), + "EXPERIMENT": config.experiment_path(ctx.paths.project).relative_to(ctx.paths.project), + } + + +def review(ctx, wave: int, kind: str, entries: list[dict], items_file: Path) -> str | None: + """Run a critic step on the entries that passed the code checks, and set their status. + Return the critic's summary, or None if the critic step failed.""" + if not entries: + return "" + write_proposed(items_file, entries) + values = {"W": wave, "PROPOSED": rel(ctx, items_file), "STATE": rel(ctx, ctx.paths.state), + "EXPERIMENT": config.experiment_path(ctx.paths.project).relative_to(ctx.paths.project)} + result = run_role(ctx, kind, render(ctx.paths.project / "prompts" / "multi" / "critic.md", values), + ctx.schemas["review"], wave) + output = structured(result, json.loads(ctx.schemas["review"])) if result else None + if output is None: + if result is not None: + ctx.log(f"wave {wave}: the {kind} step returned no valid verdicts") + return None + verdicts = {verdict["item"]: verdict for verdict in output["verdicts"]} + for entry in entries: + verdict = verdicts.get(entry["item"]) + if verdict is None: + entry.update(status="rejected", reasons=["failed: the critic gave no verdict"]) + elif verdict["accept"]: + entry.update(status="accepted", verdict=verdict["reason"]) + else: + entry.update(status="rejected", reasons=[f"critic: {verdict['reason']}"]) + return output["summary"] + + +def accept_or_review(ctx, wave: int, kind: str, entries: list[dict], items_file: Path, reviews: list[str]) -> bool: + """Send the entries that passed the code checks to the critic, or accept them if the critic is off. + Return False if the critic step failed.""" + passed = [entry for entry in entries if entry["status"] == "proposed"] + if not critic_on(ctx): + for entry in passed: + entry["status"] = "accepted" + return True + summary = review(ctx, wave, kind, passed, items_file) + if summary is None: + return False + if summary: + reviews.append(summary) + return True + + +def plan_wave(ctx, wave: int, count: int) -> list[dict] | None: + """Make the plan of a wave: planner step, code checks, critic step, and at most one revision. + + Return the accepted entries, or None if the planner step or the first critic step failed. + """ + folder = wave_dir(ctx, wave) + folder.mkdir(parents=True, exist_ok=True) + project = ctx.paths.project + values = prompt_values(ctx, wave, count) + schema = wave_schema(ctx, count) + result = run_role(ctx, "planner", render(project / "prompts" / "multi" / "planner.md", values), + json.dumps(schema), wave) + plan = structured(result, schema) if result else None + if plan is None: + ctx.log(f"wave {wave}: the planner returned no valid plan") + return None + + entries = new_entries(plan["configurations"], 1, "plan") + write_proposed(folder / "proposed.json", entries) + earlier, in_wave, reviews = earlier_keys(ctx), set(), [] + code_check(ctx, entries, earlier, in_wave) + if not accept_or_review(ctx, wave, "critic", entries, folder / "review_items.json", reviews): + write_plan(ctx, wave, plan["summary"], entries, reviews, failed="critic") + return None + + rejected = [entry for entry in entries if entry["status"] == "rejected"] + if rejected: + lines = [f"- Item {entry['item']}: {describe_config(entry['config'])}. Rejected: {'; '.join(entry['reasons'])}" + for entry in rejected] + values.update(K=len(rejected), REJECTED="\n".join(lines), PROPOSED=rel(ctx, folder / "proposed.json")) + schema = wave_schema(ctx, len(rejected)) + result = run_role(ctx, "revision", render(project / "prompts" / "multi" / "revision.md", values), + json.dumps(schema), wave) + revision = structured(result, schema) if result else None + if revision is None: + for entry in rejected: + entry["status"] = "dropped" + else: + replacements = new_entries(revision["configurations"], len(entries) + 1, "revision") + write_proposed(folder / "proposed_2.json", replacements) + code_check(ctx, replacements, earlier, in_wave) + if not accept_or_review(ctx, wave, "critic_2", replacements, folder / "review_items_2.json", reviews): + for entry in replacements: + if entry["status"] == "proposed": + entry.update(status="dropped", reasons=["failed: the second critic step failed"]) + for entry in replacements: + if entry["status"] == "rejected": + entry["status"] = "dropped" + for entry in rejected: + entry["status"] = "replaced" + entries += replacements + accepted = [entry for entry in entries if entry["status"] == "accepted"] + write_plan(ctx, wave, plan["summary"], entries, reviews) + records.write_json(folder / "accepted.json", accepted) + return accepted + + +def write_plan(ctx, wave: int, summary: str, entries: list[dict], reviews: list[str], + failed: str | None = None) -> None: + """Write plan.json for the records and plan.md for you.""" + folder = wave_dir(ctx, wave) + records.write_json(folder / "plan.json", {"wave": wave, "summary": summary, "reviews": reviews, + "failed_step": failed, "critic": critic_on(ctx), "entries": entries}) + lines = [f"# Wave {wave}", "", f"Planner: {summary}", ""] + lines += [f"Critic: {text}" for text in reviews] + ([""] if reviews else []) + lines += ["| Item | Stage | Sparsity | Steps | Epochs | LR | Weight decay | Excluded | Predicted val | " + "Pass probability | Status |", + "|---|---|---|---|---|---|---|---|---|---|---|"] + for entry in entries: + cfg, extras = entry["config"], entry["extras"] + lines.append(f"| {entry['item']} | {entry['stage']} | {cfg['sparsity']} | {cfg['steps']} | " + f"{cfg['finetune_epochs']} | {cfg['lr']} | {cfg['weight_decay']} | " + f"{', '.join(cfg['exclude_layers']) or 'none'} | {extras['predicted_val_accuracy']} | " + f"{extras['pass_probability']} | {entry['status']} |") + for entry in entries: + lines += ["", f"## Item {entry['item']} ({entry['status']})", "", + f"Hypothesis: {entry['extras']['hypothesis']}", "", f"Reason: {entry['config']['reason']}"] + if entry.get("verdict"): + lines += ["", f"Critic: {entry['verdict']}"] + for reason in entry["reasons"]: + source, _, text = reason.partition(":") + lines += ["", f"Rejected ({source}): {text.strip()}"] + if failed: + lines += ["", f"The {failed} step failed, so the wave has no accepted configuration."] + lines += ["", "To run this plan, start the loop job: sbatch slurm/multi_loop.sbatch", + f"To plan again, remove {rel(ctx, folder)} and run the --plan-only job again.", ""] + (folder / "plan.md").write_text("\n".join(lines)) + + +# Rounds and records + +def next_round_number(ctx) -> int: + """A round number above each round in history.tsv, each round folder, and each submitted task.""" + state = ctx.paths.state + numbers = [int(row["round"]) for row in records.read_history(state)] + numbers += [int(folder.name.split("_")[1]) for folder in (state / "rounds").glob("round_*") + if folder.name.split("_")[1].isdigit()] + for tasks in (state / "waves").glob("wave_*/tasks.tsv"): + numbers += [int(line.split("\t")[1]) for line in tasks.read_text().splitlines() if line.count("\t") >= 3] + return max(numbers, default=0) + 1 + + +def history_rounds(state: Path) -> set[int]: + return {int(row["round"]) for row in records.read_history(state)} + + +def ok_row(state: Path, round_no: int) -> dict: + """The history row of a complete round, with the job ID of its round job.""" + record = records.record_dir(state, "round", round_no) + return records.history_row(round_no, "ok", records.read_json(record / "config.json"), + records.read_json(record / "metrics.json"), + job=records.read_json(record / "run_meta.json")["job_id"]) + + +def write_missing_rows(ctx) -> int: + """Write the history row of each submitted round that has metrics but no row. A stop between the end + of a task and its row then loses no round.""" + state = ctx.paths.state + have = history_rounds(state) + added = 0 + for tasks in sorted((state / "waves").glob("wave_*/tasks.tsv")): + for line in tasks.read_text().splitlines(): + parts = line.split("\t") + if len(parts) < 4 or parts[2] != "round": + continue + round_no = int(parts[1]) + if round_no in have or not (records.record_dir(state, "round", round_no) / "metrics.json").exists(): + continue + records.append_history(state, ok_row(state, round_no)) + have.add(round_no) + added += 1 + if added: + ctx.log(f"wrote {added} missing history row(s)") + return added + + +def write_outcome(ctx, wave: int, complete: int, note: str, **more) -> None: + records.write_json(wave_dir(ctx, wave) / "outcome.json", {"complete_rounds": complete, "note": note, **more}) + + +def add_round_records(ctx, wave: int, rounds: list[tuple[int, dict]], job: str, outcome: str) -> int: + """Write the ledger rows and the history rows of the rounds of an array job, in round order. + Return the number of complete rounds.""" + state = ctx.paths.state + usage = ctx.slurm.array_accounting(job, len(rounds)) + have = history_rounds(state) + complete = 0 + for index, (round_no, round_config) in enumerate(rounds): + task = usage.get(index) + ctx.ledger.add(round=round_no, kind="round", status=outcome, wave=wave, + job_id=task.job_id if task else f"{job}_{index}", + # without accounting, count the time limit, so that the GPU budget stays safe + gpu_minutes=task.gpu_minutes if task else ctx.experiment["round"]["time_minutes"], + slurm_state=task.state if task else "UNKNOWN") + if (records.record_dir(state, "round", round_no) / "metrics.json").exists(): + complete += 1 + if round_no not in have: + records.append_history(state, ok_row(state, round_no)) + continue + if round_no in have: + continue + never_ran = task is None or task.runs == 0 + status = "queue_limit" if outcome == "queue_limit" and never_ran else "job_failed" + note = f"task {index} of array job {job} ended as {task.state if task else 'UNKNOWN'} without metrics" + records.append_history(state, records.history_row(round_no, status, round_config, None, note, + job=task.job_id if task else job)) + return complete + + +def close_interrupted_wave(ctx, wave: int) -> None: + """Finish the records of a wave whose loop job stopped before the wave ended.""" + folder = wave_dir(ctx, wave) + rounds = [] + for line in (folder / "tasks.tsv").read_text().splitlines(): + parts = line.split("\t") + if len(parts) >= 4: + rounds.append((int(parts[1]), records.read_json(ctx.paths.project / parts[0]))) + job = (folder / "array_job").read_text().strip() if (folder / "array_job").exists() else None + complete = add_round_records(ctx, wave, rounds, job, "interrupted") if job else 0 + write_outcome(ctx, wave, complete, "the loop job stopped during the wave", array_job=job) + ctx.log(f"wave {wave}: closed after an interruption, {complete} complete round(s)") + + +def write_notes(ctx, wave: int, entries: list[dict]) -> None: + lines = [f"## Wave {wave}", ""] + for entry in entries: + if entry.get("round"): + extras = entry["extras"] + lines += [f"### Round {entry['round']} (item {entry['item']})", "", + f"Hypothesis: {extras['hypothesis']}", + f"Predicted validation accuracy {extras['predicted_val_accuracy']}, " + f"pass probability {extras['pass_probability']}.", "", entry["config"]["reason"], ""] + elif entry["reasons"]: + lines += [f"### Item {entry['item']}, not run ({entry['status']})", "", "; ".join(entry["reasons"]), ""] + with open(ctx.paths.state / "notes.md", "a") as notes: + notes.write("\n".join(lines) + "\n") + + +def submit_wave(ctx, wave: int, accepted: list[dict]) -> int: + """Submit the accepted configurations of a wave as one array job, wait, and write the records. + Return the number of complete rounds.""" + state, project, folder = ctx.paths.state, ctx.paths.project, wave_dir(ctx, wave) + accepted = accepted[:wave_size(ctx, records.read_history(state))] # a limit can have changed since the plan + first = next_round_number(ctx) + lines, rounds = [], [] + for offset, entry in enumerate(accepted): + entry["round"] = first + offset + path = state / "configs" / f"round_{entry['round']:02d}.json" + records.write_json(path, entry["config"]) + lines.append(f"{rel(ctx, path)}\t{entry['round']}\tround\t0") + rounds.append((entry["round"], entry["config"])) + (folder / "tasks.tsv").write_text("\n".join(lines) + "\n") + records.write_json(folder / "submitted.json", [{"item": entry["item"], "round": entry["round"]} for entry in accepted]) + plan = records.read_json(folder / "plan.json") + submitted = {entry["item"]: entry["round"] for entry in accepted} + write_notes(ctx, wave, [{**entry, "round": submitted.get(entry["item"])} for entry in plan["entries"]]) + try: + job = ctx.slurm.submit_array("slurm/round_array.sbatch", str(rel(ctx, folder / "tasks.tsv")), len(lines), + ctx.experiment["wave"]["max_parallel"], project, + options=finish.round_job_options(ctx)) + except (SlurmError, subprocess.TimeoutExpired) as error: + for round_no, round_config in rounds: + records.append_history(state, records.history_row(round_no, "sbatch_failed", round_config, None, + str(error)[:200])) + write_outcome(ctx, wave, 0, f"sbatch failed: {error}"[:300]) + write_hypotheses(ctx) + ctx.log(f"wave {wave}: sbatch failed: {error}") + return 0 + (folder / "array_job").write_text(job + "\n") + (state / "current_job").write_text(job + "\n") + ctx.log(f"wave {wave}: submitted array job {job}, rounds {first} to {first + len(lines) - 1}") + outcome = ctx.slurm.wait_array(job, 60 * ctx.budget["queue_limit_minutes"]) + complete = add_round_records(ctx, wave, rounds, job, outcome) + (state / "current_job").unlink(missing_ok=True) + write_outcome(ctx, wave, complete, outcome, array_job=job) + write_hypotheses(ctx) + log_wave(ctx, wave, [round_no for round_no, _ in rounds]) + if outcome == "queue_limit": + raise StopLoop("queue_limit") + return complete + + +def log_wave(ctx, wave: int, round_numbers: list[int]) -> None: + rows = {int(row["round"]): row for row in records.read_history(ctx.paths.state)} + parts = [f"round {n}: val_accuracy={rows[n]['val_accuracy'] or '-'} weights_left={rows[n]['weights_left'] or '-'}" + for n in round_numbers if n in rows] + by_kind = {kind: int(sum(float(row["tokens"] or 0) for row in ctx.ledger.rows + if row["kind"] == kind and row.get("wave") == str(wave))) for kind in WAVE_KINDS} + tokens = ", ".join(f"{kind} {count}" for kind, count in by_kind.items() if count) + ctx.log(f"wave {wave}: {'; '.join(parts)}; tokens {tokens} " + f"(totals: {ctx.ledger.tokens} tokens, {ctx.ledger.gpu_minutes:.1f} GPU minutes)") + + +# Hypotheses and prediction scores + +def write_hypotheses(ctx) -> list[dict]: + """Rebuild hypotheses.tsv from the plans of the submitted waves and the round results.""" + state, target = ctx.paths.state, ctx.experiment["target"] + rows = [] + for wave in wave_numbers(ctx): + folder = wave_dir(ctx, wave) + if not (folder / "submitted.json").exists(): + continue + submitted = {item["item"]: item["round"] for item in records.read_json(folder / "submitted.json")} + for entry in records.read_json(folder / "plan.json")["entries"]: + round_no = submitted.get(entry["item"]) + metrics_file = records.record_dir(state, "round", round_no) / "metrics.json" if round_no else None + result = records.read_json(metrics_file) if metrics_file and metrics_file.exists() else None + predicted = entry["extras"]["predicted_val_accuracy"] + row = {"wave": wave, "item": entry["item"], "round": round_no or "", + "status": "run" if round_no else entry["status"], "verdict": entry.get("verdict", ""), + "hypothesis": " ".join(str(entry["extras"]["hypothesis"]).split()), + "predicted_val_accuracy": predicted, "pass_probability": entry["extras"]["pass_probability"], + "val_accuracy": "", "passes": "", "error": ""} + if result: + row.update(val_accuracy=f"{result['val_accuracy']:.4f}", + passes="yes" if result["val_accuracy"] >= target["val_accuracy"] else "no", + error=f"{result['val_accuracy'] - float(predicted):+.4f}") + rows.append(row) + path = state / "hypotheses.tsv" + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=HYPOTHESIS_FIELDS, delimiter="\t") + writer.writeheader() + writer.writerows(rows) + os.replace(tmp, path) + return rows + + +def prediction_scores(ctx, rows: list[dict]) -> dict: + """Score the planner's predictions next to simple baselines. + + Mean absolute error: next to the accuracy of the earlier round whose sparsity is nearest. + Brier score of the pass probabilities: next to always forecasting the pass rate of this run. + Noise in 5,000 validation images sets a floor of about 0.004 for the mean absolute error. + """ + state = ctx.paths.state + done = [row for row in rows if row["status"] == "run" and row["val_accuracy"] != ""] + if not done: + return {"rounds": 0} + history = [row for row in records.read_history(state) if row["status"] == "ok"] + errors, baseline_errors, brier, right = [], [], [], 0 + passes = [row["passes"] == "yes" for row in done] + rate = sum(passes) / len(passes) + for row, passed in zip(done, passes): + actual, predicted = float(row["val_accuracy"]), float(row["predicted_val_accuracy"]) + errors.append(abs(actual - predicted)) + probability = float(row["pass_probability"]) + brier.append((probability - passed) ** 2) + right += (probability >= 0.5) == passed + sparsity = float(records.read_json(records.record_dir(state, "round", int(row["round"])) / "config.json")["sparsity"]) + earlier = [item for item in history if int(item["round"]) < int(row["round"]) + and not same_wave(rows, item["round"], row["wave"])] + if earlier: + nearest = min(earlier, key=lambda item: abs(math.log(1 - float(item["sparsity"])) - math.log(1 - sparsity))) + baseline_errors.append(abs(actual - float(nearest["val_accuracy"]))) + return { + "rounds": len(done), + "mean_absolute_error": round(sum(errors) / len(errors), 4), + "baseline_mean_absolute_error": round(sum(baseline_errors) / len(baseline_errors), 4) if baseline_errors else None, + "brier_score": round(sum(brier) / len(brier), 4), + "base_rate_brier_score": round(rate * (1 - rate), 4), + "pass_decisions_right": right, + } + + +def same_wave(rows: list[dict], round_no: str, wave) -> bool: + """Rounds of one wave ran at the same time, so their results were not known when the wave was planned.""" + return any(str(row["round"]) == str(round_no) and str(row["wave"]) == str(wave) for row in rows) + + +# Final steps + +def final_role_step(ctx, kind: str, prompt: str, schema_text: str | None, failures: dict) -> dict | None: + """Run the writer step or the auditor step. A failure, or a step that used another model, is recorded + in failures and does not stop the code checks that follow.""" + try: + result = run_role(ctx, kind, prompt, schema_text) + except (AgentFormatError, ModelCheckError) as error: + failures[kind] = str(error) + ctx.log(f"{kind} step: {error}") + return None + if result is None: + failures[kind] = "the step failed" + return result + + +def report_steps(ctx) -> None: + """The writer step and the auditor step, then the code checks. The role steps run if the total is below + the token budget; the search limit keeps the final reserve free for them.""" + write_hypotheses(ctx) + state = ctx.paths.state + report = state / "final_report.md" + failures = {} + if ctx.ledger.tokens >= ctx.budget["tokens"]: + ctx.log("writer and auditor skipped: the token budget is used up") + failures = {kind: "skipped: the token budget is used up" for kind in ("writer", "auditor")} + else: + prompt = render(ctx.paths.project / "prompts" / "multi" / "writer.md", {"STATE": rel(ctx, state)}) + result = final_role_step(ctx, "writer", prompt, None, failures) + if result is not None: + report.write_text(str(result.get("result", "")).strip() + "\n") + run_audit(ctx, report, failures) + + +def run_audit(ctx, report: Path, failures: dict | None = None) -> dict: + """Run the auditor step, check_report, and the check_round checks of the selected round. Write audit.json. + The code checks run also when a role step failed or was skipped.""" + state = ctx.paths.state + failures = dict(failures or {}) + schema = json.loads(ctx.schemas["audit"]) + auditor = None + if not report.exists(): + failures.setdefault("auditor", "skipped: there is no report") + elif "auditor" not in failures: + prompt = render(ctx.paths.project / "prompts" / "multi" / "auditor.md", + {"STATE": rel(ctx, state), "REPORT": rel(ctx, report)}) + result = final_role_step(ctx, "auditor", prompt, ctx.schemas["audit"], failures) + auditor = structured(result, schema) + if auditor is None and "auditor" not in failures: # for example error_max_turns: no structured output + failures["auditor"] = f"no valid audit (subtype {result.get('subtype')})" + audit = { + "report": str(rel(ctx, report)), + "auditor": auditor, + "auditor_step_failed": auditor is None, + "failures": failures, + "unmatched_numbers": check_report.unmatched_numbers( + report.read_text(), state, ctx.paths.project / config.BASELINE_METRICS) if report.exists() else None, + "check_round": selected_round_checks(ctx), + } + records.write_json(state / "audit.json", audit) + return audit + + +def selected_round_checks(ctx) -> dict | None: + """Run the checks of check_round on the selected round.""" + from . import check_round # it loads PyTorch, so import it only here + + selection = ctx.paths.state / finish.SELECTION_FILE + round_no = records.read_json(selection)["round"] if selection.exists() else None + if round_no is None: + return None + try: + checks = check_round.run_checks(round_no) + except (OSError, KeyError, ValueError, RuntimeError) as error: + return {"round": round_no, "ok": False, "error": str(error)[:300]} + return {"round": round_no, "ok": all(ok for _, ok, _ in checks), + "checks": [{"name": name, "ok": ok, "detail": detail} for name, ok, detail in checks]} + + +def audit_summary(ctx) -> dict | None: + path = ctx.paths.state / "audit.json" + if not path.exists(): + return None + audit = records.read_json(path) + claims = (audit.get("auditor") or {}).get("claims", []) + return { + "claims": len(claims), + "wrong": sum(claim["verdict"] == "wrong" for claim in claims), + "unsupported": sum(claim["verdict"] == "unsupported" for claim in claims), + "limits_stated": (audit.get("auditor") or {}).get("limits_stated"), + "auditor_step_failed": audit.get("auditor_step_failed"), + "failures": audit.get("failures", {}), + "unmatched_numbers": None if audit.get("unmatched_numbers") is None else len(audit["unmatched_numbers"]), + "check_round_ok": (audit.get("check_round") or {}).get("ok"), + } + + +def summary_extra(ctx) -> dict: + """The fields of the summary that only the multi-agent loop has.""" + role_info = {} + for role in ROLES: + kinds = {kind for kind, owner in ROLE_OF.items() if owner == role} + settings = roles(ctx)[role] + role_info[role] = { + "model": settings["model"], "effort": settings["effort"], + "enabled": critic_on(ctx) if role == "critic" else True, + "steps": sum(row["kind"] in kinds for row in ctx.ledger.rows), + "failed_steps": sum(row["kind"] in kinds and row["status"] == "failed" for row in ctx.ledger.rows), + "tokens": int(ctx.ledger.total("tokens", kinds)), + "estimated_cost_usd": round(ctx.ledger.total("cost_usd", kinds), 2), + } + # Each planned configuration counts one time. "not_reviewed": a critic step failed before its verdict. + # "not_revised": rejected, but the revision did not run. + counts = {"planned": 0, "accepted": 0, "replaced": 0, "dropped": 0, "not_reviewed": 0, "not_revised": 0, + "rejected_by_code": 0, "rejected_by_critic": 0} + final_status = {"accepted": "accepted", "replaced": "replaced", "dropped": "dropped", + "proposed": "not_reviewed", "rejected": "not_revised"} + waves = [] + for wave in wave_numbers(ctx): + folder = wave_dir(ctx, wave) + if (folder / "plan.json").exists(): + for entry in records.read_json(folder / "plan.json")["entries"]: + counts["planned"] += 1 + counts[final_status[entry["status"]]] += 1 + counts["rejected_by_code"] += any(reason.startswith("code:") for reason in entry["reasons"]) + counts["rejected_by_critic"] += any(reason.startswith("critic:") for reason in entry["reasons"]) + if (folder / "outcome.json").exists(): + waves.append({"wave": wave, "tokens": wave_tokens(ctx, wave), **records.read_json(folder / "outcome.json")}) + rows = write_hypotheses(ctx) if (ctx.paths.state / "waves").exists() else [] + return { + "design": "planner and critic, waves" if critic_on(ctx) else "one agent, waves", + "roles": role_info, + "waves": waves, + "configurations": counts, + "predictions": prediction_scores(ctx, rows), + "audit": audit_summary(ctx), + } + + +# Start, plan-only, audit-only, and the loop + +def copy_round_one(ctx) -> None: + """Copy round 1 of chapter 28 from state/ to the state folder of this loop, one time.""" + state, source = ctx.paths.state, ctx.paths.project / "state" + target = records.record_dir(state, "round", 1) + if state == source or (target / "metrics.json").exists(): + return + origin = records.record_dir(source, "round", 1) + if not (origin / "metrics.json").exists(): + return + for name in ("config.json", "run_meta.json", "metrics.json"): + records.write_json(target / name, records.read_json(origin / name)) + rows = [row for row in records.read_history(source) if row["round"] == "1"] + if rows and 1 not in history_rounds(state): + records.append_history(state, rows[0]) + notes = state / "notes.md" + if not notes.exists(): + notes.write_text("# Notes\n\nThe loop script adds the reasons of the planner and the critic here.\n\n" + f"## Round 1\n\n{records.read_json(target / 'config.json').get('reason', '')}\n\n") + write_hypotheses(ctx) + ctx.log(f"copied round 1 from state/ to {rel(ctx, state)}/") + + +def check_start(ctx, plan_only: bool = False) -> None: + """Refuse to start without the baseline, round 1, and a planned wave 1, with settings that do not agree, + or after the final steps.""" + project = ctx.paths.project + finish.refuse_after_final_steps(ctx) + if not (project / "state" / "baseline" / "metrics.json").exists(): + raise SystemExit("There is no baseline. Run 'sbatch slurm/baseline.sbatch' first.") + copy_round_one(ctx) + if not (records.record_dir(ctx.paths.state, "round", 1) / "metrics.json").exists(): + raise SystemExit("There is no round 1. Run the supervised round of the single-agent example first.") + check_time_lines(ctx, ROUND_SCRIPTS) + check_roles(ctx.experiment) + records.check_integrity(project) + accepted = wave_dir(ctx, 1) / "accepted.json" + if not plan_only and not (accepted.exists() and records.read_json(accepted)): + raise SystemExit("Wave 1 has no accepted configuration. Run 'sbatch --time=01:30:00 slurm/multi_loop.sbatch " + "--plan-only', and read the plan first.") + clean_up_after_a_stop(ctx) + + +def clean_up_after_a_stop(ctx) -> None: + """Cancel the job of an earlier loop job that stopped without a cancel, close its open wave, + count the GPU minutes of its confirmation jobs, and write the missing history rows.""" + leftover = ctx.paths.state / "current_job" + if leftover.exists(): + job = leftover.read_text().strip() + ctx.slurm.cancel(job, wait_seconds=15) + leftover.unlink(missing_ok=True) + ctx.log(f"canceled the job {job} of an earlier loop job") + for wave in wave_numbers(ctx): + folder = wave_dir(ctx, wave) + if (folder / "tasks.tsv").exists() and not (folder / "outcome.json").exists(): + close_interrupted_wave(ctx, wave) + finish.count_pending_jobs(ctx) + write_missing_rows(ctx) + + +def plan_first_wave(ctx) -> int: + """Plan wave 1 and stop, so that you can read the plan before any GPU job runs.""" + check_start(ctx, plan_only=True) + folder = wave_dir(ctx, 1) + if (folder / "accepted.json").exists(): + if not records.read_json(folder / "accepted.json"): + ctx.log(f"Wave 1 has no accepted configuration. Remove {rel(ctx, folder)}, and plan again.") + return 1 + ctx.log(f"Wave 1 already has a plan: read {rel(ctx, folder / 'plan.md')}. " + f"To plan again, remove {rel(ctx, folder)}.") + return 0 + reason = stop_reason(ctx) + if reason: + ctx.log(f"cannot plan wave 1: {reason}") + return 1 + try: + accepted = plan_wave(ctx, 1, wave_size(ctx, records.read_history(ctx.paths.state))) + except StopLoop as stop: + ctx.log(f"plan of wave 1 stopped: {stop.reason}") + return 1 + except (AgentFormatError, ModelCheckError) as error: + ctx.log(f"plan of wave 1 stopped: {error}") + return 1 + if accepted is None: + ctx.log("plan of wave 1 stopped: a role step failed; see the lines above. Run the --plan-only job again.") + return 1 + if not accepted: + ctx.log(f"Wave 1 has no accepted configuration. Read {rel(ctx, folder / 'plan.md')}, " + f"remove {rel(ctx, folder)}, and plan again.") + return 1 + ctx.log(f"wave 1 planned: {len(accepted)} configuration(s) accepted. Read {rel(ctx, folder / 'plan.md')}, " + "then start the loop job with 'sbatch slurm/multi_loop.sbatch'.") + return 0 + + +def audit_only(ctx) -> int: + """Run the auditor step and check_report on the report of the state folder, for example of the single-agent run.""" + check_roles(ctx.experiment) + report = ctx.paths.state / "final_report.md" + if not report.exists(): + raise SystemExit(f"There is no report: {rel(ctx, report)}") + try: + audit = run_audit(ctx, report) + except (AgentFormatError, ModelCheckError) as error: + ctx.log(f"audit stopped: {error}") + return 1 + ctx.log(f"audit written to {rel(ctx, ctx.paths.state / 'audit.json')}: " + f"{len(audit['unmatched_numbers'])} number(s) not found in the records") + return 0 + + +def next_wave(ctx) -> int: + """The first wave with no outcome. A wave that the loop job left during its array job is closed first.""" + for wave in wave_numbers(ctx): + folder = wave_dir(ctx, wave) + if (folder / "outcome.json").exists(): + continue + if (folder / "tasks.tsv").exists(): + close_interrupted_wave(ctx, wave) + continue + return wave + numbers = wave_numbers(ctx) + return numbers[-1] + 1 if numbers else 1 + + +def run_waves(ctx) -> str: + """Plan and run waves until a stop rule applies. Return the stop reason.""" + while True: + wave = next_wave(ctx) + folder = wave_dir(ctx, wave) + planned = (folder / "accepted.json").exists() + reason = stop_reason(ctx, planned) + if reason: + return reason + try: + records.check_integrity(ctx.paths.project) + if planned: + accepted = records.read_json(folder / "accepted.json") + else: + accepted = plan_wave(ctx, wave, wave_size(ctx, records.read_history(ctx.paths.state))) + if not accepted: + write_outcome(ctx, wave, 0, "no accepted configuration" if accepted == [] else "a role step failed") + continue + if not time_for(ctx, array_minutes(ctx)): + write_outcome(ctx, wave, 0, "not run: not enough time for the array job and the final steps") + return "time_left" + submit_wave(ctx, wave, accepted) + except StopLoop as stop: + return stop.reason + except IntegrityError as error: + ctx.log(str(error)) + return "integrity" + except AgentFormatError as error: + ctx.log(f"{error}. Check the Claude Code version.") + return "agent_format" + except ModelCheckError as error: + ctx.log(f"{error}. Check which models your organization allows.") + return "model_check" + + +def run(ctx, mode: str = "loop") -> int: + """Run the waves and the final steps. On SIGTERM, cancel the running array job and stop.""" + reason = None # the stop reason of the waves, once they end + ctx.log(f"experiment file: {config.experiment_path(ctx.paths.project).relative_to(ctx.paths.project)}, " + f"state folder: {rel(ctx, ctx.paths.state)}, mode: {mode}") + try: + if mode == "plan_only": + return plan_first_wave(ctx) + if mode == "audit_only": + return audit_only(ctx) + if mode == "finish": + clean_up_after_a_stop(ctx) + reason = finish.previous_stop_reason(ctx) or "finish_only" + else: + check_start(ctx) + reason = run_waves(ctx) + ctx.log(f"stop: {reason}") + write_missing_rows(ctx) + finish.run_final_steps(ctx, reason, report_step=report_steps, extra=summary_extra) + ctx.log(f"final steps done: see {ctx.paths.state.name}/summary.json, final_report.md, and audit.json") + return 0 + except Terminated: + # SLURM sends SIGKILL after KillWait (30 seconds on the cluster), so keep each step short. + signal.signal(signal.SIGTERM, signal.SIG_IGN) + ctx.agent.stop() + if mode in ("plan_only", "audit_only"): + ctx.log("SIGTERM: stopped.") + return 1 + # Cancel first: the records below read many files, and SIGKILL can come before they are done. + current = ctx.paths.state / "current_job" + if current.exists(): + job = current.read_text().strip() + answered = ctx.slurm.cancel(job, wait_seconds=15) + current.unlink(missing_ok=True) + ctx.log(f"SIGTERM: asked SLURM to cancel array job {job}" + f"{'' if answered else ' (scancel did not answer in 15 seconds)'}.") + write_missing_rows(ctx) + # Keep the stop reason of the waves if they ended, so that --finish can use it. + finish.write_summary(ctx, reason or "terminated", final=False, terminated=True, extra=summary_extra(ctx)) + ctx.log("SIGTERM: the loop script stopped. Run 'sbatch slurm/multi_loop.sbatch --finish' for the final steps.") + return 1 + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description="Run the multi-agent loop, then the final steps.") + mode_group = parser.add_mutually_exclusive_group() + mode_group.add_argument("--plan-only", action="store_true", help="plan wave 1, then stop") + mode_group.add_argument("--finish", action="store_true", help="run only the final steps") + mode_group.add_argument("--audit-only", action="store_true", help="audit the report of the state folder") + args = parser.parse_args(argv) + mode = "plan_only" if args.plan_only else "finish" if args.finish else "audit_only" if args.audit_only else "loop" + paths = config.paths_from_env() + project = paths.project + # An audit of another state folder must not add rows to its ledger, which can have other columns. + ledger_file = "audit_ledger.tsv" if mode == "audit_only" else "ledger.tsv" + ctx = MultiContext( + paths=paths, + experiment=config.load_experiment(project), + schema_text=(project / config.SCHEMA_PATH).read_text(), + agent=Agent(project), + slurm=Slurm(), + ledger=Ledger(paths.state / ledger_file, MULTI_LEDGER_FIELDS, ROLE_OF.keys(), WAVE_KINDS), + job_end=job_end_time(), + schemas={name: (project / path).read_text() for name, path in SCHEMA_FILES.items()}, + plan_only=mode == "plan_only", + ) + signal.signal(signal.SIGTERM, on_sigterm) + return run(ctx, mode) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/pruning_example/records.py b/src/pruning_example/records.py index 4666f1d..8170b45 100644 --- a/src/pruning_example/records.py +++ b/src/pruning_example/records.py @@ -137,9 +137,31 @@ def next_round(state: Path) -> int: return max((int(row["round"]) for row in read_history(state)), default=0) + 1 -def history_row(round_no: int, status: str, round_config: dict | None, result: dict | None, note: str = "") -> dict: - """Return one row of history.tsv from a configuration and a result.""" - row = {"round": round_no, "status": status, "job_id": job_id(), "note": note} +def record_dir(state: Path, kind: str, round_no: int, seed: int | None = None) -> Path: + """Return the record folder of a round or of a confirmation round. + + A confirmation round with seed 1 keeps the path of a single confirmation round. Seed S of 2 or more + adds _seedS, so that confirmation rounds that run at the same time write to different folders. + """ + if kind == "round": + return state / "rounds" / f"round_{round_no:02d}" + suffix = f"_seed{seed}" if seed is not None and seed >= 2 else "" + return state / "confirm" / f"round_{round_no:02d}{suffix}" + + +def test_eval_path(state: Path, kind: str, round_no: int, seed: int | None = None) -> Path: + """Return the file of the one-time test evaluation of a model.""" + suffix = f"_seed{seed}" if kind == "confirm" and seed is not None and seed >= 2 else "" + return state / "test_eval" / f"{kind}_{round_no:02d}{suffix}.json" + + +def history_row(round_no: int, status: str, round_config: dict | None, result: dict | None, note: str = "", + job: str | None = None) -> dict: + """Return one row of history.tsv from a configuration and a result. + + job is the ID of the round job. Without it, the row gets the ID of the job that writes the row. + """ + row = {"round": round_no, "status": status, "job_id": job or job_id(), "note": note} if round_config: row.update({key: round_config.get(key, "") for key in ("sparsity", "steps", "finetune_epochs", "lr", "weight_decay")}) row["exclude_layers"] = ",".join(round_config.get("exclude_layers", [])) or "none" diff --git a/src/pruning_example/run_round.py b/src/pruning_example/run_round.py index 579a012..9aaec54 100644 --- a/src/pruning_example/run_round.py +++ b/src/pruning_example/run_round.py @@ -1,14 +1,15 @@ """Entry point of the round job: prune, fine-tune, and evaluate one configuration. -Usage: python -m pruning_example.run_round CONFIG [--round N] [--kind round|confirm] [--seed S] +Usage: python -m pruning_example.run_round CONFIG [--round N] [--kind round|confirm] [--seed S] [--no-history] The job saves a checkpoint after each epoch. After a requeue, the job continues from the checkpoint. +With --no-history, the job does not write its row of history.tsv. The multi-agent loop script writes +the rows of round jobs that run at the same time, so that their rows cannot mix. """ from __future__ import annotations import argparse -import shutil import time from pathlib import Path @@ -28,12 +29,13 @@ def parse(argv=None): parser.add_argument("--round", type=int, default=None, help="round number (default: the next round)") parser.add_argument("--kind", choices=["round", "confirm"], default="round") parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--no-history", action="store_true", help="do not write the row of history.tsv") return parser.parse_args(argv) def finish_records(paths, args, round_no, record_dir, round_config) -> None: """Make sure that history.tsv has the row of a complete round. A requeue can stop a job between steps.""" - if args.kind != "round": + if args.kind != "round" or args.no_history: return if any(int(row["round"]) == round_no for row in records.read_history(paths.state)): return @@ -57,7 +59,7 @@ def main(argv=None) -> None: round_no = args.round or records.next_round(paths.state) records.write_json(folder / "round.json", {"round": round_no, "kind": args.kind}) records.write_json(folder / "config.json", round_config) - record_dir = paths.state / ("rounds" if args.kind == "round" else "confirm") / f"round_{round_no:02d}" + record_dir = records.record_dir(paths.state, args.kind, round_no, args.seed) if (record_dir / "metrics.json").exists(): done = records.read_json(record_dir / "run_meta.json") if done["job_id"] != records.job_id(): @@ -144,9 +146,8 @@ def new_optimizer(block: int): meta["end"] = time.strftime("%Y-%m-%dT%H:%M:%S") records.write_json(folder / "metrics.json", result) records.write_json(folder / "run_meta.json", meta) - record_dir.mkdir(parents=True, exist_ok=True) - for name in ("config.json", "run_meta.json", "metrics.json"): - shutil.copyfile(folder / name, record_dir / name) + for name in ("config.json", "run_meta.json", "metrics.json"): # metrics.json last: it marks a complete round + records.write_json(record_dir / name, records.read_json(folder / name)) finish_records(paths, args, round_no, record_dir, round_config) checkpoint.unlink(missing_ok=True) print(f"round {round_no}: val_accuracy={val_accuracy:.4f} weights_left={counts['weights_left']:.6f} " diff --git a/src/pruning_example/slurm.py b/src/pruning_example/slurm.py index 213e4c2..6a39ace 100644 --- a/src/pruning_example/slurm.py +++ b/src/pruning_example/slurm.py @@ -1,4 +1,4 @@ -"""Small wrappers around sbatch, squeue, sacct, and scancel, for the loop script.""" +"""Small wrappers around sbatch, squeue, sacct, and scancel, for the loop scripts.""" from __future__ import annotations @@ -30,6 +30,17 @@ class Accounting: runs: int +@dataclass +class TaskAccounting: + """The accounting of one task of an array job: all of its runs, requeues included.""" + + task: int + job_id: str # the job ID of the task (JobIDRaw), or the array ID for a task that never started + state: str + gpu_minutes: float + runs: int + + def clean_env() -> dict: """Return the environment without the SLURM variables of the loop job. @@ -60,8 +71,60 @@ def parse_sacct(text: str) -> Accounting | None: return Accounting(state=rows[-1][1].split()[0], gpu_minutes=round(minutes, 2), runs=len(rows)) +def task_indexes(job_id: str) -> list[int]: + """Return the task indexes in a JobID of an array job: 123_4, or a range of waiting tasks such as 123_[1-2%3].""" + match = re.fullmatch(r"\d+_(\d+)", job_id) + if match: + return [int(match.group(1))] + match = re.fullmatch(r"\d+_\[(.*)\]", job_id) + if not match: + return [] + indexes = [] + for part in match.group(1).split("%")[0].split(","): + if "-" in part: + first, last = part.split("-") + indexes.extend(range(int(first), int(last) + 1)) + elif part: + indexes.append(int(part)) + return indexes + + +def parse_array_sacct(text: str) -> dict[int, TaskAccounting]: + """Read the output of sacct -X -D -n -P -o JobID,JobIDRaw,State,Start,ElapsedRaw,AllocTRES for an array job. + + -D gives one line for each run of a requeued task. Tasks that never started can share one line + with a range of indexes; they have no GPU time. + """ + runs: dict[int, list[list[str]]] = {} + for line in text.strip().splitlines(): + row = line.split("|") + if len(row) < 6: + continue + for index in task_indexes(row[0]): + runs.setdefault(index, []).append(row) + tasks = {} + for index, rows in runs.items(): + # Order the runs by Start. A run that never started (a waiting run that was canceled) comes last. + rows.sort(key=lambda row: (0, row[3]) if row[3][:1].isdigit() else (1, "")) + minutes = sum(int(row[4] or 0) * gpu_count(row[5]) for row in rows) / 60 + last = rows[-1] + tasks[index] = TaskAccounting(task=index, job_id=last[1], state=last[2].split()[0], + gpu_minutes=round(minutes, 2), runs=sum(int(row[4] or 0) > 0 for row in rows)) + return tasks + + +def parse_array_squeue(text: str) -> dict[int, str]: + """Read the output of squeue -h -r -o "%K %T": one line for each task of an array job.""" + states = {} + for line in text.strip().splitlines(): + parts = line.split() + if len(parts) == 2 and parts[0].isdigit(): + states[int(parts[0])] = parts[1] + return states + + class Slurm: - """The SLURM commands of the loop script. The tests replace this class with a fake.""" + """The SLURM commands of the loop scripts. The tests replace this class with a fake.""" def __init__(self, poll_seconds: int = 60, sleep=time.sleep, clock=time.monotonic): self.poll_seconds = poll_seconds @@ -71,10 +134,10 @@ def __init__(self, poll_seconds: int = 60, sleep=time.sleep, clock=time.monotoni def run(self, args: list[str], timeout: int = 120, **kwargs) -> subprocess.CompletedProcess: return subprocess.run(args, capture_output=True, text=True, timeout=timeout, **kwargs) - def submit(self, script: str, args: list[str], cwd, attempts: int = 3) -> str: + def submit(self, script: str, args: list[str], cwd, attempts: int = 3, options: tuple = ()) -> str: """Submit a job and return its ID. Try again after a timeout of the busy SLURM controller.""" for attempt in range(1, attempts + 1): - result = self.run(["sbatch", "--parsable", script, *args], cwd=cwd, env=clean_env()) + result = self.run(["sbatch", "--parsable", *options, script, *args], cwd=cwd, env=clean_env()) if result.returncode == 0: return result.stdout.strip().split(";")[0] error = result.stderr.strip() @@ -83,6 +146,11 @@ def submit(self, script: str, args: list[str], cwd, attempts: int = 3) -> str: self.sleep(30) raise SlurmError("sbatch failed") + def submit_array(self, script: str, tasks_file: str, count: int, max_parallel: int, cwd, + options: tuple = ()) -> str: + """Submit an array job with one task for each line of tasks_file, at most max_parallel at a time.""" + return self.submit(script, [tasks_file], cwd, options=(f"--array=0-{count - 1}%{max_parallel}", *options)) + def queue_state(self, job_id: str) -> str | None: """Return the state of the job in the queue, or None if the job left the queue.""" result = self.run(["squeue", "-h", "-j", job_id, "-o", "%T"], timeout=60) @@ -92,6 +160,16 @@ def queue_state(self, job_id: str) -> str | None: raise SlurmError(f"squeue failed: {result.stderr.strip()}") return result.stdout.strip() or None + def array_states(self, array_id: str) -> dict[int, str]: + """Return the state of each task of an array job that is still in the queue. -r gives one line for + each task, also for tasks that have not started.""" + result = self.run(["squeue", "-h", "-r", "-j", array_id, "-o", "%K %T"], timeout=60) + if result.returncode != 0: + if "Invalid job id" in result.stderr: + return {} + raise SlurmError(f"squeue failed: {result.stderr.strip()}") + return parse_array_squeue(result.stdout) + def wait(self, job_id: str, queue_limit_seconds: float) -> str: """Wait until the job leaves the queue. Return "ended" or "queue_limit". @@ -119,6 +197,38 @@ def wait(self, job_id: str, queue_limit_seconds: float) -> str: return "queue_limit" self.sleep(self.poll_seconds) + def wait_array(self, array_id: str, queue_limit_seconds: float) -> str: + """Wait until all tasks of an array job leave the queue. Return "ended" or "queue_limit". + + Queue time counts while at least one task waits. After the queue limit, the waiting tasks are + canceled, also a task that a preemption put back in the queue, and the running tasks can end. + """ + queued = 0.0 + last = self.clock() + errors = 0 + over_limit = False + while True: + try: + states = self.array_states(array_id) + errors = 0 + except (SlurmError, subprocess.TimeoutExpired): + errors += 1 + if errors > 10: + raise + states = None + now = self.clock() + if states == {}: + return "queue_limit" if over_limit else "ended" + waiting = bool(states) and any(state in QUEUE_WAIT_STATES for state in states.values()) + if waiting and not over_limit: + queued += now - last + last = now + if queued > queue_limit_seconds: + over_limit = True + if over_limit and waiting: + self.cancel_waiting(array_id) + self.sleep(self.poll_seconds) + def accounting(self, job_id: str, attempts: int = 10) -> Accounting: """Return the final state and the GPU minutes of all runs of the job.""" found = None @@ -133,13 +243,39 @@ def accounting(self, job_id: str, attempts: int = 10) -> Accounting: return found raise SlurmError(f"sacct has no record of job {job_id}") + def array_accounting(self, array_id: str, count: int, attempts: int = 10) -> dict[int, TaskAccounting]: + """Return the accounting of each task of an array job. Wait until each task has a final state. + + A task that sacct does not show is left out; the caller counts its reserved GPU minutes. + """ + found: dict[int, TaskAccounting] = {} + for _ in range(attempts): + try: + result = self.run(["sacct", "-j", array_id, "-X", "-D", "-n", "-P", + "-o", "JobID,JobIDRaw,State,Start,ElapsedRaw,AllocTRES"], timeout=60) + if result.returncode == 0: + found = parse_array_sacct(result.stdout) + except subprocess.TimeoutExpired: + pass + if all(index in found and found[index].state in FINAL_STATES for index in range(count)): + return found + self.sleep(30) # the accounting database can be late after the job ends + return found + def cancel(self, job_id: str, wait_seconds: float = 60) -> bool: - """Ask SLURM to cancel the job. Return True if scancel answered in time. + """Ask SLURM to cancel the job, or all tasks of an array job. Return True if scancel answered in time. Do not raise: on a busy cluster, scancel can answer late, and SLURM cancels the job anyway. """ + return self._scancel([job_id], wait_seconds) + + def cancel_waiting(self, array_id: str, wait_seconds: float = 60) -> bool: + """Cancel only the waiting tasks of an array job. The running tasks continue.""" + return self._scancel(["--state=PENDING", array_id], wait_seconds) + + def _scancel(self, args: list[str], wait_seconds: float) -> bool: try: - process = subprocess.Popen(["scancel", job_id], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + process = subprocess.Popen(["scancel", *args], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) process.wait(timeout=wait_seconds) return process.returncode == 0 except (subprocess.TimeoutExpired, OSError): diff --git a/tests/conftest.py b/tests/conftest.py index 08048c9..03755a1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -20,7 +20,7 @@ def project(tmp_path, monkeypatch) -> Path: """Return a temporary project folder: a Git repository with the committed settings and DATA_DIR=fake.""" folder = tmp_path / "project" folder.mkdir() - for name in ("experiment.yaml", ".gitignore"): + for name in ("experiment.yaml", "experiment_multi.yaml", ".gitignore"): shutil.copy(REPO / name, folder / name) for name in ("configs", "prompts", "slurm"): shutil.copytree(REPO / name, folder / name) @@ -31,6 +31,6 @@ def project(tmp_path, monkeypatch) -> Path: monkeypatch.setenv("PROJECT_DIR", str(folder)) monkeypatch.setenv("RUN_ROOT", str(tmp_path / "runs")) monkeypatch.setenv("DATA_DIR", "fake") - for name in ("STATE_DIR", "SLURM_JOB_ID", "SLURM_RESTART_COUNT"): + for name in ("STATE_DIR", "EXPERIMENT_FILE", "SLURM_JOB_ID", "SLURM_RESTART_COUNT"): monkeypatch.delenv(name, raising=False) return folder diff --git a/tests/test_array.py b/tests/test_array.py new file mode 100644 index 0000000..f575410 --- /dev/null +++ b/tests/test_array.py @@ -0,0 +1,127 @@ +"""Tests of the SLURM array-job wrappers: parsing, waiting, the queue limit, and accounting. No jobs.""" + +from __future__ import annotations + +from pruning_example.slurm import Slurm, parse_array_sacct, parse_array_squeue, task_indexes + +TRES = "billing=8,cpu=8,gres/gpu=1,mem=32G,node=1" + + +class Result: + def __init__(self, code, out="", err=""): + self.returncode, self.stdout, self.stderr = code, out, err + + +def test_task_indexes_of_array_job_ids(): + assert task_indexes("500_4") == [4] + assert task_indexes("500_[1-2%3]") == [1, 2] + assert task_indexes("500_[0,2]") == [0, 2] + assert task_indexes("500") == [] + + +def test_sacct_adds_the_runs_of_a_requeued_task_and_keeps_tasks_that_never_started(): + text = (f"500_0|501|PREEMPTED|2026-10-21T10:00:00|300|{TRES}\n" + f"500_0|501|COMPLETED|2026-10-21T10:06:00|240|{TRES}\n" + f"500_1|502|COMPLETED|2026-10-21T10:00:05|120|{TRES}\n" + "500_[2]|500|CANCELLED by 123|Unknown|0|\n") + tasks = parse_array_sacct(text) + assert tasks[0].state == "COMPLETED" + assert tasks[0].gpu_minutes == 9.0 + assert tasks[0].runs == 2 + assert tasks[0].job_id == "501" + assert tasks[1].gpu_minutes == 2.0 + assert tasks[2].state == "CANCELLED" + assert tasks[2].gpu_minutes == 0.0 + assert tasks[2].runs == 0 + + +def test_squeue_gives_one_state_for_each_task(): + assert parse_array_squeue("0 RUNNING\n1 PENDING\n 2 PENDING\n") == {0: "RUNNING", 1: "PENDING", 2: "PENDING"} + assert parse_array_squeue("") == {} + + +def fake_slurm(states, step: int = 60): + times = iter(range(0, 100_000, step)) + slurm = Slurm(sleep=lambda seconds: None, clock=lambda: next(times)) + answers = iter(states) + slurm.array_states = lambda array_id: next(answers) + slurm.waiting_cancels = 0 + + def cancel_waiting(array_id, wait_seconds=60): + slurm.waiting_cancels += 1 + return True + + slurm.cancel_waiting = cancel_waiting + slurm.cancel = lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("must not cancel all tasks")) + return slurm + + +def test_wait_ends_when_no_task_is_left(): + slurm = fake_slurm([{0: "PENDING", 1: "PENDING"}, {0: "RUNNING", 1: "RUNNING"}, {}]) + assert slurm.wait_array("500", queue_limit_seconds=3600) == "ended" + assert slurm.waiting_cancels == 0 + + +def test_queue_limit_cancels_only_waiting_tasks_also_after_a_preemption(): + states = [ + {0: "RUNNING", 1: "PENDING"}, # 60 seconds of queue time + {0: "RUNNING", 1: "PENDING"}, # 120 seconds: past the limit, so the waiting task 1 is canceled + {0: "RUNNING"}, + {0: "PENDING"}, # a preemption put task 0 back in the queue: cancel it too + {}, + ] + slurm = fake_slurm(states) + assert slurm.wait_array("500", queue_limit_seconds=100) == "queue_limit" + assert slurm.waiting_cancels == 2 + + +def test_submit_array_passes_the_range_and_the_throttle(): + seen = [] + + def run(args, timeout=120, **kwargs): + seen.append(args) + return Result(0, out="777\n") + + slurm = Slurm(sleep=lambda seconds: None) + slurm.run = run + assert slurm.submit_array("slurm/round_array.sbatch", "state_multi/waves/wave_01/tasks.tsv", 3, 2, cwd=".") == "777" + assert seen[0] == ["sbatch", "--parsable", "--array=0-2%2", "slurm/round_array.sbatch", + "state_multi/waves/wave_01/tasks.tsv"] + + +def test_array_accounting_waits_for_final_states(): + answers = iter([Result(0, out=f"500_0|501|RUNNING|2026-10-21T10:00:00|60|{TRES}\n"), + Result(0, out=f"500_0|501|COMPLETED|2026-10-21T10:00:00|90|{TRES}\n")]) + slurm = Slurm(sleep=lambda seconds: None) + slurm.run = lambda args, timeout=120, **kwargs: next(answers) + tasks = slurm.array_accounting("500", 1) + assert tasks[0].state == "COMPLETED" + assert tasks[0].gpu_minutes == 1.5 + + +def test_array_accounting_does_not_raise_without_records(): + slurm = Slurm(sleep=lambda seconds: None) + slurm.run = lambda args, timeout=120, **kwargs: Result(1, err="sacct: error") + assert slurm.array_accounting("500", 2, attempts=3) == {} + + +def test_a_canceled_waiting_run_after_a_requeue_is_the_last_run(): + text = (f"500_0|501|REQUEUED|2026-10-21T10:00:00|300|{TRES}\n" + "500_0|501|CANCELLED by 123|Unknown|0|\n") # the queue limit canceled the requeued, waiting run + task = parse_array_sacct(text)[0] + assert task.state == "CANCELLED" + assert task.gpu_minutes == 5.0 and task.runs == 1 + + +def test_submit_array_adds_more_options_after_the_range(): + seen = [] + + def run(args, timeout=120, **kwargs): + seen.append(args) + return Result(0, out="778\n") + + slurm = Slurm(sleep=lambda seconds: None) + slurm.run = run + slurm.submit_array("slurm/round_array.sbatch", "tasks.tsv", 2, 2, cwd=".", options=("--constraint=a100-mig",)) + assert seen[0] == ["sbatch", "--parsable", "--array=0-1%2", "--constraint=a100-mig", "slurm/round_array.sbatch", + "tasks.tsv"] diff --git a/tests/test_config.py b/tests/test_config.py index b1f86ac..9609f19 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -49,9 +49,9 @@ def test_finetune_epochs_must_be_at_least_steps(): def test_experiment_file_can_point_to_a_test_copy(project, monkeypatch): - test_copy = project / "calibration" / "small_limits.yaml" + test_copy = project / "local" / "small_limits.yaml" test_copy.parent.mkdir() test_copy.write_text((project / "experiment.yaml").read_text().replace("tokens: 1000000", "tokens: 100000")) assert load_experiment(project)["budget"]["tokens"] == 1_000_000 - monkeypatch.setenv("EXPERIMENT_FILE", "calibration/small_limits.yaml") + monkeypatch.setenv("EXPERIMENT_FILE", "local/small_limits.yaml") assert load_experiment(project)["budget"]["tokens"] == 100_000 diff --git a/tests/test_loop.py b/tests/test_loop.py index 575b095..79a7b2a 100644 --- a/tests/test_loop.py +++ b/tests/test_loop.py @@ -2,10 +2,14 @@ from __future__ import annotations +import json + import pytest from pruning_example import config, finish, loop, metrics, records -from pruning_example.agent import AgentFormatError, tokens_used +from pruning_example.agent import ( + AgentFormatError, ModelCheckError, UNATTENDED_SETTINGS, check_model, command, step_env, tokens_used, +) from pruning_example.slurm import Accounting, Slurm, SlurmError, parse_sacct GOOD = {"sparsity": 0.98, "steps": 4, "finetune_epochs": 16, "lr": 0.02, "weight_decay": 0.0005, @@ -63,8 +67,10 @@ def __init__(self, state, results, gpu_minutes: float = 5.0): self.gpu_minutes = gpu_minutes self.submitted = [] self.canceled = [] + self.options = [] - def submit(self, script, args, cwd): + def submit(self, script, args, cwd, options=()): + self.options.append(options) round_no = int(args[args.index("--round") + 1]) kind = args[args.index("--kind") + 1] if "--kind" in args else "round" self.submitted.append((kind, round_no)) @@ -98,6 +104,8 @@ def context(project, agent, slurm, plateau_rounds: int = 100, **budget) -> loop. experiment = config.load_experiment(project) experiment["target"] = {**TARGET, "plateau_rounds": plateau_rounds} experiment["budget"].update(budget) + for step in ("agent_step", "report_step"): # the fake agent reports model-a; test_model_check covers the check + experiment[step].pop("model", None) return loop.Context(paths=paths, experiment=experiment, schema_text=(project / config.SCHEMA_PATH).read_text(), agent=agent, slurm=slurm, ledger=loop.Ledger(paths.state / "ledger.tsv"), log=lambda message: None) @@ -308,3 +316,150 @@ def wait(self, job_id, queue_limit_seconds): assert summary["stop_reason"] == "plateau" # --finish can use it later assert summary["terminated"] is True assert summary["final_steps_done"] is False + + +def test_model_check_stops_the_loop(started): + ctx = context(started, FakeAgent(lambda call: GOOD), FakeSlurm(started / "state", lambda n: (0.91, SMALL))) + ctx.experiment["agent_step"]["model"] = "claude-opus-5-5" # the fake agent reports model-a + assert loop.run_rounds(ctx) == "model_check" + assert ctx.ledger.tokens == 50_000 # the step still counts + + +def test_check_model_skips_aliases_and_accepts_suffixes(): + result = {"modelUsage": {"claude-opus-5-5[1m]": {}, "claude-haiku-4-5": {}}} + check_model(result, "claude-opus-5-5") + check_model(result, "opus") # an alias names no fixed model + check_model(result, None) + with pytest.raises(ModelCheckError, match="not the configured model claude-sonnet-5-5"): + check_model(result, "claude-sonnet-5-5") + check_model({"modelUsage": {"claude-opus-4-5-20251101": {}}}, "claude-opus-4-5") # a dated name + with pytest.raises(ModelCheckError): + check_model({"modelUsage": {"claude-opus-5-5": {}}}, "claude-opus-5") # another model, not a dated name + + +def test_a_job_without_accounting_counts_its_time_limit(started): + class NoAccounting(FakeSlurm): + def accounting(self, job_id): + raise SlurmError(f"sacct has no record of job {job_id}") + + ctx = context(started, FakeAgent(lambda call: GOOD), NoAccounting(started / "state", lambda n: (0.89, SMALL)), + max_round_attempts=2) + assert loop.run_rounds(ctx) == "round_limit" + assert ctx.ledger.gpu_minutes == 10.0 # the time limit of the round job, so the GPU budget stays safe + + +def test_gpu_reservation_starts_a_round_only_if_its_time_limit_fits(started): + slurm = FakeSlurm(started / "state", lambda n: (0.89, SMALL), gpu_minutes=8.0) + ctx = context(started, FakeAgent(lambda call: GOOD), slurm, gpu_minutes=60, gpu_reservation=True) + assert loop.run_rounds(ctx) == "gpu_budget" + assert ctx.ledger.gpu_minutes == 56.0 # 48 + 10 fitted in 60; 56 + 10 does not + ctx = context(started, FakeAgent(lambda call: GOOD), slurm, gpu_minutes=60) + assert loop.stop_reason(ctx) is None # the rule of chapter 28 starts one more round, because 56 is below 60 + + +def test_the_ledger_of_chapter_28_keeps_its_columns(tmp_path): + ledger = loop.Ledger(tmp_path / "ledger.tsv") + ledger.add(round=2, kind="agent_step", tokens=100, wave=1, model="m") + assert (tmp_path / "ledger.tsv").read_text().splitlines()[0].split("\t") == loop.LEDGER_FIELDS + assert ledger.tokens == 100 + + +def test_command_adds_the_model_the_effort_the_role_prompt_and_the_readme_rule(): + caps = {"max_turns": 8, "max_budget_usd": 1.0, "timeout_minutes": 15} + args = command("prompt", caps) + assert "--model" not in args and "--effort" not in args and "--append-system-prompt-file" not in args + assert json.loads(args[args.index("--settings") + 1]) == UNATTENDED_SETTINGS + assert "Read(./README.md)" in UNATTENDED_SETTINGS["permissions"]["deny"] + args = command("prompt", {**caps, "model": "claude-opus-5-5", "effort": "xhigh", + "role_prompt": "prompts/multi/roles/planner.md"}, schema_text="{}") + assert args[args.index("--model") + 1] == "claude-opus-5-5" + assert args[args.index("--effort") + 1] == "xhigh" + assert args[args.index("--append-system-prompt-file") + 1] == "prompts/multi/roles/planner.md" + assert args[-2:] == ["--json-schema", "{}"] + + +def test_the_effort_variable_is_removed_only_when_the_step_sets_an_effort(monkeypatch): + monkeypatch.setenv("CLAUDE_CODE_EFFORT_LEVEL", "low") # it would override --effort + monkeypatch.setenv("ANTHROPIC_API_KEY", "not-a-real-key") + assert "CLAUDE_CODE_EFFORT_LEVEL" not in step_env({"effort": "high"}) + assert step_env({})["CLAUDE_CODE_EFFORT_LEVEL"] == "low" + assert "ANTHROPIC_API_KEY" not in step_env({}) and "ANTHROPIC_API_KEY" not in step_env({"effort": "high"}) + + +def test_record_paths_of_confirmation_seeds(tmp_path): + assert records.record_dir(tmp_path, "round", 7) == tmp_path / "rounds" / "round_07" + assert records.record_dir(tmp_path, "confirm", 7) == tmp_path / "confirm" / "round_07" + assert records.record_dir(tmp_path, "confirm", 7, 1) == tmp_path / "confirm" / "round_07" + assert records.record_dir(tmp_path, "confirm", 7, 2) == tmp_path / "confirm" / "round_07_seed2" + assert records.test_eval_path(tmp_path, "confirm", 7, 2) == tmp_path / "test_eval" / "confirm_07_seed2.json" + assert records.test_eval_path(tmp_path, "round", 7) == tmp_path / "test_eval" / "round_07.json" + + +# GPU minutes of jobs that a stopped loop job did not count + +def test_the_next_loop_job_counts_the_round_job_that_a_stopped_one_canceled(started, monkeypatch): + class StoppedSlurm(FakeSlurm): + def wait(self, job_id, queue_limit_seconds): + raise loop.Terminated() + + monkeypatch.setattr(loop.signal, "signal", lambda *args: None) # keep the signal handlers of pytest + state = started / "state" + assert loop.run(context(started, FakeAgent(lambda call: GOOD), StoppedSlurm(state, lambda n: None))) == 1 + assert (state / finish.PENDING_FILE).read_text() == "1001\tround\t2\t0\n" # not in the ledger yet + + monkeypatch.setattr(finish, "evaluate_on_test_set", lambda ctx, kind, n: {"test_accuracy": 0.902, "kind": kind}) + slurm = FakeSlurm(state, lambda n: (0.93, 1_000_000), gpu_minutes=1.5) + ctx = context(started, FakeAgent(lambda call: GOOD), slurm, max_round_attempts=1) + assert loop.run(ctx) == 0 + assert slurm.canceled == ["1001"] # in case it still runs + rows = [(row["round"], row["kind"], row["status"]) for row in ctx.ledger.rows if row["kind"] in ("round", "confirm")] + assert rows == [("2", "round", "interrupted"), ("1", "confirm", "ended")] + assert ctx.ledger.gpu_minutes == 1.5 # the canceled round counts against the GPU budget + assert not (state / finish.PENDING_FILE).exists() + + +def test_finish_counts_the_confirmation_job_that_a_stopped_loop_job_canceled(started, monkeypatch): + class StopInConfirmation(FakeSlurm): + def submit(self, script, args, cwd, options=()): + if "confirm" in args: # the job starts, but writes no records before the stop + self.submitted.append(("confirm", int(args[args.index("--round") + 1]))) + return str(1000 + len(self.submitted)) + return super().submit(script, args, cwd, options) + + def wait(self, job_id, queue_limit_seconds): + if self.submitted[-1][0] == "confirm": + raise loop.Terminated() + return "ended" + + monkeypatch.setattr(loop.signal, "signal", lambda *args: None) + state = started / "state" + stopped = StopInConfirmation(state, lambda n: (0.91, SMALL)) + assert loop.run(context(started, FakeAgent(lambda call: GOOD), stopped, plateau_rounds=1)) == 1 + job = stopped.canceled[0] + assert (state / finish.PENDING_FILE).read_text().startswith(f"{job}\tconfirm\t") + + monkeypatch.setattr(finish, "evaluate_on_test_set", lambda ctx, kind, n: {"test_accuracy": 0.902, "kind": kind}) + ctx = context(started, FakeAgent(lambda call: GOOD), FakeSlurm(state, lambda n: (0.91, SMALL))) + assert loop.run(ctx, finish_only=True) == 0 + rows = [(row["status"], row["job_id"]) for row in ctx.ledger.rows if row["kind"] == "confirm"] + assert rows[0] == ("interrupted", job) + assert len(rows) == 2 # then the new confirmation round + assert not (state / finish.PENDING_FILE).exists() + + +def test_round_constraint_in_the_experiment_file_reaches_each_round_job(started): + slurm = FakeSlurm(started / "state", lambda n: (0.91, SMALL)) + ctx = context(started, FakeAgent(lambda call: GOOD), slurm, plateau_rounds=1) + assert finish.round_job_options(ctx) == () # unset: the --constraint line of the job script applies + ctx.experiment["round"]["constraint"] = "a100-mig" + assert loop.run_rounds(ctx) == "plateau" + assert slurm.options and all(options == ("--constraint=a100-mig",) for options in slurm.options) + + +def test_a_state_folder_with_final_steps_done_refuses_a_new_search(started, monkeypatch): + monkeypatch.setattr(finish, "evaluate_on_test_set", lambda ctx, kind, n: {"test_accuracy": 0.902, "kind": kind}) + slurm = FakeSlurm(started / "state", lambda n: (0.91, SMALL)) + assert loop.run(context(started, FakeAgent(lambda call: GOOD), slurm, plateau_rounds=1)) == 0 + with pytest.raises(SystemExit, match="final steps of state/ are done"): + loop.run(context(started, FakeAgent(lambda call: GOOD), slurm, plateau_rounds=1)) + assert loop.run(context(started, FakeAgent(lambda call: GOOD), slurm), finish_only=True) == 0 # --finish still works diff --git a/tests/test_multi_loop.py b/tests/test_multi_loop.py new file mode 100644 index 0000000..46fc09e --- /dev/null +++ b/tests/test_multi_loop.py @@ -0,0 +1,700 @@ +"""Tests of the multi-agent loop with fake roles and fake SLURM array jobs. No API calls and no jobs.""" + +from __future__ import annotations + +import csv +import json +import re +from pathlib import Path + +import pytest +import yaml + +from pruning_example import check_report, config, finish, metrics, multi_loop, records +from pruning_example.agent import AgentError +from pruning_example.loop import Ledger, Terminated +from pruning_example.slurm import TaskAccounting +from conftest import REPO + +SMALL = 223_287 + + +def usage(model: str, tokens: int) -> dict: + return {model: {"inputTokens": 10, "outputTokens": 20, "cacheReadInputTokens": tokens - 40, + "cacheCreationInputTokens": 10}} + + +def item(number: int, **changes) -> dict: + planned = {"sparsity": round(0.9 + 0.001 * number, 4), "steps": 2, "finetune_epochs": 4, "lr": 0.01, + "weight_decay": 0.0005, "exclude_layers": [], "reason": f"test {number}", + "hypothesis": f"hypothesis {number}", "predicted_val_accuracy": 0.91, "pass_probability": 0.6} + return {**planned, **changes} + + +class FakeRoles: + """Answers each role step like a real role, and keeps the calls. No API calls.""" + + def __init__(self, tokens: int = 40_000, reject=None, fail=(), plan=None, models=None): + self.tokens = tokens + self.reject = reject or (lambda kind, entry: None) # a reason to reject, or None + self.fail = set(fail) + self.plan = plan + self.models = models or {} # kind -> the model that the step reports + self.calls = [] + self.counter = 0 + + def run(self, prompt, caps, schema_text=None): + role = Path(caps["role_prompt"]).stem + kind = role + if role == "planner" and "rejected" in prompt: + kind = "revision" + if role == "critic" and "review_items_2" in prompt: + kind = "critic_2" + self.calls.append((kind, prompt, caps)) + if kind in self.fail: + raise AgentError("fake failure") + result = {"subtype": "success", "modelUsage": usage(self.models.get(kind, caps["model"]), self.tokens), + "total_cost_usd": 0.1, "result": ""} + if role == "planner": + count = json.loads(schema_text)["properties"]["configurations"]["minItems"] + planned = self.plan(kind, count, self) if self.plan else [item(self.next()) for _ in range(count)] + result["structured_output"] = {"configurations": planned, "summary": f"{kind} summary"} + elif role == "critic": + items = json.loads(Path(re.search(r"They are in (\S+?),", prompt).group(1)).read_text()) + verdicts = [] + for entry in items: + reason = self.reject(kind, entry) + verdicts.append({"item": entry["item"], "accept": reason is None, "reason": reason or "fine"}) + result["structured_output"] = {"verdicts": verdicts, "summary": f"{kind} summary"} + elif role == "writer": + result["result"] = "The selected round keeps 0.9100 validation accuracy. One number is wrong: 0.9555." + elif role == "auditor": + result["structured_output"] = { + "claims": [{"claim": "0.9100 validation accuracy", "verdict": "supported", "evidence": "summary.json", + "correction": ""}], + "limits_stated": False, "summary": "one claim checked"} + return result + + def next(self) -> int: + self.counter += 1 + return self.counter + + def kinds(self) -> list[str]: + return [call[0] for call in self.calls] + + def stop(self): + pass + + +def write_task_records(state: Path, round_config: dict, round_no: int, kind: str, seed: int, + val_accuracy: float, nonzero: int, job_id: str) -> None: + """Write the records that a real round job writes with --no-history.""" + result = {"round": round_no, "kind": kind, "seed": seed, "val_accuracy": val_accuracy, + "val_interval": metrics.wilson_interval(val_accuracy, 5000), "nonzero": nonzero, + "weights_left": nonzero / 11_164_352, "sparse_size_mb": 1.0, "passes": val_accuracy >= 0.905} + folder = records.record_dir(state, kind, round_no, seed) + records.write_json(folder / "config.json", round_config) + records.write_json(folder / "run_meta.json", {"job_id": job_id, "run_folder": str(state / "no_run_folder")}) + records.write_json(folder / "metrics.json", result) + + +class FakeArraySlurm: + """Runs no jobs. Each task of a submitted array job writes the records that a real task writes.""" + + def __init__(self, state, results, gpu_minutes: float = 2.0, outcome: str = "ended", stop_on_wave=None): + self.state = state + self.results = results # (round, kind, seed) -> (val_accuracy, nonzero), or None for a task without metrics + self.gpu_minutes = gpu_minutes + self.outcome = outcome + self.stop_on_wave = stop_on_wave + self.arrays = [] + self.canceled = [] + self.options = [] + + def submit_array(self, script, tasks_file, count, max_parallel, cwd, options=()): + self.options.append(options) + lines = [line.split("\t") for line in (Path(cwd) / tasks_file).read_text().splitlines()] + assert len(lines) == count and script == "slurm/round_array.sbatch" + job = str(900 + len(self.arrays)) + self.arrays.append((job, lines, max_parallel)) + for index, (config_path, round_no, kind, seed) in enumerate(lines): + if kind == "confirm": + assert (self.state / finish.SELECTION_FILE).exists() # the selection is saved before any confirmation + outcome = self.results(int(round_no), kind, int(seed)) + if outcome is not None: + write_task_records(self.state, records.read_json(Path(cwd) / config_path), int(round_no), kind, + int(seed), *outcome, job_id=f"{job}{index}") + return job + + def wait_array(self, job, queue_limit_seconds): + if self.stop_on_wave is not None and len(self.arrays) == self.stop_on_wave: + raise Terminated() + return self.outcome + + def array_accounting(self, job, count): + lines = dict((entry[0], entry[1]) for entry in self.arrays)[job] + tasks = {} + for index, (_, round_no, kind, seed) in enumerate(lines): + ran = self.results(int(round_no), kind, int(seed)) is not None or self.outcome != "queue_limit" + tasks[index] = TaskAccounting(index, f"{job}{index}", "COMPLETED" if ran else "CANCELLED", + self.gpu_minutes if ran else 0.0, 1 if ran else 0) + return tasks + + def cancel(self, job, wait_seconds=60): + self.canceled.append(job) + return True + + def cancel_waiting(self, job, wait_seconds=60): + return True + + def rounds(self) -> list[int]: + return [int(line[1]) for _, lines, _ in self.arrays for line in lines if line[2] == "round"] + + +@pytest.fixture +def multi(project, monkeypatch): + """A project after the baseline and the supervised round 1 of chapter 28, set up for the multi-agent loop.""" + state = project / "state" + records.write_json(state / "baseline" / "metrics.json", {"val_accuracy": 0.94}) + first = {"sparsity": 0.5, "steps": 1, "finetune_epochs": 2, "lr": 0.01, "weight_decay": 0.0005, + "exclude_layers": [], "reason": "round 1"} + write_task_records(state, first, 1, "round", 0, 0.93, 5_582_176, "round-one") + records.append_history(state, records.history_row(1, "ok", first, records.read_json( + state / "rounds" / "round_01" / "metrics.json"), job="round-one")) + monkeypatch.setenv("STATE_DIR", "state_multi") + monkeypatch.setenv("EXPERIMENT_FILE", "experiment_multi.yaml") + monkeypatch.setattr(finish, "evaluate_on_test_set", + lambda ctx, kind, n, seed=None: {"kind": kind, "round": n, "test_accuracy": 0.9234}) + return project + + +def context(project, agent, slurm, critic: bool = True, plateau_rounds: int | None = None, **budget): + paths = config.paths_from_env() + experiment = config.load_experiment(project) + experiment["budget"].update(budget) + experiment["roles"]["critic"]["enabled"] = critic + if plateau_rounds is not None: + experiment["target"]["plateau_rounds"] = plateau_rounds + schemas = {name: (project / path).read_text() for name, path in multi_loop.SCHEMA_FILES.items()} + return multi_loop.MultiContext( + paths=paths, experiment=experiment, schema_text=(project / config.SCHEMA_PATH).read_text(), agent=agent, + slurm=slurm, ledger=Ledger(paths.state / "ledger.tsv", multi_loop.MULTI_LEDGER_FIELDS, + multi_loop.ROLE_OF.keys(), multi_loop.WAVE_KINDS), + log=lambda message: None, schemas=schemas) + + +def plan_then_run(project, agent, slurm, **settings): + """Run the --plan-only job, then the loop job, as a person does.""" + planner = context(project, agent, slurm, **settings) + planner.plan_only = True + assert multi_loop.run(planner, "plan_only") == 0 + ctx = context(project, agent, slurm, **settings) + return ctx, multi_loop.run(ctx, "loop") + + +def state_of(project) -> Path: + return project / "state_multi" + + +# The wave flow + +def test_plan_only_writes_the_plan_and_runs_no_job(multi): + agent = FakeRoles() + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL)) + ctx = context(multi, agent, slurm) + ctx.plan_only = True + assert multi_loop.run(ctx, "plan_only") == 0 + folder = state_of(multi) / "waves" / "wave_01" + assert len(records.read_json(folder / "accepted.json")) == 3 + assert "To run this plan, start the loop job" in (folder / "plan.md").read_text() + assert not (folder / "tasks.tsv").exists() + assert slurm.arrays == [] + assert agent.kinds() == ["planner", "critic"] + assert (state_of(multi) / "rounds" / "round_01" / "metrics.json").exists() # round 1 copied from state/ + assert [row["round"] for row in records.read_history(state_of(multi))] == ["1"] + + +def test_the_loop_refuses_to_start_without_a_planned_wave_1(multi): + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + with pytest.raises(SystemExit, match="Wave 1 has no accepted configuration"): + multi_loop.run(ctx, "loop") + + +def test_waves_run_as_array_jobs_and_the_loop_writes_the_rows_in_round_order(multi): + agent = FakeRoles() + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if kind == "round" else (0.906, SMALL)) + ctx, code = plan_then_run(multi, agent, slurm, max_round_attempts=7) + assert code == 0 + state = state_of(multi) + history = records.read_history(state) + assert [row["round"] for row in history] == [str(n) for n in range(1, 8)] + assert [row["job_id"] for row in history[1:4]] == ["9000", "9001", "9002"] # the task job IDs, not the loop job + assert slurm.rounds() == [2, 3, 4, 5, 6, 7] + assert [max_parallel for _, _, max_parallel in slurm.arrays] == [3, 3, 3] # two waves, then the confirmations + tasks = (state / "waves" / "wave_01" / "tasks.tsv").read_text().splitlines() + assert tasks[0].split("\t")[1:] == ["2", "round", "0"] + assert "hypothesis" not in records.read_json(state / "configs" / "round_02.json") + assert agent.kinds()[:4] == ["planner", "critic", "planner", "critic"] # wave 1 is not planned again + summary = records.read_json(state / "summary.json") + assert summary["stop_reason"] == "round_limit" + assert summary["selected"]["round"] == 7 + assert [entry["seed"] for entry in summary["confirmations"]] == [1, 2] + assert summary["means"]["holds"] is True + assert summary["roles"]["planner"]["model"] == "claude-opus-5-5" + assert summary["roles"]["critic"]["steps"] == 2 + assert summary["design"] == "planner and critic, waves" + ledger = (state / "ledger.tsv").read_text().splitlines() + assert ledger[0].split("\t") == multi_loop.MULTI_LEDGER_FIELDS + with open(state / "hypotheses.tsv", newline="") as f: + assert [row["status"] for row in csv.DictReader(f, delimiter="\t")] == ["run"] * 6 + assert "## Wave 2" in (state / "notes.md").read_text() + assert (state / "final_report.md").exists() + audit = records.read_json(state / "audit.json") + assert audit["auditor"]["claims"][0]["verdict"] == "supported" + assert audit["unmatched_numbers"] == [{"number": "0.9555", "line": 1}] # 0.9100 is in the records + + +def test_code_checks_send_a_repeat_and_an_invalid_configuration_to_one_revision(multi): + def plan(kind, count, fake): + if kind == "planner": + repeat = item(fake.next(), sparsity=0.5, steps=1, finetune_epochs=2) # the configuration of round 1 + invalid = item(fake.next(), steps=5, finetune_epochs=4) # fewer epochs than steps + return [repeat, invalid, item(fake.next())][:count] + return [item(fake.next()) for _ in range(count)] + + agent = FakeRoles(plan=plan) + ctx = context(multi, agent, FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + multi_loop.copy_round_one(ctx) + accepted = multi_loop.plan_wave(ctx, 1, 3) + plan_record = records.read_json(state_of(multi) / "waves" / "wave_01" / "plan.json") + assert [entry["status"] for entry in plan_record["entries"]] == ["replaced", "replaced", "accepted", + "accepted", "accepted"] + assert "repeats an earlier configuration" in plan_record["entries"][0]["reasons"][0] + assert "at least steps" in plan_record["entries"][1]["reasons"][0] + assert len(accepted) == 3 + assert agent.kinds() == ["planner", "critic", "revision", "critic_2"] + revision_prompt = agent.calls[2][1] + assert "The checks rejected 2 of the configurations" in revision_prompt + + +def test_a_configuration_rejected_twice_is_dropped(multi): + agent = FakeRoles(reject=lambda kind, entry: "too far" if entry["item"] in (1, 4) else None) + ctx = context(multi, agent, FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + accepted = multi_loop.plan_wave(ctx, 1, 3) + statuses = [entry["status"] for entry in records.read_json(state_of(multi) / "waves" / "wave_01" / "plan.json")["entries"]] + assert statuses == ["replaced", "accepted", "accepted", "dropped"] + assert len(accepted) == 2 + + +def test_the_critic_off_gives_one_agent_with_waves(multi): + agent = FakeRoles() + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n)) + ctx, code = plan_then_run(multi, agent, slurm, critic=False, max_round_attempts=4) + assert code == 0 + assert "critic" not in agent.kinds() and "critic_2" not in agent.kinds() + assert records.read_json(state_of(multi) / "summary.json")["design"] == "one agent, waves" + + +def test_failed_planner_steps_are_failed_waves_and_two_stop_the_loop(multi): + agent = FakeRoles(fail={"planner"}) + ctx = context(multi, agent, FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + multi_loop.copy_round_one(ctx) + assert multi_loop.run_waves(ctx) == "failed_attempts" + assert agent.kinds() == ["planner", "planner"] + assert [row["kind"] for row in ctx.ledger.rows] == ["planner", "planner"] + + +def test_a_step_with_another_model_stops_the_loop(multi): + agent = FakeRoles(models={"planner": "claude-sonnet-5-5"}) # an organization can replace a model + ctx = context(multi, agent, FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + multi_loop.copy_round_one(ctx) + assert multi_loop.run_waves(ctx) == "model_check" + + +def test_each_role_step_gets_its_model_effort_and_role_prompt(multi): + agent = FakeRoles() + ctx = context(multi, agent, FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + multi_loop.copy_round_one(ctx) + multi_loop.plan_wave(ctx, 1, 3) + planner_caps, critic_caps = agent.calls[0][2], agent.calls[1][2] + assert (planner_caps["model"], planner_caps["effort"]) == ("claude-opus-5-5", "xhigh") + assert (critic_caps["model"], critic_caps["effort"]) == ("claude-sonnet-5-5", "high") + assert planner_caps["role_prompt"] == "prompts/multi/roles/planner.md" + assert "enabled" not in critic_caps + row = ctx.ledger.rows[0] + assert (row["kind"], row["wave"], row["model"], row["effort"]) == ("planner", "1", "claude-opus-5-5", "xhigh") + assert "GPU minutes left: 150" in agent.calls[0][1] and "Round attempts left: 44" in agent.calls[0][1] + + +# Limits + +def test_wave_size_follows_all_four_limits(multi): + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None), gpu_minutes=60) + multi_loop.copy_round_one(ctx) + history = records.read_history(ctx.paths.state) + assert multi_loop.wave_size(ctx, history) == 3 + ctx.ledger.add(round=2, kind="round", gpu_minutes=45) + assert multi_loop.wave_size(ctx, history) == 1 # room for one 10-minute round in 60 GPU minutes + ctx.ledger.rows.pop() + ctx.experiment["budget"]["max_round_attempts"] = 3 + assert multi_loop.wave_size(ctx, history) == 2 # 2 round attempts left + ctx.experiment["budget"]["max_round_attempts"] = 45 + ctx.experiment["target"]["plateau_rounds"] = 2 + stale = {"round": "2", "status": "ok", "val_accuracy": "0.8900", "nonzero": "100"} + assert multi_loop.wave_size(ctx, history + [stale]) == 1 # one round left before the plateau rule + + +def test_stop_rules_before_a_wave(multi): + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None), gpu_minutes=60) + multi_loop.copy_round_one(ctx) + assert multi_loop.stop_reason(ctx) is None + ctx.ledger.add(round="", kind="planner", tokens=700_000, wave=1) + assert multi_loop.stop_reason(ctx) == "token_budget" # 700,000 + 150,000 is above the search limit + assert multi_loop.stop_reason(ctx, planned=True) is None # a planned wave needs no tokens for its plan + ctx.ledger.rows.clear() + ctx.ledger.add(round=2, kind="round", gpu_minutes=55) + assert multi_loop.stop_reason(ctx) == "gpu_budget" # no room for a 10-minute round + ctx.ledger.rows.clear() + ctx.job_end = ctx.clock() + 60 * 120 # two hours left; a plan needs 70 minutes and the final steps 110 + assert multi_loop.stop_reason(ctx) == "time_left" + ctx.plan_only = True + assert multi_loop.stop_reason(ctx) is None # the --plan-only job checks the time before each step instead + + +def test_plan_only_stops_without_time_for_a_step(multi): + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + ctx.plan_only = True + ctx.job_end = ctx.clock() + 60 * 10 # 10 minutes left, the planner needs 20 + assert multi_loop.run(ctx, "plan_only") == 1 + + +def test_queue_limit_stops_after_the_wave(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL) if n == 2 else None, + outcome="queue_limit") + ctx = context(multi, FakeRoles(), slurm) + multi_loop.copy_round_one(ctx) + assert multi_loop.run_waves(ctx) == "queue_limit" + statuses = [row["status"] for row in records.read_history(ctx.paths.state)] + assert statuses == ["ok", "ok", "queue_limit", "queue_limit"] + + +# Stops and missing rows + +def test_missing_rows_are_written_from_the_round_folders(multi): + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + multi_loop.copy_round_one(ctx) + state = ctx.paths.state + folder = state / "waves" / "wave_01" + folder.mkdir(parents=True) + (folder / "tasks.tsv").write_text("state_multi/configs/round_02.json\t2\tround\t0\n" + "state_multi/configs/round_03.json\t3\tround\t0\n") + write_task_records(state, item(1), 2, "round", 0, 0.91, SMALL, "task-2") + assert multi_loop.write_missing_rows(ctx) == 1 + rows = records.read_history(state) + assert [(row["round"], row["job_id"]) for row in rows] == [("1", "round-one"), ("2", "task-2")] + assert multi_loop.write_missing_rows(ctx) == 0 + assert multi_loop.next_round_number(ctx) == 4 # round 3 has no row, but its number is taken + + +def test_sigterm_during_a_wave_cancels_the_array_job_and_keeps_finished_rounds(multi, monkeypatch): + monkeypatch.setattr(multi_loop.signal, "signal", lambda *args: None) # keep the signal handlers of pytest + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if n in (2, 3) else None, + stop_on_wave=1) + ctx, code = plan_then_run(multi, FakeRoles(), slurm) + assert code == 1 + assert slurm.canceled == ["900"] + state = state_of(multi) + assert [row["round"] for row in records.read_history(state)] == ["1", "2", "3"] # round 4 has no metrics + summary = records.read_json(state / "summary.json") + assert summary["terminated"] is True and summary["final_steps_done"] is False + assert not (state / "current_job").exists() + + # The next loop job closes the interrupted wave, then plans wave 2. + slurm.stop_on_wave = None + ctx = context(multi, FakeRoles(), slurm, max_round_attempts=5) + assert multi_loop.run(ctx, "loop") == 0 + outcome = records.read_json(state / "waves" / "wave_01" / "outcome.json") + assert outcome["complete_rounds"] == 2 + assert [row["status"] for row in records.read_history(state)][:4] == ["ok", "ok", "ok", "job_failed"] + + +def test_finish_uses_the_saved_selection(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL)) + ctx = context(multi, FakeRoles(), slurm) + multi_loop.copy_round_one(ctx) + state = ctx.paths.state + for round_no, nonzero in ((2, 300_000), (3, 200_000)): + write_task_records(state, item(round_no), round_no, "round", 0, 0.91, nonzero, f"task-{round_no}") + records.append_history(state, multi_loop.ok_row(state, round_no)) + records.write_json(state / finish.SELECTION_FILE, {"round": 2}) # saved by an earlier, stopped loop job + summary = finish.run_final_steps(ctx, "plateau", report_step=multi_loop.report_steps, + extra=multi_loop.summary_extra) + assert summary["selected"]["round"] == 2 + assert (state / "confirm" / "round_02" / "metrics.json").exists() + assert (state / "confirm" / "round_02_seed2" / "metrics.json").exists() + + +def test_audit_only_checks_a_report_without_the_ledger_of_that_folder(multi, monkeypatch): + state = multi / "state" # the state folder of the single-agent run + records.write_json(state / "summary.json", {"selected": {"val_accuracy": 0.9076}}) + (state / "final_report.md").write_text("Validation accuracy 0.9076, test accuracy 0.9999.\n") + monkeypatch.setenv("STATE_DIR", "state") + paths = config.paths_from_env() + agent = FakeRoles() + ctx = context(multi, agent, FakeArraySlurm(state, lambda n, kind, seed: None)) + ctx.ledger = Ledger(paths.state / "audit_ledger.tsv", multi_loop.MULTI_LEDGER_FIELDS, + multi_loop.ROLE_OF.keys(), multi_loop.WAVE_KINDS) + assert multi_loop.run(ctx, "audit_only") == 0 + audit = records.read_json(state / "audit.json") + assert audit["unmatched_numbers"] == [{"number": "0.9999", "line": 1}] + assert agent.kinds() == ["auditor"] + assert not (state / "ledger.tsv").exists() + assert (state / "audit_ledger.tsv").exists() + + +# Files and scores + +def test_wave_schema_has_the_round_keys_plus_three(): + round_schema = json.loads((REPO / "configs" / "round.schema.json").read_text()) + wave_item = json.loads((REPO / "configs" / "wave.schema.json").read_text())["properties"]["configurations"]["items"] + extras = {"hypothesis", "predicted_val_accuracy", "pass_probability"} + assert set(wave_item["properties"]) == set(round_schema["properties"]) | extras + for key, value in round_schema["properties"].items(): + assert wave_item["properties"][key] == value + assert set(wave_item["required"]) == set(round_schema["required"]) | extras + + +def test_experiment_files_agree(): + single = yaml.safe_load((REPO / "experiment.yaml").read_text()) + multi = yaml.safe_load((REPO / "experiment_multi.yaml").read_text()) + assert multi["target"]["val_accuracy"] == single["target"]["val_accuracy"] + assert multi["baseline"] == single["baseline"] + for experiment in (single, multi): + experiment["round"].pop("constraint", None) # the GPUs can differ; the other round settings cannot + assert multi["round"] == single["round"] + multi_loop.check_roles(multi) + assert multi["roles"]["planner"]["model"] != multi["roles"]["critic"]["model"] + assert multi["roles"]["writer"]["model"] != multi["roles"]["auditor"]["model"] + + +def test_check_roles_refuses_a_bad_effort_level(): + experiment = yaml.safe_load((REPO / "experiment_multi.yaml").read_text()) + experiment["roles"]["critic"]["effort"] = "huge" + with pytest.raises(SystemExit, match="effort must be one of"): + multi_loop.check_roles(experiment) + + +def test_check_report_lists_only_the_numbers_that_the_records_do_not_show(tmp_path): + records.write_json(tmp_path / "summary.json", {"selected": {"val_accuracy": 0.9076, "weights_left": 0.007453, + "nonzero": 83215}, "tokens": {"total": 323072}}) + report = ("Round 15 keeps 0.9076 validation accuracy with 83,215 weights (0.745% left).\n" + "It used 323,072 tokens with Claude Code 2.1.302 in 2026, with 95% intervals.\n" + "The test accuracy is 0.9123.\n") + assert check_report.unmatched_numbers(report, tmp_path) == [{"number": "0.9123", "line": 3}] + + +def test_prediction_scores_next_to_baselines(multi): + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + multi_loop.copy_round_one(ctx) + state = ctx.paths.state + for round_no, val in ((2, 0.92), (3, 0.90)): + write_task_records(state, item(round_no), round_no, "round", 0, val, SMALL, f"task-{round_no}") + records.append_history(state, multi_loop.ok_row(state, round_no)) + rows = [{"wave": 1, "round": 2, "status": "run", "val_accuracy": "0.9200", "predicted_val_accuracy": 0.91, + "pass_probability": 0.9, "passes": "yes"}, + {"wave": 1, "round": 3, "status": "run", "val_accuracy": "0.9000", "predicted_val_accuracy": 0.91, + "pass_probability": 0.6, "passes": "no"}] + scores = multi_loop.prediction_scores(ctx, rows) + assert scores["rounds"] == 2 + assert scores["mean_absolute_error"] == 0.01 + assert scores["baseline_mean_absolute_error"] == 0.02 # round 1 (0.93) is the only earlier round + assert scores["brier_score"] == round((0.1 ** 2 + 0.6 ** 2) / 2, 4) + assert scores["base_rate_brier_score"] == 0.25 + assert scores["pass_decisions_right"] == 1 + + +def test_a_failed_sbatch_is_a_failed_wave(multi): + class NoSubmit(FakeArraySlurm): + def submit_array(self, script, tasks_file, count, max_parallel, cwd, options=()): + raise multi_loop.SlurmError("sbatch failed: invalid partition") + + ctx = context(multi, FakeRoles(), NoSubmit(state_of(multi), lambda n, kind, seed: (0.91, SMALL))) + multi_loop.copy_round_one(ctx) + assert multi_loop.run_waves(ctx) == "failed_attempts" # two waves in a row with no complete round + statuses = [row["status"] for row in records.read_history(ctx.paths.state)] + assert statuses == ["ok"] + ["sbatch_failed"] * 6 + + +def test_finish_reuses_the_records_of_an_earlier_run(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if kind == "round" else (0.906, SMALL)) + plan_then_run(multi, FakeRoles(), slurm, max_round_attempts=4) + arrays = len(slurm.arrays) + ctx = context(multi, FakeRoles(), slurm) + assert multi_loop.run(ctx, "finish") == 0 + assert len(slurm.arrays) == arrays # no confirmation round runs again + summary = records.read_json(state_of(multi) / "summary.json") + assert summary["stop_reason"] == "round_limit" and summary["final_steps_done"] is True + + +def test_a_new_loop_job_cancels_the_array_job_of_a_stopped_one(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL)) + planner = context(multi, FakeRoles(), slurm) + planner.plan_only = True + assert multi_loop.run(planner, "plan_only") == 0 + (state_of(multi) / "current_job").write_text("812\n") # the loop job got SIGKILL during a wave + ctx = context(multi, FakeRoles(), slurm, max_round_attempts=4) + multi_loop.check_start(ctx) + assert slurm.canceled == ["812"] + assert not (state_of(multi) / "current_job").exists() + + +# Fixes after the code review + +def test_the_loop_refuses_an_empty_plan_of_wave_1(multi): + ctx = context(multi, FakeRoles(reject=lambda kind, entry: "no"), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + ctx.plan_only = True + assert multi_loop.run(ctx, "plan_only") == 1 # the critic rejects each configuration, also after the revision + assert records.read_json(state_of(multi) / "waves" / "wave_01" / "accepted.json") == [] + assert multi_loop.run(ctx, "plan_only") == 1 # a second --plan-only job says to plan again + ctx = context(multi, FakeRoles(), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + with pytest.raises(SystemExit, match="no accepted configuration"): + multi_loop.run(ctx, "loop") + + +def test_finish_closes_the_wave_of_a_stopped_loop_job(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if n in (2, 3) or kind == "confirm" else None) + planner = context(multi, FakeRoles(), slurm) + planner.plan_only = True + assert multi_loop.run(planner, "plan_only") == 0 + ctx = context(multi, FakeRoles(), slurm) + multi_loop.check_start(ctx) + accepted = records.read_json(state_of(multi) / "waves" / "wave_01" / "accepted.json") + # The loop job submits wave 1 and gets SIGKILL: no handler runs, so current_job stays. + original_wait = slurm.wait_array + slurm.wait_array = lambda job, limit: (_ for _ in ()).throw(KeyboardInterrupt()) + with pytest.raises(KeyboardInterrupt): + multi_loop.submit_wave(ctx, 1, accepted) + slurm.wait_array = original_wait + assert (state_of(multi) / "current_job").exists() + ctx = context(multi, FakeRoles(), slurm) + assert multi_loop.run(ctx, "finish") == 0 + assert slurm.canceled == ["900"] # the array job of the stopped loop job + outcome = records.read_json(state_of(multi) / "waves" / "wave_01" / "outcome.json") + assert outcome["complete_rounds"] == 2 + statuses = [row["status"] for row in records.read_history(state_of(multi))] + assert statuses == ["ok", "ok", "ok", "job_failed"] + assert [row["kind"] for row in ctx.ledger.rows if row["kind"] == "round"] == ["round"] * 3 # GPU minutes counted + + +def test_a_writer_with_another_model_does_not_stop_the_code_checks(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if kind == "round" else (0.906, SMALL)) + agent = FakeRoles(models={"writer": "claude-haiku-4-5"}) + ctx, code = plan_then_run(multi, agent, slurm, max_round_attempts=4) + assert code == 0 + audit = records.read_json(state_of(multi) / "audit.json") + assert "not the configured model" in audit["failures"]["writer"] + assert audit["failures"]["auditor"] == "skipped: there is no report" + assert audit["auditor_step_failed"] is True + assert audit["check_round"]["round"] == 4 # the code checks still ran + assert records.read_json(state_of(multi) / "summary.json")["audit"]["failures"]["writer"] + + +def test_an_audit_without_structured_output_is_a_failed_audit(multi): + class NoAudit(FakeRoles): + def run(self, prompt, caps, schema_text=None): + result = super().run(prompt, caps, schema_text) + if Path(caps["role_prompt"]).stem == "auditor": + result.pop("structured_output") + result["subtype"] = "error_max_turns" + return result + + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if kind == "round" else (0.906, SMALL)) + plan_then_run(multi, NoAudit(), slurm, max_round_attempts=4) + audit = records.read_json(state_of(multi) / "audit.json") + assert audit["auditor_step_failed"] is True + assert audit["failures"]["auditor"] == "no valid audit (subtype error_max_turns)" + + +def test_critic_counts_count_each_configuration_one_time(multi): + agent = FakeRoles(fail={"critic"}) + ctx = context(multi, agent, FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + multi_loop.copy_round_one(ctx) + assert multi_loop.plan_wave(ctx, 1, 3) is None + multi_loop.write_outcome(ctx, 1, 0, "a role step failed") + counts = multi_loop.summary_extra(ctx)["configurations"] + assert counts["planned"] == 3 and counts["not_reviewed"] == 3 + assert counts["rejected_by_critic"] == 0 and counts["accepted"] == 0 + + +def test_check_report_uses_the_factor_100_only_for_percentages_and_keeps_signs(tmp_path): + records.write_json(tmp_path / "summary.json", {"weights_left": 0.0075, "val_accuracy": 0.9076, "error": -0.0123, + "rounds": 45}) + report = "0.75% right, 0.75 wrong, 0.9076% wrong, -0.0123 right, 0.0123 wrong, 2,048 checked, 4500 checked\n" + found = [item["number"] for item in check_report.unmatched_numbers(report, tmp_path)] + assert found == ["0.75", "0.9076%", "0.0123", "2,048", "4500"] + + +# Fixes after the cluster tests + +def test_finish_counts_the_confirmation_array_job_that_a_stopped_loop_job_canceled(multi, monkeypatch): + monkeypatch.setattr(multi_loop.signal, "signal", lambda *args: None) # keep the signal handlers of pytest + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n) if kind == "round" else None, + stop_on_wave=2) # array job 2 is the confirmation of the selected round + ctx, code = plan_then_run(multi, FakeRoles(), slurm, max_round_attempts=4) + assert code == 1 + state = state_of(multi) + job = slurm.arrays[1][0] + assert slurm.canceled == [job] + selected = records.read_json(state / finish.SELECTION_FILE)["round"] + assert (state / finish.PENDING_FILE).read_text() == f"{job}\tconfirm\t{selected}\t2\n" + + slurm.stop_on_wave = None + slurm.results = lambda n, kind, seed: (0.91, SMALL - n) if kind == "round" else (0.906, SMALL) + ctx = context(multi, FakeRoles(), slurm) + assert multi_loop.run(ctx, "finish") == 0 + rows = [(row["status"], row["job_id"]) for row in ctx.ledger.rows if row["kind"] == "confirm"] + assert rows[:2] == [("interrupted", f"{job}0"), ("interrupted", f"{job}1")] + assert len(rows) == 4 # then the two new confirmation rounds + assert not (state / finish.PENDING_FILE).exists() + + +def test_check_report_reads_the_baseline_metrics(tmp_path): + state = tmp_path / "state_multi" + records.write_json(state / "summary.json", {"selected": {"val_accuracy": 0.9182}}) + records.write_json(tmp_path / config.BASELINE_METRICS, {"val_accuracy": 0.937}) + report = "Round 4 keeps 0.9182 validation accuracy. The baseline has 0.937.\n" + assert check_report.unmatched_numbers(report, state) == [{"number": "0.937", "line": 1}] + assert check_report.unmatched_numbers(report, state, tmp_path / config.BASELINE_METRICS) == [] + + +def test_round_constraint_reaches_the_array_jobs(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n)) + planner = context(multi, FakeRoles(), slurm) + planner.plan_only = True + assert multi_loop.run(planner, "plan_only") == 0 + ctx = context(multi, FakeRoles(), slurm, max_round_attempts=4) + ctx.experiment["round"]["constraint"] = "a100-mig" + assert multi_loop.run(ctx, "loop") == 0 + assert slurm.options == [("--constraint=a100-mig",)] * 2 # the wave and the confirmation rounds + + +def test_a_state_folder_with_final_steps_done_refuses_a_new_search(multi): + slurm = FakeArraySlurm(state_of(multi), lambda n, kind, seed: (0.91, SMALL - n)) + ctx, code = plan_then_run(multi, FakeRoles(), slurm, max_round_attempts=4) + assert code == 0 + for mode in ("loop", "plan_only"): + again = context(multi, FakeRoles(), slurm) + again.plan_only = mode == "plan_only" + with pytest.raises(SystemExit, match="final steps of state_multi/ are done"): + multi_loop.run(again, mode) + assert multi_loop.run(context(multi, FakeRoles(), slurm), "finish") == 0 # --finish still works + + +def test_plan_only_says_when_a_role_step_failed(multi): + messages = [] + ctx = context(multi, FakeRoles(fail={"planner"}), FakeArraySlurm(state_of(multi), lambda n, kind, seed: None)) + ctx.plan_only = True + ctx.log = messages.append + assert multi_loop.run(ctx, "plan_only") == 1 + assert any("a role step failed" in message for message in messages) + assert not any("Read state_multi/waves/wave_01/plan.md" in message for message in messages) diff --git a/tests/test_round_smoke.py b/tests/test_round_smoke.py index 894a851..b55ed49 100644 --- a/tests/test_round_smoke.py +++ b/tests/test_round_smoke.py @@ -62,3 +62,19 @@ def test_check_round_finds_a_changed_record(project, capsys): records.write_json(project / "state" / "rounds" / "round_01" / "metrics.json", result) assert check_round.main(["1"]) == 1 assert "DIFF" in capsys.readouterr().out + + +def test_a_round_without_history_and_a_confirmation_with_seed_2(project): + (project / BASELINE_MODEL).parent.mkdir(parents=True) + torch.save(build_model().state_dict(), project / BASELINE_MODEL) + (project / "agent").mkdir() + (project / "agent" / "next_config.json").write_text(json.dumps(CONFIG)) + + run_round.main(["agent/next_config.json", "--round", "3", "--no-history"]) # the multi-agent loop writes the row + assert (project / "state" / "rounds" / "round_03" / "metrics.json").exists() + assert records.read_history(project / "state") == [] + + run_round.main(["state/rounds/round_03/config.json", "--round", "3", "--kind", "confirm", "--seed", "2"]) + result = records.read_json(project / "state" / "confirm" / "round_03_seed2" / "metrics.json") + assert result["seed"] == 2 and result["kind"] == "confirm" + assert not (project / "state" / "confirm" / "round_03").exists() # seed 1 keeps its own folder diff --git a/uv.lock b/uv.lock index abe41e5..acf9bf7 100644 --- a/uv.lock +++ b/uv.lock @@ -784,7 +784,7 @@ wheels = [ [[package]] name = "pruning-example" -version = "0.1.0" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "jsonschema", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-15-pruning-example-cpu' and extra == 'extra-15-pruning-example-gpu')" },