Compute PolicyEngine provenance once per chunked or retry invocation

[MaxGhenis]

Summary

eval-no-tools-chunked and retry-failed-responses start one python -m policybench.cli eval-no-tools subprocess per chunk (up to --parallel at a time, per model). Before this PR, each chunk worker built its resume sidecar's policyengine_bundles itself. For US chunks that means import policyengine, the cost #181 removed from supervised workers. This PR applies the same handoff to the chunked and retry paths. It is the third follow-up listed in #181.

What a chunk worker was doing (read and probed this session)

• run_chunk spawns eval-no-tools --num-scenarios --country --program-set --scenario-manifest --scenario-start --scenario-end --model -o [--no-explanations] [--single-output]. The subprocess inherited os.environ, with no provenance file.
• The only PolicyEngine import that argv reaches is _build_resume_metadata, then resolve_policyengine_bundles (falling back to computing), then policyengine_bundles_for_countries, then _load_policyengine_manifest. It happens once per worker, before the first LLM request, under both runners, with or without explanations, for every program set.
• A read-only probe of the exact run_chunk argv (fake LLM, heavy imports blocked and recorded) covered 28 combinations: us/uk × every program set × batch/single-output × explanations on/off. With a provenance file, no combination loaded or attempted a heavy module.
• Without the file, US workers attempted import policyengine. UK workers in this lockfile did not, because policyengine pins policyengine-uk 2.88.x and 2.89.0 is installed. That depends on the environment, so the UK path is handed the file too.
• _load_eval_scenarios rejects any manifest scenario whose country differs from --country. Every chunk that reaches its sidecar is therefore single-country, which Compute PolicyEngine provenance once per supervised run #181's single-country guard requires.

Change

• PolicyEngineProvenanceHandoff (policybench/chunked_eval.py). The first worker_env(countries) call writes <output_dir>/policyengine_provenance.json from a fresh interpreter, using the workers' sys.executable and environment. It reads the file back and, only if that succeeds, returns an env with POLICYBENCH_POLICYENGINE_PROVENANCE. It holds a lock, so model threads share one write; a failure is not retried per chunk.
• Once per invocation. run_chunked_eval creates one handoff and passes it to every run_model_chunks. run_retry_models creates one at the round's root rather than one per model_runs/<model>. run_model_chunks called on its own creates its own.
• Lazy. A handoff is requested only when a model has pending chunks, so a pass with nothing pending imports nothing.
• Same guards as the supervisor:
  • single-country only;
  • the one-hour writer timeout;
  • read-back before handing the file on;
  • any inherited POLICYBENCH_POLICYENGINE_PROVENANCE is dropped;
  • workers still reuse the file only when their fingerprint matches.
• One deliberate difference: no delete before writing. The runbook runs five provider groups at once on the same --output-dir "$RUN_DIR/$country". The writer replaces the file atomically (pid-named temp file plus os.replace), and this invocation hands the path on only after its own writer succeeded. Another invocation can still replace or remove the file while this one's workers run: Compute PolicyEngine provenance once per supervised run #181's writer removes whatever is at the path when its own round-trip check fails. A worker that then finds the file missing, or written for a different environment, computes the bundles itself. That is safe, but it loses the speedup.
• Shared code. PROVENANCE_WRITER, PROVENANCE_WRITER_TIMEOUT_SECONDS, run_policyengine_provenance_writer and read_written_policyengine_bundles move into policyengine_runtime.py. The supervisor's _compute_policyengine_provenance and _write_policyengine_provenance now call them. Supervisor behavior is unchanged, and its tests pass unmodified apart from using the shared probe.
• run_chunk takes env= and passes it to subprocess.run. With no env, the chunk inherits os.environ as before.
• Unchanged: the chunk argv, the sidecar format and RESUME_METADATA_VERSION, and chunk completeness and merging.
• Not added: an orchestrator-side check of each chunk sidecar against the read-back bundles, which the supervisor does. chunk_is_complete never reads sidecars, and adding that check would make resuming chunks written under older packages fail. Workers that fall back print PolicyEngine provenance file not reused to the invocation's terminal; nothing counts them, since chunk output isn't captured.

Output-dir scanners (audited this session)

Nothing that scans chunked or retry output dirs picks up the new file. Every scan falls into one of these groups:

