Admit UK datasets without area codes only when local-area runs are routed, and guard the weight-matrix fallback

[juaristi22]

Fixes #725

Summary

policyengine.py may one day publish a UK default without area codes: a Microcosm national file whose local-area data come from a separate dataset (PolicyEngine/policyengine.py#553). Today every UK run on such a dataset fails, because the executor requires la_code_oa when it loads a UK dataset and when a UK worker starts.

This PR lets the executor accept such a dataset only when it is declared national-only, meaning the policyengine.py bundle routes both constituency and local-authority regions to another dataset. A missing column on its own is never read as "national-only", so a local-area dataset that lost its area codes still fails. The PR also stops the weight-matrix fallback for constituency and local-authority runs early, with a clear message, when the matrix does not fit the dataset.

With today's bundle (policyengine 6.2.1) nothing is routed, so the executor admits exactly what main admits.

Changes

• Routing declaration (uk_local_authority_metadata.py):
  • uk_area_regions_routed() is true when policyengine.py's UK region_datasets lists both constituency and local_authority. Its regional_dataset_defaults produce those entries (Use ACS-local data for US regional simulations policyengine.py#552).
  • uk_area_codes_required(region_code) is true for every constituency or local-authority run. For any other run it is true unless both area region types are routed.
• Per-run check: detect_uk_local_authority_metadata takes the run's region code.
  • It raises on a missing la_code_oa whenever codes are required, as main does.
  • It returns None only for a national or nation-level run with routing declared.
  • Codes that are present are validated as before.
• Worker start-up check (stage12_worker_validation.py): admits an installed default without la_code_oa only when routing is declared; otherwise it runs main's strict detection. detect_uk_local_authority_metadata_from_hdf is unchanged from main, so the smoke apps still read a boundary version off it.
• Breakdowns (simulation_output_geographic.py): build_uk_constituency_impact and build_uk_local_authority_impact raise a clear error when the simulated household output lacks constituency_code_oa or la_code_oa. They never drop the breakdown, and the constituency builder skips its GCS lookup in that case.
• Stage 12 (stage12_runtime/aggregation.py, coordination.py, stage12_qualification.py): validate_uk_local_authority_metadata now takes the report's region.
  • It accepts a baseline/reform pair that both lack metadata only for a national or nation-level request whose bundle routes both constituency and local-authority regions to another dataset. That is the same evidence per-run detection requires.
  • A missing region, metadata on one side only, or mismatched boundary versions still raise.
• Weight matrices bound to the certified release (simulation_runtime.py):
  • The constituency and local-authority regions the executor builds from the enhanced-FRS matrices no longer download the unversioned bucket copy (download_missing_assets=False). Every uk-data release replaces that object, so it is not tied to the dataset being run.
  • _require_certified_uk_weight_matrix runs after the dataset loads, in both runtimes. It takes the matrix the policyengine.py bundle certifies in the same release (repository and revision) as the selected dataset, and materializes it through policyengine.py's digest-checked materializer where the strategy looks first. It then confirms the strategy will read exactly that file.
  • It then checks the run year and the household dimension, and stops with a clear message if either is off. Before, these runs failed inside WeightReplacementStrategy.apply with a KeyError or an out-of-date-matrix error.

Behaviour changes to expect

• Constituency and local-authority results can move. Until now the executor read the matrices from the bucket's unversioned objects, which each uk-data release overwrites. The latest release, 1.57.4, has different matrices from the certified 1.56.16 ones: the same 650 × 52,846 shape, with households in a different order. If production has been reading the 1.57.4 objects, its constituency and local-authority results will change when this deploys, to the weights certified with the dataset it runs.
• These runs only work for 2025. The certified matrices hold weights for 2025 only, as do the 1.57.4 ones, so constituency and local-authority runs for any other year stop with a clear message. With those matrices they failed before too, with a KeyError inside WeightReplacementStrategy.apply.
• The first such run per worker downloads the matrix from Hugging Face with the worker's UK data credential: 137 MB for constituencies, 76 MB for local authorities. Later runs reuse the digest-checked file.

Not in this PR

These move to #728:

• Serving local results from the routed dataset: constituency and local-authority runs, and the breakdowns of national reports, from the enhanced FRS in the interim or from the Microcosm local-area dataset. Until then a routed national default would raise on those breakdowns rather than drop them, so the default does not switch before UK: keep constituency and local-authority results available through a switch to Microcosm, then serve them from its local-area dataset #728.
• Checking the routed dataset at worker start-up: today it is checked when a run loads it, because workers install only the default dataset.
• The loading cost of a large local-area file.

Testing

• Tests touched:
  • Updated: test_uk_local_authority_metadata.py, test_stage12_worker_validation.py, test_simulation_output_geographic.py, test_simulation_output_builder.py and test_stage12_runtime.py.
  • New: test_uk_weight_matrix_guard.py.
• Full executor suite (uv run pytest tests/, Python 3.13): 692 passed, 22 skipped. The 9 failures in test_policyengine_package_update_scripts.py also fail on main locally, because the system python3 used by those scripts lacks tomllib.
• Formatting and lint: ruff format --check passes on the changed src files, and ruff check adds no findings over main.
• End to end against the real bundle: the guard materialized the certified 1.56.16 constituency matrix (sha256 6870c360…) from Hugging Face and passed for the 52,846-household enhanced FRS. A 2026 run stopped with the missing-year message.
• Not run: the Docker build and the integration tests.

🤖 Generated with Claude Code

[anth-volk]

@juaristi22, could you confirm the intended UK Microcosm release sequence? My understanding from PolicyEngine/policyengine.py#553 and PolicyEngine/microcosm#1114 is:

1. First, release a national dataset with no la_code_oa or constituency_code_oa identifiers.
2. Later, release a separate local-area dataset containing those identifiers and add the routing needed for constituency and local-authority simulations.

If that is correct, switching the UK default after this PR but before the local-area dataset and routing are available would make constituency and local-authority simulations unavailable and omit those breakdowns from national reports. We do not want any interval in which those existing subnational capabilities are deactivated.

Could we revise the release plan so that subnational simulations remain available throughout—for example, by continuing to route those requests to the enhanced FRS until the Microcosm local-area release is ready, or by releasing the national dataset, local-area dataset, and routing together? I do not think silently returning no geographic output is an acceptable interim state.

[juaristi22]

Thanks @anth-volk. We don't necessarily intend to publish the national dataset separately from the local-area one. This PR is about having the flexibility in case that becomes necessary: with it, the executor no longer fails outright on a UK dataset without area codes. Nothing changes while the enhanced FRS is the default.

If that eventuality did arise, constituency and local-authority simulations, and the breakdowns in national reports, would route to the enhanced FRS local-area dataset in the meantime, so those subnational capabilities stay available throughout. We wouldn't switch the UK default to a dataset without area codes before that routing is in place. I've filed #728 to make this explicit, alongside the later move to the Microcosm local-area dataset.

[juaristi22]

@anth-volk, on your concern that letting a national dataset pass without local-area values would also let a local-area dataset with broken local-area handling through undetected: agreed. The first version inferred "this dataset is national-only" from the missing column, so a file that lost its area codes by mistake looked the same. I've pushed a rework that makes it a declaration instead:

• A UK dataset may lack area codes only when the policyengine.py bundle routes both constituency and local-authority regions to another dataset. Without that routing, the worker start-up check and every run reject it exactly as main does.
• Constituency and local-authority runs always require area codes on the dataset they run on, so a local-area dataset that lost its columns fails loudly. Codes that are present are still checked for unsupported or mixed boundary versions.
• Breakdowns are never dropped. Building a constituency or local-authority breakdown over a dataset without area codes raises a clear error. Building them from the routed dataset is part of UK: keep constituency and local-authority results available through a switch to Microcosm, then serve them from its local-area dataset #728.
• Today's bundle (policyengine 6.2.1) routes nothing, so the executor admits exactly what main admits.

One check is still open, and #728 tracks it: the routed dataset is validated when a run loads it, not yet at worker start-up, because workers only install the default dataset.

[anth-volk]

Two inline findings from the code review.

[anth-volk]

Matching the matrix year and household dimension does not establish that this matrix belongs to the selected dataset. WeightReplacementStrategy.apply assigns a matrix row positionally, so a matrix from another dataset revision or household ordering will pass this check whenever it has the same number of households and can attach weights to the wrong records. Please bind the matrix to the certified dataset revision or digest, or validate equivalent provenance that covers household ordering.

[juaristi22]

Agreed: the shape check couldn't tell releases apart. uk-data 1.57.4's constituency matrix is 650 × 52,846 like the certified 1.56.16 one, and its enhanced FRS has the same 52,846 households in a different order. Fixed in 1442197:

• No more bucket downloads. The fallback regions no longer download the unversioned bucket copy (download_missing_assets=False). Every uk-data release overwrites that object, so it was never bound to the certified dataset.
• The matrix comes from the certified release. _require_certified_uk_weight_matrix takes the matrix the bundle certifies in the same release (repository and revision) as the selected dataset. It materializes that matrix through policyengine.py's digest-checked materializer, where the strategy looks first, and confirms the strategy resolves exactly that file before the run. The year and household checks stay as a last sanity check.
• Checked end to end against the real bundle: it fetched the certified 1.56.16 constituency matrix (sha256 6870c360…) and passed for the 52,846-household enhanced FRS.

This surfaced two things, now noted in the PR description:

• Results may move on deploy. If production has been reading the 1.57.4 bucket objects (constituency matrix sha256 73944f69…), its constituency and local-authority results will change when this deploys. Someone with bucket access can confirm by comparing the object's digest with 6870c360….
• These runs only work for 2025. The certified matrices only cover 2025, so these runs fail for any other year. Before, that failure was a KeyError inside policyengine.py; now it's a clear message.

[anth-volk]

This accepts every pair of UK artifacts whose boundary metadata is absent, but the coordinator can deserialize retained artifacts without rerunning detect_uk_local_authority_metadata. An invalid local-area artifact pair therefore passes this validation and fails only later during output construction. Please require evidence here that the request is national or nation-level and that both local-area dataset mappings exist, either by passing that context into this validator or encoding it in the artifact contract.

[juaristi22]

Done in 1442197.

• The check now sees the report's region. validate_uk_local_authority_metadata takes report.baseline.geography.region, passed in by the coordinator and by the qualification path.
• What a UK pair without metadata needs: the region must be national or nation-level, and the bundle must route both constituency and local-authority regions to another dataset. That is the same rule detect_uk_local_authority_metadata applies when the simulation runs.
• A missing region doesn't count as evidence.