From f1fa9f465d2a8844cc538d0643dbe4de66b821d7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 24 Sep 2026 19:00:33 -0400 Subject: [PATCH 1/5] Score the post-export release gates over household batches (Route A PR-4) Route A's post-export stages (reform-coverage smoke, reform validation, demographics) built whole-file Microsimulations of the written H5 per reform pass, which the dense export cannot afford, and a mixed period order raised the MD CCS ParameterNotFoundError. - A household-batched scorer loads the written H5 once and scores each consumer's recorded baseline plan in one batch-outer pass, in ascending period order, with one reform system per reform and every engine released (#456). The scored file's sha must equal the manifest's dataset sha. - Additivity is enforced at run time: a multi-batch engine refuses population aggregates (Medicaid SLCSP state sums, income deciles), a reform engine refuses baseline-branch readers, and a reform may not move behavioral-response parameters. Watched lists are pinned against the installed engine. - The calibrated target frame is dropped once the export frame exists. - tools/sweep_us_post_export_scoring.py pushes every shipped smoke and validation reform (107 passes, OBBBA stacks included) through the scorer in batches and compares against whole-file engines. Refs #956, #456. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 6 + .../us-batched-post-export-scoring.fixed.md | 1 + docs/us-release-build-rule.md | 126 +- .../build/us_runtime/reform_coverage_smoke.py | 6 +- .../build/us_runtime/reform_validation.py | 22 +- .../us_runtime/release_gate_preflight.py | 12 +- .../build/us_runtime/spm_composition.py | 20 +- packages/microcosm-build/tests/conftest.py | 7 +- .../tests/test_us_fiscal_refresh_builder.py | 34 +- .../tests/test_us_post_export_scoring.py | 2397 +++++++++++++++++ tools/build_us_fiscal_refresh_release.py | 1136 +++++++- tools/sweep_us_post_export_scoring.py | 279 ++ 12 files changed, 3902 insertions(+), 144 deletions(-) create mode 100644 changelog.d/us-batched-post-export-scoring.fixed.md create mode 100644 packages/microcosm-build/tests/test_us_post_export_scoring.py create mode 100644 tools/sweep_us_post_export_scoring.py diff --git a/CLAUDE.md b/CLAUDE.md index 49dc882329..66624bc18b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -118,6 +118,12 @@ to the private repository; when you run `tools/build_uk_rowwise_candidate.py` yourself, pass `--staging-local-only` unless the operator asked for a staged upload. +The US fiscal-refresh builder scores its written H5 in household batches. +Before a release rerun, run the small-H5 guard sweep described in +[the release build rule](docs/us-release-build-rule.md#post-export-scoring). +That fixture check is separate from full-export timing and release +certification. + The US native-SPM-role source-enrichment lane is a separate release type: `tools/build_us_spm_role_enrichment.py` creates a local candidate from the exact reviewed BuildP parent, preserving original variables and inherited schema-5 diff --git a/changelog.d/us-batched-post-export-scoring.fixed.md b/changelog.d/us-batched-post-export-scoring.fixed.md new file mode 100644 index 0000000000..78be14167c --- /dev/null +++ b/changelog.d/us-batched-post-export-scoring.fixed.md @@ -0,0 +1 @@ +The US release's reform-coverage smoke, reform validation and demographics score the written H5 in household batches, bind its SHA-256 to the release manifest, and evaluate each baseline plan in ascending period order. Each reform shares one tax-benefit system across its batch engines. Guards refuse watched population-aggregate formulas in multi-batch passes, watched baseline-branch readers on reform engines, and changes to behavioral-response parameters. Real-engine tests exercise the refusals and compare shipped baseline plans and selected reform totals with whole-file engines; a separate sequential-worker sweep checks every shipped smoke and out-of-sample validation reform, including OBBBA stacks and in-sample fallback simulations. The MD CCS request-order premise is observed with a warning if the upstream error disappears, not asserted to persist. Writer tests pin export-before-scorer ordering and calibrated in-sample JCT estimates. After constructing the export frame, the build drops calibrated target frames; exact-k retains the clean export frame for its receipt's household-count check. Poverty and SPM outcomes remain record-only. diff --git a/docs/us-release-build-rule.md b/docs/us-release-build-rule.md index 8ee87102cf..83a263a228 100644 --- a/docs/us-release-build-rule.md +++ b/docs/us-release-build-rule.md @@ -79,16 +79,19 @@ waits on that fix. In policyengine-us 2.2.1 one SPM unit with no classified adult (age 18 or over, or age 15 or over with `is_spm_independent_minor_role`) refuses the SPM -measurement for the whole population. The first whole-dataset simulation of the -written release file is the reform-coverage smoke gate -(`tools/build_us_fiscal_refresh_release.py:11657`), which runs after the export -write (`:11641`) and before the calibration NPZ (`:11723`); its LIHEAP probe -reaches SPM composition through `spm_unit_benefits` -(`tools/build_us_release_input_coverage_manifest.py:513-520`), so a population -with such a unit refuses there. The reform-validation diagnostics that follow -the NPZ (`_write_reform_validation`, near `:11725`; 104 `in_poverty` rows from -`us/state_spm_poverty_levels.json` on one whole-dataset simulation) would -refuse the same way; they record, they do not gate. +measurement for every unit in the simulation that holds it. The first +simulation of the written release file is the reform-coverage smoke gate +(`us_reform_coverage_smoke_gate` in `tools/build_us_fiscal_refresh_release.py`), +which runs after the export write and before the calibration NPZ; its LIHEAP +probe reaches SPM composition through `spm_unit_benefits` (the probe list in +`tools/build_us_release_input_coverage_manifest.py`), so a population with such +a unit refuses there. The reform-validation diagnostics that follow the NPZ +(`_write_reform_validation`; 104 `in_poverty` rows from +`us/state_spm_poverty_levels.json`) would refuse the same way; they record, they +do not gate. These stages now score the written file in household batches, not +one whole-dataset simulation each (see +[Post-export scoring](#post-export-scoring)); the batch that holds such a unit +refuses, and the refusal still aborts the stage. The certified default passes because `tools/build_us_spm_role_enrichment.py` added the role afterwards, and that tool accepts only Build P's exact bytes @@ -176,8 +179,14 @@ stable source IDs, and their outputs are compared with or merged into the original assembled table. This lets historical source kernels consume the projection without treating intentionally removed provenance as corruption. -`--dense-default-dataset` is diagnostic only. A release build leaves it unset, -so the default is the sparse dataset that runs on standard machines. +Until 23 September 2026, `--dense-default-dataset` was diagnostic only and a +release build left it unset, so the default was the sparse dataset that runs on +standard machines. Decided 23 September 2026 (microcosm#956): route A's first +certified release, which is not the default, calibrates the whole base dense +(`--dense-default-dataset`) on the national and state target surface +(`--target-surface national_state`), and sparse L0 tuning is pursued after it. +The standing principle for the graph era is the same target surface for every +dataset size, with L0 the only difference across sizes. ### 4. Three raw inputs exist in one untracked directory @@ -194,6 +203,99 @@ and size and fetches them as it fetches the SIPP donor; the base stage verifies the digest it is given. `us/spec/sources.yaml` is generated from a sealed six-role receipt and does not change until a real build re-cuts it. +## Post-export scoring + +The reform-coverage smoke, reform validation and demographics score the written +release H5 through one household-batched scorer +(`_HouseholdBatchedPostExportScorer` in +`tools/build_us_fiscal_refresh_release.py`, microcosm#956). Each used to build +one Microsimulation over the whole file. + +- The scorer loads the written H5 once, binds its SHA-256, and refuses to let + the release manifest pin different bytes. +- It refuses a file whose group units do not each nest in one household, or + whose batches do not partition every entity, before any engine exists. +- Batches hold at most `--maximum-microsim-batch-size` households (default + 5,000; route A runs 2,000). Every engine declares the release SPM selection + (`US_RELEASE_SPM_SELECTION`, county geography) and is released before the + next batch builds. +- Each stage's baseline requests are learned before export by an engine-free + dry run and scored in one pass over the batches, in ascending period order. + One engine refuses a period earlier than one it has already computed. + The real-engine test probes the MD CCS request-order premise by requesting + 2025 state taxes after 2027 keys. It checks for the + `md.msde.ccs.payment.informal.rates` error and warns if that error is absent; + it does not require the upstream bug to persist. +- Each reform builds one tax-benefit system and scores its measure batch by + batch. Each batch engine gets that system alone, without `reform=`: given + both, policyengine-us creates a baseline branch, which clones the system. +- Every statistic reads full-length values with their engine weights. + +The totals equal a whole-file simulation's up to floating-point summation order +only if every scored measure is additive across households and never reads the +`baseline` branch that a batch reform engine lacks. The scorer rejects these +known violations at run time: + +- In a multi-batch pass, an engine refuses once it has computed a formula that + aggregates over its whole simulation. In policyengine-us 2.2.1 these are the + Medicaid SLCSP state sums behind `medicaid_cost` and the weighted household + and SPM-unit income deciles (`POST_EXPORT_POPULATION_AGGREGATE_VARIABLES`). + Each batch would compute them over itself alone. +- A reform engine refuses once it has computed a formula that reads the + baseline branch (`POST_EXPORT_BASELINE_BRANCH_READERS`). These are the + Medicaid denominator a reform holds at baseline, and the behavioral-response + measurements. +- A reform pass refuses a reform that sets the labor-supply or capital-gains + response parameters off their baseline values. Without a baseline branch, + those responses score 0. A whole-file engine also scores 0 at the engine's + zero default elasticities. + +Tests pin the watched lists against the installed engine's formula sources, +exercise each guard on a real engine, and compare the shipped baseline plans +(18 smoke keys, 75 validation keys and 1 demographics key) against whole-file +values on a small written H5 in three batches. + +On 2026-09-24, the sweep below completed on a written fixture containing eight +households across five states, with Medicaid enrollees, under policyengine-us +2.2.1. All 107 reform passes cleared the guards in three batches: 41 smoke +passes and 66 validation passes, covering 52 out-of-sample specs, 12 in-sample +fallbacks and two OBBBA pre-baselines. The 18 OBBBA stacked provision states +were included. All three baseline plans and four selected reform passes +matched whole-file values at `rtol=1e-12, atol=1e-6`; their weights matched +exactly. This is evidence for the fixture's calculation paths. + +The complete reform sweep is a separate pre-rerun step: + +```bash +uv run --no-sync python tools/sweep_us_post_export_scoring.py \ + --dataset-path /path/to/small-written-export.h5 --batch-size 3 \ + --output /path/to/guard-sweep.json +``` + +The input must contain enough households for at least two batches. The tool +records the shipped consumers' requests, including every out-of-sample +validation reform and the OBBBA pre-baseline and stacked states. It also +exercises the fallback simulations used when an in-sample estimate is absent. +It scores at most four reform passes per worker, serially, and compares all +baseline values and four reform passes with whole-file engines. Workers stop +at a sampled 7 GiB RSS threshold and report their peak RSS, elapsed time, +counts and the scored H5 digest. +Run this fixture sweep before R4's full-export timing steps, and repeat it +after a probe, spec or engine-lock change before starting a release. +R4's existing full-export rehearsal times the smoke and validation +baseline plans plus three selected reform passes; it does not cover every +reform and does not replace this sweep. + +`calibration_diagnostics.json` records the plan under `build.post_export_scoring` +(method, batch size, batch count, each stage's baseline keys and periods), and +`reform_coverage_smoke.json` records what the smoke scored under +`post_export_scoring`. Poverty rows stay record-only. Before the export write +the build frees the target frame and the calibration result's target tables: +`Frame.with_weights` copies every table, target columns included. Dense and +L0 results drop those frames; exact-k substitutes the clean export frame for +its receipt's later household-count check. The fixture sweep does not measure +memory or wall time of the batched stages at full size. + ## Measured stage times From the September attempt, on one 128 GiB machine, under policyengine-us diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/reform_coverage_smoke.py b/packages/microcosm-build/src/microcosm/build/us_runtime/reform_coverage_smoke.py index 1fce87c756..30048e088a 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/reform_coverage_smoke.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/reform_coverage_smoke.py @@ -16,8 +16,10 @@ The gate takes an injected ``simulate(reform) -> simulation`` (the same seam as :mod:`microcosm.build.us_runtime.reform_validation`), so it unit-tests without -policyengine-us and runs live against the written release H5 in the build via -:func:`microcosm.build.us_runtime.reform_validation.default_simulate_factory`. +policyengine-us. In the build it runs live against the written release H5 +through the release tool's household-batched post-export scorer +(``tools/build_us_fiscal_refresh_release.py``, microcosm#956), which serves the +same seam from one engine per household batch. """ from __future__ import annotations diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/reform_validation.py b/packages/microcosm-build/src/microcosm/build/us_runtime/reform_validation.py index 2ee2456daf..36b1f30b53 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/reform_validation.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/reform_validation.py @@ -684,6 +684,19 @@ def _capped_weighted_total( US_RELEASE_SPM_SELECTION: dict[str, object] = {"geography_kind": "county"} +def _released_engine_state(simulation: Any) -> bool: + """Release a finished simulation; True if it held engine state to sweep. + + Callers run a full collection after releasing an object with a populations + dict (microcosm#456). The release tool's household-batched adapter releases + and collects its own batch engines; it and the engine-free recorder expose + no populations dict requiring another collection here. + """ + holds_engine_state = isinstance(getattr(simulation, "populations", None), dict) + release_engine_simulation(simulation) + return holds_engine_state + + def default_simulate_factory(dataset_path: Path) -> SimulateFn: """Build a simulate() that runs a Microsimulation over the release H5.""" @@ -768,8 +781,8 @@ def released_weighted_total( try: return _weighted_total(transient_simulation, measure, at_period) finally: - release_engine_simulation(transient_simulation) - gc.collect() + if _released_engine_state(transient_simulation): + gc.collect() def stacked_obbba_effects() -> dict[str, tuple[float, float, float]]: """Score the OBBBA provisions *stacked* in their JCX-35-25 order. @@ -1023,9 +1036,10 @@ def _level_total(level: BaselineLevelSpec) -> float | None: # microcosm#456: the shared baseline simulation has served every reform row # and baseline-level row by now; release it before assembling the payload. if baseline is not None: - release_engine_simulation(baseline) + sweep = _released_engine_state(baseline) baseline = None - gc.collect() + if sweep: + gc.collect() has_out_of_sample = any(not spec.in_sample for spec in specs) or bool( baseline_levels diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/release_gate_preflight.py b/packages/microcosm-build/src/microcosm/build/us_runtime/release_gate_preflight.py index 10d0962e27..f1e2978a42 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/release_gate_preflight.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/release_gate_preflight.py @@ -49,14 +49,14 @@ 5. **SPM measurement composition** — does every SPM unit have a classified adult? In ``spm-calculator`` 1.0.0 a *single* SPM unit whose members are all under 18, with none aged 15-17 carrying an SPM independence role, raises - ``SPMInputError("SPM_COMPOSITION_REQUIRED")`` for the **whole population's** - SPM measurement (``spm_calculator/policyengine_adapter.py`` + ``SPMInputError("SPM_COMPOSITION_REQUIRED")`` for the SPM measurement of + **every unit in its simulation** (``spm_calculator/policyengine_adapter.py`` ``policyengine_amount``: ``if np.any(adults < 1): raise``). The release's 104 - state SPM poverty levels run on one whole-dataset ``Microsimulation`` - (:func:`~microcosm.build.us_runtime.reform_validation.default_simulate_factory`) + state SPM poverty levels are measured on the written H5, one engine per + household batch (the release tool's post-export scorer, microcosm#956), *after* calibration, export and the NPZ write, so that engine traceback — - which names neither the unit nor a remedy — is what an operator gets hours - into a build. This reproduces the engine rule + which names neither the unit nor a remedy, and aborts the stage — is what an + operator gets hours into a build. This reproduces the engine rule (``adult = (age >= 18) | ((age >= 15) & role)``) over frame columns and names the offending units and the remedy in seconds. diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/spm_composition.py b/packages/microcosm-build/src/microcosm/build/us_runtime/spm_composition.py index 088682ca43..17e7911ea4 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/spm_composition.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/spm_composition.py @@ -366,16 +366,18 @@ def check_spm_composition( ) -> CheckResult: """Does every SPM unit have a classified adult (no solve, no engine)? - The release's 104 state SPM poverty levels are measured on one whole-dataset - ``Microsimulation`` (``reform_validation.default_simulate_factory``), and in - ``spm-calculator`` 1.0.0 a *single* SPM unit with no classified adult raises - ``SPMInputError("SPM_COMPOSITION_REQUIRED")`` for the **whole population's** - measurement (``spm_calculator/policyengine_adapter.py`` ``policyengine_amount``: + The release's 104 state SPM poverty levels are measured on the written H5, + one engine per household batch (the release tool's post-export scorer, + microcosm#956), and in ``spm-calculator`` 1.0.0 a *single* SPM unit with no + classified adult raises ``SPMInputError("SPM_COMPOSITION_REQUIRED")`` for the + measurement of **every unit in its simulation** + (``spm_calculator/policyengine_adapter.py`` ``policyengine_amount``: ``if np.any(adults < 1): raise``), naming neither the offending unit nor a - remedy. The release tool runs this same classification on its calibrated - export frame as a batched pre-export gate, so a build refuses by name before - the H5 and NPZ writes — but only after paying for the calibration that frame - comes from. Run pre-solve, this is the same verdict in seconds. + remedy, and aborting the stage. The release tool runs this same + classification on its calibrated export frame as a batched pre-export gate, + so a build refuses by name before the H5 and NPZ writes — but only after + paying for the calibration that frame comes from. Run pre-solve, this is the + same verdict in seconds. This reproduces the engine's classification exactly — ``adult = (age >= 18) | ((age >= 15) & role)`` with ``role`` resolved by diff --git a/packages/microcosm-build/tests/conftest.py b/packages/microcosm-build/tests/conftest.py index 3f874491a9..8de2e3dc35 100644 --- a/packages/microcosm-build/tests/conftest.py +++ b/packages/microcosm-build/tests/conftest.py @@ -9,7 +9,6 @@ import pandas as pd import pytest -from microcosm.build.uk_runtime.ods_tables import ODS_MIME_TYPE from microcosm.frame import EntitySchema, Frame, WeightKind, Weights @@ -77,10 +76,14 @@ def _write_ods( path: Path, tables_xml: str, *, - mimetype: str = ODS_MIME_TYPE, + mimetype: str | None = None, store_mimetype: bool = True, include_content: bool = True, ) -> Path: + if mimetype is None: + from microcosm.build.uk_runtime.ods_tables import ODS_MIME_TYPE + + mimetype = ODS_MIME_TYPE with ZipFile(path, "w") as archive: archive.writestr( "mimetype", diff --git a/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py b/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py index fe6a0638be..b8ea501d68 100644 --- a/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py +++ b/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py @@ -6247,6 +6247,9 @@ def test_main_writes_diagnostics_before_post_calibration_gate_failure( result = SimpleNamespace( skipped=(), diagnostics=(), + # A calibration result always carries its compiled problem; the + # pre-export post-export-scoring plan reads its targets (#956). + problem=SimpleNamespace(targets=()), initial_loss=2.0, final_loss=1.0, l0_lambda=0.2, @@ -7996,7 +7999,9 @@ def fake_write_calibration_diagnostics(result, path, *, target_registry, build): monkeypatch.setattr(builder, "calibrate_l0_refit", fake_calibrate_l0_refit) if terminal_mode == "puf_tail": - ladder_outcome = SimpleNamespace( + # The real receipt type, so ``_main``'s exact-k branch that drops the + # calibration frames before the export (microcosm#956) runs here. + ladder_outcome = builder.ExactKLadderCalibration( result=result, support=np.asarray([0, 1], dtype=np.int64), selected_inclusion_probabilities=np.ones(2), @@ -8167,6 +8172,24 @@ def fake_release_gate_failures(*args, **kwargs): "write_calibration_diagnostics", fake_write_calibration_diagnostics, ) + # The pre-export post-export-scoring plan (microcosm#956) dry-runs the real + # smoke and reform-validation consumers without an engine; only the reform + # objects they build need policyengine-core, which the engine-free lane + # does not install. + import microcosm.build.us_runtime.reform_coverage_smoke as smoke_module + import microcosm.build.us_runtime.reform_validation as reform_validation_module + + monkeypatch.setattr(smoke_module, "_build_reform", lambda probe: probe.id) + monkeypatch.setattr( + reform_validation_module.ReformValidationSpec, + "build_reform", + lambda spec: spec.id, + ) + monkeypatch.setattr( + reform_validation_module, + "_build_parameter_reform", + lambda changes: tuple(sorted(changes)), + ) if green_run: _run_green_register_release( @@ -8521,6 +8544,15 @@ def fake_release_gate_failures(*args, **kwargs): "social_security_components": ss_repair_payload, "non_sch_d_capital_gains": cgd_repair_payload, } + # The diagnostics carry a complete post-export scoring plan, never an + # error record, so no plan line joins the terminal batch (#956). + post_export_scoring = captured["diagnostics"]["post_export_scoring"] + assert "error" not in post_export_scoring + assert list(post_export_scoring["consumers"]) == [ + "reform_coverage_smoke", + "reform_validation", + "demographics", + ] assert captured["diagnostics"]["default_dataset"] == { "method": "l0_refit", "sparse": True, diff --git a/packages/microcosm-build/tests/test_us_post_export_scoring.py b/packages/microcosm-build/tests/test_us_post_export_scoring.py new file mode 100644 index 0000000000..baf1f205fb --- /dev/null +++ b/packages/microcosm-build/tests/test_us_post_export_scoring.py @@ -0,0 +1,2397 @@ +"""Household-batched post-export scoring (route A remediation, microcosm#956). + +The reform-coverage smoke, reform_validation and demographics score the +written release H5 through ``_HouseholdBatchedPostExportScorer`` instead of one +whole-pool Microsimulation each. These tests run the shipped consumers (the +41 smoke probes, the shipped reform-validation specs and baseline levels, and +the age distribution) against a fake engine on a small nested frame and pin +the plan's contract: output identical to the unbatched path, batches that +partition the pool, one live engine at a time, the SPM selection on every +construction, one reform system per reform (handed to each batch engine +without ``reform=``), one baseline pass per plan, ascending periods per engine, +engine-free recording, the ``_main`` and writer wiring, the scored-sha +binding, the nesting, order and weight premises, the additivity guards +(population aggregates, baseline-branch readers, behavioral-response +parameters, pinned against the installed engine's sources), the calibration +result's target frames dropped before the export, and a plan that cannot be +built joining the terminal batch instead of raising. Real-engine tests cover +whole-file equality for selected reforms, the additivity refusals, and all +shipped baseline plans. The MD CCS request-order premise is observed, with a +warning if the upstream error disappears. The complete reform sweep runs in +sequential worker processes through ``tools/sweep_us_post_export_scoring.py``; +fake-engine tests check its recording and worker plumbing here. +""" + +from __future__ import annotations + +import ast +import gc +import importlib.util +import inspect +import json +import re +import sys +import warnings +import zlib +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + +import microcosm.build.us_runtime.reform_coverage_smoke as smoke_module +import microcosm.build.us_runtime.reform_validation as reform_validation_module +from microcosm.build.us_runtime.demographics import demographics_payload +from microcosm.build.us_runtime.reform_validation import ( + US_RELEASE_SPM_SELECTION, + ReformValidationSpec, +) +from microcosm.build.us_runtime.release_input_coverage import ( + ReformCoverageProbe, + us_release_reform_coverage_probes, +) +from microcosm.calibrate import TargetRegistry, TargetSpec, calibrate +from microcosm.calibrate.geography_constants import US_STATE_NUMERIC_FIPS_TO_POSTAL +from microcosm.frame import EntitySchema, Frame, WeightKind, Weights +from microcosm.frame.units import US_SCHEMA + + +def _load_builder_module(): + root = Path(__file__).resolve().parents[3] + path = root / "tools" / "build_us_fiscal_refresh_release.py" + spec = importlib.util.spec_from_file_location( + "build_us_fiscal_refresh_release", path + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def builder(): + return _load_builder_module() + + +# --------------------------------------------------------------------------- +# A small nested US frame and a fake engine over it +# --------------------------------------------------------------------------- + +# (state_fips, [ages]) per household; households 2 and 5 hold two tax units. +_HOUSEHOLDS = ( + (24, (41, 39, 6)), + (6, (29, 31)), + (36, (72,)), + (24, (35, 8, 3, 1)), + (48, (55, 52, 19)), + (6, (24,)), + (36, (47, 16)), +) +_SECOND_TAX_UNIT = {2: 1, 5: 2} # household -> member index starting unit 2 + + +def _nested_frame() -> Frame: + rows = [] + person_id = 0 + for household_id, (_, ages) in enumerate(_HOUSEHOLDS, start=1): + split = _SECOND_TAX_UNIT.get(household_id) + for index, age in enumerate(ages): + person_id += 1 + tax_unit = household_id * 10 + ( + 1 if split is not None and index >= split else 0 + ) + rows.append( + { + "person_id": person_id, + "person_household_id": household_id, + "person_tax_unit_id": tax_unit, + "person_spm_unit_id": household_id * 100, + "person_family_id": household_id * 1000, + "person_marital_unit_id": person_id * 10_000, + "age": float(age), + } + ) + person = pd.DataFrame(rows) + household_ids = np.arange(1, len(_HOUSEHOLDS) + 1, dtype=np.int64) + tables = { + "person": person, + "household": pd.DataFrame( + { + "household_id": household_ids, + "state_fips": np.asarray([state for state, _ in _HOUSEHOLDS]), + } + ), + "tax_unit": pd.DataFrame( + {"tax_unit_id": np.unique(person["person_tax_unit_id"].to_numpy())} + ), + "spm_unit": pd.DataFrame({"spm_unit_id": household_ids * 100}), + "family": pd.DataFrame({"family_id": household_ids * 1000}), + "marital_unit": pd.DataFrame( + {"marital_unit_id": person["person_marital_unit_id"].to_numpy()} + ), + } + # Integer-valued weights and values keep every weighted sum exact, so the + # batched and unbatched payloads can be compared with ==. + weights = np.asarray([3.0, 5.0, 7.0, 11.0, 13.0, 17.0, 19.0]) * 100.0 + return Frame( + tables, + US_SCHEMA, + {"household": Weights(weights, WeightKind.CALIBRATED)}, + ) + + +class _FakeReform: + def __init__(self, label: str, *, parameters=None) -> None: + self.label = label + self.delta = float(zlib.crc32(label.encode()) % 41 + 1) + # Parameter values this reform sets, by full parameter name. + self.parameters = dict(parameters or {}) + + +class _FakeParameter: + """A policyengine-core ``Parameter``: named, valued, callable at an instant.""" + + def __init__(self, name: str, value: float) -> None: + self.name = name + self.values_list = [value] + + def __call__(self, instant): + return self.values_list[0] + + def get_descendants(self): + return iter(()) + + +class _FakeParameterTree: + """``system.parameters``: ``get_child`` returns a node over its leaves.""" + + _BASELINE = { + "gov.simulation.capital_gains_responses.elasticity": 0.0, + "gov.simulation.labor_supply_responses.elasticities.income": 0.0, + "gov.simulation.labor_supply_responses.elasticities.substitution.all": 0.0, + "gov.simulation.labor_supply_responses.bounds.income_change": 0.5, + } + + def __init__(self, overrides) -> None: + self._leaves = tuple( + _FakeParameter(name, value) + for name, value in {**self._BASELINE, **overrides}.items() + ) + + def get_child(self, path: str): + leaves = tuple(leaf for leaf in self._leaves if leaf.name.startswith(path)) + assert leaves, path + return SimpleNamespace( + name=path, children={}, get_descendants=lambda: iter(leaves) + ) + + +class _FakeSeries: + """The slice of a policyengine MicroSeries the consumers read.""" + + def __init__(self, values, weights) -> None: + self._values = np.asarray(values) + self.weights = np.asarray(weights, dtype=np.float64) + + def __array__(self, dtype=None, copy=None): + return np.asarray(self._values, dtype=dtype) + + @property + def values(self): + return self._values + + def sum(self) -> float: + return float(pd.Series(self._values).multiply(self.weights).sum()) + + +_HOUSEHOLD_VARIABLES = frozenset({"state_code_str", "state_fips"}) +_PERSON_VARIABLES = frozenset( + {"age", "in_poverty", "is_child", "ssi", "medicare_cost", "head_start", "wic"} +) +_SPM_UNIT_VARIABLES = frozenset({"snap", "housing_assistance", "spm_unit_benefits"}) + + +def _native_entity(variable: str) -> str: + """Person, SPM-unit and household measures as the engine keys them; every + other measure (taxes, credits, SOI caps) on the tax unit.""" + if variable in _HOUSEHOLD_VARIABLES: + return "household" + if variable in _PERSON_VARIABLES: + return "person" + if variable in _SPM_UNIT_VARIABLES or variable.startswith("spm_unit"): + return "spm_unit" + return "tax_unit" + + +def _household_positions(frame: Frame, entity: str) -> np.ndarray: + household_positions = pd.Series( + np.arange(frame.n("household")), + index=frame.table("household")["household_id"].to_numpy(), + ) + person = frame.table("person") + if entity == "household": + return np.arange(frame.n("household")) + if entity == "person": + return household_positions.reindex(person["person_household_id"]).to_numpy() + first_household = person.drop_duplicates(f"person_{entity}_id").set_index( + f"person_{entity}_id" + )["person_household_id"] + unit_households = first_household.reindex( + frame.table(entity)[f"{entity}_id"].to_numpy() + ) + return household_positions.reindex(unit_households.to_numpy()).to_numpy() + + +def _native_values(frame: Frame, variable: str, period: int, delta: float): + person = frame.table("person") + state_fips = frame.table("household")["state_fips"].to_numpy() + if variable == "state_fips": + return state_fips + if variable == "state_code_str": + return np.asarray( + [US_STATE_NUMERIC_FIPS_TO_POSTAL[int(fips)] for fips in state_fips], + dtype=object, + ) + if variable == "age": + return person["age"].to_numpy(dtype=np.float64) + if variable == "is_child": + return (person["age"].to_numpy() < 18).astype(np.float64) + if variable == "in_poverty": + return person["person_id"].to_numpy() % 3 == 0 + entity = _native_entity(variable) + ids = frame.table(entity)[f"{entity}_id"].to_numpy() + base = zlib.crc32(f"{variable}@{period}".encode()) % 89 + return (base + ids % 11 + delta).astype(np.float64) + + +def _map_values(frame: Frame, values, native: str, target: str): + if target == native: + return values + person = frame.table("person") + if target == "person": + if native == "household": + return np.asarray(values)[_household_positions(frame, "person")] + unit_positions = pd.Series( + np.arange(frame.n(native)), + index=frame.table(native)[f"{native}_id"].to_numpy(), + ).reindex(person[f"person_{native}_id"].to_numpy()) + return np.asarray(values)[unit_positions.to_numpy()] + if target == "household": + out = np.zeros(frame.n("household"), dtype=np.float64) + np.add.at(out, _household_positions(frame, native), values) + return out + raise NotImplementedError(f"{native} -> {target}") + + +class _EngineLog: + def __init__(self) -> None: + self.constructions: list = [] + self.systems: list = [] + self.max_unreleased_at_construction = 0 + self.fail_on: str | None = None + # variable -> the watched formulas computing it leaves a value for. + self.computes: dict[str, tuple[str, ...]] = {} + # variable -> periods the "written H5" stores for it as an input. + self.stored_inputs: dict[str, tuple[int, ...]] = {} + + +def _fake_engine(log: _EngineLog): + class FakeSystem: + def __init__(self, reform=None) -> None: + self.reform = reform + self.parameters = _FakeParameterTree( + {} if reform is None else reform.parameters + ) + log.systems.append(self) + + class FakeHolder: + def __init__(self, periods) -> None: + self._periods = periods + + def get_known_periods(self): + return sorted(self._periods) + + class FakeMicrosimulation: + default_tax_benefit_system = FakeSystem + default_tax_benefit_system_instance = SimpleNamespace( + parameters=_FakeParameterTree({}) + ) + + def __init__(self, *, dataset, reform=None, tax_benefit_system=None, spm=None): + unreleased = sum(engine.dataset is not None for engine in log.constructions) + log.max_unreleased_at_construction = max( + log.max_unreleased_at_construction, unreleased + ) + self.dataset = dataset + # What this construction saw, kept past release_engine_simulation + # (which severs ``dataset``). + self.household_ids = tuple(dataset.table("household")["household_id"]) + self.reform = reform + self.tax_benefit_system = tax_benefit_system + self.spm = spm + self.periods: list[int] = [] + self.branches: dict = {} + self.known = { + name: set(periods) for name, periods in log.stored_inputs.items() + } + log.constructions.append(self) + + @property + def policy_reform(self): + """The reform this engine scores: passed, or carried by its system.""" + if self.reform is not None: + return self.reform + return getattr(self.tax_benefit_system, "reform", None) + + def get_holder(self, variable): + return FakeHolder(self.known.get(variable, ())) + + def calculate(self, variable, period, map_to=None): + frame = self.dataset + assert frame is not None, "calculate on a released engine" + if variable == log.fail_on: + raise RuntimeError(f"fixture engine refuses {variable}") + period = int(period) + self.periods.append(period) + for computed in (variable, *log.computes.get(variable, ())): + self.known.setdefault(computed, set()).add(period) + native = _native_entity(variable) + reform = self.policy_reform + delta = 0.0 if reform is None else reform.delta + values = _native_values(frame, variable, period, delta) + target = map_to or native + values = _map_values(frame, values, native, target) + weights = frame.weights_for("household").values[ + _household_positions(frame, target) + ] + return _FakeSeries(values, weights) + + return FakeMicrosimulation + + +@pytest.fixture +def fake_reforms(monkeypatch, builder): + """Engine-free reform objects for the shipped probes and specs. + + Full garbage collections are stubbed too; these stand-ins do not create + the engine state collected between production passes. + """ + monkeypatch.setattr(builder, "_collect_family_garbage", lambda: None) + monkeypatch.setattr( + reform_validation_module, "gc", SimpleNamespace(collect=lambda *args: 0) + ) + monkeypatch.setattr( + smoke_module, "_build_reform", lambda probe: _FakeReform(probe.id) + ) + monkeypatch.setattr( + ReformValidationSpec, "build_reform", lambda spec: _FakeReform(spec.id) + ) + monkeypatch.setattr( + reform_validation_module, + "_build_parameter_reform", + lambda changes: _FakeReform(json.dumps(changes, sort_keys=True, default=str)), + ) + + +def _empty_calibration_result(): + return SimpleNamespace(diagnostics=[], problem=SimpleNamespace(targets=[])) + + +#: Fitted (final estimate, target) per in-sample JCT reform in +#: ``_calibration_result_with_in_sample_fit``, in spec order. +_IN_SAMPLE_FIT = ((-1_250.0, -1_000.0), (-3_500.0, -4_000.0)) + + +def _calibration_result_with_in_sample_fit(period: int): + """A calibration fit covering the first two in-sample JCT reforms, plus a + target no reform reads. Reform validation takes those two rows from the + fit instead of simulating them.""" + fitted = reform_validation_module.in_sample_reform_specs(period=period)[:2] + fit = (*_IN_SAMPLE_FIT, (7.0, 8.0)) + names = (*(spec.id for spec in fitted), "fixture_other_target") + return SimpleNamespace( + diagnostics=[ + SimpleNamespace(final_estimate=estimate, target=target) + for estimate, target in fit + ], + problem=SimpleNamespace(targets=[SimpleNamespace(name=name) for name in names]), + ) + + +def _written_h5(tmp_path: Path) -> Path: + tmp_path.mkdir(parents=True, exist_ok=True) + path = tmp_path / "populace_us_2024.h5" + path.write_bytes(b"the written release H5 bytes") + return path + + +def _scorer(builder, frame, log, batch_size, tmp_path, *, loads=None): + path = _written_h5(tmp_path) + + def load_frame(load_path, *, expected_sha256): + if loads is not None: + loads.append((Path(load_path), expected_sha256)) + return frame + + return builder._HouseholdBatchedPostExportScorer( + path, + maximum_microsim_batch_size=batch_size, + microsimulation_cls=_fake_engine(log), + dataset_from_frame=lambda batch_frame: batch_frame, + load_frame=load_frame, + ) + + +def _unbatched_simulate(frame, log): + engine_cls = _fake_engine(log) + + def simulate(reform): + return engine_cls( + dataset=frame, reform=reform, spm=dict(US_RELEASE_SPM_SELECTION) + ) + + return simulate + + +def _run_consumers(builder, simulate_for): + """Run the three shipped consumers; ``simulate_for(name, consumer)``.""" + gate = builder.us_reform_coverage_smoke_gate( + simulate=simulate_for( + "reform_coverage_smoke", builder._reform_coverage_smoke_consumer + ), + period=builder.PERIOD, + ) + payload_for = builder._reform_validation_consumer( + result=_empty_calibration_result(), release_id="fixture-release" + ) + payload = payload_for(simulate_for("reform_validation", payload_for)) + ages, weights = builder._demographics_consumer( + simulate_for("demographics", builder._demographics_consumer) + ) + demographics = demographics_payload(ages, weights, period=builder.PERIOD) + return gate, payload, demographics + + +def _batched_run(builder, frame, batch_size, tmp_path): + log = _EngineLog() + scorer = _scorer(builder, frame, log, batch_size, tmp_path) + consumers: dict[str, object] = {} + + def simulate_for(name, consumer): + plan = builder._record_post_export_baseline_plan(consumer) + consumers[name] = scorer.open_consumer(name, plan) + return consumers[name].simulate + + outputs = _run_consumers(builder, simulate_for) + return outputs, log, scorer, consumers + + +# --------------------------------------------------------------------------- +# (a) output invariance +# --------------------------------------------------------------------------- + + +def test_batched_scoring_matches_the_unbatched_path( + builder, fake_reforms, tmp_path +) -> None: + """(a) The smoke's GateResult, every reform_validation row and the + demographics payload are identical for batch sizes 1, 2 and unbatched, and + identical to one whole-frame engine per simulate() call.""" + frame = _nested_frame() + reference = _run_consumers( + builder, + lambda name, consumer: _unbatched_simulate(frame, _EngineLog()), + ) + reference_gate, reference_payload, reference_demographics = reference + assert reference_payload["out_of_sample_simulated"] is True + assert reference_gate.details["probes"] == len(us_release_reform_coverage_probes()) + + for batch_size in (1, 2, None): + (gate, payload, demographics), _, scorer, _ = _batched_run( + builder, frame, batch_size, tmp_path / f"b{batch_size}" + ) + assert scorer.n_batches == (1 if batch_size is None else -(-7 // batch_size)) + assert gate.passed == reference_gate.passed + assert gate.failures == reference_gate.failures + assert gate.details == reference_gate.details + assert payload == reference_payload + assert demographics == reference_demographics + + +# --------------------------------------------------------------------------- +# (b)-(g) the engine contract +# --------------------------------------------------------------------------- + + +def test_batched_engines_partition_release_and_order_their_work( + builder, fake_reforms, tmp_path +) -> None: + """(b) no construction sees more than the batch size and every pass + partitions the pool; (c) one engine alive at a time; (d) the SPM selection + on every construction; (e) one reform system per reform, 41 for the + shipped smoke; (f) one baseline pass per plan; (g) ascending periods.""" + frame = _nested_frame() + household_ids = frame.table("household")["household_id"].tolist() + batch_size = 3 + (gate, payload, _), log, scorer, consumers = _batched_run( + builder, frame, batch_size, tmp_path + ) + n_batches = scorer.n_batches + assert n_batches == 3 + + # (b) No construction saw more than the batch size, and the constructions + # come in passes of n_batches, each covering every household exactly once + # and in the written order, with one reform (or none) per pass. + constructions = log.constructions + assert len(constructions) % n_batches == 0 + for construction in constructions: + assert construction.dataset is None # released (see (c)) + assert 0 < len(construction.household_ids) <= batch_size + for start in range(0, len(constructions), n_batches): + chunk = constructions[start : start + n_batches] + assert [id_ for engine in chunk for id_ in engine.household_ids] == ( + household_ids + ) + assert len({id(engine.tax_benefit_system) for engine in chunk}) == 1 + + # (c) At every construction, every earlier engine had been released. + assert log.max_unreleased_at_construction == 0 + + # (d) Every construction declares the release SPM selection. + assert all(engine.spm == US_RELEASE_SPM_SELECTION for engine in constructions) + + # (e) One reform system per reform, built once and shared by the batches. + # Each batch engine gets that system alone: ``reform=`` too would make + # policyengine-us clone the system on every construction. + smoke = consumers["reform_coverage_smoke"].record() + validation = consumers["reform_validation"].record() + demographics = consumers["demographics"].record() + probes = len(us_release_reform_coverage_probes()) + assert smoke["reform_systems"] == smoke["reform_passes"] == probes == 41 + assert validation["reform_systems"] == validation["reform_passes"] > 0 + assert demographics["reform_systems"] == demographics["reform_passes"] == 0 + assert all(engine.reform is None for engine in constructions) + reform_constructions = [ + engine for engine in constructions if engine.tax_benefit_system is not None + ] + assert len(log.systems) == smoke["reform_systems"] + validation["reform_systems"] + assert {id(engine.tax_benefit_system) for engine in reform_constructions} == { + id(system) for system in log.systems + } + assert all(system.reform is not None for system in log.systems) + assert len(reform_constructions) == len(log.systems) * n_batches + + # (f) One baseline pass per consumer plan. + baseline_constructions = [ + engine for engine in constructions if engine.tax_benefit_system is None + ] + assert len(baseline_constructions) == 3 * n_batches + for record in (smoke, validation, demographics): + assert record["baseline_passes"] == 1 + assert record["n_batches"] == n_batches + assert record["max_batch_households"] == batch_size + + # (g) Each engine computed its keys in non-decreasing period order, and + # the validation baseline spans several periods. + assert all(engine.periods == sorted(engine.periods) for engine in constructions) + validation_periods = validation["baseline_plan"]["period_order"] + assert validation_periods == sorted(validation_periods) + assert len(validation_periods) > 1 + assert payload["out_of_sample_simulated"] is True + assert gate.details["probes"] == probes + + +def test_engine_guard_refuses_an_earlier_period(builder) -> None: + """(g) The guard refuses a key earlier than one the engine computed.""" + calls = [] + engine = builder._AscendingPeriodEngine( + SimpleNamespace( + calculate=lambda variable, period, **kwargs: calls.append( + (variable, period, kwargs) + ) + ), + label="fixture", + ) + engine.calculate(("income_tax", 2024, None)) + engine.calculate(("state_income_tax", 2024, "household")) + engine.calculate(("ga_income_tax", 2026, None)) + with pytest.raises(RuntimeError, match="nd_income_tax@2025 after it computed"): + engine.calculate(("nd_income_tax", 2025, None)) + assert calls == [ + ("income_tax", 2024, {}), + ("state_income_tax", 2024, {"map_to": "household"}), + ("ga_income_tax", 2026, {}), + ] + + +def test_ascending_plan_is_stable_and_consumers_refuse_other_orders( + builder, tmp_path +) -> None: + keys = [ + ("income_tax", 2024, None), + ("ga_income_tax", 2026, None), + ("nd_income_tax", 2025, None), + ("state_fips", 2024, "person"), + ("income_tax", 2024, None), + ] + assert builder._ascending_period_plan(keys) == ( + ("income_tax", 2024, None), + ("state_fips", 2024, "person"), + ("nd_income_tax", 2025, None), + ("ga_income_tax", 2026, None), + ) + scorer = _scorer(builder, _nested_frame(), _EngineLog(), 3, tmp_path) + with pytest.raises(ValueError, match="ascending period order"): + scorer.open_consumer("fixture", tuple(dict.fromkeys(keys))) + + +def test_served_baseline_refuses_a_key_outside_its_plan(builder, tmp_path) -> None: + log = _EngineLog() + scorer = _scorer(builder, _nested_frame(), log, 3, tmp_path) + consumer = scorer.open_consumer("fixture", (("income_tax", 2024, None),)) + baseline = consumer.simulate(None) + assert baseline.calculate("income_tax", 2024).sum() > 0 + with pytest.raises(RuntimeError, match="does not hold"): + baseline.calculate("income_tax", 2026) + # The served baseline answered from the plan: one pass, no new engines. + assert len(log.constructions) == scorer.n_batches + + +def test_a_failing_key_still_releases_its_batch_engine(builder, tmp_path) -> None: + """(c) The engine that raised is released before the error propagates, so + a failed pass never leaves a batch engine's arrays alive.""" + log = _EngineLog() + log.fail_on = "state_income_tax" + scorer = _scorer(builder, _nested_frame(), log, 3, tmp_path) + with pytest.raises(RuntimeError, match="fixture engine refuses"): + scorer.open_consumer( + "fixture", + (("income_tax", 2024, None), ("state_income_tax", 2024, None)), + ) + assert len(log.constructions) == 1 + assert log.constructions[0].dataset is None + + +def test_post_export_values_sum_matches_microseries_sum(builder) -> None: + """``.sum()`` equals policyengine-core's ``MicroSeries.sum`` (NaN skipped) + and the arrays are read-only.""" + microdf = pytest.importorskip("microdf") + values = np.asarray([1.5, np.nan, 2.0, 4.0]) + weights = np.asarray([10.0, 20.0, np.nan, 0.5]) + result = builder._PostExportValues(values, weights) + assert result.sum() == float(microdf.MicroSeries(values, weights=weights).sum()) + flags = np.asarray([True, False, True, True]) + assert builder._PostExportValues(flags, weights).sum() == float( + microdf.MicroSeries(flags, weights=weights).sum() + ) + with pytest.raises(ValueError): + np.asarray(result)[0] = 9.0 + with pytest.raises(ValueError, match="one weight per value"): + builder._PostExportValues(values, weights[:2]) + + +# --------------------------------------------------------------------------- +# Additivity: population aggregates and the missing baseline branch +# --------------------------------------------------------------------------- + + +def test_a_multi_batch_pass_refuses_a_population_aggregate(builder, tmp_path) -> None: + """A measure that reaches a formula aggregating over its whole simulation + (here ``medicaid_cost`` -> ``medicaid_slcsp_state_denominator``) would sum + batch-local aggregates, so every engine of a multi-batch pass refuses it, + after scoring and before its release. One batch is the whole file.""" + plan = (("medicaid_cost", 2024, None),) + log = _EngineLog() + log.computes["medicaid_cost"] = ("medicaid_slcsp_state_denominator",) + scorer = _scorer(builder, _nested_frame(), log, 3, tmp_path / "batched") + with pytest.raises( + RuntimeError, + match=( + r"not batch-invariant: a batch engine computed " + r"medicaid_slcsp_state_denominator@2024 \(aggregates over" + ), + ): + scorer.open_consumer("fixture", plan) + assert len(log.constructions) == 1 + assert log.constructions[0].dataset is None + + whole = _scorer(builder, _nested_frame(), log, None, tmp_path / "whole") + assert whole.n_batches == 1 + consumer = whole.open_consumer("fixture", plan) + assert consumer.simulate(None).calculate("medicaid_cost", 2024).sum() > 0 + + # A value the written H5 stores is read, not recomputed per batch; the same + # variable computed for another period is refused. + stored = _EngineLog() + stored.stored_inputs["household_income_decile"] = (2024,) + scorer = _scorer(builder, _nested_frame(), stored, 3, tmp_path / "stored") + scorer.open_consumer("fixture", (("household_income_decile", 2024, None),)) + assert len(stored.constructions) == scorer.n_batches + with pytest.raises(RuntimeError, match="household_income_decile@2025"): + scorer.open_consumer("fixture", (("household_income_decile", 2025, None),)) + + +def test_known_periods_include_live_branches(builder) -> None: + def engine(known): + return SimpleNamespace( + get_holder=lambda name: SimpleNamespace( + get_known_periods=lambda: known.get(name, []) + ), + branches={}, + ) + + root = engine({"medicaid_slcsp_state_denominator": ["2024"]}) + branch = engine({"household_income_decile": ["2025"]}) + root.branches["mtr"] = branch + branch.branches["loop"] = root # a cycle is walked once + assert builder._post_export_known_periods( + root, builder.POST_EXPORT_POPULATION_AGGREGATE_VARIABLES + ) == { + ("medicaid_slcsp_state_denominator", "2024"), + ("household_income_decile", "2025"), + } + + +def test_a_reform_engine_refuses_a_baseline_branch_reader(builder, tmp_path) -> None: + """A batch reform engine is built from the reform's system alone, so it + has no baseline branch: a reform pass refuses a formula that reads one, + even in a single batch. A baseline pass has no baseline branch in a + whole-file simulation either and scores it, and a single-batch reform pass + may compute a population aggregate, which then spans the whole file.""" + log = _EngineLog() + log.computes["fixture_hours"] = ("relative_wage_change",) + log.computes["fixture_decile"] = ("household_income_decile",) + scorer = _scorer(builder, _nested_frame(), log, None, tmp_path) + consumer = scorer.open_consumer("fixture", (("fixture_hours", 2024, None),)) + assert consumer.simulate(None).calculate("fixture_hours", 2024).sum() > 0 + reformed = consumer.simulate(_FakeReform("fixture")) + assert reformed.calculate("fixture_decile", 2024).sum() > 0 + with pytest.raises( + RuntimeError, + match=( + r"relative_wage_change@2024 \(reads the engine's baseline branch, " + r"which a batch reform engine does not carry\)" + ), + ): + reformed.calculate("fixture_hours", 2024) + assert all(engine.dataset is None for engine in log.constructions) + assert builder._post_export_watched_variables(batched=False, reform=False) == () + assert set(builder._post_export_watched_variables(batched=True, reform=True)) == { + *builder.POST_EXPORT_POPULATION_AGGREGATE_VARIABLES, + *builder.POST_EXPORT_BASELINE_BRANCH_READERS, + } + + +def test_a_reform_that_moves_a_behavioral_parameter_is_refused( + builder, tmp_path +) -> None: + """Without a baseline branch policyengine-us scores labor-supply and + capital-gains responses as 0, which a whole-file reform engine matches + only while the reform leaves their parameters at baseline. A reform that + moves one is refused before any engine is built.""" + log = _EngineLog() + scorer = _scorer(builder, _nested_frame(), log, 3, tmp_path) + consumer = scorer.open_consumer("fixture", ()) + dynamic = _FakeReform( + "dynamic", + parameters={"gov.simulation.labor_supply_responses.elasticities.income": -0.05}, + ) + with pytest.raises( + RuntimeError, + match=( + r"off baseline at 2024: " + r"gov\.simulation\.labor_supply_responses\.elasticities\.income\. " + ), + ): + consumer.simulate(dynamic).calculate("income_tax", 2024) + assert log.constructions == [] + assert consumer.record()["reform_passes"] == 0 + # Setting a behavioral parameter to its baseline value moves nothing. + static = _FakeReform( + "static", + parameters={"gov.simulation.capital_gains_responses.elasticity": 0.0}, + ) + assert consumer.simulate(static).calculate("income_tax", 2024).sum() > 0 + assert len(log.constructions) == scorer.n_batches + + +_POPULATION_AGGREGATE_MARKER = re.compile( + r"\bsum_by_state\(|\bMicroSeries\(|\.decile_rank\(|\bquantile\(|\bpercentile\(" + r"|\bnp\.(?:sum|mean|median|average|nansum|nanmean|nanmedian|sort|argsort" + r"|cumsum)\((?![^\n]*\baxis\s*=)" +) +_WEIGHT_READ = re.compile(r"[\"'](\w+_weight)[\"']") +_BASELINE_BRANCH_MARKER = re.compile(r"\.baseline\b|get_branch\(\s*[\"']baseline[\"']") + + +def _is_variable_class(node: ast.AST) -> bool: + return isinstance(node, ast.ClassDef) and any( + getattr(base, "id", None) == "Variable" for base in node.bases + ) + + +def _engine_formula_modules() -> list[tuple[str, ast.Module]]: + """Every policyengine-us module that can define a formula, parsed but never + imported: ``variables`` and the contrib ``reforms``.""" + spec = importlib.util.find_spec("policyengine_us") + root = Path(next(iter(spec.submodule_search_locations))) + modules = [] + for package in ("variables", "reforms"): + for path in sorted((root / package).rglob("*.py")): + source = path.read_text() + modules.append((source, ast.parse(source))) + return modules + + +def _variables_reaching(modules, marker, *, weight_reads: bool = False) -> dict: + """Variable classes whose source matches ``marker``, directly or through + a helper function (in a module that defines no Variable) that does.""" + helpers: dict[str, str] = {} + for source, tree in modules: + if not any(_is_variable_class(node) for node in ast.walk(tree)): + helpers.update( + (node.name, ast.get_source_segment(source, node)) + for node in tree.body + if isinstance(node, ast.FunctionDef) + ) + reaching = {name for name, body in helpers.items() if marker.search(body)} + grew = True + while grew: + grew = False + for name, body in helpers.items(): + if name not in reaching and any( + re.search(rf"\b{other}\(", body) for other in reaching + ): + reaching.add(name) + grew = True + found: dict[str, list[str]] = {} + for source, tree in modules: + for node in ast.walk(tree): + if not _is_variable_class(node): + continue + body = ast.get_source_segment(source, node) + evidence = [match.group(0) for match in marker.finditer(body)] + evidence += [ + f"{helper}()" + for helper in sorted(reaching) + if re.search(rf"\b{helper}\(", body) + ] + if weight_reads and not node.name.endswith("_weight"): + # A weight variable's own formula broadcasts the household + # weight; any other formula reading a weight aggregates. + evidence += _WEIGHT_READ.findall(body) + if evidence: + found[node.name] = evidence + return found + + +@pytest.mark.requires_us +def test_watched_formula_lists_match_the_installed_engine(builder) -> None: + """Pin the watched lists to the installed engine's known source patterns. + + The scan matches weight reads, selected population-aggregate operations + and baseline-branch references in ``variables`` and contrib ``reforms``. + It follows matching helper names too. Changes detected by these patterns + require the lists to be revisited; this is not a general additivity proof. + """ + modules = _engine_formula_modules() + aggregates = _variables_reaching( + modules, _POPULATION_AGGREGATE_MARKER, weight_reads=True + ) + assert sorted(aggregates) == sorted( + builder.POST_EXPORT_POPULATION_AGGREGATE_VARIABLES + ), aggregates + readers = _variables_reaching(modules, _BASELINE_BRANCH_MARKER) + assert sorted(readers) == sorted( + { + *builder.POST_EXPORT_BASELINE_BRANCH_READERS, + *builder.POST_EXPORT_BEHAVIORAL_RESPONSE_PARAMETERS, + } + ), readers + # The two responses left to the parameter check read their own subtree. + sources = { + node.name: ast.get_source_segment(source, node) + for source, tree in modules + for node in ast.walk(tree) + if _is_variable_class(node) + } + for response, subtree in builder.POST_EXPORT_BEHAVIORAL_RESPONSE_PARAMETERS.items(): + assert f"parameters(period).{subtree}" in sources[response], response + + +# --------------------------------------------------------------------------- +# (h) engine-free recording +# --------------------------------------------------------------------------- + + +def test_recording_dry_runs_construct_no_engine( + builder, fake_reforms, monkeypatch +) -> None: + """(h) Learning the plans never constructs an engine, and the plans hold + exactly the baseline keys the consumers ask for.""" + + class Refused: + def __init__(self, *args, **kwargs): + raise AssertionError("a recording dry run constructed an engine") + + @staticmethod + def default_tax_benefit_system(*args, **kwargs): + raise AssertionError("a recording dry run built a reform system") + + monkeypatch.setitem( + sys.modules, "policyengine_us", SimpleNamespace(Microsimulation=Refused) + ) + args = SimpleNamespace( + skip_reform_coverage_smoke=False, + skip_reform_validation=False, + skip_out_of_sample_reforms=False, + skip_demographics=False, + ) + plan = builder._post_export_scoring_plan( + n_households=352_932, + maximum_microsim_batch_size=2_000, + consumers=builder._post_export_consumers( + args, result=_empty_calibration_result(), release_id="fixture" + ), + ) + record = plan.record() + assert record["n_batches"] == 177 + assert record["spm"] == US_RELEASE_SPM_SELECTION + assert list(record["consumers"]) == [ + "reform_coverage_smoke", + "reform_validation", + "demographics", + ] + smoke_keys = { + (probe.budget_measure, int(probe.period or builder.PERIOD), None) + for probe in us_release_reform_coverage_probes() + } + assert set(plan.baseline_plan("reform_coverage_smoke")) == smoke_keys + assert len(plan.baseline_plan("reform_coverage_smoke")) == len(smoke_keys) + assert plan.baseline_plan("demographics") == (("age", builder.PERIOD, None),) + validation = plan.baseline_plan("reform_validation") + assert ("income_tax", builder.PERIOD, None) in validation + assert ("state_code_str", builder.PERIOD, None) in validation + for consumer in record["consumers"].values(): + periods = [key["period"] for key in consumer["baseline_plan"]["keys"]] + assert periods == sorted(periods) + + # The skip flags drop their consumers from the plan. + skipped = builder._post_export_consumers( + SimpleNamespace( + skip_reform_coverage_smoke=True, + skip_reform_validation=False, + skip_out_of_sample_reforms=True, + skip_demographics=False, + ), + result=_empty_calibration_result(), + release_id="fixture", + ) + assert list(skipped) == ["demographics"] + + +def test_reform_validation_sweeps_only_released_engines( + builder, fake_reforms, monkeypatch +) -> None: + """The dry run's recorders hold no engine state, so reform_validation runs + no full collection for them; an engine simulation, which carries a + ``populations`` map, is still swept + after its reform row and after the shared baseline.""" + collections: list = [] + monkeypatch.setattr( + reform_validation_module, + "gc", + SimpleNamespace(collect=lambda *args: collections.append(args) or 0), + ) + plan = builder._record_post_export_baseline_plan( + builder._reform_validation_consumer( + result=_empty_calibration_result(), release_id="fixture" + ) + ) + assert plan and collections == [] + + class EngineLike: + def __init__(self) -> None: + self.populations: dict = {} + + def calculate(self, variable, period, map_to=None): + return _FakeSeries(np.ones(2), np.ones(2)) + + spec = ReformValidationSpec( + id="fixture_repeal", + name="fixture repeal", + category="fixture", + in_sample=False, + period=2024, + jct_score=None, + jct_window="", + jct_source="", + jct_source_url="", + neutralized_variable="fixture_credit", + ) + reform_validation_module.reform_validation_payload( + (spec,), period=2024, simulate=lambda reform: EngineLike() + ) + assert len(collections) == 2 + + +def test_an_unbuildable_plan_is_recorded_and_refused_not_raised( + builder, fake_reforms, monkeypatch, tmp_path +) -> None: + """microcosm#547: building the plan never raises before the diagnostics. + The error rides the plan record, becomes one terminal-batch line, and the + scorer refuses the plan if a run ever reached the export.""" + args = SimpleNamespace( + skip_reform_coverage_smoke=False, + skip_reform_validation=False, + skip_out_of_sample_reforms=False, + skip_demographics=False, + maximum_microsim_batch_size=3, + ) + # A calibration result without its compiled problem cannot feed the + # in-sample reform-validation rows. + failed = builder._record_post_export_scoring_plan( + args, + n_households=7, + result=SimpleNamespace(diagnostics=()), + release_id="fixture", + ) + assert failed.error is not None and failed.error.startswith("AttributeError") + assert failed.record() == { + "method": builder.POST_EXPORT_SCORING_METHOD, + "error": failed.error, + } + (line,) = failed.terminal_failures() + assert line.startswith("Post-export scoring plan could not be built") + assert failed.error in line + with pytest.raises(RuntimeError, match="not built before export: Attribute"): + failed.baseline_plan("demographics") + with pytest.raises(RuntimeError, match="not built before export: Attribute"): + builder._open_post_export_scorer(failed, _written_h5(tmp_path)) + + built = builder._record_post_export_scoring_plan( + args, + n_households=7, + result=_empty_calibration_result(), + release_id="fixture", + ) + assert built.error is None and built.terminal_failures() == [] + assert list(built.baseline_plans) == [ + "reform_coverage_smoke", + "reform_validation", + "demographics", + ] + empty = builder._PostExportScoringPlan( + n_households=7, maximum_batch_size=3, baseline_plans={} + ) + assert builder._open_post_export_scorer(empty, _written_h5(tmp_path)) is None + + +def test_calibration_diagnostics_carry_the_post_export_plan( + builder, monkeypatch, tmp_path +) -> None: + captured: dict[str, object] = {} + + def fake_write_calibration_diagnostics(result, path, *, target_registry, build): + captured["build"] = build + return path + + monkeypatch.setattr( + builder, "write_calibration_diagnostics", fake_write_calibration_diagnostics + ) + gate = SimpleNamespace(passed=True, failures=(), details={}) + plan = builder._PostExportScoringPlan( + n_households=7, + maximum_batch_size=3, + baseline_plans={ + "demographics": (("age", builder.PERIOD, None),), + }, + ) + common = dict( + result=SimpleNamespace(), + release_dir=tmp_path, + registry=TargetRegistry((), country="us"), + base_dataset_sha256="base-sha", + compilation={"dropped_target_names": []}, + target_profile_gate=gate, + health_input_gate=gate, + base_population_gate=gate, + support_value_repairs={}, + audit_export_targets=False, + gate_failures=[], + ) + builder._write_release_calibration_diagnostics( + **common, post_export_scoring=plan.record() + ) + assert captured["build"]["post_export_scoring"] == { + "method": "household_batched_written_h5", + "maximum_batch_size": 3, + "n_households": 7, + "n_batches": 3, + "spm": {"geography_kind": "county"}, + "period_order": "ascending", + "consumers": { + "demographics": { + "baseline_plan": { + "keys": [ + {"variable": "age", "period": builder.PERIOD, "map_to": None} + ], + "period_order": [builder.PERIOD], + } + } + }, + } + builder._write_release_calibration_diagnostics(**common) + assert "post_export_scoring" not in captured["build"] + + +# --------------------------------------------------------------------------- +# (i) wiring in _main +# --------------------------------------------------------------------------- + + +def _function_source(builder, name: str) -> tuple[str, ast.AST]: + source = inspect.getsource(getattr(builder, name)) + return source, ast.parse(source) + + +def _called_names(tree: ast.AST) -> list[str]: + names = [] + for node in ast.walk(tree): + if isinstance(node, ast.Call): + func = node.func + names.append( + func.attr + if isinstance(func, ast.Attribute) + else getattr(func, "id", None) + ) + return names + + +def test_main_and_writers_never_build_a_whole_pool_simulation(builder) -> None: + """(i) Neither _main nor the reform-validation and demographics writers + call default_simulate_factory or construct a Microsimulation, and _main + feeds the smoke from the batched scorer.""" + for name in ("_main", "_write_reform_validation", "_write_demographics"): + _, tree = _function_source(builder, name) + called = _called_names(tree) + assert "default_simulate_factory" not in called, name + assert "Microsimulation" not in called, name + assert "USSingleYearDataset" not in called, name + + source, tree = _function_source(builder, "_main") + calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call)] + + def calls_to(name: str) -> list[ast.Call]: + return [ + call + for call in calls + if name + in (getattr(call.func, "id", None), getattr(call.func, "attr", None)) + ] + + # The scorer opens on the written H5 with the plan recorded before export. + assert [ast.unparse(call) for call in calls_to("_open_post_export_scorer")] == [ + "_open_post_export_scorer(post_export_scoring_plan, dataset_path)" + ] + # It opens the file this run has just written, never a previous run's + # bytes at the same path (microcosm#443), and before the smoke scores. + (write_call,) = calls_to("write_dataset") + (open_call,) = calls_to("_open_post_export_scorer") + (smoke_gate_call,) = calls_to("us_reform_coverage_smoke_gate") + assert write_call.lineno < open_call.lineno < smoke_gate_call.lineno + # The smoke's served baseline is its own recorded plan, and the gate reads + # that consumer's seam. + assert [ast.unparse(call) for call in calls_to("open_consumer")] == [ + "post_export_scorer.open_consumer('reform_coverage_smoke', " + "post_export_scoring_plan.baseline_plan('reform_coverage_smoke'))" + ] + (smoke_call,) = calls_to("us_reform_coverage_smoke_gate") + simulate = next(kw for kw in smoke_call.keywords if kw.arg == "simulate") + assert ast.unparse(simulate.value) == "smoke_scoring.simulate" + # Each writer shares the scorer and is served its own consumer's plan. + for writer, consumer in ( + ("_write_reform_validation", "reform_validation"), + ("_write_demographics", "demographics"), + ): + (call,) = calls_to(writer) + keywords = {kw.arg: kw.value for kw in call.keywords} + assert ast.unparse(keywords["post_export_scorer"]) == "post_export_scorer" + plan_calls = [ + ast.unparse(node) + for node in ast.walk(keywords["baseline_plan"]) + if isinstance(node, ast.Call) + ] + assert plan_calls == [f"post_export_scoring_plan.baseline_plan({consumer!r})"] + (validation_call,) = calls_to("_write_reform_validation") + assert {kw.arg: ast.unparse(kw.value) for kw in validation_call.keywords}[ + "simulate_out_of_sample" + ] == "not args.skip_out_of_sample_reforms" + # The manifest binds the sha of the bytes the closed scorer scored, and + # nothing else assigns it. + sha_assignments = [ + node + for node in ast.walk(tree) + if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)) + and "scored_dataset_sha256" + in { + target.id + for target in ( + node.targets if isinstance(node, ast.Assign) else [node.target] + ) + if isinstance(target, ast.Name) + } + ] + assert [ast.unparse(node.value) for node in sha_assignments] == [ + "_close_post_export_scorer(post_export_scorer)" + ] + (demographics_call,) = calls_to("_write_demographics") + (manifest_call,) = calls_to("_build_manifests") + assert demographics_call.lineno < sha_assignments[0].lineno < manifest_call.lineno + diagnostics_call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and getattr(node.func, "id", None) == "_write_release_calibration_diagnostics" + ) + assert {kw.arg: ast.unparse(kw.value) for kw in diagnostics_call.keywords}[ + "post_export_scoring" + ] == "post_export_scoring_plan.record()" + assert '"post_export_scoring": smoke_scoring.record()' in source + # The plan is recorded (never raised) before the diagnostics, and its own + # failure joins the terminal batch ahead of the export write. + recorded = source.index( + "post_export_scoring_plan = _record_post_export_scoring_plan(" + ) + joined = source.index( + "terminal_gate_failures.extend(post_export_scoring_plan.terminal_failures())" + ) + assert recorded < source.index("_write_release_calibration_diagnostics(") < joined + assert joined < source.index("release_engine.write_dataset(") + + +def _plan_args(**overrides) -> SimpleNamespace: + """The ``_main`` arguments the post-export plan reads.""" + return SimpleNamespace( + **{ + "skip_reform_coverage_smoke": False, + "skip_reform_validation": False, + "skip_out_of_sample_reforms": False, + "skip_demographics": False, + "maximum_microsim_batch_size": 3, + **overrides, + } + ) + + +@pytest.fixture +def engine_free_loader(builder, monkeypatch): + """Serve the default scorer construction without an engine or an H5. + + ``policyengine_us.Microsimulation`` is the fake engine, the written H5 + loads as the nested frame, and each batch frame is its own dataset. Yields + the fake engine's log and the load calls. + """ + log = _EngineLog() + loads: list = [] + + def load_frame(path, *, expected_sha256=None): + loads.append((Path(path), expected_sha256)) + return _nested_frame() + + monkeypatch.setitem( + sys.modules, + "policyengine_us", + SimpleNamespace(Microsimulation=_fake_engine(log)), + ) + monkeypatch.setattr(builder, "_load_frame", load_frame) + monkeypatch.setattr(builder, "_dataset_from_frame", lambda frame, **_: frame) + return log, loads + + +def test_main_opens_the_scorer_its_plan_recorded( + builder, fake_reforms, engine_free_loader, tmp_path +) -> None: + """(i) The plan ``_main`` records before export opens the scorer + ``_main`` scores with: the written H5, bound to its sha, split by the + plan's batch size, and refused if its household count is not the plan's. + Closing it returns the scored sha the manifest binds.""" + log, loads = engine_free_loader + plan = builder._record_post_export_scoring_plan( + _plan_args(), + n_households=7, + result=_empty_calibration_result(), + release_id="fixture-release", + ) + assert plan.error is None and plan.n_batches == 3 + path = _written_h5(tmp_path) + scorer = builder._open_post_export_scorer(plan, path) + assert loads == [(path, builder._sha256(path))] + assert scorer.dataset_path == path + assert scorer.n_batches == 3 and scorer.max_batch_households == 3 + # Served through the engine's own Microsimulation, one batch at a time. + scorer.open_consumer("demographics", plan.baseline_plan("demographics")) + assert [len(engine.household_ids) for engine in log.constructions] == [3, 3, 1] + assert builder._close_post_export_scorer(scorer) == builder._sha256(path) + with pytest.raises(RuntimeError, match="scorer is closed"): + scorer.open_consumer("demographics", plan.baseline_plan("demographics")) + assert builder._close_post_export_scorer(None) is None + + other = builder._record_post_export_scoring_plan( + _plan_args(), + n_households=8, + result=_empty_calibration_result(), + release_id="fixture-release", + ) + with pytest.raises(ValueError, match="post-export scoring plan was built for 8"): + builder._open_post_export_scorer(other, path) + + +def test_writers_score_through_the_shared_scorer( + builder, fake_reforms, monkeypatch, tmp_path +) -> None: + """(h) The writers ``_main`` calls write what the unbatched path computes, + with ``out_of_sample_simulated`` true (the publisher refuses false), and + score through the scorer ``_main`` hands them: opening a second one fails + this test, and demographics costs exactly one baseline pass. The + calibration's in-sample JCT estimates and targets reach their rows.""" + frame = _nested_frame() + log = _EngineLog() + scorer = _scorer(builder, frame, log, 3, tmp_path / "h5") + fit = _calibration_result_with_in_sample_fit(builder.PERIOD) + plan = builder._record_post_export_scoring_plan( + _plan_args(), + n_households=7, + result=fit, + release_id="fixture-release", + ) + + def second_scorer(*args, **kwargs): + raise AssertionError("a post-export writer opened a second scorer") + + monkeypatch.setattr(builder, "_HouseholdBatchedPostExportScorer", second_scorer) + monkeypatch.setattr( + builder, "geography_coverage_payload", lambda path: {"h5": Path(path).name} + ) + release_dir = tmp_path / "release" + release_dir.mkdir() + + builder._write_reform_validation( + release_dir=release_dir, + dataset_path=scorer.dataset_path, + result=fit, + registry=TargetRegistry((), country="us"), + release_id="fixture-release", + simulate_out_of_sample=True, + post_export_scorer=scorer, + baseline_plan=plan.baseline_plan("reform_validation"), + ) + written = json.loads((release_dir / "reform_validation.json").read_text()) + assert written["out_of_sample_simulated"] is True + # The reference calls the payload directly, with the fit spelled out, + # rather than through the consumer the writer uses. + fitted = reform_validation_module.in_sample_reform_specs(period=builder.PERIOD)[:2] + reference = reform_validation_module.reform_validation_payload( + reform_validation_module.load_default_reform_specs(period=builder.PERIOD), + period=builder.PERIOD, + simulate=_unbatched_simulate(frame, _EngineLog()), + in_sample_estimates={ + fitted[0].id: _IN_SAMPLE_FIT[0][0], + fitted[1].id: _IN_SAMPLE_FIT[1][0], + "fixture_other_target": 7.0, + }, + in_sample_targets={ + fitted[0].id: _IN_SAMPLE_FIT[0][1], + fitted[1].id: _IN_SAMPLE_FIT[1][1], + "fixture_other_target": 8.0, + }, + baseline_levels=reform_validation_module.default_baseline_level_specs(), + release_id="fixture-release", + ) + assert written == json.loads(json.dumps(reference, allow_nan=False)) + rows = {row["id"]: row for row in written["reforms"]} + for spec, (estimate, target) in zip(fitted, _IN_SAMPLE_FIT, strict=True): + row = rows[spec.id] + assert row["in_sample"] is True + assert row["microcosm"]["budget_effect"] == estimate + assert row["microcosm"]["baseline_total"] is None + assert row["jct"]["score"] == target + # The in-sample reforms the fit does not cover are simulated. + unfitted = reform_validation_module.in_sample_reform_specs(period=builder.PERIOD)[2] + assert rows[unfitted.id]["microcosm"]["baseline_total"] is not None + validation_constructions = len(log.constructions) + assert validation_constructions > scorer.n_batches + + builder._write_demographics( + release_dir=release_dir, + dataset_path=scorer.dataset_path, + release_id="fixture-release", + post_export_scorer=scorer, + baseline_plan=plan.baseline_plan("demographics"), + ) + demographics_engines = log.constructions[validation_constructions:] + assert len(demographics_engines) == scorer.n_batches + assert all(engine.tax_benefit_system is None for engine in demographics_engines) + ages, weights = builder._demographics_consumer( + _unbatched_simulate(frame, _EngineLog()) + ) + expected = demographics_payload( + ages, weights, period=builder.PERIOD, release_id="fixture-release" + ) + expected["geography_coverage"] = {"h5": "populace_us_2024.h5"} + written = json.loads((release_dir / "demographics.json").read_text()) + assert written == json.loads(json.dumps(expected, allow_nan=False)) + + +def test_a_direct_writer_call_opens_and_closes_its_own_scorer( + builder, fake_reforms, engine_free_loader, monkeypatch, tmp_path +) -> None: + """A writer called without ``_main``'s scorer records its plan by the + same dry run and scores through a scorer of its own, which it closes.""" + log, loads = engine_free_loader + closed = [] + scorer_cls = builder._HouseholdBatchedPostExportScorer + close = scorer_cls.close + monkeypatch.setattr( + scorer_cls, "close", lambda self: closed.append(self) or close(self) + ) + monkeypatch.setattr(builder, "geography_coverage_payload", lambda path: {}) + path = _written_h5(tmp_path / "h5") + release_dir = tmp_path / "release" + release_dir.mkdir() + builder._write_demographics( + release_dir=release_dir, + dataset_path=path, + release_id="fixture-release", + maximum_microsim_batch_size=3, + ) + assert loads == [(path, builder._sha256(path))] + assert [len(engine.household_ids) for engine in log.constructions] == [3, 3, 1] + assert len(closed) == 1 and closed[0].dataset_path == path + written = json.loads((release_dir / "demographics.json").read_text()) + assert written["period"] == builder.PERIOD + assert written["total_population"] > 0 + + +def test_main_frees_the_target_frame_before_the_export(builder) -> None: + """``del target_frame`` follows its last reader and precedes the export + write, and nothing reads the name afterwards.""" + _, tree = _function_source(builder, "_main") + deletes = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Delete) + and any( + isinstance(target, ast.Name) and target.id == "target_frame" + for target in node.targets + ) + ] + assert len(deletes) == 1 + deleted_at = deletes[0].lineno + writes = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "write_dataset" + ] + assert writes and deleted_at < min(writes) + reads = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Name) + and node.id == "target_frame" + and isinstance(node.ctx, ast.Load) + ] + assert reads and max(reads) < deleted_at + + +def _small_calibration_problem(n: int = 40): + """A one-person-per-household frame and one household-total target.""" + rng = np.random.default_rng(0) + person = pd.DataFrame( + { + "person_id": np.arange(n), + "person_household_id": np.arange(n), + "income": rng.uniform(0.0, 100.0, n), + } + ) + household = pd.DataFrame( + {"household_id": np.arange(n), "x": rng.uniform(0.0, 10.0, n)} + ) + frame = Frame( + {"person": person, "household": household}, + EntitySchema(group_entities=("household",)), + {"household": Weights(np.ones(n), WeightKind.DESIGN)}, + ) + target = TargetSpec( + name="x_total@2024", + entity="household", + measure="x", + value=float(household["x"].sum() * 1.1), + source="fixture", + family="fixture", + ) + return frame, TargetRegistry([target], country="us").to_target_set() + + +def test_calibration_result_holds_no_reference_to_the_target_frame() -> None: + """The premise of ``del target_frame``: ``calibrate`` keeps no reference + to its input frame, and ``result.frame`` shares no memory with it.""" + frame, targets = _small_calibration_problem() + result = calibrate(frame, targets, epochs=3, mass="conserve", seed=0) + gc.collect() + # Nothing the calibration returned (or cached) refers to the input frame. + referrers = [ + type(referrer).__name__ + for referrer in gc.get_referrers(frame) + if not inspect.isframe(referrer) + ] + assert referrers == [] + assert result.frame is not frame + for entity in ("person", "household"): + for column in frame.table(entity).columns: + assert not np.shares_memory( + frame.table(entity)[column].to_numpy(), + result.frame.table(entity)[column].to_numpy(), + ) + + +def test_calibration_results_drop_their_frames_and_keep_everything_else( + builder, +) -> None: + """Dropping target tables preserves diagnostics, estimates and weights.""" + from microcosm.build.us_runtime.exact_k_ladder import ( + ExactKLadderCalibration, + assert_exact_k_realized_count, + ) + from microcosm.calibrate import calibrate_l0_refit + from microcosm.calibrate.diagnostics import diagnostics_payload + + frame, targets = _small_calibration_problem() + dense = calibrate(frame, targets, epochs=3, mass="conserve", seed=0) + dropped = builder._without_calibrated_frames(dense) + assert dropped.frame is None + assert dropped.weights is dense.weights and dropped.problem is dense.problem + assert diagnostics_payload(dropped) == diagnostics_payload(dense) + assert builder._in_sample_estimates(dropped) == builder._in_sample_estimates(dense) + # Once the full result goes, nothing but this test holds its frame. + held = dense.frame + del dense + gc.collect() + assert [ + type(referrer).__name__ + for referrer in gc.get_referrers(held) + if not inspect.isframe(referrer) + ] == [] + del held + + refit = calibrate_l0_refit( + frame, targets, epochs=5, refit_epochs=5, target_records=20, seed=0 + ) + dropped = builder._without_calibrated_frames(refit) + assert dropped.selection.frame is None and dropped.refit.frame is None + assert diagnostics_payload(dropped) == diagnostics_payload(refit) + np.testing.assert_array_equal( + dropped.selected_entity_ids, refit.selected_entity_ids + ) + exported = builder._with_l0_refit_weights(frame, refit) + np.testing.assert_array_equal( + builder._with_l0_refit_weights(frame, dropped).weights_for("household").values, + exported.weights_for("household").values, + ) + kept = builder._without_calibrated_frames(refit, export_frame=exported) + assert kept.selection.frame is None and kept.refit.frame is exported + assert kept.refit.frame is not refit.refit.frame + k = int(refit.refit.frame.n("household")) + outcome = ExactKLadderCalibration( + result=kept, + support=np.arange(k), + selected_inclusion_probabilities=np.ones(k), + selection_receipt={}, + refit_baseline_diagnostics={}, + ) + assert assert_exact_k_realized_count(outcome, k) == k + # A test double carries no calibrated frame and passes through. + double = _empty_calibration_result() + assert builder._without_calibrated_frames(double) is double + + +def test_main_drops_the_calibration_frames_before_the_export(builder) -> None: + """Cleanup follows export construction and precedes the H5 write.""" + _, tree = _function_source(builder, "_main") + (deleted,) = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Delete) + and any( + getattr(target, "id", None) == "target_frame" for target in node.targets + ) + ] + frame_reads = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) + and node.attr == "frame" + and getattr(node.value, "id", None) == "result" + ] + later_bindings = sorted( + (node.lineno, ast.unparse(node.value)) + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any(getattr(target, "id", None) == "result" for target in node.targets) + and node.lineno > deleted + ) + (write_call,) = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and getattr(node.func, "attr", None) == "write_dataset" + ] + assert frame_reads and max(frame_reads) < deleted + assert [value for _, value in later_bindings] == [ + "_without_calibrated_frames(result, " + "export_frame=export_frame if ladder_outcome is not None else None)", + ] + assert all(deleted < lineno < write_call for lineno, _ in later_bindings) + export_bindings = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any( + getattr(target, "id", None) == "export_frame" for target in node.targets + ) + ] + assert export_bindings and max(export_bindings) < later_bindings[0][0] + assert "dataclasses.replace(ladder_outcome, result=result)" in ast.unparse(tree) + + +@pytest.mark.parametrize("full_pool", [True, False]) +@pytest.mark.parametrize("exact_k", [True, False]) +def test_main_has_no_calibrated_target_frame_at_the_export_write( + builder, full_pool, exact_k +) -> None: + """Execute the export/cleanup statements with real calibration results.""" + from microcosm.calibrate import calibrate_l0_refit + + target_frame, targets = _small_calibration_problem() + base_frame = Frame( + { + "person": target_frame.table("person"), + "household": target_frame.table("household").drop(columns="x"), + }, + target_frame.schema, + {"household": target_frame.weights_for("household")}, + ) + if full_pool: + result = calibrate(target_frame, targets, epochs=3, mass="conserve", seed=0) + else: + result = calibrate_l0_refit( + target_frame, targets, epochs=5, refit_epochs=5, target_records=20, seed=0 + ) + ladder_outcome = ( + builder.ExactKLadderCalibration( + result=result, + support=np.arange(len(result.weights)), + selected_inclusion_probabilities=np.ones(len(result.weights)), + selection_receipt={}, + refit_baseline_diagnostics={}, + ) + if exact_k + else None + ) + _, tree = _function_source(builder, "_main") + statements = tree.body[0].body + export_start = next( + index + for index, node in enumerate(statements) + if isinstance(node, ast.If) + and ast.unparse(node.test) == "full_pool_calibration" + and any( + isinstance(child, ast.Assign) + and any( + getattr(target, "id", None) == "export_frame" + for target in child.targets + ) + for child in ast.walk(node) + ) + ) + cleanup_end = next( + index + for index, node in enumerate(statements) + if index > export_start + and isinstance(node, ast.If) + and ast.unparse(node.test) + == "isinstance(ladder_outcome, ExactKLadderCalibration)" + ) + (write,) = [ + node + for node in statements + if isinstance(node, ast.Expr) + and isinstance(node.value, ast.Call) + and getattr(node.value.func, "attr", None) == "write_dataset" + ] + writes = [] + + def check_write(frame, path, *, period): + writes.append(frame) + assert "x" not in frame.table("household") + cleaned = namespace["result"] + if not full_pool: + assert cleaned.selection.frame is None + assert cleaned.frame is (frame if exact_k else None) + if exact_k: + outcome = namespace["ladder_outcome"] + assert outcome.result is cleaned + assert builder.assert_exact_k_realized_count( + outcome, len(result.weights) + ) == len(result.weights) + + namespace = { + **vars(builder), + "base_frame": base_frame, + "result": result, + "ladder_outcome": ladder_outcome, + "full_pool_calibration": full_pool, + "release_engine": SimpleNamespace(write_dataset=check_write), + "dataset_path": Path("fixture.h5"), + } + module = ast.Module( + body=[*statements[export_start : cleanup_end + 1], write], type_ignores=[] + ) + exec(compile(module, "
", "exec"), namespace) + assert len(writes) == 1 + + +# --------------------------------------------------------------------------- +# (j) the scored bytes are the manifest's bytes +# --------------------------------------------------------------------------- + + +def test_scored_sha_is_bound_at_load_and_checked_by_the_manifest( + builder, tmp_path +) -> None: + """(j) The scorer loads the written H5 against the sha it records, and + _build_manifests refuses a dataset whose sha differs from the scored one.""" + loads: list = [] + scorer = _scorer(builder, _nested_frame(), _EngineLog(), 3, tmp_path, loads=loads) + path = tmp_path / "populace_us_2024.h5" + assert scorer.dataset_sha256 == builder._sha256(path) + assert loads == [(path, scorer.dataset_sha256)] + record = scorer.open_consumer("fixture", ()).record() + assert record["dataset_sha256"] == builder._sha256(path) + assert record["baseline_passes"] == 0 + + # The manifest check runs before anything else reads the arguments. + required = { + name: None + for name, parameter in inspect.signature( + builder._build_manifests + ).parameters.items() + if parameter.default is inspect.Parameter.empty + } + required.update( + release_dir=tmp_path, + artifact_root=tmp_path, + ) + with pytest.raises(RuntimeError, match="never scored"): + builder._build_manifests( + **required, + dataset_filename=path.name, + scored_dataset_sha256="0" * 64, + ) + + _, tree = _function_source(builder, "_main") + manifest_call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and getattr(node.func, "id", None) == "_build_manifests" + ) + keywords = {kw.arg: ast.unparse(kw.value) for kw in manifest_call.keywords} + assert keywords["scored_dataset_sha256"] == "scored_dataset_sha256" + + +# --------------------------------------------------------------------------- +# (k) the nesting premise +# --------------------------------------------------------------------------- + + +def test_scorer_refuses_a_unit_that_spans_households(builder, tmp_path) -> None: + """(k) Frame.select does not check that group units nest in households, + so the scorer does, before any engine exists.""" + frame = _nested_frame() + tables = {entity: frame.table(entity).copy() for entity in frame.entities} + person = tables["person"] + # Move the second member of household 1 into household 2's tax unit. + person.loc[person["person_id"] == 2, "person_tax_unit_id"] = 20 + tables["tax_unit"] = pd.DataFrame( + {"tax_unit_id": np.unique(person["person_tax_unit_id"].to_numpy())} + ) + spanning = Frame(tables, US_SCHEMA, {"household": frame.weights_for("household")}) + log = _EngineLog() + with pytest.raises(ValueError, match="nested in households"): + _scorer(builder, spanning, log, 1, tmp_path) + assert log.constructions == [] + + # The partition check is a second line: a batch set that misses or + # repeats a household is refused. + batches = [ + builder._select_households_by_position(frame, np.asarray([0, 1, 2])), + builder._select_households_by_position(frame, np.asarray([2, 3, 4, 5, 6])), + ] + with pytest.raises(ValueError, match="do not partition the pool"): + builder._assert_post_export_batching_premises(frame, batches) + + # Batches that partition the pool must also keep the written household + # order and household weights, or the concatenated arrays would misalign. + batches = [ + builder._select_households_by_position(frame, np.asarray(positions)) + for positions in ([0, 1, 2], [3, 4, 5, 6]) + ] + builder._assert_post_export_batching_premises(frame, batches) + with pytest.raises(ValueError, match="household order"): + builder._assert_post_export_batching_premises(frame, batches[::-1]) + last = batches[1] + reweighted = Frame( + {entity: last.table(entity) for entity in last.entities}, + US_SCHEMA, + { + "household": Weights( + last.weights_for("household").values * 2.0, WeightKind.CALIBRATED + ) + }, + ) + with pytest.raises(ValueError, match="household weights"): + builder._assert_post_export_batching_premises(frame, [batches[0], reweighted]) + + +def test_scorer_refuses_a_household_count_other_than_planned(builder, tmp_path) -> None: + with pytest.raises(ValueError, match="built for 8"): + builder._HouseholdBatchedPostExportScorer( + _written_h5(tmp_path), + maximum_microsim_batch_size=3, + expected_n_households=8, + microsimulation_cls=_fake_engine(_EngineLog()), + dataset_from_frame=lambda batch_frame: batch_frame, + load_frame=lambda path, *, expected_sha256: _nested_frame(), + ) + + +# --------------------------------------------------------------------------- +# (l) the same invariance against policyengine-us +# --------------------------------------------------------------------------- + +# (state_fips, county_fips, [(age, pre-LSR wages, weekly hours), ...]). The +# first two households put North Dakota and Maryland in one batch at size 2. +# Test (l) probes the MD CCS request-order premise on that batch and warns if +# the engine no longer raises the expected error. +_ENGINE_HOUSEHOLDS = ( + (38, "38105", ((32, 47_000.0, 40.0),)), + ( + 24, + "24043", + ( + (38, 50_000.0, 20.0), + (40, 75_000.0, 40.0), + (11, 0.0, 0.0), + (7, 0.0, 0.0), + (5, 0.0, 0.0), + (2, 0.0, 0.0), + ), + ), + (6, "06037", ((58, 150_000.0, 40.0), (55, 0.0, 0.0))), + (24, "24031", ((29, 31_000.0, 30.0), (2, 0.0, 0.0))), +) + + +#: Four low-income households with children or a senior, in California, New +#: York and Texas. The guard test checks that the combined fixture has +#: Medicaid enrollees and reaches a watched SLCSP state aggregate. +_MEDICAID_HOUSEHOLDS = ( + (6, "06001", ((26, 6_000.0, 10.0), (4, 0.0, 0.0), (1, 0.0, 0.0))), + (36, "36047", ((67, 0.0, 0.0),)), + ( + 48, + "48201", + ((45, 12_000.0, 25.0), (44, 0.0, 0.0), (15, 0.0, 0.0), (9, 0.0, 0.0)), + ), + (6, "06073", ((31, 18_000.0, 30.0), (3, 0.0, 0.0))), +) + + +def _write_engine_h5(builder, directory: Path, households) -> Path: + """Write ``households`` as a release H5 through the engine's own writer. + + One tax unit, SPM unit and family per household; household ``i`` weighs + ``1,000 * i``. + """ + from microcosm.frame.adapters.policyengine_us import PolicyEngineUSEngine + + rows = [] + person_id = 0 + for household_id, (_, _, members) in enumerate(households, start=1): + for age, wages, hours in members: + person_id += 1 + rows.append( + { + "person_id": person_id, + "person_household_id": household_id, + "person_tax_unit_id": household_id * 10, + "person_spm_unit_id": household_id * 100, + "person_family_id": household_id * 1000, + "person_marital_unit_id": person_id * 10_000, + "age": age, + "employment_income_before_lsr": wages, + "weekly_hours_worked_before_lsr": hours, + } + ) + person = pd.DataFrame(rows) + household_ids = np.arange(1, len(households) + 1, dtype=np.int64) + frame = Frame( + { + "person": person, + "household": pd.DataFrame( + { + "household_id": household_ids, + "state_fips": [state for state, _, _ in households], + "county_fips": [county for _, county, _ in households], + } + ), + "tax_unit": pd.DataFrame({"tax_unit_id": household_ids * 10}), + "spm_unit": pd.DataFrame({"spm_unit_id": household_ids * 100}), + "family": pd.DataFrame({"family_id": household_ids * 1000}), + "marital_unit": pd.DataFrame( + {"marital_unit_id": person["person_marital_unit_id"].to_numpy()} + ), + }, + US_SCHEMA, + { + "household": Weights( + household_ids.astype(np.float64) * 1_000.0, WeightKind.CALIBRATED + ) + }, + ) + directory.mkdir(parents=True, exist_ok=True) + path = directory / "populace_us_2024.h5" + PolicyEngineUSEngine().write_dataset(frame, path, period=builder.PERIOD) + return path + + +def _engine_probe(probe_id, budget_measure, period, **reform) -> ReformCoverageProbe: + return ReformCoverageProbe( + id=probe_id, + name=probe_id, + parameter_changes=reform.get("parameter_changes", {}), + neutralized_variable=reform.get("neutralized_variable"), + budget_measure=budget_measure, + period=period, + binding_inputs=("employment_income_before_lsr",), + min_abs_effect=1.0, + effect_direction="baseline_minus_reform", + expected_sign="either", + reason="fixture", + issue="PolicyEngine/microcosm#956", + ) + + +@pytest.mark.slow +@pytest.mark.requires_us +def test_batched_scoring_matches_whole_pool_policyengine_us(builder, tmp_path) -> None: + """(l) On a small written H5 with Maryland households, three probes and + two validation reforms score the same totals on batch engines (each given + the reform's system alone) as on a whole-file engine given ``reform=``, a + whole-file reform engine scores no behavioral response at the engine's + defaults, and the 75-key reform-validation baseline plan runs in ascending + period order. The request-order MD CCS premise is observed, with a warning + if the engine no longer raises the expected error.""" + from policyengine_us import Microsimulation + + path = _write_engine_h5(builder, tmp_path, _ENGINE_HOUSEHOLDS) + + probes = ( + _engine_probe( + "wages_income_tax_2024", + "income_tax", + 2024, + neutralized_variable="employment_income_before_lsr", + ), + _engine_probe( + "wages_state_income_tax_2025", + "state_income_tax", + 2025, + neutralized_variable="employment_income_before_lsr", + ), + _engine_probe( + "zero_single_standard_deduction_2026", + "income_tax", + 2026, + parameter_changes={ + "gov.irs.deductions.standard.amount.SINGLE": {"2026-01-01": 0} + }, + ), + ) + whole = builder.us_reform_coverage_smoke_gate( + simulate=reform_validation_module.default_simulate_factory(path), + probes=probes, + period=builder.PERIOD, + ) + scorer = builder._HouseholdBatchedPostExportScorer( + path, maximum_microsim_batch_size=2 + ) + assert scorer.n_batches == 2 + + def smoke(simulate): + return builder.us_reform_coverage_smoke_gate( + simulate=simulate, probes=probes, period=builder.PERIOD + ) + + plan = builder._record_post_export_baseline_plan(smoke) + assert [key[1] for key in plan] == [2024, 2025, 2026] + scoring = scorer.open_consumer("reform_coverage_smoke", plan) + batched = smoke(scoring.simulate) + for probe in probes: + expected = whole.details["results"][probe.id] + observed = batched.details["results"][probe.id] + for field in ("baseline_total", "reform_total", "effect"): + assert observed[field] == pytest.approx( + expected[field], rel=1e-12, abs=1e-6 + ), (probe.id, field) + assert expected["effect"] != 0.0, probe.id + assert scoring.record()["reform_systems"] == len(probes) + + # Validation reforms, a credit repeal and a structural contrib reform: the + # batch engines (the reform's system alone, no baseline branch) score what + # a whole-file engine given ``reform=`` scores. + specs = { + spec.id: spec + for spec in reform_validation_module.load_default_reform_specs( + period=builder.PERIOD + ) + } + for spec_id in ("state_repeal_md_eitc", "federal.ubi_mechanical"): + spec = specs[spec_id] + reform = spec.build_reform() + whole_engine = reform_validation_module.default_simulate_factory(path)(reform) + expected = float(whole_engine.calculate(spec.budget_measure, spec.period).sum()) + builder.release_engine_simulation(whole_engine) + del whole_engine + observed = ( + scoring.simulate(reform).calculate(spec.budget_measure, spec.period).sum() + ) + assert expected != 0.0, spec_id + assert observed == pytest.approx(expected, rel=1e-12, abs=1e-6), spec_id + + # The premise of dropping the baseline branch: at the installed engine's + # defaults, a whole-file reform engine (baseline branch present) scores no + # labor-supply or capital-gains response, as the batch engines do. + dynamic_probe = probes[2] + whole_engine = reform_validation_module.default_simulate_factory(path)( + smoke_module._build_reform(dynamic_probe) + ) + assert whole_engine.baseline is not None + for response in builder.POST_EXPORT_BEHAVIORAL_RESPONSE_PARAMETERS: + values = whole_engine.calculate(response, dynamic_probe.period) + assert float(np.abs(np.asarray(values)).sum()) == 0.0, response + builder.release_engine_simulation(whole_engine) + del whole_engine + + validation_plan = builder._record_post_export_baseline_plan( + builder._reform_validation_consumer( + result=_empty_calibration_result(), release_id="fixture" + ) + ) + periods = [key[1] for key in validation_plan] + assert len(validation_plan) == 75 + assert periods == sorted(periods) and max(periods) > 2025 + # Probe the ascending-order premise on the first batch by requesting a + # 2025 key after 2024 and 2027 keys. The warning permits an upstream fix + # without requiring an unrelated lock bump to preserve the error. + (first_batch, _) = scorer._batches() + engine = Microsimulation( + dataset=builder._dataset_from_frame( + first_batch, assert_no_formula_owned_columns=False + ), + spm=dict(US_RELEASE_SPM_SELECTION), + ) + engine.calculate("income_tax", 2024) + engine.calculate("md_income_tax", 2027) + try: + engine.calculate("nd_income_tax", 2025) + except Exception as error: # noqa: BLE001 - the engine's own error type + assert "md.msde.ccs.payment.informal.rates" in str(error), error + else: + warnings.warn( + "policyengine-us no longer raises the MD CCS ParameterNotFoundError " + "for request-order periods; revisit the ascending-period guard " + "(microcosm#956).", + stacklevel=1, + ) + finally: + builder.release_engine_simulation(engine) + del engine + # The same batch scores every reform-validation baseline key in ascending + # period order without error. + validation = scorer.open_consumer("reform_validation", validation_plan) + income_tax = validation.simulate(None).calculate("income_tax", builder.PERIOD) + assert income_tax.sum() == pytest.approx( + whole.details["results"]["wages_income_tax_2024"]["baseline_total"], + rel=1e-12, + ) + scorer.close() + + +# --------------------------------------------------------------------------- +# The guards and the shipped request surface against policyengine-us +# --------------------------------------------------------------------------- + +#: The (l) households (two in Maryland) and the four Medicaid ones. +_GUARD_HOUSEHOLDS = (*_ENGINE_HOUSEHOLDS, *_MEDICAID_HOUSEHOLDS) + + +@pytest.mark.slow +@pytest.mark.requires_us +def test_additivity_guards_fire_on_the_installed_engine(builder, tmp_path) -> None: + """The three additivity guards refuse on policyengine-us itself. + + The fake-engine tests above pin the guards' logic; these read what a real + engine computed (its holders' known periods) and what a real reform + system sets (its parameter leaves), covering API behavior that the fake + engine does not establish. + ``medicaid_cost`` reaches the Medicaid SLCSP state sums: two batches + refuse it, one batch scores the whole-file total, and a reform engine + refuses it even in one batch, because the state denominator reads the + baseline branch. A reform that moves the behavioral-response elasticities + is refused by name before a batch engine for that reform is built.""" + from policyengine_core.reforms import Reform + from policyengine_us import Microsimulation + + path = _write_engine_h5(builder, tmp_path, _GUARD_HOUSEHOLDS) + key = ("medicaid_cost", 2024, None) + whole = reform_validation_module.default_simulate_factory(path)(None) + expected = float(whole.calculate("medicaid_cost", 2024).sum()) + enrolled = int(np.asarray(whole.calculate("medicaid_enrolled", 2024)).sum()) + builder.release_engine_simulation(whole) + del whole + assert expected > 0.0 and enrolled > 0 + + batched = builder._HouseholdBatchedPostExportScorer( + path, maximum_microsim_batch_size=4 + ) + assert batched.n_batches == 2 + with pytest.raises( + RuntimeError, + match=( + r"not batch-invariant: a batch engine computed .*" + r"medicaid_slcsp_state_denominator@2024 \(aggregates over" + ), + ): + batched.open_consumer("guard", (key,)) + batched.close() + + scorer = builder._HouseholdBatchedPostExportScorer( + path, maximum_microsim_batch_size=None + ) + assert scorer.n_batches == 1 + consumer = scorer.open_consumer("guard", (key,)) + scored = consumer.simulate(None).calculate("medicaid_cost", 2024).sum() + assert scored == pytest.approx(expected, rel=1e-12) + + # A reform that leaves the behavioral parameters alone passes that check + # and reaches an engine, which refuses the denominator. + static = Reform.from_dict( + {"gov.irs.deductions.standard.amount.SINGLE": {"2024-01-01": 0}}, + country_id="us", + ) + with pytest.raises( + RuntimeError, + match=( + r"computed medicaid_slcsp_state_denominator@2024 \(reads the " + r"engine's baseline branch, which a batch reform engine does not " + r"carry\)" + ), + ): + consumer.simulate(static).calculate("medicaid_cost", 2024) + income_elasticity = "gov.simulation.labor_supply_responses.elasticities.income" + dynamic = Reform.from_dict( + {income_elasticity: {"2024-01-01": -0.05}}, + country_id="us", + ) + assert builder._moved_behavioral_response_parameters( + Microsimulation.default_tax_benefit_system(reform=dynamic), + Microsimulation.default_tax_benefit_system_instance, + 2024, + ) == [income_elasticity] + with pytest.raises( + RuntimeError, + match=( + r"off baseline at 2024: " + r"gov\.simulation\.labor_supply_responses\.elasticities\.income\. " + ), + ): + consumer.simulate(dynamic).calculate("income_tax", 2024) + record = consumer.record() + assert (record["reform_systems"], record["reform_passes"]) == (2, 1) + scorer.close() + + +def _load_sweep_module(): + path = ( + Path(__file__).resolve().parents[3] + / "tools" + / "sweep_us_post_export_scoring.py" + ) + spec = importlib.util.spec_from_file_location("post_export_guard_sweep", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_guard_sweep_records_and_scores_shipped_reform_passes( + builder, fake_reforms, monkeypatch, tmp_path +) -> None: + sweep = _load_sweep_module() + plans, requests, counts = sweep._record_requests(builder) + assert counts["smoke_probes"] == len(us_release_reform_coverage_probes()) == 41 + specs = reform_validation_module.load_default_reform_specs(period=builder.PERIOD) + assert counts["validation_specs"] == sum(not spec.in_sample for spec in specs) == 52 + assert ( + counts["in_sample_fallback_specs"] + == sum(spec.in_sample for spec in specs) + == 12 + ) + assert counts["baseline_keys"] == { + "reform_coverage_smoke": 18, + "reform_validation": 75, + "demographics": 1, + } + smoke_requests = [ + request for request in requests if request[0] == "reform_coverage_smoke" + ] + assert [reform.label for _, reform, _ in smoke_requests] == [ + probe.id for probe in us_release_reform_coverage_probes() + ] + validation_requests = [ + request for request in requests if request[0] == "reform_validation" + ] + ordinary_specs = [ + spec for spec in specs if not reform_validation_module._is_obbba_spec(spec) + ] + assert {spec.id for spec in ordinary_specs}.issubset( + {reform.label for _, reform, _ in validation_requests} + ) + obbba_specs = [ + spec for spec in specs if reform_validation_module._is_obbba_spec(spec) + ] + pre_baseline = reform_validation_module._merged_parameter_changes(obbba_specs) + pre_baseline_requests = [ + key + for _, reform, key in validation_requests + if reform.label == json.dumps(pre_baseline, sort_keys=True, default=str) + ] + obbba_groups = {(spec.budget_measure, spec.period) for spec in obbba_specs} + assert {(key[0], key[1]) for key in pre_baseline_requests} == obbba_groups + assert len(pre_baseline_requests) == len(obbba_groups) == 2 + assert counts["validation_reform_passes"] == len(specs) + len(obbba_groups) == 66 + assert ( + len(requests) + == counts["smoke_probes"] + counts["validation_reform_passes"] + == 107 + ) + + frame, log = _nested_frame(), _EngineLog() + monkeypatch.setattr( + reform_validation_module, + "default_simulate_factory", + lambda path: _unbatched_simulate(frame, _EngineLog()), + ) + scorer = _scorer(builder, frame, log, 3, tmp_path) + sweep._score_baselines(builder, scorer, tmp_path / "fixture.h5", plans) + records = [ + sweep._score_reforms( + builder, scorer, tmp_path / "fixture.h5", requests, start, 4 + ) + for start in range(0, len(requests), 4) + ] + assert sum(record["reform_passes"] for record in records) == len(requests) + assert sum(len(record["whole_file_comparisons"]) for record in records) == 4 + assert len(log.systems) == len(requests) + scorer.close() + + +def test_guard_sweep_runs_every_chunk_and_propagates_worker_failure(tmp_path) -> None: + import subprocess + + sweep = _load_sweep_module() + starts = [] + + def run(command, *, check): + assert check is True + output = Path(command[command.index("--output") + 1]) + if "--worker-start" in command: + start = int(command[command.index("--worker-start") + 1]) + starts.append(start) + record = { + "reform_passes": min(4, 9 - start), + "whole_file_comparisons": [start] if start == 0 else [], + } + else: + record = {"total_reform_passes": 9} + output.write_text( + json.dumps( + { + **record, + "peak_rss_bytes": 1234, + "elapsed_seconds": 1.0, + } + ) + ) + + args = SimpleNamespace( + dataset_path=tmp_path / "fixture.h5", + batch_size=3, + reforms_per_worker=4, + output=tmp_path / "report.json", + ) + report = sweep._run_workers(args, run=run) + assert starts == [0, 4, 8] + assert report["scored_reform_passes"] == 9 + assert report["whole_file_comparisons"] == [0] + assert len(report["workers"]) == 4 + assert report["peak_rss_bytes"] == 1234 + + def fail(command, *, check): + raise subprocess.CalledProcessError(2, command) + + with pytest.raises(subprocess.CalledProcessError): + sweep._run_workers(args, run=fail) + + +@pytest.mark.slow +@pytest.mark.requires_us +def test_shipped_baseline_plans_match_whole_file_policyengine_us( + builder, tmp_path +) -> None: + """The smoke, validation and demographics baseline values match in 3 batches. + + The complete reform sweep runs separately through + ``tools/sweep_us_post_export_scoring.py`` in bounded worker processes. + """ + sweep = _load_sweep_module() + path = _write_engine_h5(builder, tmp_path, _GUARD_HOUSEHOLDS) + plans, _, counts = sweep._record_requests(builder) + assert counts["baseline_keys"] == { + "reform_coverage_smoke": 18, + "reform_validation": 75, + "demographics": 1, + } + scorer = builder._HouseholdBatchedPostExportScorer( + path, maximum_microsim_batch_size=3 + ) + assert scorer.n_batches == 3 + try: + sweep._score_baselines(builder, scorer, path, plans) + finally: + scorer.close() diff --git a/tools/build_us_fiscal_refresh_release.py b/tools/build_us_fiscal_refresh_release.py index 327aac94c5..1fce8e3a6a 100644 --- a/tools/build_us_fiscal_refresh_release.py +++ b/tools/build_us_fiscal_refresh_release.py @@ -30,6 +30,7 @@ os.environ.setdefault(_thread_pool_variable, _THREAD_POOL_DEFAULT) import argparse +import dataclasses import gc import hashlib import importlib.metadata @@ -267,8 +268,8 @@ assert_puf_capital_gains_tail_survives_selection, ) from microcosm.build.us_runtime.reform_validation import ( + US_RELEASE_SPM_SELECTION, default_baseline_level_specs, - default_simulate_factory, load_default_reform_specs, reform_validation_payload, write_reform_validation, @@ -294,6 +295,8 @@ select_frozen_support, ) from microcosm.calibrate import ( + CalibrationResult, + L0RefitResult, TargetRegistry, TargetSpec, calibrate, @@ -4280,119 +4283,892 @@ def _select_households_by_position(frame: Frame, positions: np.ndarray) -> Frame return frame.select(person_mask) -class _BatchedScalarTotal: - def __init__(self, value: float): - self.value = float(value) +# --------------------------------------------------------------------------- +# Household-batched post-export scoring (route A remediation, microcosm#956) +# --------------------------------------------------------------------------- +# +# The reform-coverage smoke, reform_validation and demographics each used to +# build one Microsimulation over the whole written H5 +# (``reform_validation.default_simulate_factory``). Reform validation keeps +# one shared baseline for all of its keys. +# +# This scorer keeps the consumers' ``simulate(reform) -> simulation`` seam and +# changes only how the seam is served: +# +# * the WRITTEN H5 is loaded once (its sha256 bound to the bytes read) and +# split once into household batches; every group entity must nest in one +# household and the batches must partition every entity, or the scorer +# refuses before any engine exists; +# * every construction declares ``spm=US_RELEASE_SPM_SELECTION``; +# * a consumer's baseline requests are learned by an engine-free recording +# dry run and scored in ONE batch-outer pass, in ascending period order; +# the served baseline refuses any key outside that recorded plan; +# * each reform builds ONE tax-benefit system (microcosm#456) and scores its +# keys batch by batch; every batch engine is released before the next; +# * a batch reform engine gets that system alone, not ``reform=`` too (see +# ``_HouseholdBatchedPostExportScorer._construct``), so it carries no +# ``baseline`` branch; it refuses the formulas that read one +# (``POST_EXPORT_BASELINE_BRANCH_READERS``), and a reform may not move the +# behavioral-response parameters that decide whether the rest do; +# * ``calculate`` returns full-length concatenated values with their engine +# weights, so ``.sum()`` (== ``MicroSeries.sum``), ``np.asarray`` and +# ``.weights`` feed every existing statistic unchanged. +# +# Totals then differ from one whole-pool simulation only in floating-point +# summation order, PROVIDED every scored measure is additive across +# households: a formula that aggregates over its simulation's whole population +# (``POST_EXPORT_POPULATION_AGGREGATE_VARIABLES``: the Medicaid SLCSP state +# sums and the weighted income deciles) sees one batch, not the file. Every +# batch engine of a multi-batch pass refuses once it has computed one of them, +# so a measure that reaches such a formula fails the stage instead of scoring +# a batch-local aggregate. + +#: Recorded in ``post_export_scoring`` blocks; bump with any change to how the +#: written H5 is scored. +POST_EXPORT_SCORING_METHOD = "household_batched_written_h5" + +#: One scored request: (variable, period, map_to). +PostExportKey = tuple[str, int, str | None] + +#: policyengine-us formulas that aggregate over their simulation's whole +#: population instead of within one household, so a batch engine computes them +#: over its own batch only: ``medicaid_slcsp_state_average_cost_index`` and +#: ``medicaid_slcsp_state_denominator`` sum person weights by state +#: (``sum_by_state``) and feed ``medicaid_cost`` (and through it ``medicaid`` +#: and ``household_health_benefits``), and the two income deciles are weighted +#: ranks over every household or SPM unit. A multi-batch pass refuses an engine +#: that computed any of them. ``test_us_post_export_scoring.py`` pins this list +#: against the installed engine's variable sources and exercises the Medicaid +#: refusal on a written fixture H5. +POST_EXPORT_POPULATION_AGGREGATE_VARIABLES = ( + "household_income_decile", + "medicaid_slcsp_state_average_cost_index", + "medicaid_slcsp_state_denominator", + "spm_unit_income_decile", +) + +#: policyengine-us formulas whose value reads the engine's ``baseline`` branch, +#: which a batch reform engine does not have. ``medicaid_slcsp_state_denominator`` +#: holds a reform at the baseline's denominator; the others reach +#: ``get_behavioral_response_measurements``, which measures a reform against +#: ``get_branch("baseline")`` and would measure it against itself here. A reform +#: pass refuses an engine that computed any of them, whatever its batch count. +#: Pinned with the aggregates above. +POST_EXPORT_BASELINE_BRANCH_READERS = ( + "income_elasticity_lsr", + "medicaid_slcsp_state_denominator", + "relative_capital_gains_mtr_change", + "relative_income_change", + "relative_wage_change", + "substitution_elasticity_lsr", +) + +#: The other two formulas that read ``simulation.baseline``, and the parameter +#: subtree each response is computed from. Without a baseline branch both +#: return 0; with one they return 0 at zero elasticities. Every reform pass +#: requires these parameters to remain at their baseline values before it +#: builds a batch engine. +POST_EXPORT_BEHAVIORAL_RESPONSE_PARAMETERS = { + "capital_gains_behavioral_response": "gov.simulation.capital_gains_responses", + "labor_supply_behavioral_response": "gov.simulation.labor_supply_responses", +} + + +def _post_export_key( + variable: str, period: object, map_to: str | None = None +) -> PostExportKey: + if period is None: + raise ValueError( + f"Post-export scoring needs an explicit period for {variable!r}; the " + "engine default period is not part of the recorded plan." + ) + return (str(variable), int(str(period)), None if map_to is None else str(map_to)) + + +def _post_export_key_record(key: PostExportKey) -> dict[str, object]: + variable, period, map_to = key + return {"variable": variable, "period": period, "map_to": map_to} + + +def _ascending_period_plan(keys: Iterable[PostExportKey]) -> tuple[PostExportKey, ...]: + """Order a baseline plan by period, keeping request order within a period. + + Baseline consumers can request interleaved periods. Sort those requests + before engine calls; the engine guard refuses a return to an earlier + period. The real-engine test probes the MD CCS request-order error and + warns if the upstream engine no longer reproduces it. + """ + return tuple(sorted(dict.fromkeys(keys), key=lambda key: key[1])) + + +def _post_export_baseline_plan_record( + plan: Sequence[PostExportKey], +) -> dict[str, object]: + return { + "keys": [_post_export_key_record(key) for key in plan], + "period_order": sorted({key[1] for key in plan}), + } + + +class _PostExportValues: + """Full-length values of one scored key, with the engine's weights. + + The duck type the post-export consumers read from a Microsimulation + result: ``np.asarray(values)``, ``values.values``, ``values.weights`` and + ``values.sum()``. Both arrays are read-only, so no consumer can mutate a + baseline result another consumer statistic still reads. + """ + + __slots__ = ("_values", "_weights") + + def __init__(self, values: np.ndarray, weights: np.ndarray) -> None: + values = np.asarray(values) + weights = np.asarray(weights, dtype=np.float64) + if values.ndim != 1 or weights.shape != values.shape: + raise ValueError( + "Post-export scoring needs one weight per value; got values " + f"of shape {values.shape} and weights of shape {weights.shape}." + ) + values.flags.writeable = False + weights.flags.writeable = False + self._values = values + self._weights = weights + + def __array__(self, dtype=None, copy=None) -> np.ndarray: + if dtype is None and not copy: + return self._values + return np.array(self._values, dtype=dtype, copy=True) + + def __len__(self) -> int: + return len(self._values) + + @property + def values(self) -> np.ndarray: + return self._values + + @property + def weights(self) -> np.ndarray: + return self._weights def sum(self) -> float: - return self.value + """Weighted total, computed as policyengine-core's ``MicroSeries.sum``. + + ``MicroSeries.sum`` is ``self.multiply(self.weights).sum()``: a pandas + product summed with NaN skipped. Doing the same over the concatenated + arrays gives the whole-pool total up to summation order. + """ + return float(pd.Series(self._values).multiply(self._weights).sum()) + + +def _concatenate_post_export_parts( + parts: Sequence[tuple[np.ndarray, np.ndarray]], +) -> _PostExportValues: + return _PostExportValues( + np.concatenate([values for values, _ in parts]), + np.concatenate([weights for _, weights in parts]), + ) + + +def _post_export_values_and_weights(result, key: PostExportKey): + weights = getattr(result, "weights", None) + if weights is None: + raise RuntimeError( + f"Post-export scoring of {key} got an unweighted engine result; " + "every scored key must carry the engine's entity weights." + ) + values = np.asarray(result) + weights = np.asarray(weights, dtype=np.float64) + if values.ndim != 1 or weights.shape != values.shape: + raise RuntimeError( + f"Post-export scoring of {key} got {values.shape} values against " + f"{weights.shape} weights from one batch engine." + ) + return values, weights + + +class _RecordingPostExportSimulation: + """Engine-free stand-in that records the baseline keys a consumer asks for.""" + + def __init__(self, keys: dict[PostExportKey, None] | None) -> None: + self._keys = keys + + def calculate(self, variable, period=None, map_to=None) -> _PostExportValues: + key = _post_export_key(variable, period, map_to) + if self._keys is not None: + self._keys.setdefault(key, None) + return _PostExportValues(np.zeros(0), np.zeros(0)) -class _BatchedReformValidationSimulation: +def _record_post_export_baseline_plan( + consumer: Callable[[Callable[[Any], Any]], object], +) -> tuple[PostExportKey, ...]: + """Learn a consumer's baseline keys by running it against recorders. + + No engine is constructed: ``simulate(None)`` returns a recorder that + logs every request, ``simulate(reform)`` a recorder that logs nothing + (reforms are scored key by key when the consumer runs for real). The + consumers ask for the same keys whatever values they get back, and the + served baseline refuses any key this dry run did not see, so a + value-dependent request cannot slip past the plan. + """ + keys: dict[PostExportKey, None] = {} + + def simulate(reform): + return _RecordingPostExportSimulation(keys if reform is None else None) + + consumer(simulate) + return _ascending_period_plan(keys) + + +class _AscendingPeriodEngine: + """One batch engine that refuses a period earlier than one it computed.""" + + def __init__(self, simulation, *, label: str) -> None: + self.simulation = simulation + self._label = label + self._latest_period: int | None = None + + def calculate(self, key: PostExportKey): + variable, period, map_to = key + if self._latest_period is not None and period < self._latest_period: + raise RuntimeError( + f"Post-export scoring ({self._label}) asked one engine for " + f"{variable}@{period} after it computed period " + f"{self._latest_period}. policyengine-us 2.2.1 can fail to " + "compute an earlier period after a later one on the same " + "engine (the MD CCS ParameterNotFoundError; see " + "_ascending_period_plan), so each engine scores its keys in " + "ascending period order." + ) + self._latest_period = period + if map_to is None: + return self.simulation.calculate(variable, period) + return self.simulation.calculate(variable, period, map_to=map_to) + + +def _post_export_watched_variables(*, batched: bool, reform: bool) -> tuple[str, ...]: + """The formulas a batch engine may not compute (see the block comment).""" + watched = POST_EXPORT_POPULATION_AGGREGATE_VARIABLES if batched else () + if reform: + watched = (*watched, *POST_EXPORT_BASELINE_BRANCH_READERS) + return tuple(dict.fromkeys(watched)) + + +def _post_export_known_periods( + simulation, variables: Sequence[str] +) -> set[tuple[str, str]]: + """The (variable, period) values an engine or any live branch of it holds. + + The holder's ``get_known_periods`` includes memory and disk storage. + Branches keep their own holders; the ones still attached are read too. + The real-engine guard test checks that the watched Medicaid computation + leaves a known period visible through this API. + """ + known: set[tuple[str, str]] = set() + stack = [simulation] + seen: set[int] = set() + while stack: + current = stack.pop() + if id(current) in seen: + continue + seen.add(id(current)) + for variable in variables: + known.update( + (variable, str(period)) + for period in current.get_holder(variable).get_known_periods() + ) + branches = getattr(current, "branches", None) + if isinstance(branches, dict): + stack.extend(branches.values()) + return known + + +def _assert_post_export_scoring_is_batch_invariant( + simulation, + watched: Sequence[str], + known_at_load: set[tuple[str, str]], + *, + label: str, + batched: bool, + reform: bool, +) -> None: + """Refuse an engine that computed a formula batching would score wrongly. + + ``known_at_load`` holds what the engine held before scoring (values the + written H5 stores as inputs, which batching does not change). + """ + computed = sorted(_post_export_known_periods(simulation, watched) - known_at_load) + if not computed: + return + reasons = [] + for variable, period in computed: + why = [] + if batched and variable in POST_EXPORT_POPULATION_AGGREGATE_VARIABLES: + why.append( + "aggregates over its simulation's whole population, here one " + "household batch" + ) + if reform and variable in POST_EXPORT_BASELINE_BRANCH_READERS: + why.append( + "reads the engine's baseline branch, which a batch reform engine " + "does not carry" + ) + reasons.append(f"{variable}@{period} ({'; '.join(why)})") + raise RuntimeError( + f"Post-export scoring ({label}) is not batch-invariant: a batch engine " + f"computed {', '.join(reasons)}. A measure that reaches these formulas " + "cannot be scored in household batches (microcosm#956)." + ) + + +def _parameter_leaf_values(node, instant: str) -> dict[str, object]: + """Every parameter under ``node`` (itself included) at ``instant``.""" + return { + leaf.name: leaf(instant) + for leaf in (node, *node.get_descendants()) + if hasattr(leaf, "values_list") + } + + +def _moved_behavioral_response_parameters( + reform_system, baseline_system, period: int +) -> list[str]: + """Behavioral-response parameters a reform system sets off baseline.""" + instant = f"{int(period)}-01-01" + moved: list[str] = [] + for subtree in POST_EXPORT_BEHAVIORAL_RESPONSE_PARAMETERS.values(): + reformed = _parameter_leaf_values( + reform_system.parameters.get_child(subtree), instant + ) + baseline = _parameter_leaf_values( + baseline_system.parameters.get_child(subtree), instant + ) + moved.extend( + name + for name in sorted(reformed.keys() | baseline.keys()) + if name not in reformed + or name not in baseline + or reformed[name] != baseline[name] + ) + return moved + + +def _assert_post_export_batching_premises( + frame: Frame, batch_frames: Sequence[Frame] +) -> None: + """Refuse a written frame that household batching would score wrongly. + + Batching is output-invariant only if every group unit sits wholly inside + one household, so that each batch holds whole units and every entity row + lands in exactly one batch. ``Frame.select`` keeps every group row any + selected person references without checking that, so a unit spanning two + batches would be scored twice, each time with part of its members. (The + Frame constructor already refuses group rows no person references.) + """ + for entity in US_SCHEMA.group_entities: + if entity != "household": + # Raises when a unit's persons sit in more than one household. + _group_to_household_positions(frame, entity) + for entity in frame.entities: + batched = sum(batch.n(entity) for batch in batch_frames) + if batched != frame.n(entity): + raise ValueError( + f"Post-export household batches hold {batched} {entity} rows " + f"against {frame.n(entity)} in the written H5; the batches do " + "not partition the pool." + ) + household_ids = frame.table("household")["household_id"].to_numpy() + batched_ids = np.concatenate( + [batch.table("household")["household_id"].to_numpy() for batch in batch_frames] + ) + if not np.array_equal(batched_ids, household_ids): + raise ValueError( + "Post-export household batches do not reproduce the written H5's " + "household order." + ) + batched_weights = np.concatenate( + [batch.weights_for("household").values for batch in batch_frames] + ) + if not np.array_equal(batched_weights, frame.weights_for("household").values): + raise ValueError( + "Post-export household batches do not carry the written H5's " + "household weights." + ) + + +class _PostExportScoringPlan: + """The pre-export record of how the written H5 will be scored. + + Built before ``calibration_diagnostics.json`` is written, so that artifact + carries the plan; the post-export stages then score exactly these + baseline plans. (A plain class: this tool is loaded by file path in the + tests, outside ``sys.modules``, where ``dataclasses`` cannot resolve the + module's string annotations.) + """ + def __init__( self, - frame: Frame, *, - reform, + n_households: int, + maximum_batch_size: int | None, + baseline_plans: Mapping[str, tuple[PostExportKey, ...]], + error: str | None = None, + ) -> None: + self.n_households = int(n_households) + self.maximum_batch_size = maximum_batch_size + self.baseline_plans = { + consumer: tuple(plan) for consumer, plan in baseline_plans.items() + } + self.error = error + + @property + def n_batches(self) -> int: + return sum( + 1 + for _ in _household_position_batches( + self.n_households, self.maximum_batch_size + ) + ) + + def baseline_plan(self, consumer: str) -> tuple[PostExportKey, ...]: + if self.error is not None: + raise RuntimeError( + "The post-export scoring plan was not built before export: " + f"{self.error}" + ) + return self.baseline_plans[consumer] + + def terminal_failures(self) -> list[str]: + """The plan's own failure, as a line of the batched terminal raise.""" + if self.error is None: + return [] + return [ + "Post-export scoring plan could not be built before export " + f"(microcosm#956): {self.error}" + ] + + def record(self) -> dict[str, object]: + if self.error is not None: + return {"method": POST_EXPORT_SCORING_METHOD, "error": self.error} + return { + "method": POST_EXPORT_SCORING_METHOD, + "maximum_batch_size": self.maximum_batch_size, + "n_households": self.n_households, + "n_batches": self.n_batches, + "spm": dict(US_RELEASE_SPM_SELECTION), + "period_order": "ascending", + "consumers": { + consumer: {"baseline_plan": _post_export_baseline_plan_record(plan)} + for consumer, plan in self.baseline_plans.items() + }, + } + + +def _post_export_scoring_plan( + *, + n_households: int, + maximum_microsim_batch_size: int | None, + consumers: Mapping[str, Callable[[Callable[[Any], Any]], object]], +) -> _PostExportScoringPlan: + return _PostExportScoringPlan( + n_households=int(n_households), + maximum_batch_size=maximum_microsim_batch_size, + baseline_plans={ + name: _record_post_export_baseline_plan(consumer) + for name, consumer in consumers.items() + }, + ) + + +def _record_post_export_scoring_plan( + args: argparse.Namespace, *, n_households: int, result, release_id: str +) -> _PostExportScoringPlan: + """The plan ``_main`` records before the calibration diagnostics. + + Never raises (microcosm#547): a crash while building the plan rides the + plan's record into the diagnostics and joins the batched terminal raise + (``terminal_failures``), so it cannot pre-empt the diagnostics, the QRF + tail evidence or any later terminal group, and the run still refuses + before the export write. + """ + try: + return _post_export_scoring_plan( + n_households=n_households, + maximum_microsim_batch_size=args.maximum_microsim_batch_size, + consumers=_post_export_consumers( + args, result=result, release_id=release_id + ), + ) + except Exception as error: + return _PostExportScoringPlan( + n_households=n_households, + maximum_batch_size=args.maximum_microsim_batch_size, + baseline_plans={}, + error=f"{type(error).__name__}: {error}", + ) + + +def _reform_coverage_smoke_consumer(simulate) -> GateResult: + """The smoke exactly as ``_main`` runs it, for the plan's dry run.""" + return us_reform_coverage_smoke_gate(simulate=simulate, period=PERIOD) + + +def _demographics_consumer(simulate) -> tuple[np.ndarray, np.ndarray]: + return population_by_age_from_sim(simulate(None), PERIOD) + + +def _reform_validation_consumer( + *, result, release_id: str +) -> Callable[[Callable[[Any], Any] | None], dict[str, Any]]: + """reform_validation.json's payload as a function of the simulate seam.""" + specs = load_default_reform_specs(period=PERIOD) + in_sample_estimates = _in_sample_estimates(result) + in_sample_targets = _in_sample_targets(result) + baseline_levels = default_baseline_level_specs() + + def payload_for(simulate) -> dict[str, Any]: + return reform_validation_payload( + specs, + period=PERIOD, + simulate=simulate, + in_sample_estimates=in_sample_estimates, + in_sample_targets=in_sample_targets, + baseline_levels=baseline_levels, + release_id=release_id, + ) + + return payload_for + + +def _post_export_consumers( + args: argparse.Namespace, *, result, release_id: str +) -> dict[str, Callable[[Callable[[Any], Any]], object]]: + """The post-export stages this build will run, in run order.""" + consumers: dict[str, Callable[[Callable[[Any], Any]], object]] = {} + if not args.skip_reform_coverage_smoke: + consumers["reform_coverage_smoke"] = _reform_coverage_smoke_consumer + if not args.skip_reform_validation and not args.skip_out_of_sample_reforms: + consumers["reform_validation"] = _reform_validation_consumer( + result=result, release_id=release_id + ) + if not args.skip_demographics: + consumers["demographics"] = _demographics_consumer + return consumers + + +class _HouseholdBatchedPostExportScorer: + """Scores the written release H5 in household batches (see the block + comment above). One instance serves every post-export consumer.""" + + def __init__( + self, + dataset_path: Path, + *, maximum_microsim_batch_size: int | None, - microsimulation_cls, - dataset_from_frame, - ): - self._frame = frame - self._reform = reform - self._maximum_microsim_batch_size = maximum_microsim_batch_size + expected_n_households: int | None = None, + microsimulation_cls=None, + dataset_from_frame=None, + load_frame=None, + ) -> None: + self.dataset_path = Path(dataset_path) + self.maximum_batch_size = maximum_microsim_batch_size + # The scored bytes are bound to this digest: the loader refuses a file + # whose bytes differ from it, and the manifest's dataset sha must equal + # it (``_build_manifests(scored_dataset_sha256=...)``). + self.dataset_sha256 = _sha256(self.dataset_path) + frame = (load_frame or _load_frame)( + self.dataset_path, expected_sha256=self.dataset_sha256 + ) + self.n_households = int(frame.n("household")) + if ( + expected_n_households is not None + and self.n_households != expected_n_households + ): + raise ValueError( + f"The written H5 carries {self.n_households} households; the " + f"post-export scoring plan was built for {expected_n_households}." + ) + batch_frames = tuple( + frame + if len(positions) == self.n_households + else _select_households_by_position(frame, positions) + for positions in _household_position_batches( + self.n_households, maximum_microsim_batch_size + ) + ) + _assert_post_export_batching_premises(frame, batch_frames) + del frame + self._batch_frames: tuple[Frame, ...] | None = batch_frames + self.n_batches = len(batch_frames) + self.max_batch_households = max( + (batch.n("household") for batch in batch_frames), default=0 + ) + if microsimulation_cls is None: + from policyengine_us import Microsimulation + + microsimulation_cls = Microsimulation self._microsimulation_cls = microsimulation_cls + if dataset_from_frame is None: + + def dataset_from_frame(batch_frame: Frame): + return _dataset_from_frame( + batch_frame, + assert_no_formula_owned_columns=False, + ) + self._dataset_from_frame = dataset_from_frame - self._cache: dict[tuple[str, int], float] = {} - self._reform_system = None - def calculate(self, measure: str, period: int) -> _BatchedScalarTotal: - key = (str(measure), int(period)) - if key not in self._cache: - self._cache[key] = self._calculate_total(str(measure), int(period)) - return _BatchedScalarTotal(self._cache[key]) + def open_consumer( + self, name: str, baseline_plan: Sequence[PostExportKey] + ) -> _PostExportConsumer: + """Score ``baseline_plan`` now; return the consumer's simulate seam.""" + return _PostExportConsumer(self, name, tuple(baseline_plan)) + + def close(self) -> None: + self._batch_frames = None + _collect_family_garbage() - def _calculate_total(self, measure: str, period: int) -> float: - n_households = self._frame.n("household") - batches = tuple( - _household_position_batches( - n_households, - self._maximum_microsim_batch_size, + def _batches(self) -> tuple[Frame, ...]: + if self._batch_frames is None: + raise RuntimeError("The post-export scorer is closed.") + return self._batch_frames + + def _construct(self, batch_frame: Frame, *, reform_system): + dataset = self._dataset_from_frame(batch_frame) + if reform_system is None: + return self._microsimulation_cls( + dataset=dataset, spm=dict(US_RELEASE_SPM_SELECTION) + ) + # The reform's system alone, not ``reform=`` as well. Given both, + # policyengine-us 2.2.1 clones the supplied system on every + # construction (``SPMSimulationMixin._prepare_spm_system`` -> + # ``clone_spm_system``: the whole parameter tree rebuilt and every + # variable deep-copied, so each batch starts on cold caches), core + # applies the reform again to the clone, and core adds a ``baseline`` + # branch whose holders copy every input array. Given the system alone + # it shares the system's policy (``share_spm_policy``). + # ``POST_EXPORT_BASELINE_BRANCH_READERS`` and + # ``POST_EXPORT_BEHAVIORAL_RESPONSE_PARAMETERS`` list formulas that + # read the baseline branch, and the scorer guards those uses. + return self._microsimulation_cls( + tax_benefit_system=reform_system, + dataset=dataset, + spm=dict(US_RELEASE_SPM_SELECTION), + ) + + def _assert_reform_keeps_baseline_responses( + self, reform_system, key: PostExportKey, *, label: str + ) -> None: + """Refuse a reform that moves a behavioral-response parameter. + + A batch reform engine has no ``baseline`` branch, so policyengine-us + scores its labor-supply and capital-gains responses as 0. A whole-file + reform engine has one and scores them from these parameters, as 0 at + the engine's zero default elasticities. The guard requires these + parameters to stay at baseline before building a batch reform engine. + """ + moved = _moved_behavioral_response_parameters( + reform_system, + self._microsimulation_cls.default_tax_benefit_system_instance, + key[1], + ) + if moved: + raise RuntimeError( + f"Post-export scoring ({label}) refuses a reform that sets " + f"behavioral-response parameters off baseline at {key[1]}: " + f"{', '.join(moved)}. Batch reform engines carry no baseline " + "branch, so policyengine-us would score these responses as zero " + "(microcosm#956)." ) - ) - if len(batches) > 1: + + def _score( + self, + keys: Sequence[PostExportKey], + *, + label: str, + reform_system=None, + ) -> dict[PostExportKey, _PostExportValues]: + """One batch-outer pass: one engine per batch scores every key.""" + batches = self._batches() + batched = len(batches) > 1 + watched = _post_export_watched_variables( + batched=batched, reform=reform_system is not None + ) + if batched: print( - "Scoring reform validation measure " - f"{measure} in {len(batches)} batches of up to " - f"{self._maximum_microsim_batch_size:,} households.", + f"Scoring post-export {label} ({len(keys)} key(s)) in " + f"{len(batches)} batches of up to {self.max_batch_households:,} " + "households.", flush=True, ) - total = 0.0 - for household_positions in batches: + parts: dict[PostExportKey, list[tuple[np.ndarray, np.ndarray]]] = { + key: [] for key in keys + } + for batch_frame in batches: with _automatic_gc_suspended(): - full_batch = len(household_positions) == n_households - batch_frame = ( - self._frame - if full_batch - else _select_households_by_position( - self._frame, household_positions - ) - ) - dataset = self._dataset_from_frame(batch_frame) - if self._reform is None: - simulation = self._microsimulation_cls(dataset=dataset) - else: - # microcosm#456: one reform system per scored reform, not - # one per batch (each engine build permanently leaks - # ~5,600 sys.modules entries). - if self._reform_system is None: - self._reform_system = ( - self._microsimulation_cls.default_tax_benefit_system( - reform=self._reform - ) + simulation = self._construct(batch_frame, reform_system=reform_system) + try: + known_at_load = _post_export_known_periods(simulation, watched) + engine = _AscendingPeriodEngine(simulation, label=label) + for key in keys: + parts[key].append( + _post_export_values_and_weights(engine.calculate(key), key) ) - simulation = self._microsimulation_cls( - tax_benefit_system=self._reform_system, - dataset=dataset, - reform=self._reform, + # Before the release below drops the engine's holders. + _assert_post_export_scoring_is_batch_invariant( + simulation, + watched, + known_at_load, + label=label, + batched=batched, + reform=reform_system is not None, ) - total += float(simulation.calculate(measure, period).sum()) - release_engine_simulation(simulation) - del simulation, dataset, batch_frame + finally: + # microcosm#456: free this batch engine's array mass by + # refcount before the next batch builds its own, also + # when a key raises. + release_engine_simulation(simulation) + del simulation, engine _collect_batch_garbage() _collect_family_garbage() - return total + return { + key: _concatenate_post_export_parts(key_parts) + for key, key_parts in parts.items() + } -def _batched_reform_validation_simulate_factory_from_frame( - frame: Frame, - *, - maximum_microsim_batch_size: int | None = DEFAULT_MAXIMUM_MICROSIM_BATCH_SIZE, - microsimulation_cls=None, - dataset_from_frame=None, -): - if microsimulation_cls is None: - from policyengine_us import Microsimulation +def _open_post_export_scorer( + plan: _PostExportScoringPlan, dataset_path: Path +) -> _HouseholdBatchedPostExportScorer | None: + """Open the scorer the plan recorded before export, or None if no stage runs.""" + if plan.error is not None: + raise RuntimeError( + f"The post-export scoring plan was not built before export: {plan.error}" + ) + if not plan.baseline_plans: + return None + return _HouseholdBatchedPostExportScorer( + dataset_path, + maximum_microsim_batch_size=plan.maximum_batch_size, + expected_n_households=plan.n_households, + ) - microsimulation_cls = Microsimulation - if dataset_from_frame is None: - def dataset_from_frame(batch_frame: Frame): - return _dataset_from_frame( - batch_frame, - assert_no_formula_owned_columns=False, - ) +def _close_post_export_scorer( + scorer: _HouseholdBatchedPostExportScorer | None, +) -> str | None: + """Close the scorer; return the sha256 of the H5 it scored (None: no stage).""" + if scorer is None: + return None + scorer.close() + return scorer.dataset_sha256 - def simulate(reform): - return _BatchedReformValidationSimulation( - frame, - reform=reform, - maximum_microsim_batch_size=maximum_microsim_batch_size, - microsimulation_cls=microsimulation_cls, - dataset_from_frame=dataset_from_frame, + +class _PostExportConsumer: + """One consumer's simulate seam over a scored baseline plan.""" + + def __init__( + self, + scorer: _HouseholdBatchedPostExportScorer, + name: str, + baseline_plan: tuple[PostExportKey, ...], + ) -> None: + ordered = _ascending_period_plan(baseline_plan) + if ordered != baseline_plan: + raise ValueError( + f"{name}: a post-export baseline plan must be unique keys in " + "ascending period order." + ) + self._scorer = scorer + self.name = name + self.baseline_plan = baseline_plan + self.baseline_passes = 1 if baseline_plan else 0 + self.reform_passes = 0 + self.reform_systems = 0 + self._baseline = ( + scorer._score(baseline_plan, label=f"{name} baseline plan") + if baseline_plan + else {} ) - return simulate + def simulate(self, reform): + if reform is None: + return _PlannedPostExportBaseline(self) + return _BatchedPostExportReform(self, reform) + + def record(self) -> dict[str, object]: + scorer = self._scorer + return { + "method": POST_EXPORT_SCORING_METHOD, + "dataset_sha256": scorer.dataset_sha256, + "maximum_batch_size": scorer.maximum_batch_size, + "n_households": scorer.n_households, + "n_batches": scorer.n_batches, + "max_batch_households": scorer.max_batch_households, + "spm": dict(US_RELEASE_SPM_SELECTION), + "baseline_plan": _post_export_baseline_plan_record(self.baseline_plan), + "baseline_passes": self.baseline_passes, + "reform_passes": self.reform_passes, + "reform_systems": self.reform_systems, + } + + def _baseline_values(self, key: PostExportKey) -> _PostExportValues: + try: + return self._baseline[key] + except KeyError: + raise RuntimeError( + f"{self.name} asked the baseline for {key}, which its recorded " + "plan does not hold; the engine-free dry run must see every " + "baseline request, so nothing is scored off-plan." + ) from None + + def _reform_values(self, reform_system, key: PostExportKey): + label = f"{self.name} reform {key[0]}@{key[1]}" + self._scorer._assert_reform_keeps_baseline_responses( + reform_system, key, label=label + ) + self.reform_passes += 1 + scored = self._scorer._score((key,), label=label, reform_system=reform_system) + return scored[key] + + +class _PlannedPostExportBaseline: + """The baseline simulation a consumer sees: served from its scored plan.""" + + def __init__(self, consumer: _PostExportConsumer) -> None: + self._consumer = consumer + + def calculate(self, variable, period=None, map_to=None) -> _PostExportValues: + return self._consumer._baseline_values( + _post_export_key(variable, period, map_to) + ) + + +class _BatchedPostExportReform: + """A reformed simulation a consumer sees: each key is one batched pass. + + The reform's tax-benefit system is built once, on the first key, and + handed to every batch engine (microcosm#456). The engines get that system + alone (see ``_HouseholdBatchedPostExportScorer._construct``). + """ + + def __init__(self, consumer: _PostExportConsumer, reform) -> None: + self._consumer = consumer + self._reform = reform + self._system = None + self._results: dict[PostExportKey, _PostExportValues] = {} + + def calculate(self, variable, period=None, map_to=None) -> _PostExportValues: + key = _post_export_key(variable, period, map_to) + if key not in self._results: + if self._system is None: + scorer = self._consumer._scorer + self._system = scorer._microsimulation_cls.default_tax_benefit_system( + reform=self._reform + ) + self._consumer.reform_systems += 1 + self._results[key] = self._consumer._reform_values(self._system, key) + return self._results[key] def _reform_household_income_tax( @@ -6141,6 +6917,23 @@ def _with_l0_refit_weights(base_frame: Frame, result) -> Frame: ) +def _without_calibrated_frames(result, *, export_frame: Frame | None = None): + """Drop calibrated target tables while keeping weights and diagnostics. + + The exact-k receipt reads ``result.frame.n("household")``; give it the + clean export frame instead of retaining the calibrated target tables. + """ + if isinstance(result, L0RefitResult): + return dataclasses.replace( + result, + selection=_without_calibrated_frames(result.selection), + refit=_without_calibrated_frames(result.refit, export_frame=export_frame), + ) + if isinstance(result, CalibrationResult): + return dataclasses.replace(result, frame=export_frame) + return result + + def _selected_plan_ratio_bucket(values: np.ndarray) -> dict[str, object]: finite = np.asarray(values, dtype=np.float64) finite = finite[np.isfinite(finite)] @@ -8544,6 +9337,7 @@ def _write_release_calibration_diagnostics( target_loss_basis: Mapping[str, object] | None = None, exact_k_ladder: Mapping[str, object] | None = None, calibration_runtime: Mapping[str, object] | None = None, + post_export_scoring: Mapping[str, object] | None = None, ) -> None: """Write calibration diagnostics even when hard release gates fail.""" failures = list(gate_failures) @@ -8737,6 +9531,14 @@ def _write_release_calibration_diagnostics( }, "incumbent_diagnostics": incumbent_payload, "post_export_target_audit": bool(audit_export_targets), + # How the written H5 will be scored after export (route A + # remediation, microcosm#956): household batches, their count, + # and each post-export stage's baseline plan in period order. + **( + {"post_export_scoring": dict(post_export_scoring)} + if post_export_scoring is not None + else {} + ), }, ) @@ -8989,6 +9791,51 @@ def _in_sample_targets(result) -> dict[str, float]: return targets +def _score_post_export_consumer( + name: str, + consumer: Callable[[Callable[[Any], Any]], object], + *, + dataset_path: Path, + post_export_scorer: _HouseholdBatchedPostExportScorer | None, + baseline_plan: Sequence[PostExportKey] | None, + maximum_microsim_batch_size: int | None, +): + """Run one post-export consumer through the household-batched scorer. + + ``_main`` passes the scorer it opened on the written H5 and the plan it + recorded in the calibration diagnostics. A direct caller may pass neither: + the scorer is then opened (and closed) here, and the plan is recorded by + the same engine-free dry run. + """ + scorer = post_export_scorer + if scorer is None: + scorer = _HouseholdBatchedPostExportScorer( + dataset_path, + maximum_microsim_batch_size=maximum_microsim_batch_size, + ) + try: + plan = ( + _record_post_export_baseline_plan(consumer) + if baseline_plan is None + else tuple(baseline_plan) + ) + scoring = scorer.open_consumer(name, plan) + output = consumer(scoring.simulate) + record = scoring.record() + print( + f"Post-export {name}: {record['baseline_passes']} baseline pass " + f"({len(plan)} keys), {record['reform_passes']} reform passes, " + f"{record['reform_systems']} reform systems, " + f"{record['n_batches']} batches of up to " + f"{record['max_batch_households']:,} households.", + flush=True, + ) + return output + finally: + if post_export_scorer is None: + scorer.close() + + def _write_reform_validation( *, release_dir: Path, @@ -8997,15 +9844,19 @@ def _write_reform_validation( registry: TargetRegistry, release_id: str, simulate_out_of_sample: bool, + post_export_scorer: _HouseholdBatchedPostExportScorer | None = None, + baseline_plan: Sequence[PostExportKey] | None = None, + maximum_microsim_batch_size: int | None = DEFAULT_MAXIMUM_MICROSIM_BATCH_SIZE, ) -> None: """Emit reform_validation.json: microcosm budget effects vs JCT scores. In-sample JCT tax-expenditure reforms come straight from the calibration fit; out-of-sample OBBBA provisions are simulated on the freshly written - release H5 (skipped if ``simulate_out_of_sample`` is False, e.g. for a fast - diagnostics-only build). + release H5 in household batches (skipped if ``simulate_out_of_sample`` is + False, e.g. for a fast diagnostics-only build). The shared baseline is one + batch-outer pass over the plan's keys in ascending period order. """ - specs = load_default_reform_specs(period=PERIOD) + payload_for = _reform_validation_consumer(result=result, release_id=release_id) if not simulate_out_of_sample: print( "\n".join( @@ -9023,18 +9874,16 @@ def _write_reform_validation( ), file=sys.stderr, ) - simulate = ( - default_simulate_factory(dataset_path) if simulate_out_of_sample else None - ) - payload = reform_validation_payload( - specs, - period=PERIOD, - simulate=simulate, - in_sample_estimates=_in_sample_estimates(result), - in_sample_targets=_in_sample_targets(result), - baseline_levels=default_baseline_level_specs(), - release_id=release_id, - ) + payload = payload_for(None) + else: + payload = _score_post_export_consumer( + "reform_validation", + payload_for, + dataset_path=dataset_path, + post_export_scorer=post_export_scorer, + baseline_plan=baseline_plan, + maximum_microsim_batch_size=maximum_microsim_batch_size, + ) write_reform_validation(payload, release_dir / "reform_validation.json") @@ -9043,17 +9892,24 @@ def _write_demographics( release_dir: Path, dataset_path: Path, release_id: str, + post_export_scorer: _HouseholdBatchedPostExportScorer | None = None, + baseline_plan: Sequence[PostExportKey] | None = None, + maximum_microsim_batch_size: int | None = DEFAULT_MAXIMUM_MICROSIM_BATCH_SIZE, ) -> None: """Emit demographics.json: the dataset's weighted population by age band. The fiscal-refresh release calibrates source-backed Census PEP age targets; - this file remains a compact summary diagnostic for release consumers. + this file remains a compact summary diagnostic for release consumers. Ages + and person weights are scored on the written H5 in household batches. """ - from policyengine_us import Microsimulation - from policyengine_us.data import USSingleYearDataset - - sim = Microsimulation(dataset=USSingleYearDataset(file_path=str(dataset_path))) - ages, weights = population_by_age_from_sim(sim, PERIOD) + ages, weights = _score_post_export_consumer( + "demographics", + _demographics_consumer, + dataset_path=dataset_path, + post_export_scorer=post_export_scorer, + baseline_plan=baseline_plan, + maximum_microsim_batch_size=maximum_microsim_batch_size, + ) payload = demographics_payload(ages, weights, period=PERIOD, release_id=release_id) # Household-record counts by state and congressional district: the # release's sub-national resolution floor, surfaced on the dashboard. @@ -9103,12 +9959,22 @@ def _build_manifests( export_input_mass_reference: Mapping[str, object] | None = None, calibration_runtime: Mapping[str, object] | None = None, skipped_gates: Iterable[str] = (), + scored_dataset_sha256: str | None = None, ) -> None: dataset_path = artifact_root / dataset_filename calibration_path = artifact_root / calibration_filename diagnostics_path = release_dir / "calibration_diagnostics.json" coverage_path = release_dir / "us_source_coverage.json" dataset_sha = _sha256(dataset_path) + if scored_dataset_sha256 is not None and scored_dataset_sha256 != dataset_sha: + # The post-export smoke, reform validation and demographics scored the + # H5 whose sha256 the scorer bound at load; the manifest must pin the + # same bytes, or those verdicts describe a different file. + raise RuntimeError( + f"The post-export stages scored {dataset_path} at sha256 " + f"{scored_dataset_sha256}, but the file now hashes to {dataset_sha}; " + "the manifest would pin bytes the release gates never scored." + ) calibration_sha = _sha256(calibration_path) diagnostics_sha = _sha256(diagnostics_path) coverage_sha = _sha256(coverage_path) @@ -12615,6 +13481,10 @@ def _main(argv: Sequence[str] | None = None) -> None: "refit_initial_loss": _finite_or_none(result.initial_loss), "refit_final_loss": _finite_or_none(result.final_loss), } + # Frame.with_weights copies the target tables into each calibrated frame. + # Drop the input now and replace those result frames after building the + # clean export from base_frame and the calibrated weights below. + del target_frame if telemetry is not None: telemetry.stage( "take_up_final_diagnostics", @@ -12635,6 +13505,12 @@ def _main(argv: Sequence[str] | None = None) -> None: ) else: export_frame = _with_l0_refit_weights(base_frame, result) + result = _without_calibrated_frames( + result, + export_frame=export_frame if ladder_outcome is not None else None, + ) + if isinstance(ladder_outcome, ExactKLadderCalibration): + ladder_outcome = dataclasses.replace(ladder_outcome, result=result) compilation = dict(compilation) final_uncapped_ssi = _ssi_person_uncapped_amount( export_frame, @@ -12962,6 +13838,17 @@ def _main(argv: Sequence[str] | None = None) -> None: *exact_k_fit_failures, *gate_failures, ] + # microcosm#956: how the written H5 will be scored. The baseline plans come + # from engine-free dry runs of the post-export consumers, so they are + # fixed here and recorded in the calibration diagnostics; the post-export + # stages then score exactly these plans. A plan that cannot be built is + # recorded, not raised, and joins the terminal batch below. + post_export_scoring_plan = _record_post_export_scoring_plan( + args, + n_households=int(export_frame.n("household")), + result=result, + release_id=release_id, + ) _write_release_calibration_diagnostics( result=result, release_dir=release_dir, @@ -12999,6 +13886,7 @@ def _main(argv: Sequence[str] | None = None) -> None: target_loss_basis=target_loss_basis, exact_k_ladder=exact_k_ladder_provenance, calibration_runtime=calibration_runtime, + post_export_scoring=post_export_scoring_plan.record(), ) # Terminal-gate batching: evaluate EVERY terminal gate # group and raise once with the full failure list, instead of aborting at @@ -13009,6 +13897,9 @@ def _main(argv: Sequence[str] | None = None) -> None: # an evaluation crash in degraded mode records a line rather than masking # the earlier failures) and certification manifests are never written. terminal_gate_failures: list[str] = list(gate_failures) + # An unbuildable post-export plan refuses here, before the export write, + # with every other terminal group still evaluated (microcosm#956). + terminal_gate_failures.extend(post_export_scoring_plan.terminal_failures()) terminal_batch_telemetry = _TerminalBatchTelemetry( telemetry, terminal_gate_failures, @@ -13373,6 +14264,14 @@ def _main(argv: Sequence[str] | None = None) -> None: # microcosm#443: #437 dropped this call while inserting the batched raise, # so attempts 13/14 smoke-scored a stale artifact from a prior run. release_engine.write_dataset(export_frame, dataset_path, period=PERIOD) + # Route A remediation (microcosm#956): every post-export stage below (the + # smoke, reform validation, demographics) scores THIS file through one + # household-batched scorer instead of one whole-pool Microsimulation each. + # It loads the written H5 once, binds its sha256 (the manifest must pin + # the same bytes), and serves the baseline plans recorded above. + post_export_scorer = _open_post_export_scorer( + post_export_scoring_plan, dataset_path + ) # microcosm#368: reform-coverage smoke on the WRITTEN release H5. The column # gate above proves the required keys exist and carry signal; this is the # end-to-end backstop: each pinned probe (first: SSI asset limits at @@ -13388,8 +14287,12 @@ def _main(argv: Sequence[str] | None = None) -> None: "reform_coverage_smoke", message="Scoring pinned bound-reform probes on the written H5.", ) + smoke_scoring = post_export_scorer.open_consumer( + "reform_coverage_smoke", + post_export_scoring_plan.baseline_plan("reform_coverage_smoke"), + ) reform_coverage_smoke_gate = us_reform_coverage_smoke_gate( - simulate=default_simulate_factory(dataset_path), + simulate=smoke_scoring.simulate, period=PERIOD, ) reform_coverage_smoke_path = release_dir / "reform_coverage_smoke.json" @@ -13403,12 +14306,16 @@ def _main(argv: Sequence[str] | None = None) -> None: "failures": list(reform_coverage_smoke_gate.failures), "details": dict(reform_coverage_smoke_gate.details), }, + "post_export_scoring": smoke_scoring.record(), }, indent=2, sort_keys=True, ) + "\n" ) + # The smoke's scored baseline plan is recorded; free its arrays before + # reform validation scores its own. + del smoke_scoring if telemetry is not None: telemetry.attach_artifact( "reform_coverage_smoke", reform_coverage_smoke_path @@ -13469,6 +14376,12 @@ def _main(argv: Sequence[str] | None = None) -> None: registry=registry, release_id=release_id, simulate_out_of_sample=not args.skip_out_of_sample_reforms, + post_export_scorer=post_export_scorer, + baseline_plan=( + None + if args.skip_out_of_sample_reforms + else post_export_scoring_plan.baseline_plan("reform_validation") + ), ) if telemetry is not None: telemetry.attach_artifact( @@ -13483,9 +14396,15 @@ def _main(argv: Sequence[str] | None = None) -> None: release_dir=release_dir, dataset_path=dataset_path, release_id=release_id, + post_export_scorer=post_export_scorer, + baseline_plan=post_export_scoring_plan.baseline_plan("demographics"), ) if telemetry is not None: telemetry.attach_artifact("demographics", release_dir / "demographics.json") + # Every post-export engine stage has run; free the scorer's batch frames. + # The manifest must pin the bytes those stages scored. + scored_dataset_sha256 = _close_post_export_scorer(post_export_scorer) + post_export_scorer = None if telemetry is not None: telemetry.stage( @@ -13658,6 +14577,7 @@ def _main(argv: Sequence[str] | None = None) -> None: skipped_gates=( ("reform_coverage_smoke",) if args.skip_reform_coverage_smoke else () ), + scored_dataset_sha256=scored_dataset_sha256, ) if telemetry is not None: telemetry.attach_artifact("build_manifest", release_dir / "build_manifest.json") diff --git a/tools/sweep_us_post_export_scoring.py b/tools/sweep_us_post_export_scoring.py new file mode 100644 index 0000000000..11b2bf6aae --- /dev/null +++ b/tools/sweep_us_post_export_scoring.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +"""Check shipped post-export requests on a small written US release H5. + +Run sequential workers with four reform passes each by default, so their reform +systems are discarded when each worker exits. This is a fixture guard sweep; +the release-frame memory and timing rehearsal remains a separate operation. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import resource +import subprocess +import sys +import tempfile +import threading +import time +from pathlib import Path + + +def _load_builder(): + path = Path(__file__).with_name("build_us_fiscal_refresh_release.py") + spec = importlib.util.spec_from_file_location("post_export_sweep_builder", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _record_requests(builder): + """Record OBBBA stacks and validation's missing-estimate fallback passes.""" + specs = builder.load_default_reform_specs(period=builder.PERIOD) + + def validation(simulate): + return builder.reform_validation_payload( + specs, + period=builder.PERIOD, + simulate=simulate, + baseline_levels=builder.default_baseline_level_specs(), + release_id="post-export-guard-sweep", + ) + + consumers = { + "reform_coverage_smoke": builder._reform_coverage_smoke_consumer, + "reform_validation": validation, + "demographics": builder._demographics_consumer, + } + plans, requests = {}, [] + for name, consume in consumers.items(): + baseline = {} + + class RecordingSimulation: + def __init__(self, reform, *, name=name, baseline=baseline): + self.reform = reform + self.name = name + self.baseline = baseline + self.keys = set() + + def calculate(self, variable, period=None, map_to=None): + key = builder._post_export_key(variable, period, map_to) + if self.reform is None: + self.baseline.setdefault(key, None) + elif key not in self.keys: + requests.append((self.name, self.reform, key)) + self.keys.add(key) + return builder._RecordingPostExportSimulation(None).calculate(*key) + + consume(RecordingSimulation) + plans[name] = builder._ascending_period_plan(baseline) + counts = { + "smoke_probes": sum(name == "reform_coverage_smoke" for name, _, _ in requests), + "validation_specs": sum(not spec.in_sample for spec in specs), + "in_sample_fallback_specs": sum(spec.in_sample for spec in specs), + "validation_reform_passes": sum( + name == "reform_validation" for name, _, _ in requests + ), + "baseline_keys": {name: len(plan) for name, plan in plans.items()}, + } + return plans, requests, counts + + +def _assert_same_values(observed, expected, key): + import numpy as np + + np.testing.assert_array_equal(observed.weights, expected.weights, err_msg=str(key)) + values, reference = np.asarray(observed), np.asarray(expected) + if values.dtype.kind in "biuf": + np.testing.assert_allclose( + values, reference, rtol=1e-12, atol=1e-6, err_msg=str(key) + ) + else: + np.testing.assert_array_equal(values, reference, err_msg=str(key)) + + +def _score_baselines(builder, scorer, path, plans): + from microcosm.build.us_runtime.reform_validation import default_simulate_factory + + whole = default_simulate_factory(path)(None) + try: + engine = builder._AscendingPeriodEngine( + whole, label="sweep whole-file baseline" + ) + reference = { + key: builder._PostExportValues( + *builder._post_export_values_and_weights(engine.calculate(key), key) + ) + for key in builder._ascending_period_plan( + key for plan in plans.values() for key in plan + ) + } + finally: + builder.release_engine_simulation(whole) + for name, plan in plans.items(): + consumer = scorer.open_consumer(name, plan) + for key in plan: + _assert_same_values( + consumer.simulate(None).calculate(*key), reference[key], key + ) + + +def _comparison_indices(requests): + """Compare the first two reform passes of each consumer to a whole file.""" + seen, selected = {}, set() + for index, (name, _, _) in enumerate(requests): + seen[name] = seen.get(name, 0) + 1 + if seen[name] <= 2: + selected.add(index) + return selected + + +def _score_reforms(builder, scorer, path, requests, start, count): + from microcosm.build.us_runtime.reform_validation import default_simulate_factory + + consumer = scorer.open_consumer("shipped-reform-sweep", ()) + comparisons = _comparison_indices(requests) + compared = [] + for index in range(start, min(start + count, len(requests))): + name, reform, key = requests[index] + observed = consumer.simulate(reform).calculate(*key) + if index in comparisons: + whole = default_simulate_factory(path)(reform) + try: + expected = whole.calculate(key[0], key[1], map_to=key[2]) + _assert_same_values(observed, expected, (name, index, key)) + finally: + builder.release_engine_simulation(whole) + compared.append(index) + print(f"Guard sweep pass {index + 1}/{len(requests)}: {name} {key}", flush=True) + return {"reform_passes": consumer.reform_passes, "whole_file_comparisons": compared} + + +def _peak_rss_bytes(): + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return int(peak if sys.platform == "darwin" else peak * 1024) + + +def _watch_worker_rss(): + def watch(): + while True: + if _peak_rss_bytes() > 7 * 1024**3: + try: + print( + "Guard sweep worker exceeded 7 GiB RSS", + file=sys.stderr, + flush=True, + ) + finally: + os._exit(2) + time.sleep(0.1) + + threading.Thread(target=watch, daemon=True).start() + + +def _worker(args): + _watch_worker_rss() + started = time.monotonic() + builder = _load_builder() + plans, requests, counts = _record_requests(builder) + scorer = builder._HouseholdBatchedPostExportScorer( + args.dataset_path, maximum_microsim_batch_size=args.batch_size + ) + if scorer.n_batches < 2: + raise ValueError("The guard sweep needs at least two household batches") + try: + if args.worker_start is None: + _score_baselines(builder, scorer, args.dataset_path, plans) + report = {**counts, "total_reform_passes": len(requests)} + else: + report = _score_reforms( + builder, + scorer, + args.dataset_path, + requests, + args.worker_start, + args.reforms_per_worker, + ) + report["n_batches"] = scorer.n_batches + report["n_households"] = scorer.n_households + report["dataset_sha256"] = scorer.dataset_sha256 + report["peak_rss_bytes"] = _peak_rss_bytes() + report["elapsed_seconds"] = time.monotonic() - started + return report + finally: + scorer.close() + + +def _run_workers(args, *, run=subprocess.run): + started = time.monotonic() + with tempfile.TemporaryDirectory( + prefix="post-export-guard-sweep-", dir=args.output.parent + ) as temporary: + output = Path(temporary) / "worker.json" + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--dataset-path", + str(args.dataset_path), + "--batch-size", + str(args.batch_size), + "--reforms-per-worker", + str(args.reforms_per_worker), + "--worker", + "--output", + str(output), + ] + run(command, check=True) + report = json.loads(output.read_text()) + report["workers"] = [ + { + "baseline": True, + "peak_rss_bytes": report["peak_rss_bytes"], + "elapsed_seconds": report["elapsed_seconds"], + } + ] + report["scored_reform_passes"] = 0 + report["whole_file_comparisons"] = [] + for start in range(0, report["total_reform_passes"], args.reforms_per_worker): + run([*command, "--worker-start", str(start)], check=True) + worker = json.loads(output.read_text()) + report["workers"].append({"reform_start": start, **worker}) + report["scored_reform_passes"] += worker["reform_passes"] + report["whole_file_comparisons"].extend(worker["whole_file_comparisons"]) + assert report["scored_reform_passes"] == report["total_reform_passes"] + report["peak_rss_bytes"] = max( + worker["peak_rss_bytes"] for worker in report["workers"] + ) + report["elapsed_seconds"] = time.monotonic() - started + return report + + +def _positive_integer(value): + number = int(value) + if number < 1: + raise argparse.ArgumentTypeError("must be positive") + return number + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-path", type=Path, required=True) + parser.add_argument("--batch-size", type=_positive_integer, default=4) + parser.add_argument( + "--reforms-per-worker", type=_positive_integer, choices=range(1, 5), default=4 + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--worker-start", type=int, help=argparse.SUPPRESS) + args = parser.parse_args() + args.output.parent.mkdir(parents=True, exist_ok=True) + report = _worker(args) if args.worker else _run_workers(args) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, sort_keys=True), flush=True) + + +if __name__ == "__main__": + main() From 8142e8159db6bbeb834b70769810cff2ff01a152 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 24 Sep 2026 19:46:51 -0400 Subject: [PATCH 2/5] Keep the frame-cleanup tests engine-free for the fast CI groups The two frame-cleanup tests reach _assert_no_formula_owned_columns through the export helpers, which reads the installed policyengine-us metadata index; the fast CI groups run without policyengine-us, so five cases failed there with ImportError. Formula ownership is not under test here, so stub it; reproduced the failure locally with policyengine-us metadata hidden, and the fast subset passes (33) under the same condition. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_post_export_scoring.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_post_export_scoring.py b/packages/microcosm-build/tests/test_us_post_export_scoring.py index baf1f205fb..38517df1a6 100644 --- a/packages/microcosm-build/tests/test_us_post_export_scoring.py +++ b/packages/microcosm-build/tests/test_us_post_export_scoring.py @@ -1559,9 +1559,13 @@ def test_calibration_result_holds_no_reference_to_the_target_frame() -> None: def test_calibration_results_drop_their_frames_and_keep_everything_else( - builder, + builder, monkeypatch ) -> None: """Dropping target tables preserves diagnostics, estimates and weights.""" + # The export helpers check formula ownership against the installed engine's + # metadata index; that check is not under test here, and the fast CI groups + # run without policyengine-us. + monkeypatch.setattr(builder, "_assert_no_formula_owned_columns", lambda frame: None) from microcosm.build.us_runtime.exact_k_ladder import ( ExactKLadderCalibration, assert_exact_k_realized_count, @@ -1670,9 +1674,11 @@ def test_main_drops_the_calibration_frames_before_the_export(builder) -> None: @pytest.mark.parametrize("full_pool", [True, False]) @pytest.mark.parametrize("exact_k", [True, False]) def test_main_has_no_calibrated_target_frame_at_the_export_write( - builder, full_pool, exact_k + builder, full_pool, exact_k, monkeypatch ) -> None: """Execute the export/cleanup statements with real calibration results.""" + # As above: formula ownership needs the engine and is not under test. + monkeypatch.setattr(builder, "_assert_no_formula_owned_columns", lambda frame: None) from microcosm.calibrate import calibrate_l0_refit target_frame, targets = _small_calibration_problem() From c743ff1bb498961ac6dd8cdb761dbc1fb490335a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 24 Sep 2026 21:12:59 -0400 Subject: [PATCH 3/5] Score PR-3's green-register smoke through the batched scorer PR-3's green-register test stubbed default_simulate_factory, which the batched post-export scorer replaces in _main. Stand in for the scorer instead: it binds the written file's sha256, the smoke scores that file, and its post_export_scoring record pins the bytes the build manifest pins. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_fiscal_refresh_builder.py | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py b/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py index b8ea501d68..e92dc2a628 100644 --- a/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py +++ b/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py @@ -5977,9 +5977,29 @@ def fake_load_us_frame(path): return reference_frame monkeypatch.setattr(builder, "load_us_frame", fake_load_us_frame) - monkeypatch.setattr(builder, "default_simulate_factory", lambda path: path) + + class WrittenH5Scorer: + """The household-batched post-export scorer (#956) without an engine: + it binds the sha256 of the file it opens, as the real one does, and + hands each consumer a seam naming that file.""" + + def __init__(self, dataset_path, **kwargs): + self.dataset_path = Path(dataset_path) + self.dataset_sha256 = builder._sha256(self.dataset_path) + captured["scorer_opened_on"] = self.dataset_path + + def open_consumer(self, name, baseline_plan): + record = {"dataset_sha256": self.dataset_sha256, "consumer": name} + return SimpleNamespace(simulate=self.dataset_path, record=lambda: record) + + def close(self): + captured["scorer_closed"] = True + + monkeypatch.setattr(builder, "_HouseholdBatchedPostExportScorer", WrittenH5Scorer) def fake_smoke(*, simulate, period): + # The pre-export plan dry-runs the smoke on a recording seam first; + # the last call is the gate scoring the written release. captured["smoke_scored"] = simulate return builder.GateResult( name="reform_coverage_smoke", @@ -6018,12 +6038,27 @@ def fake_smoke(*, simulate, period): "estate_income": "donor tail concentrated before calibration" } if skipped_smoke: + # No post-export stage runs, so no scorer opens. assert "smoke_scored" not in captured + assert "scorer_opened_on" not in captured assert not (release_dir / "reform_coverage_smoke.json").exists() else: + assert captured["scorer_opened_on"] == captured["written_dataset"] assert captured["smoke_scored"] == captured["written_dataset"] + assert captured["scorer_closed"] is True build_manifest = json.loads((release_dir / "build_manifest.json").read_text()) release_manifest = json.loads((release_dir / "release_manifest.json").read_text()) + if not skipped_smoke: + # The smoke verdict names the bytes it scored, and they are the bytes + # the manifest pins. + smoke_scoring = json.loads( + (release_dir / "reform_coverage_smoke.json").read_text() + )["post_export_scoring"] + assert smoke_scoring == { + "dataset_sha256": hashlib.sha256(b"release h5").hexdigest(), + "consumer": "reform_coverage_smoke", + } + assert build_manifest["dataset"]["sha256"] == smoke_scoring["dataset_sha256"] artifacts = release_manifest["artifacts"] bound = dict(builder.US_RELEASE_GATE_EVIDENCE_FILES) From 37f09a824566e2ffad92f23df6f24af8b02e9d61 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 25 Sep 2026 04:20:05 -0400 Subject: [PATCH 4/5] Carry the post-export scoring block in both manifests Route A remediation PR-3 records the evidence the release gates judged in build_manifest.json and release_manifest.json's build section, not only in loose diagnostics. The batched scorer's records reached only the calibration diagnostics and the smoke evidence file. The scorer now keeps each finished consumer's record (records, not the consumers, so their arrays are still freed). Before closing it, _main takes a post_export_scoring block: the method, the scored sha256, the batching and each consumer's baseline plan and pass counts. That block rides both manifests beside calibration_runtime. _build_manifests refuses a block whose sha256 is not the pinned dataset's. With no post-export stage there is no scorer and no block. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_fiscal_refresh_builder.py | 26 ++++++- .../tests/test_us_post_export_scoring.py | 66 +++++++++++++++- tools/build_us_fiscal_refresh_release.py | 78 ++++++++++++++++++- 3 files changed, 164 insertions(+), 6 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py b/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py index e92dc2a628..ab63415104 100644 --- a/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py +++ b/packages/microcosm-build/tests/test_us_fiscal_refresh_builder.py @@ -5990,7 +5990,19 @@ def __init__(self, dataset_path, **kwargs): def open_consumer(self, name, baseline_plan): record = {"dataset_sha256": self.dataset_sha256, "consumer": name} - return SimpleNamespace(simulate=self.dataset_path, record=lambda: record) + return SimpleNamespace( + name=name, simulate=self.dataset_path, record=lambda: record + ) + + def finish_consumer(self, scoring): + captured.setdefault("finished_consumers", []).append(scoring.name) + return scoring.record() + + def manifest_record(self): + return { + "dataset_sha256": self.dataset_sha256, + "consumers": list(captured.get("finished_consumers", [])), + } def close(self): captured["scorer_closed"] = True @@ -6059,6 +6071,18 @@ def fake_smoke(*, simulate, period): "consumer": "reform_coverage_smoke", } assert build_manifest["dataset"]["sha256"] == smoke_scoring["dataset_sha256"] + # Both manifests carry how the post-export gates were scored (route A + # remediation PR-3's rule): the scorer's block, taken after every consumer + # finished and naming the bytes the manifest pins. With no post-export + # stage there is no scorer and no block. + for block in (build_manifest, release_manifest["build"]): + if skipped_smoke: + assert "post_export_scoring" not in block + else: + assert block["post_export_scoring"] == { + "dataset_sha256": build_manifest["dataset"]["sha256"], + "consumers": ["reform_coverage_smoke"], + } artifacts = release_manifest["artifacts"] bound = dict(builder.US_RELEASE_GATE_EVIDENCE_FILES) diff --git a/packages/microcosm-build/tests/test_us_post_export_scoring.py b/packages/microcosm-build/tests/test_us_post_export_scoring.py index 38517df1a6..8bf76509cc 100644 --- a/packages/microcosm-build/tests/test_us_post_export_scoring.py +++ b/packages/microcosm-build/tests/test_us_post_export_scoring.py @@ -1258,7 +1258,36 @@ def calls_to(name: str) -> list[ast.Call]: assert {kw.arg: ast.unparse(kw.value) for kw in diagnostics_call.keywords}[ "post_export_scoring" ] == "post_export_scoring_plan.record()" - assert '"post_export_scoring": smoke_scoring.record()' in source + # The smoke's record is registered with the scorer before its evidence is + # written, and the evidence carries that same record. + assert ( + "smoke_scoring_record = post_export_scorer.finish_consumer(smoke_scoring)" + in source + ) + assert '"post_export_scoring": smoke_scoring_record' in source + # Both manifests get the scorer's block, taken after every consumer ran + # and before the scorer closes. + block_assignments = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "post_export_scoring" + for target in node.targets + ) + ] + assert [ast.unparse(node.value) for node in block_assignments] == [ + "_post_export_scoring_manifest_block(post_export_scorer)" + ] + assert ( + demographics_call.lineno + < block_assignments[0].lineno + < sha_assignments[0].lineno + < manifest_call.lineno + ) + assert {kw.arg: ast.unparse(kw.value) for kw in manifest_call.keywords}[ + "post_export_scoring" + ] == "post_export_scoring" # The plan is recorded (never raised) before the diagnostics, and its own # failure joins the terminal batch ahead of the export write. recorded = source.index( @@ -1331,8 +1360,30 @@ def test_main_opens_the_scorer_its_plan_recorded( assert scorer.dataset_path == path assert scorer.n_batches == 3 and scorer.max_batch_households == 3 # Served through the engine's own Microsimulation, one batch at a time. - scorer.open_consumer("demographics", plan.baseline_plan("demographics")) + demographics = scorer.open_consumer( + "demographics", plan.baseline_plan("demographics") + ) assert [len(engine.household_ids) for engine in log.constructions] == [3, 3, 1] + # Only finished consumers reach the manifests' block, once each. + assert builder._post_export_scoring_manifest_block(scorer)["consumers"] == {} + finished = scorer.finish_consumer(demographics) + with pytest.raises(ValueError, match="already finished"): + scorer.finish_consumer(demographics) + block = builder._post_export_scoring_manifest_block(scorer) + assert block["dataset_sha256"] == builder._sha256(path) + assert (block["n_batches"], block["max_batch_households"]) == (3, 3) + assert block["consumers"] == { + "demographics": { + key: finished[key] + for key in ( + "baseline_plan", + "baseline_passes", + "reform_passes", + "reform_systems", + ) + } + } + assert builder._post_export_scoring_manifest_block(None) is None assert builder._close_post_export_scorer(scorer) == builder._sha256(path) with pytest.raises(RuntimeError, match="scorer is closed"): scorer.open_consumer("demographics", plan.baseline_plan("demographics")) @@ -1807,6 +1858,17 @@ def test_scored_sha_is_bound_at_load_and_checked_by_the_manifest( dataset_filename=path.name, scored_dataset_sha256="0" * 64, ) + # The manifests' scoring block must describe the same bytes. + with pytest.raises(RuntimeError, match="does not pin"): + builder._build_manifests( + **required, + dataset_filename=path.name, + post_export_scoring={"dataset_sha256": "0" * 64}, + ) + # A consumer opened on another scorer cannot be finished on this one. + other = _scorer(builder, _nested_frame(), _EngineLog(), 3, tmp_path / "other") + with pytest.raises(ValueError, match="not opened on this scorer"): + scorer.finish_consumer(other.open_consumer("fixture", ())) _, tree = _function_source(builder, "_main") manifest_call = next( diff --git a/tools/build_us_fiscal_refresh_release.py b/tools/build_us_fiscal_refresh_release.py index 1fce8e3a6a..056fa2e7ad 100644 --- a/tools/build_us_fiscal_refresh_release.py +++ b/tools/build_us_fiscal_refresh_release.py @@ -4917,6 +4917,9 @@ def dataset_from_frame(batch_frame: Frame): ) self._dataset_from_frame = dataset_from_frame + # Each finished consumer's scoring record, for the manifests' block. + # Records, not consumers: a finished consumer's arrays must be freed. + self.consumer_records: dict[str, dict[str, object]] = {} def open_consumer( self, name: str, baseline_plan: Sequence[PostExportKey] @@ -4924,6 +4927,46 @@ def open_consumer( """Score ``baseline_plan`` now; return the consumer's simulate seam.""" return _PostExportConsumer(self, name, tuple(baseline_plan)) + def finish_consumer(self, scoring: _PostExportConsumer) -> dict[str, object]: + """Record a consumer that has finished scoring and return its record.""" + if scoring._scorer is not self: + raise ValueError(f"{scoring.name} was not opened on this scorer.") + if scoring.name in self.consumer_records: + raise ValueError(f"{scoring.name} already finished on this scorer.") + record = scoring.record() + self.consumer_records[scoring.name] = record + return record + + def manifest_record(self) -> dict[str, object]: + """The ``post_export_scoring`` block both manifests carry. + + Route A remediation PR-3 records the evidence the gates judged in the + manifests, not only in loose diagnostics. This block names the bytes + scored, the batching, and each finished consumer's baseline plan and + pass counts, in the order the consumers finished. + """ + return { + "method": POST_EXPORT_SCORING_METHOD, + "dataset_sha256": self.dataset_sha256, + "maximum_batch_size": self.maximum_batch_size, + "n_households": self.n_households, + "n_batches": self.n_batches, + "max_batch_households": self.max_batch_households, + "spm": dict(US_RELEASE_SPM_SELECTION), + "consumers": { + name: { + key: record[key] + for key in ( + "baseline_plan", + "baseline_passes", + "reform_passes", + "reform_systems", + ) + } + for name, record in self.consumer_records.items() + }, + } + def close(self) -> None: self._batch_frames = None _collect_family_garbage() @@ -5054,6 +5097,15 @@ def _open_post_export_scorer( ) +def _post_export_scoring_manifest_block( + scorer: _HouseholdBatchedPostExportScorer | None, +) -> dict[str, object] | None: + """The manifests' ``post_export_scoring`` block (None: no stage ran).""" + if scorer is None: + return None + return scorer.manifest_record() + + def _close_post_export_scorer( scorer: _HouseholdBatchedPostExportScorer | None, ) -> str | None: @@ -9821,7 +9873,7 @@ def _score_post_export_consumer( ) scoring = scorer.open_consumer(name, plan) output = consumer(scoring.simulate) - record = scoring.record() + record = scorer.finish_consumer(scoring) print( f"Post-export {name}: {record['baseline_passes']} baseline pass " f"({len(plan)} keys), {record['reform_passes']} reform passes, " @@ -9960,6 +10012,7 @@ def _build_manifests( calibration_runtime: Mapping[str, object] | None = None, skipped_gates: Iterable[str] = (), scored_dataset_sha256: str | None = None, + post_export_scoring: Mapping[str, object] | None = None, ) -> None: dataset_path = artifact_root / dataset_filename calibration_path = artifact_root / calibration_filename @@ -9975,6 +10028,16 @@ def _build_manifests( f"{scored_dataset_sha256}, but the file now hashes to {dataset_sha}; " "the manifest would pin bytes the release gates never scored." ) + if ( + post_export_scoring is not None + and post_export_scoring.get("dataset_sha256") != dataset_sha + ): + raise RuntimeError( + "The post_export_scoring block records sha256 " + f"{post_export_scoring.get('dataset_sha256')}, but {dataset_path} " + f"hashes to {dataset_sha}; the manifest would describe scoring of " + "bytes it does not pin." + ) calibration_sha = _sha256(calibration_path) diagnostics_sha = _sha256(diagnostics_path) coverage_sha = _sha256(coverage_path) @@ -10004,6 +10067,11 @@ def _build_manifests( if calibration_runtime is not None else {} ), + **( + {"post_export_scoring": dict(post_export_scoring)} + if post_export_scoring is not None + else {} + ), "fiscal_target_exclusion_receipt": _fiscal_target_exclusion_receipt_reference( coverage_path, coverage_sha ), @@ -14295,6 +14363,7 @@ def _main(argv: Sequence[str] | None = None) -> None: simulate=smoke_scoring.simulate, period=PERIOD, ) + smoke_scoring_record = post_export_scorer.finish_consumer(smoke_scoring) reform_coverage_smoke_path = release_dir / "reform_coverage_smoke.json" reform_coverage_smoke_path.write_text( json.dumps( @@ -14306,7 +14375,7 @@ def _main(argv: Sequence[str] | None = None) -> None: "failures": list(reform_coverage_smoke_gate.failures), "details": dict(reform_coverage_smoke_gate.details), }, - "post_export_scoring": smoke_scoring.record(), + "post_export_scoring": smoke_scoring_record, }, indent=2, sort_keys=True, @@ -14402,7 +14471,9 @@ def _main(argv: Sequence[str] | None = None) -> None: if telemetry is not None: telemetry.attach_artifact("demographics", release_dir / "demographics.json") # Every post-export engine stage has run; free the scorer's batch frames. - # The manifest must pin the bytes those stages scored. + # The manifest must pin the bytes those stages scored, and both manifests + # carry how they were scored (Route A remediation PR-3). + post_export_scoring = _post_export_scoring_manifest_block(post_export_scorer) scored_dataset_sha256 = _close_post_export_scorer(post_export_scorer) post_export_scorer = None @@ -14578,6 +14649,7 @@ def _main(argv: Sequence[str] | None = None) -> None: ("reform_coverage_smoke",) if args.skip_reform_coverage_smoke else () ), scored_dataset_sha256=scored_dataset_sha256, + post_export_scoring=post_export_scoring, ) if telemetry is not None: telemetry.attach_artifact("build_manifest", release_dir / "build_manifest.json") From 5055805e3cb6f12fbc49bd220c6d4a85f3cf96ba Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 25 Sep 2026 05:05:16 -0400 Subject: [PATCH 5/5] Test that the scoring writers register with the shared scorer Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_post_export_scoring.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/packages/microcosm-build/tests/test_us_post_export_scoring.py b/packages/microcosm-build/tests/test_us_post_export_scoring.py index 8bf76509cc..bf611ae4ac 100644 --- a/packages/microcosm-build/tests/test_us_post_export_scoring.py +++ b/packages/microcosm-build/tests/test_us_post_export_scoring.py @@ -1881,6 +1881,28 @@ def test_scored_sha_is_bound_at_load_and_checked_by_the_manifest( assert keywords["scored_dataset_sha256"] == "scored_dataset_sha256" +def test_writers_register_their_consumer_with_the_shared_scorer( + builder, tmp_path +) -> None: + """(j) Reform validation and demographics score through + ``_score_post_export_consumer``; on ``_main``'s shared scorer each + registers its finished record, so the manifests' block names it.""" + scorer = _scorer(builder, _nested_frame(), _EngineLog(), 3, tmp_path) + output = builder._score_post_export_consumer( + "reform_validation", + lambda simulate: "payload", + dataset_path=scorer.dataset_path, + post_export_scorer=scorer, + baseline_plan=(), + maximum_microsim_batch_size=3, + ) + assert output == "payload" + assert list(scorer.consumer_records) == ["reform_validation"] + block = builder._post_export_scoring_manifest_block(scorer) + assert list(block["consumers"]) == ["reform_validation"] + assert block["consumers"]["reform_validation"]["reform_passes"] == 0 + + # --------------------------------------------------------------------------- # (k) the nesting premise # ---------------------------------------------------------------------------