• by_model/*.csv only: full_run_export.load_predictions, reparse_predictions.reparse_country_run.
• Fixed filenames:
  • runstore.import_run_dir: predictions.csv[.gz], scenarios.csv[.meta.json];
  • _available_countries and export_country;
  • fold_board's five named reference files;
  • freeze_snapshot's explicit copy list.
• Narrow globs: supervisor._validate_resume (scenarios/ only), load_repeated_predictions (*.csv).

The digests (reference CSV sha, runstore source_csv_sha256, retry_metadata.json, snapshot manifest) cover only named files. test_output_dir_scanners_ignore_the_provenance_file runs load_predictions, reparse_country_run and import_run_dir on a chunked dir with and without the file. It asserts identical results and that the file is left byte-for-byte unchanged.

Invariants

1. At most one write per invocation. For every request sequence, the writer runs at most once, and only when the first request is single-country.
2. A worker gets the file exactly when it should: when the file was written and read back, and the worker's countries equal the written ones. Otherwise the worker computes as before.
3. The worker environment is otherwise unchanged. It equals os.environ minus any inherited provenance variable.
4. No PolicyEngine outside the writer. Chunk workers and the orchestrator never load policyengine, policyengine_core, policyengine_us, policyengine_uk or h5py, checked in real subprocesses at every LLM request and at exit. The one-shot writer does load them (positive control).
5. Differential, byte-identical sidecars. A chunk sidecar written through the handoff is byte-identical to the sidecar the same run_chunk command writes with no file (the pre-change behavior). That baseline worker loads policyengine_us before its first request, which is the positive control that the probe sees PolicyEngine in a worker.

Invariants 1–3 are checked exhaustively rather than sampled. The test enumerates every request sequence of length 1–3 over {us}, {uk}, {us, uk} and {US}, crossed with 3 writer outcomes and with and without an inherited variable: 504 cases.

Tests

• tests/test_chunked_eval.py:
  • handoff: one write across 8 concurrent threads; an inherited variable replaced; four write-failure modes, none retried; a valid file left by another invocation not handed on when this invocation's writer fails; mixed-country skipped; other countries not handed the file; an existing file not deleted; the exhaustive invariant test;
  • run_model_chunks at --parallel 1 and 3; no write when nothing is pending;
  • run_chunked_eval at --model-parallel 1 and 2 writes once for two models;
  • the scanner differential;
  • the updated run_chunk fakes (they now take env).
• tests/test_retry_eval.py: run_retry_models with two models and --model-parallel 2 writes once, at the round root, and no file lands under model_runs/.
• End to end (real subprocesses, real writer, fake LLM via the shared tests/policyengine_probe.py sitecustomize):
  • test_chunked_run_keeps_policyengine_out_of_chunk_workers: eval-no-tools-chunked with 2 Claude models × 2 chunks. It checks one writer, a clean orchestrator, clean workers and no other probed process loading PolicyEngine, plus byte-identical sidecars against the baseline worker.
  • test_retry_round_keeps_policyengine_out_of_chunk_workers: retry-failed-responses with 2 models: one file at the root, and all 4 model-household responses retried with predictions.
• Mutation check: 9 mutants, each breaking one property. They cover writing on every call, handing the file to any country, keeping an inherited variable, dropping the single-country guard, writing with nothing pending, run_chunk dropping env, a per-model handoff in retry, a per-model handoff in chunked, and skipping read-back. The light tests caught all 9.
• Locally: ruff check, ruff format --check, and the non-PolicyEngine tests in test_chunked_eval.py and test_retry_eval.py (46 passed). The host is memory-constrained right now, so the PolicyEngine-importing tests (the three end-to-end tests, test_supervisor.py, test_policyengine_runtime.py) run in CI here, not locally.
• CI: at 8086ec2, uv run pytest -m "not slow" passed the whole suite, including both new end-to-end tests. The first CI run failed only on the retry test's last assertion, which expected the merge to accept the probe's fake answers. The fake's explanations have no value = X trailer, so _enforce_explanation_value_contract drops them and the merge rejects every group. The test now asserts that every model-household response was retried with predictions.

Adversarial review (three lenses, each finding put through a refutation pass)

• Fixed in this PR:
  • the runbook and docstring overclaimed that sharing an output dir never pulls the file from another invocation's workers (Compute PolicyEngine provenance once per supervised run #181's writer still unlinks it on a failed round-trip);
  • the Section 6 merge pass was a wrong example of a "no pending chunks" pass (it passes no --model);
  • the retry-dir list missed that a later round overwrites its copy of the source predictions;
  • two test gaps: an earlier valid file with a failing writer, and unclassified probed processes.
• Refuted:
  • "the lock test passes without the lock": 15 of 15 lock-less runs failed the test;
  • "a writer spawn OSError aborts the run": every spawn in the codebase already behaves this way, including run_chunk's own next spawn and the supervisor's writer on main, and no data is lost.
• Confirmed safe: the remaining findings describe safe fallbacks in configurations the runbook doesn't use, such as a supervisor run dir that is also a chunked output dir under a different venv.

Not measured

I have not timed chunked runs before and after. #181 measured the same eval-no-tools worker path under the supervisor: about 1 GB peak RSS and 11–14 CPU-seconds per worker for the PolicyEngine import. A chunk worker pays that once per chunk (--chunk-size scenarios, 50 by default and 10 for retries), not once per scenario.

axiom: n/a: benchmark runner infrastructure (no policy change)

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated
policybench-site	Ready	Preview	Sep 29, 2026 9:44pm UTC

[MaxGhenis]

Heads-up: #186 overlaps this PR in policyengine_runtime.py, supervisor.py and the supervisor tests, so whichever lands second rebases.

• Runtime: _load_policyengine_manifest and its dead policyengine.core.release_manifest fallback are gone. Bundles now come from package metadata plus the bundled manifest on disk. When the installed model equals the pin, the manifest is validated with policyengine's own CountryReleaseManifest. No PolicyEngine module is imported.
• Fingerprint: policyengine_provenance_inputs() drops the environment block (import flags and token presence), since those can no longer change the bundles.
• Supervisor docs: the docstring's cost paragraph and the single-country handoff comment are rewritten.
• This PR's writer: the comment that it runs in a fresh interpreter "so a single-country file records what a worker computing them itself would" stays true, but the fresh interpreter is no longer needed, because the result no longer depends on process state.

In the pinned environment the bundles are byte-identical to main's, and RESUME_METADATA_VERSION stays 7.