diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 5fc93bea0..ce4676bc5 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -16,6 +16,7 @@ - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: + - Split the artifact precompute (precompute.py, 931 lines) and the Stage 12 cache precompute (stage12_cache/precompute.py, 844 lines) each into planning, dataset-building and baseline modules, with the manifest and run_precompute left in the original module, which re-exports every public name, so the Modal precompute apps and their tests import from the same paths; the cache precompute's _scoped_planned_simulation and _baseline_identity are now public, and the worker validation imports scoped_planned_simulation from its new module; the 1,835-line precompute test file is split into four, and the cache precompute tests patch each name in the module whose code reads it - Split the 4,709-line Stage 12 runtime test file into thirteen test files by behaviour (dataset loading, provenance, the single worker and its failures, aggregates, the coordinator and its failures, runtime store retries, failure records, cancellation, child state errors and the wait window) and three support modules for the shared builders, fake stores and fake child calls, whose helpers are now public; the other Stage 12 tests import them from those modules, every test keeps its name and body, and two parametrized cases are named after the renamed bundle helper - Split Stage 12 report coordination (stage12_runtime/coordination.py, 1,091 lines) into modules for the report coordinator class, its child simulations' records, failures and cancellation, and the parent report's checks and claim, with coordinate_report and every public name still in coordination.py and child failures still logged under its logger; and turned the contract's stage12_execution.py (929 lines) into a stage12_execution package of modules for the shared field types and simulation input, report inputs and output plans, result artifacts, and comparison-run records, which re-exports every name from the same import path, with unchanged JSON schemas; models subclass the base models through the package path ruff.toml lists, so no ruff setting changes - Split the gateway's endpoints.py (1,454 lines) into modules for the shared router and request helpers, route resolution from the active routing state, Modal reads of job metadata and routing, bundle payloads, and the job submission, job status and version routes, which register their routes in the same order, so the OpenAPI document is unchanged; endpoints.py re-exports every public name, the routes keep logging under the policyengine_simulation_gateway.endpoints logger, and the private helpers the new modules share are now public; its 1,801-line test file is split into five test files, with the shared expected-response helpers in fixtures/gateway_responses.py diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 71b7fbb04..7f8323b2c 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -66,7 +66,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` -- `sim.quality.python-file-length`: 7 remaining in `rules/baselines/python-file-length.txt` +- `sim.quality.python-file-length`: 4 remaining in `rules/baselines/python-file-length.txt` ### Exceptions and review dates diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 50d0561f6..512ec9498 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -28,30 +28,30 @@ identity collection and store I/O must run inside containers: the bundle receipt that the keys digest lives in the image, not on the deploying machine. + +The pipeline's steps live in precompute_planning, precompute_datasets and +precompute_baselines. This module runs them, builds and publishes the +manifest, and re-exports their public names. """ from __future__ import annotations -import gc import logging -import tempfile -import time -from dataclasses import dataclass, field -from pathlib import Path -from typing import TYPE_CHECKING, Any - -from policyengine_simulation_executor.artifact_keys import ( - BaselineArtifactIdentity, - baseline_artifact_filename, +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from policyengine_simulation_executor.precompute_baselines import ( + compute_baseline_impl, + verify_determinism_impl, +) +from policyengine_simulation_executor.precompute_datasets import ( + build_dataset_impl, + build_state_datasets_impl, ) -from policyengine_simulation_executor.output_frames import EntityFrames from policyengine_simulation_executor.precompute_models import ( ArtifactManifest, BaselineComputeResult, - BaselinePlanEntry, - BundleVersionIdentity, DatasetBuildResult, - DatasetPlanEntry, DeterminismVerdict, ManifestArtifact, ManifestEntry, @@ -59,32 +59,46 @@ RemoteFunction, StateDatasetArtifact, StateDatasetBuildResult, - StateDatasetPlanEntry, - WorkSelection, ) -from policyengine_simulation_executor.state_datasets import ( - state_dataset_relative_folder, - year_file_reuse_problem, +from policyengine_simulation_executor.precompute_planning import ( + PRECOMPUTE_COUNTRY, + PRECOMPUTE_YEARS, + StateDatasetTarget, + cohort_identity, + cohort_params, + plan_artifacts_impl, + select_work, + state_dataset_targets, ) if TYPE_CHECKING: from collections.abc import Callable - from types import ModuleType - from microdf import MicroDataFrame - from policyengine.core import Simulation - from policyengine.core.scoping_strategy import RowFilterStrategy - - from policyengine_simulation_executor.artifact_store import ArtifactStore - from policyengine_simulation_executor.simulation_runtime import RegionResolution +__all__ = [ + "MANIFEST_DIGEST_PREFIX", + "MANIFEST_SCHEMA", + "PRECOMPUTE_COUNTRY", + "PRECOMPUTE_YEARS", + "PrecomputeFunctions", + "StateDatasetTarget", + "build_dataset_impl", + "build_manifest", + "build_state_datasets_impl", + "cohort_identity", + "cohort_params", + "compute_baseline_impl", + "logger", + "manifest_digest_line", + "plan_artifacts_impl", + "publish_manifest_impl", + "run_precompute", + "select_work", + "state_dataset_targets", + "verify_determinism_impl", +] logger = logging.getLogger(__name__) -PRECOMPUTE_COUNTRY = "us" -# The user-facing priority order (2026 first); with parallel waves it only -# orders spawn submission, but keep it intentional. -PRECOMPUTE_YEARS = [2026, 2027, 2025] - MANIFEST_SCHEMA = "mf1" # The stdout contract the CI deploy job parses; keep format changes in @@ -97,155 +111,6 @@ def manifest_digest_line(digest: str) -> str: return f"{MANIFEST_DIGEST_PREFIX}{digest}" -def cohort_params(year: int, group: list[str]) -> dict: - """Return the synthetic request a segmented-national child would receive.""" - return { - "country": PRECOMPUTE_COUNTRY, - "scope": "macro", - "region_group": list(group), - "time_period": year, - } - - -def _cohort_resolution( - year: int, group: list[str] -) -> tuple[dict[str, Any], ModuleType, RegionResolution]: - """Cohort request params, country module, and region resolution. - - The planner's identity (``cohort_identity``) and the worker's - simulation (``_prepare_cohort_baseline``) must derive from ONE - resolution or the in-container id checks abort the run; sharing this - prologue removes the drift surface. - """ - from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _country_module, - _resolve_region, - ) - - params = cohort_params(year, group) - country_module = _country_module(PRECOMPUTE_COUNTRY) - resolution = _resolve_region( - country_module=country_module, - country=PRECOMPUTE_COUNTRY, - params=params, - ) - return params, country_module, resolution - - -def cohort_identity(year: int, group: list[str]) -> BaselineArtifactIdentity: - """Identity via the executor's own resolution path (writer==reader).""" - from policyengine_simulation_executor.baseline_artifacts import ( - BaselineIdentityRequest, - qualifying_baseline_identity, - ) - from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _resolve_dataset_selection, - ) - - params, _, resolution = _cohort_resolution(year, group) - dataset_selection = _resolve_dataset_selection(params, region_resolution=resolution) - identity = qualifying_baseline_identity( - BaselineIdentityRequest( - params, - country=PRECOMPUTE_COUNTRY, - dataset_is_default=dataset_selection.is_default, - policy=None, - region_code=resolution.code, - scoping_strategy=resolution.scoping_strategy, - year=year, - ) - ) - if identity is None: - msg = ( - f"Cohort request for year {year} group {group} does not qualify " - "for a deterministic baseline id — writer and reader predicates " - "have diverged." - ) - raise RuntimeError(msg) - return identity - - -@dataclass(frozen=True) -class StateDatasetTarget: - """A US state whose requests load a non-default dataset, and its row filter. - - The precompute bakes this state's rows of ``dataset`` as their own year - file: the households whose ``filter_variable`` equals ``state_fips``. - ``scoping_strategy`` is the region's own .py strategy, the one a request - for the state applies, so a baked file holds exactly the rows the request - would keep. - """ - - region: str - state_code: str - state_fips: int - filter_variable: str - dataset: str - scoping_strategy: RowFilterStrategy | None = field( - default=None, compare=False, repr=False - ) - - -def state_dataset_targets() -> list[StateDatasetTarget]: - """Return every US state region that loads a non-default dataset, by state code. - - Each state is resolved exactly as a request for it is (``_resolve_region`` - then ``_resolve_dataset_selection``), so a bundle that sends states back to - the default dataset gets no state files. - """ - from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _country_module, - _resolve_dataset_selection, - _resolve_region, - ) - - country_module = _country_module(PRECOMPUTE_COUNTRY) - targets = [] - for region in country_module.model.region_registry.get_by_type("state"): - params = { - "country": PRECOMPUTE_COUNTRY, - "region": region.code, - "time_period": PRECOMPUTE_YEARS[0], - } - resolution = _resolve_region( - country_module=country_module, country=PRECOMPUTE_COUNTRY, params=params - ) - selection = _resolve_dataset_selection(params, region_resolution=resolution) - if selection.is_default: - continue - strategy = resolution.scoping_strategy - variable_name = getattr(strategy, "variable_name", None) - variable_value = getattr(strategy, "variable_value", None) - if not variable_name or variable_value is None or not resolution.state_code: - msg = ( - f"State region {resolution.code} loads {selection.name} but has no " - "row filter and state code to build its state file from" - ) - raise RuntimeError(msg) - targets.append( - StateDatasetTarget( - region=resolution.code, - state_code=resolution.state_code, - state_fips=int(variable_value), - filter_variable=variable_name, - dataset=selection.name, - scoping_strategy=strategy, - ) - ) - return sorted(targets, key=lambda target: target.state_code) - - -def select_work(plan: PrecomputePlan, *, force: bool) -> WorkSelection: - """Which plan entries need computing.""" - return WorkSelection( - datasets=[entry for entry in plan.datasets if force or not entry.exists], - baselines=[entry for entry in plan.baselines if force or not entry.exists], - state_datasets=[ - entry for entry in plan.state_datasets if force or not entry.exists - ], - ) - - def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: """Build the deploy manifest: the receipt and each artifact the image must fetch. @@ -294,526 +159,6 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: ) -def plan_artifacts_impl(bucket: str) -> PrecomputePlan: - """Compute every expected artifact identity and its store presence.""" - from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity - collect_dataset_identity, - collect_state_dataset_identity, - ) - from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore - ArtifactStore, - ) - from policyengine_simulation_executor.national_partition import ( # noqa: PLC0415 # resolved per call so tests can patch national_partition.national_region_groups - national_region_groups, - ) - from policyengine_simulation_executor.release_bundle import ( # noqa: PLC0415 # resolved per call so tests can patch release_bundle.get_country_release_bundle - get_country_release_bundle, - ) - - store = ArtifactStore(bucket) - groups = national_region_groups(PRECOMPUTE_COUNTRY) - if not groups: - msg = f"No national partition for {PRECOMPUTE_COUNTRY!r}" - raise RuntimeError(msg) - - targets = state_dataset_targets() - datasets: list[DatasetPlanEntry] = [] - baselines: list[BaselinePlanEntry] = [] - state_datasets: list[StateDatasetPlanEntry] = [] - for year in PRECOMPUTE_YEARS: - dataset_identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, year) - datasets.append( - DatasetPlanEntry( - year=year, - digest=dataset_identity.digest, - path=dataset_identity.store_path, - filename=dataset_identity.filename, - exists=store.exists(dataset_identity.store_path), - ) - ) - for group in groups: - identity = cohort_identity(year, group) - baselines.append( - BaselinePlanEntry( - year=year, - group=list(group), - region=identity.region, - digest=identity.digest, - path=identity.store_path, - simulation_id=identity.simulation_id, - exists=store.exists(identity.store_path), - ) - ) - for target in targets: - state_identity = collect_state_dataset_identity( - PRECOMPUTE_COUNTRY, - year, - dataset=target.dataset, - state_code=target.state_code, - state_fips=target.state_fips, - filter_variable=target.filter_variable, - ) - state_datasets.append( - StateDatasetPlanEntry( - year=year, - state_code=target.state_code, - region=target.region, - dataset=target.dataset, - digest=state_identity.digest, - path=state_identity.store_path, - filename=state_identity.filename, - folder=state_dataset_relative_folder( - target.dataset, target.state_code - ), - exists=store.exists(state_identity.store_path), - ) - ) - - bundle = get_country_release_bundle(PRECOMPUTE_COUNTRY) - receipt = BundleVersionIdentity( - policyengine_version=bundle.policyengine_version, - model_version=bundle.model_version, - data_version=bundle.data_version, - data_artifact_revision=bundle.data_artifact_revision, - default_dataset=bundle.default_dataset, - ) - return PrecomputePlan( - datasets=datasets, - baselines=baselines, - receipt=receipt, - state_datasets=state_datasets, - ) - - -def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: - """Wave 1: build one single-year dataset and upload it.""" - from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity - collect_dataset_identity, - ) - from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore - ArtifactStore, - ) - from policyengine_simulation_executor.release_bundle import ( # noqa: PLC0415 # resolved per call so tests can patch release_bundle.get_country_release_bundle - get_country_release_bundle, - ) - from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _country_module, - resolve_data_folder, - ) - - identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, expected.year) - if identity.store_path != expected.path: - msg = ( - "Planned and in-container dataset identities disagree " - f"({expected.path} != {identity.store_path}); refusing to upload " - "under a mismatched key." - ) - raise RuntimeError(msg) - - data_folder = resolve_data_folder() - country_module = _country_module(PRECOMPUTE_COUNTRY) - started = time.monotonic() - country_module.ensure_datasets( - datasets=[get_country_release_bundle(PRECOMPUTE_COUNTRY).default_dataset], - years=[expected.year], - data_folder=data_folder, - ) - build_seconds = time.monotonic() - started - - local_file = Path(data_folder) / identity.filename - if not local_file.exists(): - msg = f"ensure_datasets produced no file at {local_file}" - raise RuntimeError(msg) - uploaded = ArtifactStore(bucket).upload_file(identity.store_path, local_file) - return DatasetBuildResult( - year=expected.year, - path=identity.store_path, - uploaded=uploaded, - build_seconds=round(build_seconds, 1), - size_bytes=local_file.stat().st_size, - ) - - -def _checked_state_targets( - year: int, expected: list[StateDatasetPlanEntry] -) -> dict[str, RowFilterStrategy]: - """Return each planned state's row filter, after re-deriving its key here. - - The planner and this builder run in different containers; both derive a - state's key from the installed bundle. A disagreement means they see - different bundles, so nothing is built rather than uploading a file under - a key that does not describe it. A state without a row filter is refused - here too, before anything is prepared. - """ - from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_state_dataset_identity - collect_state_dataset_identity, - ) - - targets = {target.region: target for target in state_dataset_targets()} - datasets = set() - for entry in expected: - if entry.year != year: - msg = ( - f"State dataset entry for {entry.state_code} is for year " - f"{entry.year}, not {year}" - ) - raise ValueError(msg) - target = targets.get(entry.region) - if target is None: - msg = f"No state dataset target for {entry.region} in this container" - raise RuntimeError(msg) - if target.scoping_strategy is None: - msg = f"{entry.region} has no row filter to build its state file with" - raise RuntimeError(msg) - identity = collect_state_dataset_identity( - PRECOMPUTE_COUNTRY, - year, - dataset=target.dataset, - state_code=target.state_code, - state_fips=target.state_fips, - filter_variable=target.filter_variable, - ) - folder = state_dataset_relative_folder(target.dataset, target.state_code) - if (identity.store_path, identity.filename, folder) != ( - entry.path, - entry.filename, - entry.folder, - ): - msg = ( - f"Planned and in-container identities of {entry.region} {year} " - f"disagree ({entry.path} != {identity.store_path}); refusing to " - "build under a mismatched key." - ) - raise RuntimeError(msg) - datasets.add(target.dataset) - if len(datasets) > 1: - msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" - raise ValueError(msg) - strategies = {} - for entry in expected: - strategy = targets[entry.region].scoping_strategy - if strategy is not None: - strategies[entry.region] = strategy - return strategies - - -def _registered_district_geoids(country_module: ModuleType) -> set[int]: - """Return the district codes the US congressional-district regions filter on.""" - return { - int(region.scoping_strategy.variable_value) - for region in country_module.model.region_registry.get_by_type( - "congressional_district" - ) - } - - -def _require_districts_within_their_states( - household: MicroDataFrame, district_geoids: set[int] -) -> None: - """Refuse a dataset in which a household sits in another state's district. - - District requests read their state's file and filter it by - ``congressional_district_geoid`` (state FIPS times 100 plus the district - number). A household whose code names a district of another state would - be missing from that district's results, so the build stops instead. A - code that names no district, such as an unassigned 0, matches no district - request on either path and changes no result. Both columns are labels; - nothing here is weighted. - """ - if "congressional_district_geoid" not in household.columns: - return - district = household["congressional_district_geoid"] - in_another_state = district.isin(sorted(district_geoids)) & ( - (district // 100) != household["state_fips"] - ) - if in_another_state.any(): - msg = ( - "Some households have a congressional district code that names " - "another state's district; their state's file would not match " - "that district's results" - ) - raise RuntimeError(msg) - - -def build_state_datasets_impl( - bucket: str, - year: int, - expected: list[StateDatasetPlanEntry], - *, - work_folder: Path | None = None, -) -> list[StateDatasetBuildResult]: - """Build, check and upload one year's per-state datasets. - - Prepares the regional dataset's full-country year file once, then for - each planned state keeps the rows its region's own scoping strategy keeps, - writes them with .py's ``save()``, copies the year-file stamp, confirms - that .py's reuse checks accept the file, and uploads it. - """ - from policyengine.tax_benefit_models.us.datasets import ( - PolicyEngineUSDataset, - USYearData, - ) - from policyengine.tax_benefit_models.us.legacy_inputs import ( - read_year_file_stamp, - write_year_file_stamp, - ) - from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES - - from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore - ArtifactStore, - ) - from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _country_module, - ) - - if not expected: - return [] - started = time.monotonic() - strategies = _checked_state_targets(year, expected) - dataset = expected[0].dataset - work_folder = work_folder or Path( - tempfile.mkdtemp(prefix="policyengine-state-datasets-") - ) - source_folder = str(work_folder / "source") - country_module = _country_module(PRECOMPUTE_COUNTRY) - - # Preparing the year file runs a Microsimulation over every state, whose - # objects hold reference cycles; release them before filtering 51 copies. - country_module.ensure_datasets( - datasets=[dataset], years=[year], data_folder=source_folder - ) - gc.collect() - (source,) = country_module.load_datasets( - datasets=[dataset], years=[year], data_folder=source_folder - ).values() - stamp = read_year_file_stamp(source.filepath) - if stamp is None: - msg = f"The prepared year file {source.filepath} has no year-file stamp" - raise RuntimeError(msg) - _require_districts_within_their_states( - source.data.household, _registered_district_geoids(country_module) - ) - - store = ArtifactStore(bucket) - results = [] - for entry in expected: - rows = strategies[entry.region].apply( - entity_data=source.data.entity_data, - group_entities=US_GROUP_ENTITIES, - year=year, - ) - local_file = work_folder / entry.folder / entry.filename - PolicyEngineUSDataset( - name=f"{dataset}-{entry.state_code.lower()}-year-{year}", - description=f"{entry.region} rows of {dataset} for {year}", - filepath=str(local_file), - year=year, - metadata=dict(source.metadata), - data=USYearData(**rows), - ).save() - # save() rewrites the file without the stamp, so stamp it afterwards, - # as create_datasets does. - write_year_file_stamp(local_file, stamp) - problem = year_file_reuse_problem(local_file, dataset) - if problem is not None: - msg = f"ensure_datasets would not reuse {local_file}: {problem}" - raise RuntimeError(msg) - results.append( - StateDatasetBuildResult( - year=year, - state_code=entry.state_code, - path=entry.path, - uploaded=store.upload_file(entry.path, local_file), - size_bytes=local_file.stat().st_size, - build_seconds=round(time.monotonic() - started, 1), - ) - ) - return results - - -def _prepare_cohort_baseline( - bucket: str, expected: BaselinePlanEntry -) -> tuple[ArtifactStore, Path, Simulation]: - """Shared by compute and verify: dataset download + simulation build.""" - from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity - collect_dataset_identity, - ) - from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore - ArtifactStore, - ) - from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _build_simulation, - _load_dataset, - _resolve_dataset_selection, - resolve_data_folder, - ) - - store = ArtifactStore(bucket) - data_folder = Path(resolve_data_folder()) - - dataset_identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, expected.year) - local_dataset = data_folder / dataset_identity.filename - if not local_dataset.exists(): - store.download_file(dataset_identity.store_path, local_dataset) - - params, country_module, resolution = _cohort_resolution( - expected.year, expected.group - ) - dataset_selection = _resolve_dataset_selection(params, region_resolution=resolution) - dataset = _load_dataset( - params, country_module=country_module, selection=dataset_selection - ) - baseline = _build_simulation( - params, - dataset=dataset, - dataset_selection=dataset_selection, - policy=None, - region=resolution, - ) - if baseline.id != expected.simulation_id: - msg = ( - "Planned and in-container baseline ids disagree " - f"({expected.simulation_id} != {baseline.id}); refusing to act " - "under a mismatched key." - ) - raise RuntimeError(msg) - planned_storage_id = Path(expected.path).stem - storage_id = getattr(baseline, "storage_id", baseline.id) - if storage_id != planned_storage_id: - msg = ( - "Planned and in-container baseline storage ids disagree " - f"({planned_storage_id} != {storage_id}); refusing to act " - "under a mismatched key." - ) - raise RuntimeError(msg) - - # The extras economic_impact_analysis applies unconditionally before - # ensure(); the artifact must carry them or every request would fail - # the column guard and recompute. - from policyengine.tax_benefit_models.us.analysis import ( - configure_budgetary_impact_variables, - ) - - configure_budgetary_impact_variables(baseline, baseline) - return store, data_folder, baseline - - -def compute_baseline_impl( - bucket: str, expected: BaselinePlanEntry -) -> BaselineComputeResult: - """Wave 2: compute one cohort baseline and upload its output.""" - from policyengine_simulation_executor.baseline_artifacts import ( - ArtifactBaselineSimulation, - ) - - store, data_folder, baseline = _prepare_cohort_baseline(bucket, expected) - if not isinstance(baseline, ArtifactBaselineSimulation): - msg = ( - "Precompute built a plain Simulation — the qualifying predicate " - "rejected the cohort request shape." - ) - raise RuntimeError(msg) # noqa: TRY004 # a precompute invariant failure, not a bad argument; raised across the compute_baseline Modal boundary and asserted by test_precompute - - started = time.monotonic() - baseline.ensure() - compute_seconds = time.monotonic() - started - - artifact_file = data_folder / baseline_artifact_filename( - getattr(baseline, "storage_id", baseline.id) - ) - if not artifact_file.exists(): - msg = f"ensure() left no artifact at {artifact_file}" - raise RuntimeError(msg) - uploaded = store.upload_file(expected.path, artifact_file) - return BaselineComputeResult( - year=expected.year, - group=list(expected.group), - simulation_id=baseline.id, - outcome=baseline.artifact_outcome, - uploaded=uploaded, - compute_seconds=round(compute_seconds, 1), - size_bytes=artifact_file.stat().st_size, - ) - - -def verify_determinism_impl( - bucket: str, expected: BaselinePlanEntry -) -> DeterminismVerdict: - """Load the uploaded artifact, recompute independently, compare frames. - - Exact comparison of every entity column (values and dtypes) between the - store artifact (downloaded and loaded through ``ensure()``, so the save/ - load round-trip is under test too) and a fresh independent run — the - direct check of "same key, same bytes". File-level byte equality is - deliberately not used (HDF5 container metadata is not guaranteed - stable). - """ - from policyengine.core import ( # noqa: PLC0415 # resolved per call so tests can patch policyengine.core.Simulation - Simulation, - ) - - from policyengine_simulation_executor.baseline_artifacts import ( - OUTCOME_HIT, - ) - - store, data_folder, baseline = _prepare_cohort_baseline(bucket, expected) - artifact_file = data_folder / baseline_artifact_filename( - getattr(baseline, "storage_id", baseline.id) - ) - if not artifact_file.exists(): - store.download_file(expected.path, artifact_file) - baseline.ensure() - # A cohort baseline is an ArtifactBaselineSimulation, which records the - # outcome; the attribute isn't on the Simulation type it is declared as. - outcome = getattr(baseline, "artifact_outcome", None) - if outcome != OUTCOME_HIT: - msg = ( - "Determinism gate could not load the artifact it verifies " - f"(outcome={outcome})" - ) - raise RuntimeError(msg) - - selection = getattr(baseline, "spm_config", None) - fresh = Simulation( - **({"spm": selection} if selection is not None else {}), - dataset=baseline.dataset, - tax_benefit_model_version=baseline.tax_benefit_model_version, - policy=None, - scoping_strategy=baseline.scoping_strategy, - extra_variables=dict(baseline.extra_variables), - ) - fresh.run() - - missing_frames = "Determinism gate found a simulation with no output entity frames" - if baseline.output_dataset is None or not isinstance( - baseline.output_dataset.data, EntityFrames - ): - raise TypeError(missing_frames) - if fresh.output_dataset is None or not isinstance( - fresh.output_dataset.data, EntityFrames - ): - raise TypeError(missing_frames) - differences = [] - loaded = baseline.output_dataset.data.entity_data - recomputed = fresh.output_dataset.data.entity_data - for entity in sorted(set(loaded) | set(recomputed)): - left, right = loaded.get(entity), recomputed.get(entity) - if left is None or right is None: - differences.append(f"{entity}: missing on one side") - continue - left_cols, right_cols = set(left.columns), set(right.columns) - differences.extend( - f"{entity}.{column}: only on one side" - for column in sorted(left_cols ^ right_cols) - ) - for column in sorted(left_cols & right_cols): - if str(left[column].dtype) != str(right[column].dtype): - differences.append(f"{entity}.{column}: dtype differs") - elif not left[column].equals(right[column]): - differences.append(f"{entity}.{column}: values differ") - return DeterminismVerdict(equal=not differences, differences=differences[:50]) - - def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: """Store the plan's manifest write-once and return its content digest.""" from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_baselines.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_baselines.py new file mode 100644 index 000000000..f0a5b827f --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_baselines.py @@ -0,0 +1,207 @@ +"""Compute the precompute's cohort baselines, and check that they are deterministic.""" + +from __future__ import annotations + +import time +from pathlib import Path +from typing import TYPE_CHECKING + +from policyengine_simulation_executor.artifact_keys import baseline_artifact_filename +from policyengine_simulation_executor.output_frames import EntityFrames +from policyengine_simulation_executor.precompute_models import ( + BaselineComputeResult, + BaselinePlanEntry, + DeterminismVerdict, +) +from policyengine_simulation_executor.precompute_planning import ( + PRECOMPUTE_COUNTRY, + cohort_resolution, +) + +if TYPE_CHECKING: + from policyengine.core import Simulation + + from policyengine_simulation_executor.artifact_store import ArtifactStore + + +def _prepare_cohort_baseline( + bucket: str, expected: BaselinePlanEntry +) -> tuple[ArtifactStore, Path, Simulation]: + """Shared by compute and verify: dataset download + simulation build.""" + from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity + collect_dataset_identity, + ) + from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore + ArtifactStore, + ) + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _build_simulation, + _load_dataset, + _resolve_dataset_selection, + resolve_data_folder, + ) + + store = ArtifactStore(bucket) + data_folder = Path(resolve_data_folder()) + + dataset_identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, expected.year) + local_dataset = data_folder / dataset_identity.filename + if not local_dataset.exists(): + store.download_file(dataset_identity.store_path, local_dataset) + + params, country_module, resolution = cohort_resolution( + expected.year, expected.group + ) + dataset_selection = _resolve_dataset_selection(params, region_resolution=resolution) + dataset = _load_dataset( + params, country_module=country_module, selection=dataset_selection + ) + baseline = _build_simulation( + params, + dataset=dataset, + dataset_selection=dataset_selection, + policy=None, + region=resolution, + ) + if baseline.id != expected.simulation_id: + msg = ( + "Planned and in-container baseline ids disagree " + f"({expected.simulation_id} != {baseline.id}); refusing to act " + "under a mismatched key." + ) + raise RuntimeError(msg) + planned_storage_id = Path(expected.path).stem + storage_id = getattr(baseline, "storage_id", baseline.id) + if storage_id != planned_storage_id: + msg = ( + "Planned and in-container baseline storage ids disagree " + f"({planned_storage_id} != {storage_id}); refusing to act " + "under a mismatched key." + ) + raise RuntimeError(msg) + + # The extras economic_impact_analysis applies unconditionally before + # ensure(); the artifact must carry them or every request would fail + # the column guard and recompute. + from policyengine.tax_benefit_models.us.analysis import ( + configure_budgetary_impact_variables, + ) + + configure_budgetary_impact_variables(baseline, baseline) + return store, data_folder, baseline + + +def compute_baseline_impl( + bucket: str, expected: BaselinePlanEntry +) -> BaselineComputeResult: + """Wave 2: compute one cohort baseline and upload its output.""" + from policyengine_simulation_executor.baseline_artifacts import ( + ArtifactBaselineSimulation, + ) + + store, data_folder, baseline = _prepare_cohort_baseline(bucket, expected) + if not isinstance(baseline, ArtifactBaselineSimulation): + msg = ( + "Precompute built a plain Simulation — the qualifying predicate " + "rejected the cohort request shape." + ) + raise RuntimeError(msg) # noqa: TRY004 # a precompute invariant failure, not a bad argument; raised across the compute_baseline Modal boundary and asserted by test_precompute + + started = time.monotonic() + baseline.ensure() + compute_seconds = time.monotonic() - started + + artifact_file = data_folder / baseline_artifact_filename( + getattr(baseline, "storage_id", baseline.id) + ) + if not artifact_file.exists(): + msg = f"ensure() left no artifact at {artifact_file}" + raise RuntimeError(msg) + uploaded = store.upload_file(expected.path, artifact_file) + return BaselineComputeResult( + year=expected.year, + group=list(expected.group), + simulation_id=baseline.id, + outcome=baseline.artifact_outcome, + uploaded=uploaded, + compute_seconds=round(compute_seconds, 1), + size_bytes=artifact_file.stat().st_size, + ) + + +def verify_determinism_impl( + bucket: str, expected: BaselinePlanEntry +) -> DeterminismVerdict: + """Load the uploaded artifact, recompute independently, compare frames. + + Exact comparison of every entity column (values and dtypes) between the + store artifact (downloaded and loaded through ``ensure()``, so the save/ + load round-trip is under test too) and a fresh independent run — the + direct check of "same key, same bytes". File-level byte equality is + deliberately not used (HDF5 container metadata is not guaranteed + stable). + """ + from policyengine.core import ( # noqa: PLC0415 # resolved per call so tests can patch policyengine.core.Simulation + Simulation, + ) + + from policyengine_simulation_executor.baseline_artifacts import ( + OUTCOME_HIT, + ) + + store, data_folder, baseline = _prepare_cohort_baseline(bucket, expected) + artifact_file = data_folder / baseline_artifact_filename( + getattr(baseline, "storage_id", baseline.id) + ) + if not artifact_file.exists(): + store.download_file(expected.path, artifact_file) + baseline.ensure() + # A cohort baseline is an ArtifactBaselineSimulation, which records the + # outcome; the attribute isn't on the Simulation type it is declared as. + outcome = getattr(baseline, "artifact_outcome", None) + if outcome != OUTCOME_HIT: + msg = ( + "Determinism gate could not load the artifact it verifies " + f"(outcome={outcome})" + ) + raise RuntimeError(msg) + + selection = getattr(baseline, "spm_config", None) + fresh = Simulation( + **({"spm": selection} if selection is not None else {}), + dataset=baseline.dataset, + tax_benefit_model_version=baseline.tax_benefit_model_version, + policy=None, + scoping_strategy=baseline.scoping_strategy, + extra_variables=dict(baseline.extra_variables), + ) + fresh.run() + + missing_frames = "Determinism gate found a simulation with no output entity frames" + if baseline.output_dataset is None or not isinstance( + baseline.output_dataset.data, EntityFrames + ): + raise TypeError(missing_frames) + if fresh.output_dataset is None or not isinstance( + fresh.output_dataset.data, EntityFrames + ): + raise TypeError(missing_frames) + differences = [] + loaded = baseline.output_dataset.data.entity_data + recomputed = fresh.output_dataset.data.entity_data + for entity in sorted(set(loaded) | set(recomputed)): + left, right = loaded.get(entity), recomputed.get(entity) + if left is None or right is None: + differences.append(f"{entity}: missing on one side") + continue + left_cols, right_cols = set(left.columns), set(right.columns) + differences.extend( + f"{entity}.{column}: only on one side" + for column in sorted(left_cols ^ right_cols) + ) + for column in sorted(left_cols & right_cols): + if str(left[column].dtype) != str(right[column].dtype): + differences.append(f"{entity}.{column}: dtype differs") + elif not left[column].equals(right[column]): + differences.append(f"{entity}.{column}: values differ") + return DeterminismVerdict(equal=not differences, differences=differences[:50]) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_datasets.py new file mode 100644 index 000000000..9682fc68d --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_datasets.py @@ -0,0 +1,276 @@ +"""Build the precompute's national and per-state datasets.""" + +from __future__ import annotations + +import gc +import tempfile +import time +from pathlib import Path +from typing import TYPE_CHECKING + +from policyengine_simulation_executor.precompute_models import ( + DatasetBuildResult, + DatasetPlanEntry, + StateDatasetBuildResult, + StateDatasetPlanEntry, +) +from policyengine_simulation_executor.precompute_planning import ( + PRECOMPUTE_COUNTRY, + state_dataset_targets, +) +from policyengine_simulation_executor.state_datasets import ( + state_dataset_relative_folder, + year_file_reuse_problem, +) + +if TYPE_CHECKING: + from types import ModuleType + + from microdf import MicroDataFrame + from policyengine.core.scoping_strategy import RowFilterStrategy + + +def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: + """Wave 1: build one single-year dataset and upload it.""" + from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity + collect_dataset_identity, + ) + from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore + ArtifactStore, + ) + from policyengine_simulation_executor.release_bundle import ( # noqa: PLC0415 # resolved per call so tests can patch release_bundle.get_country_release_bundle + get_country_release_bundle, + ) + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _country_module, + resolve_data_folder, + ) + + identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, expected.year) + if identity.store_path != expected.path: + msg = ( + "Planned and in-container dataset identities disagree " + f"({expected.path} != {identity.store_path}); refusing to upload " + "under a mismatched key." + ) + raise RuntimeError(msg) + + data_folder = resolve_data_folder() + country_module = _country_module(PRECOMPUTE_COUNTRY) + started = time.monotonic() + country_module.ensure_datasets( + datasets=[get_country_release_bundle(PRECOMPUTE_COUNTRY).default_dataset], + years=[expected.year], + data_folder=data_folder, + ) + build_seconds = time.monotonic() - started + + local_file = Path(data_folder) / identity.filename + if not local_file.exists(): + msg = f"ensure_datasets produced no file at {local_file}" + raise RuntimeError(msg) + uploaded = ArtifactStore(bucket).upload_file(identity.store_path, local_file) + return DatasetBuildResult( + year=expected.year, + path=identity.store_path, + uploaded=uploaded, + build_seconds=round(build_seconds, 1), + size_bytes=local_file.stat().st_size, + ) + + +def _checked_state_targets( + year: int, expected: list[StateDatasetPlanEntry] +) -> dict[str, RowFilterStrategy]: + """Return each planned state's row filter, after re-deriving its key here. + + The planner and this builder run in different containers; both derive a + state's key from the installed bundle. A disagreement means they see + different bundles, so nothing is built rather than uploading a file under + a key that does not describe it. A state without a row filter is refused + here too, before anything is prepared. + """ + from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_state_dataset_identity + collect_state_dataset_identity, + ) + + targets = {target.region: target for target in state_dataset_targets()} + datasets = set() + for entry in expected: + if entry.year != year: + msg = ( + f"State dataset entry for {entry.state_code} is for year " + f"{entry.year}, not {year}" + ) + raise ValueError(msg) + target = targets.get(entry.region) + if target is None: + msg = f"No state dataset target for {entry.region} in this container" + raise RuntimeError(msg) + if target.scoping_strategy is None: + msg = f"{entry.region} has no row filter to build its state file with" + raise RuntimeError(msg) + identity = collect_state_dataset_identity( + PRECOMPUTE_COUNTRY, + year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + folder = state_dataset_relative_folder(target.dataset, target.state_code) + if (identity.store_path, identity.filename, folder) != ( + entry.path, + entry.filename, + entry.folder, + ): + msg = ( + f"Planned and in-container identities of {entry.region} {year} " + f"disagree ({entry.path} != {identity.store_path}); refusing to " + "build under a mismatched key." + ) + raise RuntimeError(msg) + datasets.add(target.dataset) + if len(datasets) > 1: + msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" + raise ValueError(msg) + strategies = {} + for entry in expected: + strategy = targets[entry.region].scoping_strategy + if strategy is not None: + strategies[entry.region] = strategy + return strategies + + +def _registered_district_geoids(country_module: ModuleType) -> set[int]: + """Return the district codes the US congressional-district regions filter on.""" + return { + int(region.scoping_strategy.variable_value) + for region in country_module.model.region_registry.get_by_type( + "congressional_district" + ) + } + + +def _require_districts_within_their_states( + household: MicroDataFrame, district_geoids: set[int] +) -> None: + """Refuse a dataset in which a household sits in another state's district. + + District requests read their state's file and filter it by + ``congressional_district_geoid`` (state FIPS times 100 plus the district + number). A household whose code names a district of another state would + be missing from that district's results, so the build stops instead. A + code that names no district, such as an unassigned 0, matches no district + request on either path and changes no result. Both columns are labels; + nothing here is weighted. + """ + if "congressional_district_geoid" not in household.columns: + return + district = household["congressional_district_geoid"] + in_another_state = district.isin(sorted(district_geoids)) & ( + (district // 100) != household["state_fips"] + ) + if in_another_state.any(): + msg = ( + "Some households have a congressional district code that names " + "another state's district; their state's file would not match " + "that district's results" + ) + raise RuntimeError(msg) + + +def build_state_datasets_impl( + bucket: str, + year: int, + expected: list[StateDatasetPlanEntry], + *, + work_folder: Path | None = None, +) -> list[StateDatasetBuildResult]: + """Build, check and upload one year's per-state datasets. + + Prepares the regional dataset's full-country year file once, then for + each planned state keeps the rows its region's own scoping strategy keeps, + writes them with .py's ``save()``, copies the year-file stamp, confirms + that .py's reuse checks accept the file, and uploads it. + """ + from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + read_year_file_stamp, + write_year_file_stamp, + ) + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore + ArtifactStore, + ) + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _country_module, + ) + + if not expected: + return [] + started = time.monotonic() + strategies = _checked_state_targets(year, expected) + dataset = expected[0].dataset + work_folder = work_folder or Path( + tempfile.mkdtemp(prefix="policyengine-state-datasets-") + ) + source_folder = str(work_folder / "source") + country_module = _country_module(PRECOMPUTE_COUNTRY) + + # Preparing the year file runs a Microsimulation over every state, whose + # objects hold reference cycles; release them before filtering 51 copies. + country_module.ensure_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ) + gc.collect() + (source,) = country_module.load_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ).values() + stamp = read_year_file_stamp(source.filepath) + if stamp is None: + msg = f"The prepared year file {source.filepath} has no year-file stamp" + raise RuntimeError(msg) + _require_districts_within_their_states( + source.data.household, _registered_district_geoids(country_module) + ) + + store = ArtifactStore(bucket) + results = [] + for entry in expected: + rows = strategies[entry.region].apply( + entity_data=source.data.entity_data, + group_entities=US_GROUP_ENTITIES, + year=year, + ) + local_file = work_folder / entry.folder / entry.filename + PolicyEngineUSDataset( + name=f"{dataset}-{entry.state_code.lower()}-year-{year}", + description=f"{entry.region} rows of {dataset} for {year}", + filepath=str(local_file), + year=year, + metadata=dict(source.metadata), + data=USYearData(**rows), + ).save() + # save() rewrites the file without the stamp, so stamp it afterwards, + # as create_datasets does. + write_year_file_stamp(local_file, stamp) + problem = year_file_reuse_problem(local_file, dataset) + if problem is not None: + msg = f"ensure_datasets would not reuse {local_file}: {problem}" + raise RuntimeError(msg) + results.append( + StateDatasetBuildResult( + year=year, + state_code=entry.state_code, + path=entry.path, + uploaded=store.upload_file(entry.path, local_file), + size_bytes=local_file.stat().st_size, + build_seconds=round(time.monotonic() - started, 1), + ) + ) + return results diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py new file mode 100644 index 000000000..6f5ceb332 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py @@ -0,0 +1,271 @@ +"""Plan the artifact precompute: cohorts, state dataset targets and the work to do.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_executor.precompute_models import ( + BaselinePlanEntry, + BundleVersionIdentity, + DatasetPlanEntry, + PrecomputePlan, + StateDatasetPlanEntry, + WorkSelection, +) +from policyengine_simulation_executor.state_datasets import ( + state_dataset_relative_folder, +) + +if TYPE_CHECKING: + from types import ModuleType + + from policyengine.core.scoping_strategy import RowFilterStrategy + + from policyengine_simulation_executor.artifact_keys import BaselineArtifactIdentity + from policyengine_simulation_executor.simulation_runtime import RegionResolution + +PRECOMPUTE_COUNTRY = "us" +# The user-facing priority order (2026 first); with parallel waves it only +# orders spawn submission, but keep it intentional. +PRECOMPUTE_YEARS = [2026, 2027, 2025] + + +def cohort_params(year: int, group: list[str]) -> dict: + """Return the synthetic request a segmented-national child would receive.""" + return { + "country": PRECOMPUTE_COUNTRY, + "scope": "macro", + "region_group": list(group), + "time_period": year, + } + + +def cohort_resolution( + year: int, group: list[str] +) -> tuple[dict[str, Any], ModuleType, RegionResolution]: + """Cohort request params, country module, and region resolution. + + The planner's identity (``cohort_identity``) and the worker's + simulation (``_prepare_cohort_baseline``) must derive from ONE + resolution or the in-container id checks abort the run; sharing this + prologue removes the drift surface. + """ + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _country_module, + _resolve_region, + ) + + params = cohort_params(year, group) + country_module = _country_module(PRECOMPUTE_COUNTRY) + resolution = _resolve_region( + country_module=country_module, + country=PRECOMPUTE_COUNTRY, + params=params, + ) + return params, country_module, resolution + + +def cohort_identity(year: int, group: list[str]) -> BaselineArtifactIdentity: + """Identity via the executor's own resolution path (writer==reader).""" + from policyengine_simulation_executor.baseline_artifacts import ( + BaselineIdentityRequest, + qualifying_baseline_identity, + ) + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _resolve_dataset_selection, + ) + + params, _, resolution = cohort_resolution(year, group) + dataset_selection = _resolve_dataset_selection(params, region_resolution=resolution) + identity = qualifying_baseline_identity( + BaselineIdentityRequest( + params, + country=PRECOMPUTE_COUNTRY, + dataset_is_default=dataset_selection.is_default, + policy=None, + region_code=resolution.code, + scoping_strategy=resolution.scoping_strategy, + year=year, + ) + ) + if identity is None: + msg = ( + f"Cohort request for year {year} group {group} does not qualify " + "for a deterministic baseline id — writer and reader predicates " + "have diverged." + ) + raise RuntimeError(msg) + return identity + + +@dataclass(frozen=True) +class StateDatasetTarget: + """A US state whose requests load a non-default dataset, and its row filter. + + The precompute bakes this state's rows of ``dataset`` as their own year + file: the households whose ``filter_variable`` equals ``state_fips``. + ``scoping_strategy`` is the region's own .py strategy, the one a request + for the state applies, so a baked file holds exactly the rows the request + would keep. + """ + + region: str + state_code: str + state_fips: int + filter_variable: str + dataset: str + scoping_strategy: RowFilterStrategy | None = field( + default=None, compare=False, repr=False + ) + + +def state_dataset_targets() -> list[StateDatasetTarget]: + """Return every US state region that loads a non-default dataset, by state code. + + Each state is resolved exactly as a request for it is (``_resolve_region`` + then ``_resolve_dataset_selection``), so a bundle that sends states back to + the default dataset gets no state files. + """ + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _country_module, + _resolve_dataset_selection, + _resolve_region, + ) + + country_module = _country_module(PRECOMPUTE_COUNTRY) + targets = [] + for region in country_module.model.region_registry.get_by_type("state"): + params = { + "country": PRECOMPUTE_COUNTRY, + "region": region.code, + "time_period": PRECOMPUTE_YEARS[0], + } + resolution = _resolve_region( + country_module=country_module, country=PRECOMPUTE_COUNTRY, params=params + ) + selection = _resolve_dataset_selection(params, region_resolution=resolution) + if selection.is_default: + continue + strategy = resolution.scoping_strategy + variable_name = getattr(strategy, "variable_name", None) + variable_value = getattr(strategy, "variable_value", None) + if not variable_name or variable_value is None or not resolution.state_code: + msg = ( + f"State region {resolution.code} loads {selection.name} but has no " + "row filter and state code to build its state file from" + ) + raise RuntimeError(msg) + targets.append( + StateDatasetTarget( + region=resolution.code, + state_code=resolution.state_code, + state_fips=int(variable_value), + filter_variable=variable_name, + dataset=selection.name, + scoping_strategy=strategy, + ) + ) + return sorted(targets, key=lambda target: target.state_code) + + +def select_work(plan: PrecomputePlan, *, force: bool) -> WorkSelection: + """Which plan entries need computing.""" + return WorkSelection( + datasets=[entry for entry in plan.datasets if force or not entry.exists], + baselines=[entry for entry in plan.baselines if force or not entry.exists], + state_datasets=[ + entry for entry in plan.state_datasets if force or not entry.exists + ], + ) + + +def plan_artifacts_impl(bucket: str) -> PrecomputePlan: + """Compute every expected artifact identity and its store presence.""" + from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity + collect_dataset_identity, + collect_state_dataset_identity, + ) + from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore + ArtifactStore, + ) + from policyengine_simulation_executor.national_partition import ( # noqa: PLC0415 # resolved per call so tests can patch national_partition.national_region_groups + national_region_groups, + ) + from policyengine_simulation_executor.release_bundle import ( # noqa: PLC0415 # resolved per call so tests can patch release_bundle.get_country_release_bundle + get_country_release_bundle, + ) + + store = ArtifactStore(bucket) + groups = national_region_groups(PRECOMPUTE_COUNTRY) + if not groups: + msg = f"No national partition for {PRECOMPUTE_COUNTRY!r}" + raise RuntimeError(msg) + + targets = state_dataset_targets() + datasets: list[DatasetPlanEntry] = [] + baselines: list[BaselinePlanEntry] = [] + state_datasets: list[StateDatasetPlanEntry] = [] + for year in PRECOMPUTE_YEARS: + dataset_identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, year) + datasets.append( + DatasetPlanEntry( + year=year, + digest=dataset_identity.digest, + path=dataset_identity.store_path, + filename=dataset_identity.filename, + exists=store.exists(dataset_identity.store_path), + ) + ) + for group in groups: + identity = cohort_identity(year, group) + baselines.append( + BaselinePlanEntry( + year=year, + group=list(group), + region=identity.region, + digest=identity.digest, + path=identity.store_path, + simulation_id=identity.simulation_id, + exists=store.exists(identity.store_path), + ) + ) + for target in targets: + state_identity = collect_state_dataset_identity( + PRECOMPUTE_COUNTRY, + year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + state_datasets.append( + StateDatasetPlanEntry( + year=year, + state_code=target.state_code, + region=target.region, + dataset=target.dataset, + digest=state_identity.digest, + path=state_identity.store_path, + filename=state_identity.filename, + folder=state_dataset_relative_folder( + target.dataset, target.state_code + ), + exists=store.exists(state_identity.store_path), + ) + ) + + bundle = get_country_release_bundle(PRECOMPUTE_COUNTRY) + receipt = BundleVersionIdentity( + policyengine_version=bundle.policyengine_version, + model_version=bundle.model_version, + data_version=bundle.data_version, + data_artifact_revision=bundle.data_artifact_revision, + default_dataset=bundle.default_dataset, + ) + return PrecomputePlan( + datasets=datasets, + baselines=baselines, + receipt=receipt, + state_datasets=state_datasets, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index 782647cf8..2aeca5ac5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -1,90 +1,67 @@ -"""Build and publish the independent Stage 12 US cache.""" +"""Build and publish the independent Stage 12 US cache. + +The steps live in precompute_planning, precompute_datasets and +precompute_baselines. This module runs them, builds and publishes the +manifest, and re-exports their public names. +""" from __future__ import annotations -import gc -import tempfile -import time -from dataclasses import dataclass, field -from pathlib import Path +from dataclasses import dataclass from typing import TYPE_CHECKING, Any -from uuid import NAMESPACE_URL, uuid5 - -from policyengine_simulation_contract.stage12_execution import ( - BundleProvenance, - DatasetArtifactMediaType, - DatasetArtifactReference, - DatasetPopulationInput, - DatasetProvenance, - GeographySelection, - PlannedSimulationExecutionInput, - ReportAggregate, - ReportExecutionInput, - SimulationExecutionInput, - SimulationRole, - stage12_output_plan_sha256, -) - -from policyengine_simulation_executor import simulation_runtime -from policyengine_simulation_executor.simulation_runtime import ( - _country_module, - _resolve_dataset_selection, - _resolve_region, -) -from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle -from policyengine_simulation_executor.stage12_runtime import ( - simulation as stage12_simulation, -) -from policyengine_simulation_executor.stage12_runtime.cached_simulation import CACHE_HIT -from policyengine_simulation_executor.stage12_runtime.output_planning import ( - plan_simulation_input, - resolve_report_output_plan, -) -from policyengine_simulation_executor.stage12_runtime.partition import ( - US_REGION_GROUPS, -) -from policyengine_simulation_executor.stage12_runtime.simulation import ( - simulation_runtime_params, -) -from .keys import ( - Stage12BaselineCacheIdentity, - canonical_digest, - collect_baseline_identity, - collect_dataset_identity, - collect_state_dataset_identity, -) from .models import ( BaselineBuildResult, - BaselinePlanEntry, - CacheBundleIdentity, CacheManifest, CacheManifestEntry, DatasetBuildResult, - DatasetPlanEntry, DeterminismVerdict, ManifestArtifact, PrecomputePlan, StateDatasetBuildResult, StateDatasetManifestArtifact, - StateDatasetPlanEntry, - WorkSelection, ) -from .state_datasets import state_dataset_relative_folder, year_file_reuse_problem +from .precompute_baselines import compute_baseline_impl, verify_determinism_impl +from .precompute_datasets import build_dataset_impl, build_state_datasets_impl +from .precompute_planning import ( + DETERMINISM_SCHEMA, + PRECOMPUTE_YEARS, + StateDatasetTarget, + determinism_record, + determinism_record_path, + plan_artifacts_impl, + select_work, + state_dataset_targets, +) from .store import Stage12CacheStore if TYPE_CHECKING: - from collections.abc import Callable, Mapping, Sequence - from types import ModuleType + from collections.abc import Callable import modal - import pandas as pd - from microdf import MicroDataFrame - from policyengine.core.scoping_strategy import RowFilterStrategy -PRECOMPUTE_YEARS = (2026, 2027, 2025) +__all__ = [ + "DETERMINISM_SCHEMA", + "MANIFEST_DIGEST_PREFIX", + "PRECOMPUTE_YEARS", + "CachePrecomputeFunctions", + "StateDatasetTarget", + "build_dataset_impl", + "build_manifest", + "build_state_datasets_impl", + "compute_baseline_impl", + "determinism_record", + "determinism_record_path", + "manifest_digest_line", + "plan_artifacts_impl", + "publish_manifest_impl", + "run_precompute", + "select_work", + "state_dataset_targets", + "verify_determinism_impl", +] + MANIFEST_DIGEST_PREFIX = "STAGE12_CACHE_MANIFEST_DIGEST=" -DETERMINISM_SCHEMA = "stage12-dv1" def manifest_digest_line(digest: str) -> str: @@ -92,601 +69,6 @@ def manifest_digest_line(digest: str) -> str: return f"{MANIFEST_DIGEST_PREFIX}{digest}" -def _bundle_identity() -> CacheBundleIdentity: - resolved = load_stage12_bundle() - country = next(item for item in resolved.bundle.countries if item.country == "us") - return CacheBundleIdentity( - bundle_manifest_sha256=resolved.bundle_manifest_sha256, - policyengine_version=resolved.bundle.policyengine_version, - core_package_version=resolved.bundle.core_package_version, - country_package_version=country.country_package_version, - data_package_version=country.data_package_version, - data_artifact_revision=country.data_artifact_revision, - default_dataset=country.default_dataset, - ) - - -def _planned_pair( - year: int, -) -> tuple[PlannedSimulationExecutionInput, PlannedSimulationExecutionInput]: - resolved = load_stage12_bundle() - country = next(item for item in resolved.bundle.countries if item.country == "us") - dataset = next( - item for item in country.datasets if item.identity == country.default_dataset - ) - provenance = BundleProvenance( - policyengine_version=resolved.bundle.policyengine_version, - country_package_name=country.country_package_name, - country_package_version=country.country_package_version, - dataset=DatasetProvenance( - identity=dataset.identity, - uri=dataset.uri, - artifact_revision=dataset.artifact_revision, - data_package_name=country.data_package_name, - data_package_version=country.data_package_version, - ), - bundle_manifest_sha256=resolved.bundle_manifest_sha256, - ) - population = DatasetPopulationInput( - artifact=DatasetArtifactReference( - uri=dataset.uri, - media_type=DatasetArtifactMediaType.HDF5, - content_sha256=dataset.sha256, - ) - ) - evaluation_id = uuid5(NAMESPACE_URL, f"stage12-cache:{year}") - - def simulation(role: SimulationRole) -> SimulationExecutionInput: - return SimulationExecutionInput( - evaluation_id=evaluation_id, - simulation_execution_id=uuid5(evaluation_id, role.value), - role=role, - policy={}, - population=population, - year=year, - geography=GeographySelection(country="us", region="us"), - bundle=provenance, - ) - - report = ReportExecutionInput( - evaluation_id=evaluation_id, - baseline=simulation(SimulationRole.BASELINE), - reform=simulation(SimulationRole.REFORM), - requested_aggregates=tuple(ReportAggregate), - ) - plan = resolve_report_output_plan(report) - return ( - plan_simulation_input(report.baseline, plan), - plan_simulation_input(report.reform, plan), - ) - - -def _scoped_planned_simulation( - year: int, group: tuple[str, ...], *, reform: bool = False -) -> PlannedSimulationExecutionInput: - baseline, reform_input = _planned_pair(year) - selected = reform_input if reform else baseline - return selected.model_copy(update={"options": {"region_group": list(group)}}) - - -def _scope_key(year: int, group: tuple[str, ...]) -> str: - - params = { - "country": "us", - "scope": "macro", - "time_period": str(year), - "region_group": list(group), - } - resolution = _resolve_region( - country_module=_country_module("us"), - country="us", - params=params, - ) - value = getattr(resolution.scoping_strategy, "cache_key", None) - if not isinstance(value, str) or not value: - msg = "Stage 12 region group produced no cache identity" - raise RuntimeError(msg) - return value - - -def _baseline_identity( - year: int, group: tuple[str, ...] -) -> tuple[PlannedSimulationExecutionInput, Stage12BaselineCacheIdentity]: - """Return a planned baseline and the identity its worker reads it under.""" - from policyengine_simulation_executor.spm import normalize_runtime_spm - - simulation = _scoped_planned_simulation(year, group) - identity = collect_baseline_identity( - simulation, - resolved=load_stage12_bundle(), - region_codes=group, - scope_key=_scope_key(year, group), - # The selection the worker builds the simulation with, which names - # the file .py saves and loads. - spm=normalize_runtime_spm(simulation_runtime_params(simulation)), - ) - return simulation, identity - - -def determinism_record(baselines: Sequence[BaselinePlanEntry]) -> dict[str, Any]: - """Return the record a passing determinism check writes for these baselines.""" - return { - "schema": DETERMINISM_SCHEMA, - "baselines": [entry.digest for entry in baselines], - } - - -def determinism_record_path(baselines: Sequence[BaselinePlanEntry]) -> str: - """Return the bucket path of the determinism record for these baselines. - - The path is named by the record's digest, so any change to the planned - baselines needs a new passing check. - """ - return f"determinism/us/{canonical_digest(determinism_record(baselines))}.json" - - -@dataclass(frozen=True) -class StateDatasetTarget: - """A US state whose requests load a non-default dataset, and its row filter. - - Stage 12 bakes this state's rows of ``dataset`` as their own year file: - the households whose ``filter_variable`` equals ``state_fips``. - ``scoping_strategy`` is the region's own .py strategy, the one a request - for the state applies. - """ - - region: str - state_code: str - state_fips: int - filter_variable: str - dataset: str - scoping_strategy: RowFilterStrategy | None = field( - default=None, compare=False, repr=False - ) - - -def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: - """Return every US state region that loads a non-default dataset, by state code. - - Each state is resolved as a Stage 12 worker resolves a request for it - (``_resolve_region`` then ``_resolve_dataset_selection``), so a bundle - that sends states back to the default dataset gets no state files. - """ - country_module = _country_module("us") - targets = [] - for region in country_module.model.region_registry.get_by_type("state"): - params = { - "country": "us", - "region": region.code, - "time_period": PRECOMPUTE_YEARS[0], - } - resolution = _resolve_region( - country_module=country_module, country="us", params=params - ) - selection = _resolve_dataset_selection(params, region_resolution=resolution) - if selection.is_default: - continue - strategy = resolution.scoping_strategy - variable_name = getattr(strategy, "variable_name", None) - variable_value = getattr(strategy, "variable_value", None) - if not variable_name or variable_value is None or not resolution.state_code: - msg = ( - f"State region {resolution.code} loads {selection.name} but has no " - "row filter and state code to build its state file from" - ) - raise RuntimeError(msg) - targets.append( - StateDatasetTarget( - region=resolution.code, - state_code=resolution.state_code, - state_fips=int(variable_value), - filter_variable=variable_name, - dataset=selection.name, - scoping_strategy=strategy, - ) - ) - return tuple(sorted(targets, key=lambda target: target.state_code)) - - -def plan_artifacts_impl(bucket: str) -> PrecomputePlan: - """Plan every dataset and baseline file and record which the bucket holds. - - The plan covers each precompute year and US region group, and records - whether a passing determinism check exists for these baselines. - """ - resolved = load_stage12_bundle() - store = Stage12CacheStore(bucket) - targets = state_dataset_targets() - datasets = [] - baselines = [] - state_datasets = [] - for year in PRECOMPUTE_YEARS: - dataset = collect_dataset_identity(resolved, country="us", year=year) - datasets.append( - DatasetPlanEntry( - year=year, - digest=dataset.digest, - path=dataset.store_path, - filename=dataset.filename, - exists=store.exists(dataset.store_path), - ) - ) - for index, group in enumerate(US_REGION_GROUPS): - simulation, identity = _baseline_identity(year, group) - baselines.append( - BaselinePlanEntry( - year=year, - segment_index=index, - region_codes=group, - digest=identity.digest, - path=identity.store_path, - filename=identity.filename, - simulation_id=identity.simulation_id, - output_plan_sha256=stage12_output_plan_sha256( - simulation.output_plan - ), - exists=store.exists(identity.store_path), - ) - ) - for target in targets: - state_identity = collect_state_dataset_identity( - resolved, - country="us", - year=year, - dataset=target.dataset, - state_code=target.state_code, - state_fips=target.state_fips, - filter_variable=target.filter_variable, - ) - state_datasets.append( - StateDatasetPlanEntry( - year=year, - state_code=target.state_code, - region=target.region, - dataset=target.dataset, - digest=state_identity.digest, - path=state_identity.store_path, - filename=state_identity.filename, - folder=state_dataset_relative_folder( - target.dataset, target.state_code - ), - exists=store.exists(state_identity.store_path), - ) - ) - return PrecomputePlan( - years=PRECOMPUTE_YEARS, - datasets=tuple(datasets), - baselines=tuple(baselines), - state_datasets=tuple(state_datasets), - bundle=_bundle_identity(), - partition_sha256=canonical_digest( - {"groups": [list(group) for group in US_REGION_GROUPS]} - ), - determinism_verified=store.exists(determinism_record_path(baselines)), - ) - - -def select_work(plan: PrecomputePlan) -> WorkSelection: - """Select only missing files: a stored cache file is never replaced.""" - return WorkSelection( - datasets=tuple(entry for entry in plan.datasets if not entry.exists), - baselines=tuple(entry for entry in plan.baselines if not entry.exists), - state_datasets=tuple( - entry for entry in plan.state_datasets if not entry.exists - ), - ) - - -def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: - """Have .py put one planned dataset year in the data folder, then upload it. - - Raises ``RuntimeError`` when the installed bundle gives another identity - than the plan, or when the cache path already holds other data. - """ - resolved = load_stage12_bundle() - identity = collect_dataset_identity(resolved, country="us", year=expected.year) - if identity.store_path != expected.path: - msg = "planned and runtime Stage 12 dataset identities differ" - raise RuntimeError(msg) - started = time.monotonic() - _country_module("us").ensure_datasets( - datasets=[identity.dataset], - years=[expected.year], - data_folder=simulation_runtime.resolve_data_folder(), - ) - local_path = Path(simulation_runtime.resolve_data_folder()) / identity.filename - uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) - return DatasetBuildResult( - **stored.model_dump(), - year=expected.year, - uploaded=uploaded, - build_seconds=round(time.monotonic() - started, 1), - ) - - -def _checked_state_targets( - year: int, expected: Sequence[StateDatasetPlanEntry] -) -> dict[str, RowFilterStrategy]: - """Return each planned state's row filter, after re-deriving its key here. - - Raises ``RuntimeError`` when the installed bundle gives another identity - than the plan, so nothing is built under a key that does not describe it, - and when a state has no row filter, before anything is prepared. - """ - resolved = load_stage12_bundle() - targets = {target.region: target for target in state_dataset_targets()} - datasets = set() - for entry in expected: - if entry.year != year: - msg = ( - f"State dataset entry for {entry.state_code} is for year " - f"{entry.year}, not {year}" - ) - raise ValueError(msg) - target = targets.get(entry.region) - if target is None: - msg = f"No state dataset target for {entry.region} in this container" - raise RuntimeError(msg) - if target.scoping_strategy is None: - msg = f"{entry.region} has no row filter to build its state file with" - raise RuntimeError(msg) - identity = collect_state_dataset_identity( - resolved, - country="us", - year=year, - dataset=target.dataset, - state_code=target.state_code, - state_fips=target.state_fips, - filter_variable=target.filter_variable, - ) - folder = state_dataset_relative_folder(target.dataset, target.state_code) - if (identity.store_path, identity.filename, folder) != ( - entry.path, - entry.filename, - entry.folder, - ): - msg = ( - f"planned and runtime Stage 12 state dataset identities differ " - f"for {entry.region} {year}" - ) - raise RuntimeError(msg) - datasets.add(target.dataset) - if len(datasets) > 1: - msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" - raise ValueError(msg) - strategies = {} - for entry in expected: - strategy = targets[entry.region].scoping_strategy - if strategy is not None: - strategies[entry.region] = strategy - return strategies - - -def _registered_district_geoids(country_module: ModuleType) -> set[int]: - """Return the district codes the US congressional-district regions filter on.""" - return { - int(region.scoping_strategy.variable_value) - for region in country_module.model.region_registry.get_by_type( - "congressional_district" - ) - } - - -def _require_districts_within_their_states( - household: MicroDataFrame, district_geoids: set[int] -) -> None: - """Refuse a dataset in which a household sits in another state's district. - - District requests read their state's file and filter it by - ``congressional_district_geoid`` (state FIPS times 100 plus the district - number). A code that names no district, such as an unassigned 0, matches - no district request and changes no result. Both columns are labels; - nothing here is weighted. - """ - if "congressional_district_geoid" not in household.columns: - return - district = household["congressional_district_geoid"] - in_another_state = district.isin(sorted(district_geoids)) & ( - (district // 100) != household["state_fips"] - ) - if in_another_state.any(): - msg = ( - "Some households have a congressional district code that names " - "another state's district; their state's file would not match " - "that district's results" - ) - raise RuntimeError(msg) - - -def build_state_datasets_impl( - bucket: str, - year: int, - expected: Sequence[StateDatasetPlanEntry], - *, - work_folder: Path | None = None, -) -> tuple[StateDatasetBuildResult, ...]: - """Build, check and upload one year's Stage 12 state datasets. - - Has .py prepare the regional dataset's full-country year file once, then - for each planned state keeps the rows its region's own scoping strategy - keeps, writes them with .py's ``save()``, copies the year-file stamp, - confirms that .py's reuse checks accept the file, and uploads it. - """ - from policyengine.tax_benefit_models.us.datasets import ( - PolicyEngineUSDataset, - USYearData, - ) - from policyengine.tax_benefit_models.us.legacy_inputs import ( - read_year_file_stamp, - write_year_file_stamp, - ) - from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES - - if not expected: - return () - started = time.monotonic() - strategies = _checked_state_targets(year, expected) - dataset = expected[0].dataset - work_folder = work_folder or Path( - tempfile.mkdtemp(prefix="policyengine-stage12-state-datasets-") - ) - source_folder = str(work_folder / "source") - country_module = _country_module("us") - - # Preparing the year file runs a Microsimulation over every state, whose - # objects hold reference cycles; release them before filtering 51 copies. - country_module.ensure_datasets( - datasets=[dataset], years=[year], data_folder=source_folder - ) - gc.collect() - (source,) = country_module.load_datasets( - datasets=[dataset], years=[year], data_folder=source_folder - ).values() - stamp = read_year_file_stamp(source.filepath) - if stamp is None: - msg = f"The prepared year file {source.filepath} has no year-file stamp" - raise RuntimeError(msg) - _require_districts_within_their_states( - source.data.household, _registered_district_geoids(country_module) - ) - - store = Stage12CacheStore(bucket) - results = [] - for entry in expected: - rows = strategies[entry.region].apply( - entity_data=source.data.entity_data, - group_entities=US_GROUP_ENTITIES, - year=year, - ) - local_file = work_folder / entry.folder / entry.filename - PolicyEngineUSDataset( - name=f"{dataset}-{entry.state_code.lower()}-year-{year}", - description=f"{entry.region} rows of {dataset} for {year}", - filepath=str(local_file), - year=year, - metadata=dict(source.metadata), - data=USYearData(**rows), - ).save() - # save() rewrites the file without the stamp, so stamp it afterwards, - # as create_datasets does. - write_year_file_stamp(local_file, stamp) - problem = year_file_reuse_problem(local_file, dataset) - if problem is not None: - msg = f"ensure_datasets would not reuse {local_file}: {problem}" - raise RuntimeError(msg) - uploaded, stored = store.upload_file(entry.path, local_file) - results.append( - StateDatasetBuildResult( - **stored.model_dump(), - year=year, - state_code=entry.state_code, - uploaded=uploaded, - build_seconds=round(time.monotonic() - started, 1), - ) - ) - return tuple(results) - - -def _ensure_dataset_file(bucket: str, year: int) -> None: - - identity = collect_dataset_identity(load_stage12_bundle(), country="us", year=year) - local_path = Path(simulation_runtime.resolve_data_folder()) / identity.filename - if not local_path.exists(): - Stage12CacheStore(bucket).download_file(identity.store_path, local_path) - - -def compute_baseline_impl( - bucket: str, expected: BaselinePlanEntry -) -> BaselineBuildResult: - """Calculate one planned baseline and upload the file .py saved. - - Raises ``RuntimeError`` when the runtime identity differs from the plan, - when .py saved no file under the planned name, or when the cache path - already holds other data. - """ - simulation, identity = _baseline_identity( - expected.year, tuple(expected.region_codes) - ) - if identity.store_path != expected.path: - msg = "planned and runtime Stage 12 baseline identities differ" - raise RuntimeError(msg) - _ensure_dataset_file(bucket, expected.year) - started = time.monotonic() - stage12_simulation.calculate_simulation_frames(simulation) - local_path = Path(simulation_runtime.resolve_data_folder()) / identity.filename - if not local_path.exists(): - msg = ( - f"Stage 12 baseline calculation produced no cache file {identity.filename}" - ) - raise RuntimeError(msg) - uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) - return BaselineBuildResult( - **stored.model_dump(), - year=expected.year, - segment_index=expected.segment_index, - simulation_id=expected.simulation_id, - cache_outcome="miss", - uploaded=uploaded, - compute_seconds=round(time.monotonic() - started, 1), - ) - - -def _frame_differences( - cached: Mapping[str, pd.DataFrame], fresh: Mapping[str, pd.DataFrame] -) -> list[str]: - """Describe every entity and column where two sets of frames differ, in order.""" - differences = [] - for entity in sorted(set(cached) | set(fresh)): - left = cached.get(entity) - right = fresh.get(entity) - if left is None or right is None: - differences.append(f"{entity}: missing on one side") - continue - for column in sorted(set(left) | set(right)): - if column not in left or column not in right: - differences.append(f"{entity}.{column}: only on one side") - elif str(left[column].dtype) != str(right[column].dtype): - differences.append(f"{entity}.{column}: dtype differs") - elif not left[column].equals(right[column]): - differences.append(f"{entity}.{column}: values differ") - return differences - - -def verify_determinism_impl( - bucket: str, plan: PrecomputePlan, expected: BaselinePlanEntry -) -> DeterminismVerdict: - """Compare a stored baseline with a fresh calculation and record a pass.""" - if expected not in plan.baselines: - msg = "Stage 12 determinism probe is not a planned baseline" - raise RuntimeError(msg) - store = Stage12CacheStore(bucket) - group = tuple(expected.region_codes) - baseline = _scoped_planned_simulation(expected.year, group) - reform = _scoped_planned_simulation(expected.year, group, reform=True) - _ensure_dataset_file(bucket, expected.year) - local_path = Path(simulation_runtime.resolve_data_folder()) / expected.filename - if not local_path.exists(): - store.download_file(expected.path, local_path) - cached_calculation = stage12_simulation.calculate_simulation_frames(baseline) - if cached_calculation.cache_outcome != CACHE_HIT: - msg = ( - "Stage 12 determinism check did not load the stored baseline " - f"(cache outcome {cached_calculation.cache_outcome})" - ) - raise RuntimeError(msg) - fresh = stage12_simulation.calculate_simulation_frames(reform).frames - differences = _frame_differences(cached_calculation.frames, fresh) - if not differences: - store.write_json( - determinism_record_path(plan.baselines), - determinism_record(plan.baselines), - ) - return DeterminismVerdict( - equal=not differences, - differences=tuple(differences[:50]), - ) - - def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManifest: """List every planned file with the digest and size the bucket reports. diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_baselines.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_baselines.py new file mode 100644 index 000000000..3a0dce0ae --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_baselines.py @@ -0,0 +1,134 @@ +"""Compute the Stage 12 cache's baselines, and check that they are deterministic.""" + +from __future__ import annotations + +import time +from pathlib import Path +from typing import TYPE_CHECKING + +from policyengine_simulation_executor import simulation_runtime +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.stage12_runtime import ( + simulation as stage12_simulation, +) +from policyengine_simulation_executor.stage12_runtime.cached_simulation import CACHE_HIT + +from .keys import collect_dataset_identity +from .models import ( + BaselineBuildResult, + BaselinePlanEntry, + DeterminismVerdict, + PrecomputePlan, +) +from .precompute_planning import ( + baseline_identity, + determinism_record, + determinism_record_path, + scoped_planned_simulation, +) +from .store import Stage12CacheStore + +if TYPE_CHECKING: + from collections.abc import Mapping + + import pandas as pd + + +def _ensure_dataset_file(bucket: str, year: int) -> None: + + identity = collect_dataset_identity(load_stage12_bundle(), country="us", year=year) + local_path = Path(simulation_runtime.resolve_data_folder()) / identity.filename + if not local_path.exists(): + Stage12CacheStore(bucket).download_file(identity.store_path, local_path) + + +def compute_baseline_impl( + bucket: str, expected: BaselinePlanEntry +) -> BaselineBuildResult: + """Calculate one planned baseline and upload the file .py saved. + + Raises ``RuntimeError`` when the runtime identity differs from the plan, + when .py saved no file under the planned name, or when the cache path + already holds other data. + """ + simulation, identity = baseline_identity( + expected.year, tuple(expected.region_codes) + ) + if identity.store_path != expected.path: + msg = "planned and runtime Stage 12 baseline identities differ" + raise RuntimeError(msg) + _ensure_dataset_file(bucket, expected.year) + started = time.monotonic() + stage12_simulation.calculate_simulation_frames(simulation) + local_path = Path(simulation_runtime.resolve_data_folder()) / identity.filename + if not local_path.exists(): + msg = ( + f"Stage 12 baseline calculation produced no cache file {identity.filename}" + ) + raise RuntimeError(msg) + uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) + return BaselineBuildResult( + **stored.model_dump(), + year=expected.year, + segment_index=expected.segment_index, + simulation_id=expected.simulation_id, + cache_outcome="miss", + uploaded=uploaded, + compute_seconds=round(time.monotonic() - started, 1), + ) + + +def _frame_differences( + cached: Mapping[str, pd.DataFrame], fresh: Mapping[str, pd.DataFrame] +) -> list[str]: + """Describe every entity and column where two sets of frames differ, in order.""" + differences = [] + for entity in sorted(set(cached) | set(fresh)): + left = cached.get(entity) + right = fresh.get(entity) + if left is None or right is None: + differences.append(f"{entity}: missing on one side") + continue + for column in sorted(set(left) | set(right)): + if column not in left or column not in right: + differences.append(f"{entity}.{column}: only on one side") + elif str(left[column].dtype) != str(right[column].dtype): + differences.append(f"{entity}.{column}: dtype differs") + elif not left[column].equals(right[column]): + differences.append(f"{entity}.{column}: values differ") + return differences + + +def verify_determinism_impl( + bucket: str, plan: PrecomputePlan, expected: BaselinePlanEntry +) -> DeterminismVerdict: + """Compare a stored baseline with a fresh calculation and record a pass.""" + if expected not in plan.baselines: + msg = "Stage 12 determinism probe is not a planned baseline" + raise RuntimeError(msg) + store = Stage12CacheStore(bucket) + group = tuple(expected.region_codes) + baseline = scoped_planned_simulation(expected.year, group) + reform = scoped_planned_simulation(expected.year, group, reform=True) + _ensure_dataset_file(bucket, expected.year) + local_path = Path(simulation_runtime.resolve_data_folder()) / expected.filename + if not local_path.exists(): + store.download_file(expected.path, local_path) + cached_calculation = stage12_simulation.calculate_simulation_frames(baseline) + if cached_calculation.cache_outcome != CACHE_HIT: + msg = ( + "Stage 12 determinism check did not load the stored baseline " + f"(cache outcome {cached_calculation.cache_outcome})" + ) + raise RuntimeError(msg) + fresh = stage12_simulation.calculate_simulation_frames(reform).frames + differences = _frame_differences(cached_calculation.frames, fresh) + if not differences: + store.write_json( + determinism_record_path(plan.baselines), + determinism_record(plan.baselines), + ) + return DeterminismVerdict( + equal=not differences, + differences=tuple(differences[:50]), + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_datasets.py new file mode 100644 index 000000000..663cf6f4b --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_datasets.py @@ -0,0 +1,241 @@ +"""Build the Stage 12 cache's national and per-state datasets.""" + +from __future__ import annotations + +import gc +import tempfile +import time +from pathlib import Path +from typing import TYPE_CHECKING + +from policyengine_simulation_executor import simulation_runtime +from policyengine_simulation_executor.simulation_runtime import _country_module +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle + +from .keys import collect_dataset_identity, collect_state_dataset_identity +from .models import ( + DatasetBuildResult, + DatasetPlanEntry, + StateDatasetBuildResult, + StateDatasetPlanEntry, +) +from .precompute_planning import state_dataset_targets +from .state_datasets import state_dataset_relative_folder, year_file_reuse_problem +from .store import Stage12CacheStore + +if TYPE_CHECKING: + from collections.abc import Sequence + from types import ModuleType + + from microdf import MicroDataFrame + from policyengine.core.scoping_strategy import RowFilterStrategy + + +def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: + """Have .py put one planned dataset year in the data folder, then upload it. + + Raises ``RuntimeError`` when the installed bundle gives another identity + than the plan, or when the cache path already holds other data. + """ + resolved = load_stage12_bundle() + identity = collect_dataset_identity(resolved, country="us", year=expected.year) + if identity.store_path != expected.path: + msg = "planned and runtime Stage 12 dataset identities differ" + raise RuntimeError(msg) + started = time.monotonic() + _country_module("us").ensure_datasets( + datasets=[identity.dataset], + years=[expected.year], + data_folder=simulation_runtime.resolve_data_folder(), + ) + local_path = Path(simulation_runtime.resolve_data_folder()) / identity.filename + uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) + return DatasetBuildResult( + **stored.model_dump(), + year=expected.year, + uploaded=uploaded, + build_seconds=round(time.monotonic() - started, 1), + ) + + +def _checked_state_targets( + year: int, expected: Sequence[StateDatasetPlanEntry] +) -> dict[str, RowFilterStrategy]: + """Return each planned state's row filter, after re-deriving its key here. + + Raises ``RuntimeError`` when the installed bundle gives another identity + than the plan, so nothing is built under a key that does not describe it, + and when a state has no row filter, before anything is prepared. + """ + resolved = load_stage12_bundle() + targets = {target.region: target for target in state_dataset_targets()} + datasets = set() + for entry in expected: + if entry.year != year: + msg = ( + f"State dataset entry for {entry.state_code} is for year " + f"{entry.year}, not {year}" + ) + raise ValueError(msg) + target = targets.get(entry.region) + if target is None: + msg = f"No state dataset target for {entry.region} in this container" + raise RuntimeError(msg) + if target.scoping_strategy is None: + msg = f"{entry.region} has no row filter to build its state file with" + raise RuntimeError(msg) + identity = collect_state_dataset_identity( + resolved, + country="us", + year=year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + folder = state_dataset_relative_folder(target.dataset, target.state_code) + if (identity.store_path, identity.filename, folder) != ( + entry.path, + entry.filename, + entry.folder, + ): + msg = ( + f"planned and runtime Stage 12 state dataset identities differ " + f"for {entry.region} {year}" + ) + raise RuntimeError(msg) + datasets.add(target.dataset) + if len(datasets) > 1: + msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" + raise ValueError(msg) + strategies = {} + for entry in expected: + strategy = targets[entry.region].scoping_strategy + if strategy is not None: + strategies[entry.region] = strategy + return strategies + + +def _registered_district_geoids(country_module: ModuleType) -> set[int]: + """Return the district codes the US congressional-district regions filter on.""" + return { + int(region.scoping_strategy.variable_value) + for region in country_module.model.region_registry.get_by_type( + "congressional_district" + ) + } + + +def _require_districts_within_their_states( + household: MicroDataFrame, district_geoids: set[int] +) -> None: + """Refuse a dataset in which a household sits in another state's district. + + District requests read their state's file and filter it by + ``congressional_district_geoid`` (state FIPS times 100 plus the district + number). A code that names no district, such as an unassigned 0, matches + no district request and changes no result. Both columns are labels; + nothing here is weighted. + """ + if "congressional_district_geoid" not in household.columns: + return + district = household["congressional_district_geoid"] + in_another_state = district.isin(sorted(district_geoids)) & ( + (district // 100) != household["state_fips"] + ) + if in_another_state.any(): + msg = ( + "Some households have a congressional district code that names " + "another state's district; their state's file would not match " + "that district's results" + ) + raise RuntimeError(msg) + + +def build_state_datasets_impl( + bucket: str, + year: int, + expected: Sequence[StateDatasetPlanEntry], + *, + work_folder: Path | None = None, +) -> tuple[StateDatasetBuildResult, ...]: + """Build, check and upload one year's Stage 12 state datasets. + + Has .py prepare the regional dataset's full-country year file once, then + for each planned state keeps the rows its region's own scoping strategy + keeps, writes them with .py's ``save()``, copies the year-file stamp, + confirms that .py's reuse checks accept the file, and uploads it. + """ + from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + read_year_file_stamp, + write_year_file_stamp, + ) + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + if not expected: + return () + started = time.monotonic() + strategies = _checked_state_targets(year, expected) + dataset = expected[0].dataset + work_folder = work_folder or Path( + tempfile.mkdtemp(prefix="policyengine-stage12-state-datasets-") + ) + source_folder = str(work_folder / "source") + country_module = _country_module("us") + + # Preparing the year file runs a Microsimulation over every state, whose + # objects hold reference cycles; release them before filtering 51 copies. + country_module.ensure_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ) + gc.collect() + (source,) = country_module.load_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ).values() + stamp = read_year_file_stamp(source.filepath) + if stamp is None: + msg = f"The prepared year file {source.filepath} has no year-file stamp" + raise RuntimeError(msg) + _require_districts_within_their_states( + source.data.household, _registered_district_geoids(country_module) + ) + + store = Stage12CacheStore(bucket) + results = [] + for entry in expected: + rows = strategies[entry.region].apply( + entity_data=source.data.entity_data, + group_entities=US_GROUP_ENTITIES, + year=year, + ) + local_file = work_folder / entry.folder / entry.filename + PolicyEngineUSDataset( + name=f"{dataset}-{entry.state_code.lower()}-year-{year}", + description=f"{entry.region} rows of {dataset} for {year}", + filepath=str(local_file), + year=year, + metadata=dict(source.metadata), + data=USYearData(**rows), + ).save() + # save() rewrites the file without the stamp, so stamp it afterwards, + # as create_datasets does. + write_year_file_stamp(local_file, stamp) + problem = year_file_reuse_problem(local_file, dataset) + if problem is not None: + msg = f"ensure_datasets would not reuse {local_file}: {problem}" + raise RuntimeError(msg) + uploaded, stored = store.upload_file(entry.path, local_file) + results.append( + StateDatasetBuildResult( + **stored.model_dump(), + year=year, + state_code=entry.state_code, + uploaded=uploaded, + build_seconds=round(time.monotonic() - started, 1), + ) + ) + return tuple(results) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py new file mode 100644 index 000000000..f67362226 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py @@ -0,0 +1,350 @@ +"""Plan the Stage 12 cache: baseline identities, state dataset targets and work.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from uuid import NAMESPACE_URL, uuid5 + +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + PlannedSimulationExecutionInput, + ReportAggregate, + ReportExecutionInput, + SimulationExecutionInput, + SimulationRole, + stage12_output_plan_sha256, +) + +from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + _resolve_dataset_selection, + _resolve_region, +) +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.stage12_runtime.output_planning import ( + plan_simulation_input, + resolve_report_output_plan, +) +from policyengine_simulation_executor.stage12_runtime.partition import US_REGION_GROUPS +from policyengine_simulation_executor.stage12_runtime.simulation import ( + simulation_runtime_params, +) + +from .keys import ( + Stage12BaselineCacheIdentity, + canonical_digest, + collect_baseline_identity, + collect_dataset_identity, + collect_state_dataset_identity, +) +from .models import ( + BaselinePlanEntry, + CacheBundleIdentity, + DatasetPlanEntry, + PrecomputePlan, + StateDatasetPlanEntry, + WorkSelection, +) +from .state_datasets import state_dataset_relative_folder +from .store import Stage12CacheStore + +if TYPE_CHECKING: + from collections.abc import Sequence + + from policyengine.core.scoping_strategy import RowFilterStrategy + +PRECOMPUTE_YEARS = (2026, 2027, 2025) + +DETERMINISM_SCHEMA = "stage12-dv1" + + +def _bundle_identity() -> CacheBundleIdentity: + resolved = load_stage12_bundle() + country = next(item for item in resolved.bundle.countries if item.country == "us") + return CacheBundleIdentity( + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + policyengine_version=resolved.bundle.policyengine_version, + core_package_version=resolved.bundle.core_package_version, + country_package_version=country.country_package_version, + data_package_version=country.data_package_version, + data_artifact_revision=country.data_artifact_revision, + default_dataset=country.default_dataset, + ) + + +def _planned_pair( + year: int, +) -> tuple[PlannedSimulationExecutionInput, PlannedSimulationExecutionInput]: + resolved = load_stage12_bundle() + country = next(item for item in resolved.bundle.countries if item.country == "us") + dataset = next( + item for item in country.datasets if item.identity == country.default_dataset + ) + provenance = BundleProvenance( + policyengine_version=resolved.bundle.policyengine_version, + country_package_name=country.country_package_name, + country_package_version=country.country_package_version, + dataset=DatasetProvenance( + identity=dataset.identity, + uri=dataset.uri, + artifact_revision=dataset.artifact_revision, + data_package_name=country.data_package_name, + data_package_version=country.data_package_version, + ), + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + ) + population = DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256=dataset.sha256, + ) + ) + evaluation_id = uuid5(NAMESPACE_URL, f"stage12-cache:{year}") + + def simulation(role: SimulationRole) -> SimulationExecutionInput: + return SimulationExecutionInput( + evaluation_id=evaluation_id, + simulation_execution_id=uuid5(evaluation_id, role.value), + role=role, + policy={}, + population=population, + year=year, + geography=GeographySelection(country="us", region="us"), + bundle=provenance, + ) + + report = ReportExecutionInput( + evaluation_id=evaluation_id, + baseline=simulation(SimulationRole.BASELINE), + reform=simulation(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), + ) + plan = resolve_report_output_plan(report) + return ( + plan_simulation_input(report.baseline, plan), + plan_simulation_input(report.reform, plan), + ) + + +def scoped_planned_simulation( + year: int, group: tuple[str, ...], *, reform: bool = False +) -> PlannedSimulationExecutionInput: + """Return the year's planned baseline, or reform, scoped to one region group.""" + baseline, reform_input = _planned_pair(year) + selected = reform_input if reform else baseline + return selected.model_copy(update={"options": {"region_group": list(group)}}) + + +def _scope_key(year: int, group: tuple[str, ...]) -> str: + + params = { + "country": "us", + "scope": "macro", + "time_period": str(year), + "region_group": list(group), + } + resolution = _resolve_region( + country_module=_country_module("us"), + country="us", + params=params, + ) + value = getattr(resolution.scoping_strategy, "cache_key", None) + if not isinstance(value, str) or not value: + msg = "Stage 12 region group produced no cache identity" + raise RuntimeError(msg) + return value + + +def baseline_identity( + year: int, group: tuple[str, ...] +) -> tuple[PlannedSimulationExecutionInput, Stage12BaselineCacheIdentity]: + """Return a planned baseline and the identity its worker reads it under.""" + from policyengine_simulation_executor.spm import normalize_runtime_spm + + simulation = scoped_planned_simulation(year, group) + identity = collect_baseline_identity( + simulation, + resolved=load_stage12_bundle(), + region_codes=group, + scope_key=_scope_key(year, group), + # The selection the worker builds the simulation with, which names + # the file .py saves and loads. + spm=normalize_runtime_spm(simulation_runtime_params(simulation)), + ) + return simulation, identity + + +def determinism_record(baselines: Sequence[BaselinePlanEntry]) -> dict[str, Any]: + """Return the record a passing determinism check writes for these baselines.""" + return { + "schema": DETERMINISM_SCHEMA, + "baselines": [entry.digest for entry in baselines], + } + + +def determinism_record_path(baselines: Sequence[BaselinePlanEntry]) -> str: + """Return the bucket path of the determinism record for these baselines. + + The path is named by the record's digest, so any change to the planned + baselines needs a new passing check. + """ + return f"determinism/us/{canonical_digest(determinism_record(baselines))}.json" + + +@dataclass(frozen=True) +class StateDatasetTarget: + """A US state whose requests load a non-default dataset, and its row filter. + + Stage 12 bakes this state's rows of ``dataset`` as their own year file: + the households whose ``filter_variable`` equals ``state_fips``. + ``scoping_strategy`` is the region's own .py strategy, the one a request + for the state applies. + """ + + region: str + state_code: str + state_fips: int + filter_variable: str + dataset: str + scoping_strategy: RowFilterStrategy | None = field( + default=None, compare=False, repr=False + ) + + +def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: + """Return every US state region that loads a non-default dataset, by state code. + + Each state is resolved as a Stage 12 worker resolves a request for it + (``_resolve_region`` then ``_resolve_dataset_selection``), so a bundle + that sends states back to the default dataset gets no state files. + """ + country_module = _country_module("us") + targets = [] + for region in country_module.model.region_registry.get_by_type("state"): + params = { + "country": "us", + "region": region.code, + "time_period": PRECOMPUTE_YEARS[0], + } + resolution = _resolve_region( + country_module=country_module, country="us", params=params + ) + selection = _resolve_dataset_selection(params, region_resolution=resolution) + if selection.is_default: + continue + strategy = resolution.scoping_strategy + variable_name = getattr(strategy, "variable_name", None) + variable_value = getattr(strategy, "variable_value", None) + if not variable_name or variable_value is None or not resolution.state_code: + msg = ( + f"State region {resolution.code} loads {selection.name} but has no " + "row filter and state code to build its state file from" + ) + raise RuntimeError(msg) + targets.append( + StateDatasetTarget( + region=resolution.code, + state_code=resolution.state_code, + state_fips=int(variable_value), + filter_variable=variable_name, + dataset=selection.name, + scoping_strategy=strategy, + ) + ) + return tuple(sorted(targets, key=lambda target: target.state_code)) + + +def plan_artifacts_impl(bucket: str) -> PrecomputePlan: + """Plan every dataset and baseline file and record which the bucket holds. + + The plan covers each precompute year and US region group, and records + whether a passing determinism check exists for these baselines. + """ + resolved = load_stage12_bundle() + store = Stage12CacheStore(bucket) + targets = state_dataset_targets() + datasets = [] + baselines = [] + state_datasets = [] + for year in PRECOMPUTE_YEARS: + dataset = collect_dataset_identity(resolved, country="us", year=year) + datasets.append( + DatasetPlanEntry( + year=year, + digest=dataset.digest, + path=dataset.store_path, + filename=dataset.filename, + exists=store.exists(dataset.store_path), + ) + ) + for index, group in enumerate(US_REGION_GROUPS): + simulation, identity = baseline_identity(year, group) + baselines.append( + BaselinePlanEntry( + year=year, + segment_index=index, + region_codes=group, + digest=identity.digest, + path=identity.store_path, + filename=identity.filename, + simulation_id=identity.simulation_id, + output_plan_sha256=stage12_output_plan_sha256( + simulation.output_plan + ), + exists=store.exists(identity.store_path), + ) + ) + for target in targets: + state_identity = collect_state_dataset_identity( + resolved, + country="us", + year=year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + state_datasets.append( + StateDatasetPlanEntry( + year=year, + state_code=target.state_code, + region=target.region, + dataset=target.dataset, + digest=state_identity.digest, + path=state_identity.store_path, + filename=state_identity.filename, + folder=state_dataset_relative_folder( + target.dataset, target.state_code + ), + exists=store.exists(state_identity.store_path), + ) + ) + return PrecomputePlan( + years=PRECOMPUTE_YEARS, + datasets=tuple(datasets), + baselines=tuple(baselines), + state_datasets=tuple(state_datasets), + bundle=_bundle_identity(), + partition_sha256=canonical_digest( + {"groups": [list(group) for group in US_REGION_GROUPS]} + ), + determinism_verified=store.exists(determinism_record_path(baselines)), + ) + + +def select_work(plan: PrecomputePlan) -> WorkSelection: + """Select only missing files: a stored cache file is never replaced.""" + return WorkSelection( + datasets=tuple(entry for entry in plan.datasets if not entry.exists), + baselines=tuple(entry for entry in plan.baselines if not entry.exists), + state_datasets=tuple( + entry for entry in plan.state_datasets if not entry.exists + ), + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 102593c99..69371b11f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -22,9 +22,9 @@ precompute as cache_precompute, ) from policyengine_simulation_executor.stage12_cache.models import CacheManifest -from policyengine_simulation_executor.stage12_cache.precompute import ( - PRECOMPUTE_YEARS, - _scoped_planned_simulation, +from policyengine_simulation_executor.stage12_cache.precompute import PRECOMPUTE_YEARS +from policyengine_simulation_executor.stage12_cache.precompute_planning import ( + scoped_planned_simulation, ) from policyengine_simulation_executor.stage12_runtime.partition import ( US_REGION_GROUPS, @@ -200,7 +200,7 @@ def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: probe_path = data_folder / probe.filename _check_dataset_access(str(probe_path), probe.content_sha256) calculation = calculate_simulation_frames( - _scoped_planned_simulation(STAGE12_CACHE_PROBE_YEAR, US_REGION_GROUPS[0]) + scoped_planned_simulation(STAGE12_CACHE_PROBE_YEAR, US_REGION_GROUPS[0]) ) if calculation.cache_outcome != "hit": msg = "deployed Stage 12 cache probe did not load a cache hit" diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index aca03f6d6..29e90d210 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -7,15 +7,7 @@ serialization boundary does. """ -import sys -from pathlib import Path -from types import ModuleType, SimpleNamespace -from typing import ClassVar - -import pandas as pd import pytest -from fixtures.identity_stubs import install_identity_stubs -from fixtures.wrapper_spm import wrapper_storage_id from policyengine_simulation_executor import precompute from policyengine_simulation_executor.artifact_keys import canonical_digest from policyengine_simulation_executor.precompute_models import ( @@ -581,1255 +573,3 @@ def test_malformed_worker_result_fails_loudly(self): ) with pytest.raises(ValidationError): self._run(functions) - - -class TestWriterReaderContract: - """The writer and the runtime reader must derive identical identities.""" - - @pytest.fixture - def identity_stubs(self, monkeypatch): - from policyengine_simulation_executor import simulation_runtime as sr - - state = install_identity_stubs(monkeypatch) - monkeypatch.setattr( - sr, "get_country_release_bundle", lambda _country: state.bundle - ) - - @pytest.fixture - def fake_regions(self, monkeypatch): - from policyengine.core.scoping_strategy import RowFilterStrategy - from policyengine_simulation_executor import simulation_runtime as sr - - def get_region(code): - state = code.split("/")[-1].upper() - return SimpleNamespace( - code=code, - scoping_strategy=RowFilterStrategy( - variable_name="state_code", variable_value=state - ), - ) - - monkeypatch.setattr( - sr, - "_country_module", - lambda _country: SimpleNamespace( - model=SimpleNamespace(get_region=get_region) - ), - ) - - @pytest.mark.usefixtures("identity_stubs", "fake_regions") - def test_precompute_identity_equals_runtime_id(self): - from policyengine_simulation_executor import simulation_runtime as sr - from policyengine_simulation_executor.baseline_artifacts import ( - BaselineIdentityRequest, - deterministic_baseline_id, - ) - - group = ["state/ca", "state/wv"] - writer_identity = precompute.cohort_identity(2026, group) - - params = precompute.cohort_params(2026, group) - resolution = sr._resolve_region( - country_module=sr._country_module("us"), country="us", params=params - ) - dataset_selection = sr._resolve_dataset_selection( - params, region_resolution=resolution - ) - reader_id = deterministic_baseline_id( - BaselineIdentityRequest( - params, - country="us", - dataset_is_default=dataset_selection.is_default, - policy=None, - region_code=resolution.code, - scoping_strategy=resolution.scoping_strategy, - year=2026, - ) - ) - assert writer_identity.simulation_id == reader_id - assert writer_identity.store_path.endswith(f"/{writer_identity.storage_id}.h5") - - def test_dataset_filename_matches_runtime_stem_lookup(self, monkeypatch): - """The dataset artifact filename is the name the runtime looks for. - - Real-manifest coverage: it equals the exact name the runtime's - ensure_datasets existence check resolves. - """ - from policyengine.provenance.manifest import ( - dataset_logical_name, - resolve_dataset_reference, - ) - from policyengine_simulation_executor.artifact_keys import ( - collect_dataset_identity, - ) - from policyengine_simulation_executor.release_bundle import ( - get_country_release_bundle, - resolve_bundle_dataset_name, - ) - from policyengine_simulation_executor.simulation_runtime import DEFAULT_YEAR - - monkeypatch.delenv("POLICYENGINE_BUNDLE_RECEIPT", raising=False) - get_country_release_bundle.cache_clear() - try: - expected_stem = dataset_logical_name( - resolve_dataset_reference("us", resolve_bundle_dataset_name("us", None)) - ) - identity = collect_dataset_identity("us", DEFAULT_YEAR) - finally: - get_country_release_bundle.cache_clear() - assert identity.filename == f"{expected_stem}_year_{DEFAULT_YEAR}.h5" - - -class TestPlanArtifactsImpl: - """The 63-identity planner: enumeration, presence flags, receipt.""" - - @pytest.fixture - def planning_stubs(self, monkeypatch): - from policyengine_simulation_executor import ( - artifact_keys, - artifact_store, - national_partition, - release_bundle, - ) - - existing = { - "datasets/us/ds-2026/populace_year_2026.h5", - "baselines/us/bl-2026-1/bl1-2026-1.h5", - } - - class FakeStore: - def __init__(self, bucket): - self.bucket = bucket - - def exists(self, path): - return path in existing - - monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) - monkeypatch.setattr( - national_partition, - "national_region_groups", - lambda _country: [["state/ca"], ["state/ny", "state/wv"]], - ) - monkeypatch.setattr( - artifact_keys, - "collect_dataset_identity", - lambda _country, year: SimpleNamespace( - digest=f"ds-{year}", - store_path=f"datasets/us/ds-{year}/populace_year_{year}.h5", - filename=f"populace_year_{year}.h5", - ), - ) - - def fake_cohort_identity(year, group): - tag = f"{year}-{len(group)}" - return SimpleNamespace( - region="+".join(group), - digest=f"bl-{tag}", - store_path=f"baselines/us/bl-{tag}/bl1-{tag}.h5", - simulation_id=f"bl1-{tag}", - ) - - monkeypatch.setattr(precompute, "cohort_identity", fake_cohort_identity) - targets = [ - precompute.StateDatasetTarget( - region="state/ut", - state_code="UT", - state_fips=49, - filter_variable="state_fips", - dataset="populace_local", - ), - precompute.StateDatasetTarget( - region="state/wy", - state_code="WY", - state_fips=56, - filter_variable="state_fips", - dataset="populace_local", - ), - ] - monkeypatch.setattr(precompute, "state_dataset_targets", lambda: list(targets)) - existing.add("state-datasets/us/sd-2026-UT/populace_local_year_2026.h5") - - def fake_state_identity( - _country, year, *, dataset, state_code, state_fips, filter_variable - ): - assert (state_fips, filter_variable) in { - (49, "state_fips"), - (56, "state_fips"), - } - digest = f"sd-{year}-{state_code}" - return SimpleNamespace( - digest=digest, - store_path=f"state-datasets/us/{digest}/{dataset}_year_{year}.h5", - filename=f"{dataset}_year_{year}.h5", - ) - - monkeypatch.setattr( - artifact_keys, "collect_state_dataset_identity", fake_state_identity - ) - monkeypatch.setattr( - release_bundle, - "get_country_release_bundle", - lambda _country: SimpleNamespace( - policyengine_version="4.22.0", - model_version="1.0.0", - data_version="1.2.3", - data_artifact_revision="rev-abc", - default_dataset="populace_cps", - ), - ) - return SimpleNamespace(existing=existing, targets=targets) - - @pytest.mark.usefixtures("planning_stubs") - def test_plan_lists_every_state_dataset_year_major_with_presence(self): - plan = precompute.plan_artifacts_impl("bucket-x") - - assert [(entry.year, entry.state_code) for entry in plan.state_datasets] == [ - (2026, "UT"), - (2026, "WY"), - (2027, "UT"), - (2027, "WY"), - (2025, "UT"), - (2025, "WY"), - ] - assert [entry.exists for entry in plan.state_datasets] == [ - True, - False, - False, - False, - False, - False, - ] - first = plan.state_datasets[0] - assert first.region == "state/ut" - assert first.dataset == "populace_local" - assert first.digest == "sd-2026-UT" - assert first.path == "state-datasets/us/sd-2026-UT/populace_local_year_2026.h5" - assert first.filename == "populace_local_year_2026.h5" - assert first.folder == "state-datasets/populace_local/UT" - work = precompute.select_work(plan, force=False) - assert len(work.state_datasets) == 5 - - def test_states_on_the_default_dataset_add_no_state_datasets( - self, planning_stubs, monkeypatch - ): - monkeypatch.setattr(precompute, "state_dataset_targets", list) - - plan = precompute.plan_artifacts_impl("bucket-x") - - assert plan.state_datasets == [] - assert len(plan.datasets) == 3 - - @pytest.mark.usefixtures("planning_stubs") - def test_plan_enumerates_years_and_cohorts_with_presence(self): - plan = precompute.plan_artifacts_impl("bucket-x") - - assert [entry.year for entry in plan.datasets] == [2026, 2027, 2025] - assert [entry.exists for entry in plan.datasets] == [True, False, False] - assert plan.datasets[0].digest == "ds-2026" - assert plan.datasets[0].path == "datasets/us/ds-2026/populace_year_2026.h5" - assert plan.datasets[0].filename == "populace_year_2026.h5" - - # Year-major, partition order inside each year. - assert [entry.simulation_id for entry in plan.baselines] == [ - "bl1-2026-1", - "bl1-2026-2", - "bl1-2027-1", - "bl1-2027-2", - "bl1-2025-1", - "bl1-2025-2", - ] - assert [entry.exists for entry in plan.baselines] == [ - True, - False, - False, - False, - False, - False, - ] - assert plan.baselines[0].group == ["state/ca"] - assert plan.baselines[1].group == ["state/ny", "state/wv"] - assert plan.baselines[1].path == "baselines/us/bl-2026-2/bl1-2026-2.h5" - - assert plan.receipt == BundleVersionIdentity( - policyengine_version="4.22.0", - model_version="1.0.0", - data_version="1.2.3", - data_artifact_revision="rev-abc", - default_dataset="populace_cps", - ) - # An exists-flag inversion would make select_work recompute (or, - # worse, skip) the wrong entries — lock the wiring end to end. - work = precompute.select_work(plan, force=False) - assert [entry.year for entry in work.datasets] == [2027, 2025] - assert len(work.baselines) == 5 - - @pytest.mark.usefixtures("planning_stubs") - def test_empty_partition_fails_loudly(self, monkeypatch): - from policyengine_simulation_executor import national_partition - - monkeypatch.setattr( - national_partition, "national_region_groups", lambda _country: [] - ) - with pytest.raises(RuntimeError, match="No national partition"): - precompute.plan_artifacts_impl("bucket-x") - - -def _us_state_region(state_code: str, state_fips: int, dataset_path: str): - from policyengine.core.scoping_strategy import RowFilterStrategy - - return SimpleNamespace( - code=f"state/{state_code.lower()}", - dataset_path=dataset_path, - scoping_strategy=RowFilterStrategy( - variable_name="state_fips", variable_value=state_fips - ), - parent_code="us", - state_code=state_code, - ) - - -class TestStateDatasetTargets: - """Which states get baked files, resolved the way a request resolves them.""" - - @pytest.fixture - def regional_bundle(self): - from policyengine_simulation_executor.release_bundle import ( - get_country_release_bundle, - ) - - bundle = get_country_release_bundle("us") - regional = next( - name for name in bundle.dataset_uris if name.endswith("_acs_local") - ) - return SimpleNamespace( - default_uri=bundle.default_dataset_uri, - regional=regional, - regional_uri=bundle.dataset_uris[regional], - ) - - def _install_regions(self, monkeypatch, regions): - from policyengine_simulation_executor import simulation_runtime - - by_code = {region.code: region for region in regions} - country_module = SimpleNamespace( - model=SimpleNamespace( - get_region=by_code.get, - region_registry=SimpleNamespace( - get_by_type=lambda region_type: ( - list(regions) if region_type == "state" else [] - ) - ), - ) - ) - monkeypatch.setattr( - simulation_runtime, "_country_module", lambda _country: country_module - ) - - def test_targets_states_whose_requests_load_a_regional_dataset( - self, monkeypatch, regional_bundle - ): - self._install_regions( - monkeypatch, - [ - _us_state_region("WY", 56, regional_bundle.regional_uri), - _us_state_region("UT", 49, regional_bundle.regional_uri), - ], - ) - - targets = precompute.state_dataset_targets() - - assert targets == [ - precompute.StateDatasetTarget( - region="state/ut", - state_code="UT", - state_fips=49, - filter_variable="state_fips", - dataset=regional_bundle.regional, - ), - precompute.StateDatasetTarget( - region="state/wy", - state_code="WY", - state_fips=56, - filter_variable="state_fips", - dataset=regional_bundle.regional, - ), - ] - - def test_skips_states_whose_requests_load_the_default_dataset( - self, monkeypatch, regional_bundle - ): - self._install_regions( - monkeypatch, - [ - _us_state_region("UT", 49, regional_bundle.default_uri), - _us_state_region("WY", 56, regional_bundle.regional_uri), - ], - ) - - assert [target.state_code for target in precompute.state_dataset_targets()] == [ - "WY" - ] - - def test_refuses_a_state_without_a_row_filter(self, monkeypatch, regional_bundle): - unfiltered = _us_state_region("UT", 49, regional_bundle.regional_uri) - unfiltered.scoping_strategy = None - self._install_regions(monkeypatch, [unfiltered]) - - with pytest.raises(RuntimeError, match="state/ut"): - precompute.state_dataset_targets() - - -class TestCohortIdentityGuard: - def test_diverged_predicates_fail_loudly(self, monkeypatch): - from policyengine_simulation_executor import baseline_artifacts as ba - from policyengine_simulation_executor import simulation_runtime as sr - - monkeypatch.setattr(sr, "_country_module", lambda _country: SimpleNamespace()) - monkeypatch.setattr( - sr, - "_resolve_region", - lambda **_kwargs: SimpleNamespace( - code="x", dataset_reference=None, scoping_strategy=None - ), - ) - monkeypatch.setattr( - ba, "qualifying_baseline_identity", lambda *_args, **_kwargs: None - ) - with pytest.raises(RuntimeError, match="diverged"): - precompute.cohort_identity(2026, ["state/ca"]) - - -class TestBuildDatasetImpl: - @pytest.fixture - def dataset_stubs(self, monkeypatch, tmp_path): - from policyengine_simulation_executor import ( - artifact_keys, - artifact_store, - release_bundle, - ) - from policyengine_simulation_executor import simulation_runtime as sr - - state = SimpleNamespace( - folder=tmp_path, - uploads=[], - ensure_calls=[], - write_file=True, - identity=SimpleNamespace( - digest="ds-2026", - store_path="datasets/us/ds-2026/populace_year_2026.h5", - filename="populace_year_2026.h5", - ), - ) - - class FakeStore: - def __init__(self, bucket): - self.bucket = bucket - - def upload_file(self, path, local_path): - state.uploads.append((path, str(local_path))) - return True - - def fake_ensure_datasets(*, datasets, years, data_folder): - state.ensure_calls.append((datasets, years, data_folder)) - if state.write_file: - (tmp_path / state.identity.filename).write_bytes(b"h5-bytes") - - monkeypatch.setattr( - artifact_keys, - "collect_dataset_identity", - lambda _country, _year: state.identity, - ) - monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) - monkeypatch.setattr(sr, "resolve_data_folder", lambda: str(tmp_path)) - monkeypatch.setattr( - sr, - "_country_module", - lambda _country: SimpleNamespace(ensure_datasets=fake_ensure_datasets), - ) - monkeypatch.setattr( - release_bundle, - "get_country_release_bundle", - lambda _country: SimpleNamespace(default_dataset="populace_cps"), - ) - return state - - def _entry(self, path="datasets/us/ds-2026/populace_year_2026.h5"): - return DatasetPlanEntry( - year=2026, - digest="ds-2026", - path=path, - filename="populace_year_2026.h5", - exists=False, - ) - - def test_builds_and_uploads_under_the_planned_key(self, dataset_stubs): - result = precompute.build_dataset_impl("bucket-x", self._entry()) - assert dataset_stubs.ensure_calls == [ - (["populace_cps"], [2026], str(dataset_stubs.folder)) - ] - assert dataset_stubs.uploads == [ - ( - "datasets/us/ds-2026/populace_year_2026.h5", - str(dataset_stubs.folder / "populace_year_2026.h5"), - ) - ] - assert result.year == 2026 - assert result.uploaded is True - assert result.size_bytes == len(b"h5-bytes") - - def test_refuses_to_upload_under_a_mismatched_key(self, dataset_stubs): - with pytest.raises(RuntimeError, match="refusing to upload"): - precompute.build_dataset_impl( - "bucket-x", self._entry(path="datasets/us/OTHER/file.h5") - ) - assert dataset_stubs.uploads == [] - - def test_fails_loudly_when_no_file_is_produced(self, dataset_stubs): - dataset_stubs.write_file = False - with pytest.raises(RuntimeError, match="produced no file"): - precompute.build_dataset_impl("bucket-x", self._entry()) - assert dataset_stubs.uploads == [] - - -class TestBuildStateDatasetsImpl: - """Building, checking and uploading one year's per-state datasets.""" - - YEAR = 2026 - - @pytest.fixture - def state_build(self, monkeypatch, tmp_path): - from fixtures import state_dataset_sources as sources - from policyengine.core.scoping_strategy import RowFilterStrategy - from policyengine.tax_benefit_models import us as us_module - from policyengine_simulation_executor import ( - artifact_keys, - artifact_store, - ) - from policyengine_simulation_executor import simulation_runtime as sr - - state = SimpleNamespace( - uploads=[], - existing=set(), - ensure_calls=[], - entity_frames=sources.source_entity_frames(), - work_folder=tmp_path / "work", - ) - - def target(state_code, fips): - return precompute.StateDatasetTarget( - region=f"state/{state_code.lower()}", - state_code=state_code, - state_fips=fips, - filter_variable="state_fips", - dataset=sources.REGIONAL_DATASET, - scoping_strategy=RowFilterStrategy( - variable_name="state_fips", variable_value=fips - ), - ) - - state.targets = [ - target("UT", sources.UTAH_FIPS), - target("WY", sources.WYOMING_FIPS), - ] - monkeypatch.setattr( - precompute, "state_dataset_targets", lambda: list(state.targets) - ) - - def fake_state_identity( - _country, year, *, dataset, state_code, state_fips, filter_variable - ): - digest = f"sd-{year}-{state_code}-{state_fips}-{filter_variable}" - return SimpleNamespace( - digest=digest, - store_path=f"state-datasets/us/{digest}/{dataset}_year_{year}.h5", - filename=f"{dataset}_year_{year}.h5", - ) - - monkeypatch.setattr( - artifact_keys, "collect_state_dataset_identity", fake_state_identity - ) - - class FakeStore: - def __init__(self, bucket): - self.bucket = bucket - - def upload_file(self, path, local_path): - state.uploads.append((path, Path(local_path))) - return path not in state.existing - - monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) - - def fake_ensure_datasets(*, datasets, years, data_folder): - state.ensure_calls.append((datasets, years, data_folder)) - folder = Path(data_folder) - folder.mkdir(parents=True, exist_ok=True) - for year in years: - sources.write_year_file( - folder, year=year, entity_frames=state.entity_frames - ) - return us_module.load_datasets( - datasets=datasets, years=years, data_folder=data_folder - ) - - districts = [ - SimpleNamespace( - scoping_strategy=RowFilterStrategy( - variable_name="congressional_district_geoid", - variable_value=geoid, - ) - ) - for geoid in (4901, 4902, 4903, 4904, 5600) - ] - monkeypatch.setattr( - sr, - "_country_module", - lambda _country: SimpleNamespace( - ensure_datasets=fake_ensure_datasets, - load_datasets=us_module.load_datasets, - model=SimpleNamespace( - region_registry=SimpleNamespace( - get_by_type=lambda kind: ( - districts if kind == "congressional_district" else [] - ) - ) - ), - ), - ) - state.sources = sources - state.entry = lambda state_code, fips: self._entry( - fake_state_identity, state_code, fips - ) - return state - - def _entry(self, identity_for, state_code, fips): - from fixtures.state_dataset_sources import REGIONAL_DATASET - - identity = identity_for( - "us", - self.YEAR, - dataset=REGIONAL_DATASET, - state_code=state_code, - state_fips=fips, - filter_variable="state_fips", - ) - return StateDatasetPlanEntry( - year=self.YEAR, - state_code=state_code, - region=f"state/{state_code.lower()}", - dataset=REGIONAL_DATASET, - digest=identity.digest, - path=identity.store_path, - filename=identity.filename, - folder=f"state-datasets/{REGIONAL_DATASET}/{state_code}", - exists=False, - ) - - def _build(self, state_build, entries): - return precompute.build_state_datasets_impl( - "bucket-x", self.YEAR, entries, work_folder=state_build.work_folder - ) - - def _both_states(self, state_build): - return [ - state_build.entry("UT", state_build.sources.UTAH_FIPS), - state_build.entry("WY", state_build.sources.WYOMING_FIPS), - ] - - def test_writes_each_state_file_with_only_its_households(self, state_build): - sources = state_build.sources - results = self._build(state_build, self._both_states(state_build)) - - files = dict(state_build.uploads) - utah, wyoming = (files[entry.path] for entry in self._both_states(state_build)) - assert sources.stored_ids(utah, "household") == [1, 2] - assert sources.stored_ids(utah, "person") == [1, 2, 3] - assert sources.stored_ids(utah, "tax_unit") == [1, 2] - assert sources.stored_ids(utah, "marital_unit") == [1, 2, 3] - assert sources.stored_ids(wyoming, "household") == [3] - assert sources.stored_ids(wyoming, "person") == [4] - assert sources.stored_ids(wyoming, "spm_unit") == [3] - assert [(result.state_code, result.year) for result in results] == [ - ("UT", 2026), - ("WY", 2026), - ] - - def test_keeps_every_weight_unchanged(self, state_build): - sources = state_build.sources - self._build(state_build, self._both_states(state_build)) - - utah = state_build.uploads[0][1] - assert sources.stored_weights(utah, "household") == [10.0, 20.0] - assert sources.stored_weights(utah, "person") == [10.0, 11.0, 20.0] - assert sources.stored_weights(utah, "marital_unit") == [10.0, 11.0, 20.0] - - def test_state_files_pass_dotpy_reuse_checks(self, state_build): - from policyengine.tax_benefit_models.us.datasets import ( - _expected_legacy_spm_profile, - ) - from policyengine.tax_benefit_models.us.legacy_inputs import ( - read_year_file_stamp, - year_file_stamp, - year_file_stamp_problem, - ) - from policyengine.tax_benefit_models.us.legacy_spm.record import ( - year_file_problem, - ) - - self._build(state_build, self._both_states(state_build)) - - for _path, local in state_build.uploads: - assert read_year_file_stamp(local) == year_file_stamp() - assert year_file_stamp_problem(local) is None - assert ( - year_file_problem( - local, - _expected_legacy_spm_profile(state_build.sources.REGIONAL_DATASET), - ) - is None - ) - - def test_dotpy_reuses_a_state_file_without_rebuilding( - self, state_build, monkeypatch - ): - from policyengine.tax_benefit_models.us import datasets as us_datasets - - self._build(state_build, self._both_states(state_build)) - - def refuse_to_rebuild(**_kwargs): - msg = "ensure_datasets rebuilt a baked state file" - raise AssertionError(msg) - - monkeypatch.setattr(us_datasets, "create_datasets", refuse_to_rebuild) - utah_folder = state_build.uploads[0][1].parent - loaded = us_datasets.ensure_datasets( - datasets=[state_build.sources.REGIONAL_DATASET], - years=[self.YEAR], - data_folder=str(utah_folder), - ) - - (dataset,) = loaded.values() - assert sorted(dataset.data.household["household_id"]) == [1, 2] - - def test_each_state_file_holds_exactly_the_rows_its_filter_keeps(self, state_build): - from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES - - entries = self._both_states(state_build) - self._build(state_build, entries) - stored = dict(state_build.uploads) - - for target, entry in zip(state_build.targets, entries, strict=True): - expected = target.scoping_strategy.apply( - entity_data=state_build.sources.source_entity_frames(), - group_entities=US_GROUP_ENTITIES, - year=self.YEAR, - ) - for entity, frame in expected.items(): - assert state_build.sources.stored_ids(stored[entry.path], entity) == [ - int(value) for value in frame[f"{entity}_id"] - ], (target.state_code, entity) - - def test_prepares_the_full_country_year_file_once(self, state_build): - self._build(state_build, self._both_states(state_build)) - - assert len(state_build.ensure_calls) == 1 - datasets, years, folder = state_build.ensure_calls[0] - assert datasets == [state_build.sources.REGIONAL_DATASET] - assert years == [self.YEAR] - assert Path(folder).is_relative_to(state_build.work_folder) - - def test_uploads_each_file_under_its_planned_path(self, state_build): - entries = self._both_states(state_build) - state_build.existing.add(entries[1].path) - - results = self._build(state_build, entries) - - assert [path for path, _local in state_build.uploads] == [ - entry.path for entry in entries - ] - assert [local.name for _path, local in state_build.uploads] == [ - entry.filename for entry in entries - ] - assert [result.uploaded for result in results] == [True, False] - assert [result.path for result in results] == [entry.path for entry in entries] - assert all(result.size_bytes > 0 for result in results) - assert 0 <= results[0].build_seconds <= results[1].build_seconds - - def test_refuses_a_district_that_lies_outside_its_state(self, state_build): - state_build.entity_frames = state_build.sources.source_entity_frames( - wyoming_district_geoid=4903 - ) - - with pytest.raises(RuntimeError, match="congressional district"): - self._build(state_build, self._both_states(state_build)) - assert state_build.uploads == [] - - def test_a_household_without_a_district_code_does_not_stop_the_build( - self, state_build - ): - state_build.entity_frames = state_build.sources.source_entity_frames( - wyoming_district_geoid=0 - ) - - results = self._build(state_build, self._both_states(state_build)) - - assert [result.state_code for result in results] == ["UT", "WY"] - - def test_refuses_a_target_without_a_row_filter(self, state_build): - from dataclasses import replace - - state_build.targets[0] = replace(state_build.targets[0], scoping_strategy=None) - - with pytest.raises(RuntimeError, match="no row filter"): - self._build(state_build, self._both_states(state_build)) - assert state_build.uploads == [] - assert state_build.ensure_calls == [] - - def test_refuses_a_state_with_no_households(self, state_build): - from policyengine.core.scoping_strategy import RowFilterStrategy - - state_build.targets.append( - precompute.StateDatasetTarget( - region="state/id", - state_code="ID", - state_fips=16, - filter_variable="state_fips", - dataset=state_build.sources.REGIONAL_DATASET, - scoping_strategy=RowFilterStrategy( - variable_name="state_fips", variable_value=16 - ), - ) - ) - - with pytest.raises(ValueError, match="No households"): - self._build(state_build, [state_build.entry("ID", 16)]) - assert state_build.uploads == [] - - def test_refuses_a_plan_entry_whose_key_differs_in_the_container(self, state_build): - entry = state_build.entry("UT", state_build.sources.UTAH_FIPS) - mismatched = entry.model_copy(update={"path": "state-datasets/us/OTHER/x.h5"}) - - with pytest.raises(RuntimeError, match="refusing to build"): - self._build(state_build, [mismatched]) - assert state_build.ensure_calls == [] - - def test_refuses_a_state_the_container_does_not_target(self, state_build): - state_build.targets = state_build.targets[:1] - - with pytest.raises(RuntimeError, match="state/wy"): - self._build(state_build, self._both_states(state_build)) - assert state_build.ensure_calls == [] - - def test_refuses_entries_for_another_year(self, state_build): - entry = state_build.entry("UT", state_build.sources.UTAH_FIPS) - - with pytest.raises(ValueError, match="2027"): - precompute.build_state_datasets_impl( - "bucket-x", - 2027, - [entry], - work_folder=state_build.work_folder, - ) - - -class TestComputeBaselineImpl: - @pytest.fixture - def cohort_stubs(self, monkeypatch, tmp_path): - from policyengine_simulation_executor import artifact_keys, artifact_store - from policyengine_simulation_executor import simulation_runtime as sr - from policyengine_simulation_executor.baseline_artifacts import ( - ArtifactBaselineSimulation, - ) - - state = SimpleNamespace( - folder=tmp_path, downloads=[], uploads=[], configured=[] - ) - - class StubBaseline(ArtifactBaselineSimulation): - def ensure(self): - self._artifact_outcome = "miss" - (tmp_path / f"{self.id}.h5").write_bytes(b"artifact-bytes") - - class SilentBaseline(ArtifactBaselineSimulation): - def ensure(self): - self._artifact_outcome = "miss" - - class SPMStubBaseline(ArtifactBaselineSimulation): - """A baseline that names its artifact the way the wrapper does. - - The canonical wrapper derives ``storage_id`` from the resolved - selection and saves under it; ``compute_baseline_impl`` plans a - path from the identity's own derivation and refuses to act when - the two disagree. Only a stub that computes its side - independently can show the guard passing rather than aborting. - """ - - @property - def spm_config(self): - return state.spm - - @property - def storage_id(self): - return wrapper_storage_id(self.id, self.spm_config) - - def ensure(self): - self._artifact_outcome = "miss" - (tmp_path / f"{self.storage_id}.h5").write_bytes(b"artifact-bytes") - - state.spm = None - state.baseline = StubBaseline.model_construct(id="bl1-cohort") - state.make_silent = lambda: SilentBaseline.model_construct(id="bl1-cohort") - state.make_spm = lambda sim_id="bl1-cohort": SPMStubBaseline.model_construct( - id=sim_id - ) - - class FakeStore: - def __init__(self, bucket): - self.bucket = bucket - - def download_file(self, path, local_path): - state.downloads.append((path, str(local_path))) - Path(local_path).parent.mkdir(parents=True, exist_ok=True) - Path(local_path).write_bytes(b"dataset-bytes") - - def upload_file(self, path, local_path): - state.uploads.append((path, str(local_path))) - return True - - fake_analysis = ModuleType("policyengine.tax_benefit_models.us.analysis") - fake_analysis.configure_budgetary_impact_variables = lambda baseline, reform: ( - state.configured.append((baseline, reform)) - ) - monkeypatch.setitem( - sys.modules, "policyengine.tax_benefit_models.us.analysis", fake_analysis - ) - - monkeypatch.setattr( - artifact_keys, - "collect_dataset_identity", - lambda _country, year: SimpleNamespace( - store_path=f"datasets/us/ds-{year}/populace_year_{year}.h5", - filename=f"populace_year_{year}.h5", - ), - ) - monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) - monkeypatch.setattr(sr, "resolve_data_folder", lambda: str(tmp_path)) - monkeypatch.setattr(sr, "_country_module", lambda _country: SimpleNamespace()) - monkeypatch.setattr( - sr, - "_resolve_region", - lambda **_kwargs: SimpleNamespace( - code="state/ca+state/wv", - scoping_strategy="scoping", - dataset_reference=None, - ), - ) - monkeypatch.setattr( - sr, - "_load_dataset", - lambda _params, country_module, selection: "dataset", - ) - monkeypatch.setattr( - sr, "_build_simulation", lambda _params, **_kwargs: state.baseline - ) - return state - - def _entry(self, sim_id="bl1-cohort"): - return BaselinePlanEntry( - year=2026, - group=["state/ca", "state/wv"], - region="state/ca+state/wv", - digest="bl-d", - path=f"baselines/us/bl-d/{sim_id}.h5", - simulation_id=sim_id, - exists=False, - ) - - def test_downloads_dataset_configures_extras_and_uploads(self, cohort_stubs): - result = precompute.compute_baseline_impl("bucket-x", self._entry()) - assert cohort_stubs.downloads == [ - ( - "datasets/us/ds-2026/populace_year_2026.h5", - str(cohort_stubs.folder / "populace_year_2026.h5"), - ) - ] - # The extras configuration is load-bearing: without it every - # production request fails the column guard and recomputes. - assert cohort_stubs.configured == [ - (cohort_stubs.baseline, cohort_stubs.baseline) - ] - assert cohort_stubs.uploads == [ - ( - "baselines/us/bl-d/bl1-cohort.h5", - str(cohort_stubs.folder / "bl1-cohort.h5"), - ) - ] - assert result.simulation_id == "bl1-cohort" - assert result.outcome == "miss" - assert result.uploaded is True - assert result.size_bytes == len(b"artifact-bytes") - - def test_present_local_dataset_skips_download(self, cohort_stubs): - (cohort_stubs.folder / "populace_year_2026.h5").write_bytes(b"already-here") - precompute.compute_baseline_impl("bucket-x", self._entry()) - assert cohort_stubs.downloads == [] - - def test_refuses_to_act_under_a_mismatched_id(self, cohort_stubs): - with pytest.raises(RuntimeError, match="refusing to act"): - precompute.compute_baseline_impl("bucket-x", self._entry("bl1-other")) - assert cohort_stubs.uploads == [] - - def test_refuses_a_mismatched_planned_storage_id(self, cohort_stubs): - entry = self._entry() - entry.path = "baselines/us/bl-d/bl1-cohort-spm-planned.h5" - with pytest.raises(RuntimeError, match="storage ids disagree"): - precompute.compute_baseline_impl("bucket-x", entry) - assert cohort_stubs.configured == [] - assert cohort_stubs.uploads == [] - assert not (cohort_stubs.folder / "bl1-cohort.h5").exists() - - @staticmethod - def _planned_identity(selection): - """Return a real ``BaselineArtifactIdentity``, not a restated expression. - - The plan the deploy publishes comes from this object's properties, - so keying the test off them is what makes the write path a - writer==reader check: reformat or truncate ``storage_id``'s digest - and the container -- which derives its own the wrapper's way -- no - longer names the planned file. - """ - from policyengine_simulation_executor.artifact_keys import ( - BaselineArtifactIdentity, - DatasetArtifactIdentity, - ) - - return BaselineArtifactIdentity( - spm=selection, - country="us", - region="state/ca+state/wv", - scope_key="scoping", - dataset=DatasetArtifactIdentity( - spm=selection, - country="us", - dataset="populace_cps", - stem="populace", - year=2026, - data_version="1.2.3", - data_artifact_revision="rev-abc", - source_sha256=None, - data_build_fingerprint=None, - model_version="9.9.9", - policyengine_version="4.22.0", - ), - ) - - def test_publishes_a_selection_scoped_artifact_under_the_planned_path( - self, cohort_stubs - ): - """The canonical write path, end to end through the guard. - - Nothing else in hermetic CI runs ``compute_baseline_impl`` with a - selection: the identity's storage id was only ever compared to - itself, and the one case that reached the guard covered the abort. - Here the plan is the planner's own ``BaselineArtifactIdentity`` and - the container names its artifact the wrapper's way, so the guard - passing at all is the claim -- if the two derivations differed, no - canonical artifact could be published and the deploy would be - blocked. - """ - selection = { - "forecast_content_sha256": "a" * 64, - "scenario": "ce_trend", - "geography_kind": "national", - "geography_id": None, - "county_vintage": "2020", - "as_of": None, - } - planned = self._planned_identity(selection) - cohort_stubs.spm = selection - cohort_stubs.baseline = cohort_stubs.make_spm(planned.simulation_id) - - entry = self._entry(planned.simulation_id) - entry.digest = planned.digest - entry.path = planned.store_path - - result = precompute.compute_baseline_impl("bucket-x", entry) - - assert cohort_stubs.uploads == [ - ( - planned.store_path, - str(cohort_stubs.folder / f"{planned.storage_id}.h5"), - ) - ] - assert result.simulation_id == planned.simulation_id - assert result.uploaded is True - assert result.size_bytes == len(b"artifact-bytes") - # The digest the container computed is the one the planner planned. - assert planned.storage_id == wrapper_storage_id( - planned.simulation_id, selection - ) - # A selection-free plan for the same cohort is a different artifact. - assert planned.storage_id != planned.simulation_id - - def test_refuses_a_selection_scoped_plan_the_container_does_not_share( - self, cohort_stubs - ): - """Abort when the container resolved a different SPM selection than the plan. - - This is the abort the review worried about, under a real selection: a - container that resolved a different selection names a different - artifact, and nothing is published. - """ - cohort_stubs.spm = {"scenario": "zero_real"} - cohort_stubs.baseline = cohort_stubs.make_spm() - entry = self._entry() - entry.path = ( - f"baselines/us/bl-d/bl1-cohort-spm-" - f"{canonical_digest({'scenario': 'ce_trend'})}.h5" - ) - with pytest.raises(RuntimeError, match="storage ids disagree"): - precompute.compute_baseline_impl("bucket-x", entry) - assert cohort_stubs.configured == [] - assert cohort_stubs.uploads == [] - - def test_refuses_a_plain_simulation(self, cohort_stubs): - cohort_stubs.baseline = SimpleNamespace(id="bl1-cohort") - with pytest.raises(RuntimeError, match="plain Simulation"): - precompute.compute_baseline_impl("bucket-x", self._entry()) - assert cohort_stubs.uploads == [] - - def test_fails_loudly_when_ensure_leaves_no_artifact(self, cohort_stubs): - cohort_stubs.baseline = cohort_stubs.make_silent() - with pytest.raises(RuntimeError, match="left no artifact"): - precompute.compute_baseline_impl("bucket-x", self._entry()) - assert cohort_stubs.uploads == [] - - -class TestVerifyDeterminismImpl: - """The determinism gate's comparator must never fail open. - - It is the pipeline's only correctness check. - """ - - @pytest.fixture - def run_verification(self, monkeypatch, tmp_path): - """Run the determinism gate on stubbed loaded and recomputed frames.""" - - def verify(loaded, recomputed, *, outcome="hit", on_disk=True): - downloads = [] - captured = [] - - class FakeStore: - def download_file(self, path, local_path): - downloads.append(path) - Path(local_path).write_bytes(b"artifact") - - class FakeBaseline: - id = "bl1-verify" - dataset = "dataset" - tax_benefit_model_version = "model-version" - scoping_strategy = "scoping" - extra_variables: ClassVar[dict[str, int]] = {"extra": 1} - artifact_outcome = None - - def ensure(self): - self.artifact_outcome = outcome - self.output_dataset = SimpleNamespace( - data=SimpleNamespace(entity_data=loaded) - ) - - class FakeFresh: - def __init__(self, **kwargs): - captured.append(kwargs) - - def run(self): - self.output_dataset = SimpleNamespace( - data=SimpleNamespace(entity_data=recomputed) - ) - - monkeypatch.setattr( - precompute, - "_prepare_cohort_baseline", - lambda _bucket, _expected: (FakeStore(), tmp_path, FakeBaseline()), - ) - monkeypatch.setattr("policyengine.core.Simulation", FakeFresh) - if on_disk: - (tmp_path / "bl1-verify.h5").write_bytes(b"artifact") - entry = BaselinePlanEntry( - year=2026, - group=["state/ca"], - region="state/ca", - digest="d", - path="baselines/us/d/bl1-verify.h5", - simulation_id="bl1-verify", - exists=True, - ) - verdict = precompute.verify_determinism_impl("bucket-x", entry) - return verdict, downloads, captured - - return verify - - def _frames(self, age=(30.0, 40.0)): - return {"person": pd.DataFrame({"age": list(age)})} - - def test_identical_frames_pass(self, run_verification): - verdict, downloads, _ = run_verification(self._frames(), self._frames()) - assert verdict.equal is True - assert verdict.differences == [] - assert downloads == [] - - def test_value_difference_fails(self, run_verification): - verdict, _, _ = run_verification(self._frames(), self._frames(age=(30.0, 41.0))) - assert verdict.equal is False - assert verdict.differences == ["person.age: values differ"] - - def test_dtype_difference_fails(self, run_verification): - left = {"person": pd.DataFrame({"age": pd.array([30, 40], dtype="int64")})} - right = {"person": pd.DataFrame({"age": pd.array([30.0, 40.0])})} - verdict, _, _ = run_verification(left, right) - assert verdict.equal is False - assert verdict.differences == ["person.age: dtype differs"] - - def test_column_asymmetry_fails(self, run_verification): - right = {"person": pd.DataFrame({"age": [30.0, 40.0], "income": [1.0, 2.0]})} - verdict, _, _ = run_verification(self._frames(), right) - assert verdict.equal is False - assert verdict.differences == ["person.income: only on one side"] - - def test_missing_entity_fails(self, run_verification): - verdict, _, _ = run_verification(self._frames(), {}) - assert verdict.equal is False - assert verdict.differences == ["person: missing on one side"] - - def test_differences_are_truncated_but_verdict_stays_false(self, run_verification): - left = {"person": pd.DataFrame({f"c{i}": [1.0] for i in range(60)})} - right = {"person": pd.DataFrame({f"c{i}": [2.0] for i in range(60)})} - verdict, _, _ = run_verification(left, right) - assert verdict.equal is False - assert len(verdict.differences) == 50 - - def test_unloadable_artifact_fails_loudly(self, run_verification): - with pytest.raises(RuntimeError, match="could not load the artifact"): - run_verification( - self._frames(), - self._frames(), - outcome="incomplete", - ) - - def test_absent_artifact_is_downloaded_first(self, run_verification): - verdict, downloads, _ = run_verification( - self._frames(), self._frames(), on_disk=False - ) - assert downloads == ["baselines/us/d/bl1-verify.h5"] - assert verdict.equal is True - - def test_fresh_run_reuses_the_baseline_inputs(self, run_verification): - _, _, captured = run_verification(self._frames(), self._frames()) - assert captured == [ - { - "dataset": "dataset", - "tax_benefit_model_version": "model-version", - "policy": None, - "scoping_strategy": "scoping", - "extra_variables": {"extra": 1}, - } - ] diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py b/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py new file mode 100644 index 000000000..a5f3c393b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py @@ -0,0 +1,417 @@ +"""Tests for computing the precompute's cohort baselines and checking determinism.""" + +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import ClassVar + +import pandas as pd +import pytest +from fixtures.wrapper_spm import wrapper_storage_id +from policyengine_simulation_executor import precompute, precompute_baselines +from policyengine_simulation_executor.artifact_keys import canonical_digest +from policyengine_simulation_executor.precompute_models import BaselinePlanEntry + + +class TestComputeBaselineImpl: + @pytest.fixture + def cohort_stubs(self, monkeypatch, tmp_path): + from policyengine_simulation_executor import artifact_keys, artifact_store + from policyengine_simulation_executor import simulation_runtime as sr + from policyengine_simulation_executor.baseline_artifacts import ( + ArtifactBaselineSimulation, + ) + + state = SimpleNamespace( + folder=tmp_path, downloads=[], uploads=[], configured=[] + ) + + class StubBaseline(ArtifactBaselineSimulation): + def ensure(self): + self._artifact_outcome = "miss" + (tmp_path / f"{self.id}.h5").write_bytes(b"artifact-bytes") + + class SilentBaseline(ArtifactBaselineSimulation): + def ensure(self): + self._artifact_outcome = "miss" + + class SPMStubBaseline(ArtifactBaselineSimulation): + """A baseline that names its artifact the way the wrapper does. + + The canonical wrapper derives ``storage_id`` from the resolved + selection and saves under it; ``compute_baseline_impl`` plans a + path from the identity's own derivation and refuses to act when + the two disagree. Only a stub that computes its side + independently can show the guard passing rather than aborting. + """ + + @property + def spm_config(self): + return state.spm + + @property + def storage_id(self): + return wrapper_storage_id(self.id, self.spm_config) + + def ensure(self): + self._artifact_outcome = "miss" + (tmp_path / f"{self.storage_id}.h5").write_bytes(b"artifact-bytes") + + state.spm = None + state.baseline = StubBaseline.model_construct(id="bl1-cohort") + state.make_silent = lambda: SilentBaseline.model_construct(id="bl1-cohort") + state.make_spm = lambda sim_id="bl1-cohort": SPMStubBaseline.model_construct( + id=sim_id + ) + + class FakeStore: + def __init__(self, bucket): + self.bucket = bucket + + def download_file(self, path, local_path): + state.downloads.append((path, str(local_path))) + Path(local_path).parent.mkdir(parents=True, exist_ok=True) + Path(local_path).write_bytes(b"dataset-bytes") + + def upload_file(self, path, local_path): + state.uploads.append((path, str(local_path))) + return True + + fake_analysis = ModuleType("policyengine.tax_benefit_models.us.analysis") + fake_analysis.configure_budgetary_impact_variables = lambda baseline, reform: ( + state.configured.append((baseline, reform)) + ) + monkeypatch.setitem( + sys.modules, "policyengine.tax_benefit_models.us.analysis", fake_analysis + ) + + monkeypatch.setattr( + artifact_keys, + "collect_dataset_identity", + lambda _country, year: SimpleNamespace( + store_path=f"datasets/us/ds-{year}/populace_year_{year}.h5", + filename=f"populace_year_{year}.h5", + ), + ) + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) + monkeypatch.setattr(sr, "resolve_data_folder", lambda: str(tmp_path)) + monkeypatch.setattr(sr, "_country_module", lambda _country: SimpleNamespace()) + monkeypatch.setattr( + sr, + "_resolve_region", + lambda **_kwargs: SimpleNamespace( + code="state/ca+state/wv", + scoping_strategy="scoping", + dataset_reference=None, + ), + ) + monkeypatch.setattr( + sr, + "_load_dataset", + lambda _params, country_module, selection: "dataset", + ) + monkeypatch.setattr( + sr, "_build_simulation", lambda _params, **_kwargs: state.baseline + ) + return state + + def _entry(self, sim_id="bl1-cohort"): + return BaselinePlanEntry( + year=2026, + group=["state/ca", "state/wv"], + region="state/ca+state/wv", + digest="bl-d", + path=f"baselines/us/bl-d/{sim_id}.h5", + simulation_id=sim_id, + exists=False, + ) + + def test_downloads_dataset_configures_extras_and_uploads(self, cohort_stubs): + result = precompute.compute_baseline_impl("bucket-x", self._entry()) + assert cohort_stubs.downloads == [ + ( + "datasets/us/ds-2026/populace_year_2026.h5", + str(cohort_stubs.folder / "populace_year_2026.h5"), + ) + ] + # The extras configuration is load-bearing: without it every + # production request fails the column guard and recomputes. + assert cohort_stubs.configured == [ + (cohort_stubs.baseline, cohort_stubs.baseline) + ] + assert cohort_stubs.uploads == [ + ( + "baselines/us/bl-d/bl1-cohort.h5", + str(cohort_stubs.folder / "bl1-cohort.h5"), + ) + ] + assert result.simulation_id == "bl1-cohort" + assert result.outcome == "miss" + assert result.uploaded is True + assert result.size_bytes == len(b"artifact-bytes") + + def test_present_local_dataset_skips_download(self, cohort_stubs): + (cohort_stubs.folder / "populace_year_2026.h5").write_bytes(b"already-here") + precompute.compute_baseline_impl("bucket-x", self._entry()) + assert cohort_stubs.downloads == [] + + def test_refuses_to_act_under_a_mismatched_id(self, cohort_stubs): + with pytest.raises(RuntimeError, match="refusing to act"): + precompute.compute_baseline_impl("bucket-x", self._entry("bl1-other")) + assert cohort_stubs.uploads == [] + + def test_refuses_a_mismatched_planned_storage_id(self, cohort_stubs): + entry = self._entry() + entry.path = "baselines/us/bl-d/bl1-cohort-spm-planned.h5" + with pytest.raises(RuntimeError, match="storage ids disagree"): + precompute.compute_baseline_impl("bucket-x", entry) + assert cohort_stubs.configured == [] + assert cohort_stubs.uploads == [] + assert not (cohort_stubs.folder / "bl1-cohort.h5").exists() + + @staticmethod + def _planned_identity(selection): + """Return a real ``BaselineArtifactIdentity``, not a restated expression. + + The plan the deploy publishes comes from this object's properties, + so keying the test off them is what makes the write path a + writer==reader check: reformat or truncate ``storage_id``'s digest + and the container -- which derives its own the wrapper's way -- no + longer names the planned file. + """ + from policyengine_simulation_executor.artifact_keys import ( + BaselineArtifactIdentity, + DatasetArtifactIdentity, + ) + + return BaselineArtifactIdentity( + spm=selection, + country="us", + region="state/ca+state/wv", + scope_key="scoping", + dataset=DatasetArtifactIdentity( + spm=selection, + country="us", + dataset="populace_cps", + stem="populace", + year=2026, + data_version="1.2.3", + data_artifact_revision="rev-abc", + source_sha256=None, + data_build_fingerprint=None, + model_version="9.9.9", + policyengine_version="4.22.0", + ), + ) + + def test_publishes_a_selection_scoped_artifact_under_the_planned_path( + self, cohort_stubs + ): + """The canonical write path, end to end through the guard. + + Nothing else in hermetic CI runs ``compute_baseline_impl`` with a + selection: the identity's storage id was only ever compared to + itself, and the one case that reached the guard covered the abort. + Here the plan is the planner's own ``BaselineArtifactIdentity`` and + the container names its artifact the wrapper's way, so the guard + passing at all is the claim -- if the two derivations differed, no + canonical artifact could be published and the deploy would be + blocked. + """ + selection = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + planned = self._planned_identity(selection) + cohort_stubs.spm = selection + cohort_stubs.baseline = cohort_stubs.make_spm(planned.simulation_id) + + entry = self._entry(planned.simulation_id) + entry.digest = planned.digest + entry.path = planned.store_path + + result = precompute.compute_baseline_impl("bucket-x", entry) + + assert cohort_stubs.uploads == [ + ( + planned.store_path, + str(cohort_stubs.folder / f"{planned.storage_id}.h5"), + ) + ] + assert result.simulation_id == planned.simulation_id + assert result.uploaded is True + assert result.size_bytes == len(b"artifact-bytes") + # The digest the container computed is the one the planner planned. + assert planned.storage_id == wrapper_storage_id( + planned.simulation_id, selection + ) + # A selection-free plan for the same cohort is a different artifact. + assert planned.storage_id != planned.simulation_id + + def test_refuses_a_selection_scoped_plan_the_container_does_not_share( + self, cohort_stubs + ): + """Abort when the container resolved a different SPM selection than the plan. + + This is the abort the review worried about, under a real selection: a + container that resolved a different selection names a different + artifact, and nothing is published. + """ + cohort_stubs.spm = {"scenario": "zero_real"} + cohort_stubs.baseline = cohort_stubs.make_spm() + entry = self._entry() + entry.path = ( + f"baselines/us/bl-d/bl1-cohort-spm-" + f"{canonical_digest({'scenario': 'ce_trend'})}.h5" + ) + with pytest.raises(RuntimeError, match="storage ids disagree"): + precompute.compute_baseline_impl("bucket-x", entry) + assert cohort_stubs.configured == [] + assert cohort_stubs.uploads == [] + + def test_refuses_a_plain_simulation(self, cohort_stubs): + cohort_stubs.baseline = SimpleNamespace(id="bl1-cohort") + with pytest.raises(RuntimeError, match="plain Simulation"): + precompute.compute_baseline_impl("bucket-x", self._entry()) + assert cohort_stubs.uploads == [] + + def test_fails_loudly_when_ensure_leaves_no_artifact(self, cohort_stubs): + cohort_stubs.baseline = cohort_stubs.make_silent() + with pytest.raises(RuntimeError, match="left no artifact"): + precompute.compute_baseline_impl("bucket-x", self._entry()) + assert cohort_stubs.uploads == [] + + +class TestVerifyDeterminismImpl: + """The determinism gate's comparator must never fail open. + + It is the pipeline's only correctness check. + """ + + @pytest.fixture + def run_verification(self, monkeypatch, tmp_path): + """Run the determinism gate on stubbed loaded and recomputed frames.""" + + def verify(loaded, recomputed, *, outcome="hit", on_disk=True): + downloads = [] + captured = [] + + class FakeStore: + def download_file(self, path, local_path): + downloads.append(path) + Path(local_path).write_bytes(b"artifact") + + class FakeBaseline: + id = "bl1-verify" + dataset = "dataset" + tax_benefit_model_version = "model-version" + scoping_strategy = "scoping" + extra_variables: ClassVar[dict[str, int]] = {"extra": 1} + artifact_outcome = None + + def ensure(self): + self.artifact_outcome = outcome + self.output_dataset = SimpleNamespace( + data=SimpleNamespace(entity_data=loaded) + ) + + class FakeFresh: + def __init__(self, **kwargs): + captured.append(kwargs) + + def run(self): + self.output_dataset = SimpleNamespace( + data=SimpleNamespace(entity_data=recomputed) + ) + + monkeypatch.setattr( + precompute_baselines, + "_prepare_cohort_baseline", + lambda _bucket, _expected: (FakeStore(), tmp_path, FakeBaseline()), + ) + monkeypatch.setattr("policyengine.core.Simulation", FakeFresh) + if on_disk: + (tmp_path / "bl1-verify.h5").write_bytes(b"artifact") + entry = BaselinePlanEntry( + year=2026, + group=["state/ca"], + region="state/ca", + digest="d", + path="baselines/us/d/bl1-verify.h5", + simulation_id="bl1-verify", + exists=True, + ) + verdict = precompute.verify_determinism_impl("bucket-x", entry) + return verdict, downloads, captured + + return verify + + def _frames(self, age=(30.0, 40.0)): + return {"person": pd.DataFrame({"age": list(age)})} + + def test_identical_frames_pass(self, run_verification): + verdict, downloads, _ = run_verification(self._frames(), self._frames()) + assert verdict.equal is True + assert verdict.differences == [] + assert downloads == [] + + def test_value_difference_fails(self, run_verification): + verdict, _, _ = run_verification(self._frames(), self._frames(age=(30.0, 41.0))) + assert verdict.equal is False + assert verdict.differences == ["person.age: values differ"] + + def test_dtype_difference_fails(self, run_verification): + left = {"person": pd.DataFrame({"age": pd.array([30, 40], dtype="int64")})} + right = {"person": pd.DataFrame({"age": pd.array([30.0, 40.0])})} + verdict, _, _ = run_verification(left, right) + assert verdict.equal is False + assert verdict.differences == ["person.age: dtype differs"] + + def test_column_asymmetry_fails(self, run_verification): + right = {"person": pd.DataFrame({"age": [30.0, 40.0], "income": [1.0, 2.0]})} + verdict, _, _ = run_verification(self._frames(), right) + assert verdict.equal is False + assert verdict.differences == ["person.income: only on one side"] + + def test_missing_entity_fails(self, run_verification): + verdict, _, _ = run_verification(self._frames(), {}) + assert verdict.equal is False + assert verdict.differences == ["person: missing on one side"] + + def test_differences_are_truncated_but_verdict_stays_false(self, run_verification): + left = {"person": pd.DataFrame({f"c{i}": [1.0] for i in range(60)})} + right = {"person": pd.DataFrame({f"c{i}": [2.0] for i in range(60)})} + verdict, _, _ = run_verification(left, right) + assert verdict.equal is False + assert len(verdict.differences) == 50 + + def test_unloadable_artifact_fails_loudly(self, run_verification): + with pytest.raises(RuntimeError, match="could not load the artifact"): + run_verification( + self._frames(), + self._frames(), + outcome="incomplete", + ) + + def test_absent_artifact_is_downloaded_first(self, run_verification): + verdict, downloads, _ = run_verification( + self._frames(), self._frames(), on_disk=False + ) + assert downloads == ["baselines/us/d/bl1-verify.h5"] + assert verdict.equal is True + + def test_fresh_run_reuses_the_baseline_inputs(self, run_verification): + _, _, captured = run_verification(self._frames(), self._frames()) + assert captured == [ + { + "dataset": "dataset", + "tax_benefit_model_version": "model-version", + "policy": None, + "scoping_strategy": "scoping", + "extra_variables": {"extra": 1}, + } + ] diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_datasets.py b/projects/policyengine-simulation-executor/tests/test_precompute_datasets.py new file mode 100644 index 000000000..baad2b383 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_precompute_datasets.py @@ -0,0 +1,443 @@ +"""Tests for building the precompute's national and per-state datasets.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +from policyengine_simulation_executor import precompute, precompute_datasets +from policyengine_simulation_executor.precompute_models import ( + DatasetPlanEntry, + StateDatasetPlanEntry, +) + + +class TestBuildDatasetImpl: + @pytest.fixture + def dataset_stubs(self, monkeypatch, tmp_path): + from policyengine_simulation_executor import ( + artifact_keys, + artifact_store, + release_bundle, + ) + from policyengine_simulation_executor import simulation_runtime as sr + + state = SimpleNamespace( + folder=tmp_path, + uploads=[], + ensure_calls=[], + write_file=True, + identity=SimpleNamespace( + digest="ds-2026", + store_path="datasets/us/ds-2026/populace_year_2026.h5", + filename="populace_year_2026.h5", + ), + ) + + class FakeStore: + def __init__(self, bucket): + self.bucket = bucket + + def upload_file(self, path, local_path): + state.uploads.append((path, str(local_path))) + return True + + def fake_ensure_datasets(*, datasets, years, data_folder): + state.ensure_calls.append((datasets, years, data_folder)) + if state.write_file: + (tmp_path / state.identity.filename).write_bytes(b"h5-bytes") + + monkeypatch.setattr( + artifact_keys, + "collect_dataset_identity", + lambda _country, _year: state.identity, + ) + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) + monkeypatch.setattr(sr, "resolve_data_folder", lambda: str(tmp_path)) + monkeypatch.setattr( + sr, + "_country_module", + lambda _country: SimpleNamespace(ensure_datasets=fake_ensure_datasets), + ) + monkeypatch.setattr( + release_bundle, + "get_country_release_bundle", + lambda _country: SimpleNamespace(default_dataset="populace_cps"), + ) + return state + + def _entry(self, path="datasets/us/ds-2026/populace_year_2026.h5"): + return DatasetPlanEntry( + year=2026, + digest="ds-2026", + path=path, + filename="populace_year_2026.h5", + exists=False, + ) + + def test_builds_and_uploads_under_the_planned_key(self, dataset_stubs): + result = precompute.build_dataset_impl("bucket-x", self._entry()) + assert dataset_stubs.ensure_calls == [ + (["populace_cps"], [2026], str(dataset_stubs.folder)) + ] + assert dataset_stubs.uploads == [ + ( + "datasets/us/ds-2026/populace_year_2026.h5", + str(dataset_stubs.folder / "populace_year_2026.h5"), + ) + ] + assert result.year == 2026 + assert result.uploaded is True + assert result.size_bytes == len(b"h5-bytes") + + def test_refuses_to_upload_under_a_mismatched_key(self, dataset_stubs): + with pytest.raises(RuntimeError, match="refusing to upload"): + precompute.build_dataset_impl( + "bucket-x", self._entry(path="datasets/us/OTHER/file.h5") + ) + assert dataset_stubs.uploads == [] + + def test_fails_loudly_when_no_file_is_produced(self, dataset_stubs): + dataset_stubs.write_file = False + with pytest.raises(RuntimeError, match="produced no file"): + precompute.build_dataset_impl("bucket-x", self._entry()) + assert dataset_stubs.uploads == [] + + +class TestBuildStateDatasetsImpl: + """Building, checking and uploading one year's per-state datasets.""" + + YEAR = 2026 + + @pytest.fixture + def state_build(self, monkeypatch, tmp_path): + from fixtures import state_dataset_sources as sources + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine.tax_benefit_models import us as us_module + from policyengine_simulation_executor import ( + artifact_keys, + artifact_store, + ) + from policyengine_simulation_executor import simulation_runtime as sr + + state = SimpleNamespace( + uploads=[], + existing=set(), + ensure_calls=[], + entity_frames=sources.source_entity_frames(), + work_folder=tmp_path / "work", + ) + + def target(state_code, fips): + return precompute.StateDatasetTarget( + region=f"state/{state_code.lower()}", + state_code=state_code, + state_fips=fips, + filter_variable="state_fips", + dataset=sources.REGIONAL_DATASET, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + ) + + state.targets = [ + target("UT", sources.UTAH_FIPS), + target("WY", sources.WYOMING_FIPS), + ] + monkeypatch.setattr( + precompute_datasets, "state_dataset_targets", lambda: list(state.targets) + ) + + def fake_state_identity( + _country, year, *, dataset, state_code, state_fips, filter_variable + ): + digest = f"sd-{year}-{state_code}-{state_fips}-{filter_variable}" + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/us/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + artifact_keys, "collect_state_dataset_identity", fake_state_identity + ) + + class FakeStore: + def __init__(self, bucket): + self.bucket = bucket + + def upload_file(self, path, local_path): + state.uploads.append((path, Path(local_path))) + return path not in state.existing + + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) + + def fake_ensure_datasets(*, datasets, years, data_folder): + state.ensure_calls.append((datasets, years, data_folder)) + folder = Path(data_folder) + folder.mkdir(parents=True, exist_ok=True) + for year in years: + sources.write_year_file( + folder, year=year, entity_frames=state.entity_frames + ) + return us_module.load_datasets( + datasets=datasets, years=years, data_folder=data_folder + ) + + districts = [ + SimpleNamespace( + scoping_strategy=RowFilterStrategy( + variable_name="congressional_district_geoid", + variable_value=geoid, + ) + ) + for geoid in (4901, 4902, 4903, 4904, 5600) + ] + monkeypatch.setattr( + sr, + "_country_module", + lambda _country: SimpleNamespace( + ensure_datasets=fake_ensure_datasets, + load_datasets=us_module.load_datasets, + model=SimpleNamespace( + region_registry=SimpleNamespace( + get_by_type=lambda kind: ( + districts if kind == "congressional_district" else [] + ) + ) + ), + ), + ) + state.sources = sources + state.entry = lambda state_code, fips: self._entry( + fake_state_identity, state_code, fips + ) + return state + + def _entry(self, identity_for, state_code, fips): + from fixtures.state_dataset_sources import REGIONAL_DATASET + + identity = identity_for( + "us", + self.YEAR, + dataset=REGIONAL_DATASET, + state_code=state_code, + state_fips=fips, + filter_variable="state_fips", + ) + return StateDatasetPlanEntry( + year=self.YEAR, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset=REGIONAL_DATASET, + digest=identity.digest, + path=identity.store_path, + filename=identity.filename, + folder=f"state-datasets/{REGIONAL_DATASET}/{state_code}", + exists=False, + ) + + def _build(self, state_build, entries): + return precompute.build_state_datasets_impl( + "bucket-x", self.YEAR, entries, work_folder=state_build.work_folder + ) + + def _both_states(self, state_build): + return [ + state_build.entry("UT", state_build.sources.UTAH_FIPS), + state_build.entry("WY", state_build.sources.WYOMING_FIPS), + ] + + def test_writes_each_state_file_with_only_its_households(self, state_build): + sources = state_build.sources + results = self._build(state_build, self._both_states(state_build)) + + files = dict(state_build.uploads) + utah, wyoming = (files[entry.path] for entry in self._both_states(state_build)) + assert sources.stored_ids(utah, "household") == [1, 2] + assert sources.stored_ids(utah, "person") == [1, 2, 3] + assert sources.stored_ids(utah, "tax_unit") == [1, 2] + assert sources.stored_ids(utah, "marital_unit") == [1, 2, 3] + assert sources.stored_ids(wyoming, "household") == [3] + assert sources.stored_ids(wyoming, "person") == [4] + assert sources.stored_ids(wyoming, "spm_unit") == [3] + assert [(result.state_code, result.year) for result in results] == [ + ("UT", 2026), + ("WY", 2026), + ] + + def test_keeps_every_weight_unchanged(self, state_build): + sources = state_build.sources + self._build(state_build, self._both_states(state_build)) + + utah = state_build.uploads[0][1] + assert sources.stored_weights(utah, "household") == [10.0, 20.0] + assert sources.stored_weights(utah, "person") == [10.0, 11.0, 20.0] + assert sources.stored_weights(utah, "marital_unit") == [10.0, 11.0, 20.0] + + def test_state_files_pass_dotpy_reuse_checks(self, state_build): + from policyengine.tax_benefit_models.us.datasets import ( + _expected_legacy_spm_profile, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + read_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, + ) + from policyengine.tax_benefit_models.us.legacy_spm.record import ( + year_file_problem, + ) + + self._build(state_build, self._both_states(state_build)) + + for _path, local in state_build.uploads: + assert read_year_file_stamp(local) == year_file_stamp() + assert year_file_stamp_problem(local) is None + assert ( + year_file_problem( + local, + _expected_legacy_spm_profile(state_build.sources.REGIONAL_DATASET), + ) + is None + ) + + def test_dotpy_reuses_a_state_file_without_rebuilding( + self, state_build, monkeypatch + ): + from policyengine.tax_benefit_models.us import datasets as us_datasets + + self._build(state_build, self._both_states(state_build)) + + def refuse_to_rebuild(**_kwargs): + msg = "ensure_datasets rebuilt a baked state file" + raise AssertionError(msg) + + monkeypatch.setattr(us_datasets, "create_datasets", refuse_to_rebuild) + utah_folder = state_build.uploads[0][1].parent + loaded = us_datasets.ensure_datasets( + datasets=[state_build.sources.REGIONAL_DATASET], + years=[self.YEAR], + data_folder=str(utah_folder), + ) + + (dataset,) = loaded.values() + assert sorted(dataset.data.household["household_id"]) == [1, 2] + + def test_each_state_file_holds_exactly_the_rows_its_filter_keeps(self, state_build): + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + entries = self._both_states(state_build) + self._build(state_build, entries) + stored = dict(state_build.uploads) + + for target, entry in zip(state_build.targets, entries, strict=True): + expected = target.scoping_strategy.apply( + entity_data=state_build.sources.source_entity_frames(), + group_entities=US_GROUP_ENTITIES, + year=self.YEAR, + ) + for entity, frame in expected.items(): + assert state_build.sources.stored_ids(stored[entry.path], entity) == [ + int(value) for value in frame[f"{entity}_id"] + ], (target.state_code, entity) + + def test_prepares_the_full_country_year_file_once(self, state_build): + self._build(state_build, self._both_states(state_build)) + + assert len(state_build.ensure_calls) == 1 + datasets, years, folder = state_build.ensure_calls[0] + assert datasets == [state_build.sources.REGIONAL_DATASET] + assert years == [self.YEAR] + assert Path(folder).is_relative_to(state_build.work_folder) + + def test_uploads_each_file_under_its_planned_path(self, state_build): + entries = self._both_states(state_build) + state_build.existing.add(entries[1].path) + + results = self._build(state_build, entries) + + assert [path for path, _local in state_build.uploads] == [ + entry.path for entry in entries + ] + assert [local.name for _path, local in state_build.uploads] == [ + entry.filename for entry in entries + ] + assert [result.uploaded for result in results] == [True, False] + assert [result.path for result in results] == [entry.path for entry in entries] + assert all(result.size_bytes > 0 for result in results) + assert 0 <= results[0].build_seconds <= results[1].build_seconds + + def test_refuses_a_district_that_lies_outside_its_state(self, state_build): + state_build.entity_frames = state_build.sources.source_entity_frames( + wyoming_district_geoid=4903 + ) + + with pytest.raises(RuntimeError, match="congressional district"): + self._build(state_build, self._both_states(state_build)) + assert state_build.uploads == [] + + def test_a_household_without_a_district_code_does_not_stop_the_build( + self, state_build + ): + state_build.entity_frames = state_build.sources.source_entity_frames( + wyoming_district_geoid=0 + ) + + results = self._build(state_build, self._both_states(state_build)) + + assert [result.state_code for result in results] == ["UT", "WY"] + + def test_refuses_a_target_without_a_row_filter(self, state_build): + from dataclasses import replace + + state_build.targets[0] = replace(state_build.targets[0], scoping_strategy=None) + + with pytest.raises(RuntimeError, match="no row filter"): + self._build(state_build, self._both_states(state_build)) + assert state_build.uploads == [] + assert state_build.ensure_calls == [] + + def test_refuses_a_state_with_no_households(self, state_build): + from policyengine.core.scoping_strategy import RowFilterStrategy + + state_build.targets.append( + precompute.StateDatasetTarget( + region="state/id", + state_code="ID", + state_fips=16, + filter_variable="state_fips", + dataset=state_build.sources.REGIONAL_DATASET, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=16 + ), + ) + ) + + with pytest.raises(ValueError, match="No households"): + self._build(state_build, [state_build.entry("ID", 16)]) + assert state_build.uploads == [] + + def test_refuses_a_plan_entry_whose_key_differs_in_the_container(self, state_build): + entry = state_build.entry("UT", state_build.sources.UTAH_FIPS) + mismatched = entry.model_copy(update={"path": "state-datasets/us/OTHER/x.h5"}) + + with pytest.raises(RuntimeError, match="refusing to build"): + self._build(state_build, [mismatched]) + assert state_build.ensure_calls == [] + + def test_refuses_a_state_the_container_does_not_target(self, state_build): + state_build.targets = state_build.targets[:1] + + with pytest.raises(RuntimeError, match="state/wy"): + self._build(state_build, self._both_states(state_build)) + assert state_build.ensure_calls == [] + + def test_refuses_entries_for_another_year(self, state_build): + entry = state_build.entry("UT", state_build.sources.UTAH_FIPS) + + with pytest.raises(ValueError, match="2027"): + precompute.build_state_datasets_impl( + "bucket-x", + 2027, + [entry], + work_folder=state_build.work_folder, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_planning.py b/projects/policyengine-simulation-executor/tests/test_precompute_planning.py new file mode 100644 index 000000000..f95629a8b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_precompute_planning.py @@ -0,0 +1,428 @@ +"""Tests for planning the artifact precompute, and its writer and reader identities.""" + +from types import SimpleNamespace + +import pytest +from fixtures.identity_stubs import install_identity_stubs +from policyengine_simulation_executor import precompute, precompute_planning +from policyengine_simulation_executor.precompute_models import BundleVersionIdentity + + +class TestWriterReaderContract: + """The writer and the runtime reader must derive identical identities.""" + + @pytest.fixture + def identity_stubs(self, monkeypatch): + from policyengine_simulation_executor import simulation_runtime as sr + + state = install_identity_stubs(monkeypatch) + monkeypatch.setattr( + sr, "get_country_release_bundle", lambda _country: state.bundle + ) + + @pytest.fixture + def fake_regions(self, monkeypatch): + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine_simulation_executor import simulation_runtime as sr + + def get_region(code): + state = code.split("/")[-1].upper() + return SimpleNamespace( + code=code, + scoping_strategy=RowFilterStrategy( + variable_name="state_code", variable_value=state + ), + ) + + monkeypatch.setattr( + sr, + "_country_module", + lambda _country: SimpleNamespace( + model=SimpleNamespace(get_region=get_region) + ), + ) + + @pytest.mark.usefixtures("identity_stubs", "fake_regions") + def test_precompute_identity_equals_runtime_id(self): + from policyengine_simulation_executor import simulation_runtime as sr + from policyengine_simulation_executor.baseline_artifacts import ( + BaselineIdentityRequest, + deterministic_baseline_id, + ) + + group = ["state/ca", "state/wv"] + writer_identity = precompute.cohort_identity(2026, group) + + params = precompute.cohort_params(2026, group) + resolution = sr._resolve_region( + country_module=sr._country_module("us"), country="us", params=params + ) + dataset_selection = sr._resolve_dataset_selection( + params, region_resolution=resolution + ) + reader_id = deterministic_baseline_id( + BaselineIdentityRequest( + params, + country="us", + dataset_is_default=dataset_selection.is_default, + policy=None, + region_code=resolution.code, + scoping_strategy=resolution.scoping_strategy, + year=2026, + ) + ) + assert writer_identity.simulation_id == reader_id + assert writer_identity.store_path.endswith(f"/{writer_identity.storage_id}.h5") + + def test_dataset_filename_matches_runtime_stem_lookup(self, monkeypatch): + """The dataset artifact filename is the name the runtime looks for. + + Real-manifest coverage: it equals the exact name the runtime's + ensure_datasets existence check resolves. + """ + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + from policyengine_simulation_executor.artifact_keys import ( + collect_dataset_identity, + ) + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + resolve_bundle_dataset_name, + ) + from policyengine_simulation_executor.simulation_runtime import DEFAULT_YEAR + + monkeypatch.delenv("POLICYENGINE_BUNDLE_RECEIPT", raising=False) + get_country_release_bundle.cache_clear() + try: + expected_stem = dataset_logical_name( + resolve_dataset_reference("us", resolve_bundle_dataset_name("us", None)) + ) + identity = collect_dataset_identity("us", DEFAULT_YEAR) + finally: + get_country_release_bundle.cache_clear() + assert identity.filename == f"{expected_stem}_year_{DEFAULT_YEAR}.h5" + + +class TestPlanArtifactsImpl: + """The 63-identity planner: enumeration, presence flags, receipt.""" + + @pytest.fixture + def planning_stubs(self, monkeypatch): + from policyengine_simulation_executor import ( + artifact_keys, + artifact_store, + national_partition, + release_bundle, + ) + + existing = { + "datasets/us/ds-2026/populace_year_2026.h5", + "baselines/us/bl-2026-1/bl1-2026-1.h5", + } + + class FakeStore: + def __init__(self, bucket): + self.bucket = bucket + + def exists(self, path): + return path in existing + + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) + monkeypatch.setattr( + national_partition, + "national_region_groups", + lambda _country: [["state/ca"], ["state/ny", "state/wv"]], + ) + monkeypatch.setattr( + artifact_keys, + "collect_dataset_identity", + lambda _country, year: SimpleNamespace( + digest=f"ds-{year}", + store_path=f"datasets/us/ds-{year}/populace_year_{year}.h5", + filename=f"populace_year_{year}.h5", + ), + ) + + def fake_cohort_identity(year, group): + tag = f"{year}-{len(group)}" + return SimpleNamespace( + region="+".join(group), + digest=f"bl-{tag}", + store_path=f"baselines/us/bl-{tag}/bl1-{tag}.h5", + simulation_id=f"bl1-{tag}", + ) + + monkeypatch.setattr( + precompute_planning, "cohort_identity", fake_cohort_identity + ) + targets = [ + precompute.StateDatasetTarget( + region="state/ut", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + dataset="populace_local", + ), + precompute.StateDatasetTarget( + region="state/wy", + state_code="WY", + state_fips=56, + filter_variable="state_fips", + dataset="populace_local", + ), + ] + monkeypatch.setattr( + precompute_planning, "state_dataset_targets", lambda: list(targets) + ) + existing.add("state-datasets/us/sd-2026-UT/populace_local_year_2026.h5") + + def fake_state_identity( + _country, year, *, dataset, state_code, state_fips, filter_variable + ): + assert (state_fips, filter_variable) in { + (49, "state_fips"), + (56, "state_fips"), + } + digest = f"sd-{year}-{state_code}" + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/us/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + artifact_keys, "collect_state_dataset_identity", fake_state_identity + ) + monkeypatch.setattr( + release_bundle, + "get_country_release_bundle", + lambda _country: SimpleNamespace( + policyengine_version="4.22.0", + model_version="1.0.0", + data_version="1.2.3", + data_artifact_revision="rev-abc", + default_dataset="populace_cps", + ), + ) + return SimpleNamespace(existing=existing, targets=targets) + + @pytest.mark.usefixtures("planning_stubs") + def test_plan_lists_every_state_dataset_year_major_with_presence(self): + plan = precompute.plan_artifacts_impl("bucket-x") + + assert [(entry.year, entry.state_code) for entry in plan.state_datasets] == [ + (2026, "UT"), + (2026, "WY"), + (2027, "UT"), + (2027, "WY"), + (2025, "UT"), + (2025, "WY"), + ] + assert [entry.exists for entry in plan.state_datasets] == [ + True, + False, + False, + False, + False, + False, + ] + first = plan.state_datasets[0] + assert first.region == "state/ut" + assert first.dataset == "populace_local" + assert first.digest == "sd-2026-UT" + assert first.path == "state-datasets/us/sd-2026-UT/populace_local_year_2026.h5" + assert first.filename == "populace_local_year_2026.h5" + assert first.folder == "state-datasets/populace_local/UT" + work = precompute.select_work(plan, force=False) + assert len(work.state_datasets) == 5 + + def test_states_on_the_default_dataset_add_no_state_datasets( + self, planning_stubs, monkeypatch + ): + monkeypatch.setattr(precompute_planning, "state_dataset_targets", list) + + plan = precompute.plan_artifacts_impl("bucket-x") + + assert plan.state_datasets == [] + assert len(plan.datasets) == 3 + + @pytest.mark.usefixtures("planning_stubs") + def test_plan_enumerates_years_and_cohorts_with_presence(self): + plan = precompute.plan_artifacts_impl("bucket-x") + + assert [entry.year for entry in plan.datasets] == [2026, 2027, 2025] + assert [entry.exists for entry in plan.datasets] == [True, False, False] + assert plan.datasets[0].digest == "ds-2026" + assert plan.datasets[0].path == "datasets/us/ds-2026/populace_year_2026.h5" + assert plan.datasets[0].filename == "populace_year_2026.h5" + + # Year-major, partition order inside each year. + assert [entry.simulation_id for entry in plan.baselines] == [ + "bl1-2026-1", + "bl1-2026-2", + "bl1-2027-1", + "bl1-2027-2", + "bl1-2025-1", + "bl1-2025-2", + ] + assert [entry.exists for entry in plan.baselines] == [ + True, + False, + False, + False, + False, + False, + ] + assert plan.baselines[0].group == ["state/ca"] + assert plan.baselines[1].group == ["state/ny", "state/wv"] + assert plan.baselines[1].path == "baselines/us/bl-2026-2/bl1-2026-2.h5" + + assert plan.receipt == BundleVersionIdentity( + policyengine_version="4.22.0", + model_version="1.0.0", + data_version="1.2.3", + data_artifact_revision="rev-abc", + default_dataset="populace_cps", + ) + # An exists-flag inversion would make select_work recompute (or, + # worse, skip) the wrong entries — lock the wiring end to end. + work = precompute.select_work(plan, force=False) + assert [entry.year for entry in work.datasets] == [2027, 2025] + assert len(work.baselines) == 5 + + @pytest.mark.usefixtures("planning_stubs") + def test_empty_partition_fails_loudly(self, monkeypatch): + from policyengine_simulation_executor import national_partition + + monkeypatch.setattr( + national_partition, "national_region_groups", lambda _country: [] + ) + with pytest.raises(RuntimeError, match="No national partition"): + precompute.plan_artifacts_impl("bucket-x") + + +def _us_state_region(state_code: str, state_fips: int, dataset_path: str): + from policyengine.core.scoping_strategy import RowFilterStrategy + + return SimpleNamespace( + code=f"state/{state_code.lower()}", + dataset_path=dataset_path, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=state_fips + ), + parent_code="us", + state_code=state_code, + ) + + +class TestStateDatasetTargets: + """Which states get baked files, resolved the way a request resolves them.""" + + @pytest.fixture + def regional_bundle(self): + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + + bundle = get_country_release_bundle("us") + regional = next( + name for name in bundle.dataset_uris if name.endswith("_acs_local") + ) + return SimpleNamespace( + default_uri=bundle.default_dataset_uri, + regional=regional, + regional_uri=bundle.dataset_uris[regional], + ) + + def _install_regions(self, monkeypatch, regions): + from policyengine_simulation_executor import simulation_runtime + + by_code = {region.code: region for region in regions} + country_module = SimpleNamespace( + model=SimpleNamespace( + get_region=by_code.get, + region_registry=SimpleNamespace( + get_by_type=lambda region_type: ( + list(regions) if region_type == "state" else [] + ) + ), + ) + ) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: country_module + ) + + def test_targets_states_whose_requests_load_a_regional_dataset( + self, monkeypatch, regional_bundle + ): + self._install_regions( + monkeypatch, + [ + _us_state_region("WY", 56, regional_bundle.regional_uri), + _us_state_region("UT", 49, regional_bundle.regional_uri), + ], + ) + + targets = precompute.state_dataset_targets() + + assert targets == [ + precompute.StateDatasetTarget( + region="state/ut", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + dataset=regional_bundle.regional, + ), + precompute.StateDatasetTarget( + region="state/wy", + state_code="WY", + state_fips=56, + filter_variable="state_fips", + dataset=regional_bundle.regional, + ), + ] + + def test_skips_states_whose_requests_load_the_default_dataset( + self, monkeypatch, regional_bundle + ): + self._install_regions( + monkeypatch, + [ + _us_state_region("UT", 49, regional_bundle.default_uri), + _us_state_region("WY", 56, regional_bundle.regional_uri), + ], + ) + + assert [target.state_code for target in precompute.state_dataset_targets()] == [ + "WY" + ] + + def test_refuses_a_state_without_a_row_filter(self, monkeypatch, regional_bundle): + unfiltered = _us_state_region("UT", 49, regional_bundle.regional_uri) + unfiltered.scoping_strategy = None + self._install_regions(monkeypatch, [unfiltered]) + + with pytest.raises(RuntimeError, match="state/ut"): + precompute.state_dataset_targets() + + +class TestCohortIdentityGuard: + def test_diverged_predicates_fail_loudly(self, monkeypatch): + from policyengine_simulation_executor import baseline_artifacts as ba + from policyengine_simulation_executor import simulation_runtime as sr + + monkeypatch.setattr(sr, "_country_module", lambda _country: SimpleNamespace()) + monkeypatch.setattr( + sr, + "_resolve_region", + lambda **_kwargs: SimpleNamespace( + code="x", dataset_reference=None, scoping_strategy=None + ), + ) + monkeypatch.setattr( + ba, "qualifying_baseline_identity", lambda *_args, **_kwargs: None + ) + with pytest.raises(RuntimeError, match="diverged"): + precompute.cohort_identity(2026, ["state/ca"]) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 3291fb84d..621481fad 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -9,7 +9,12 @@ import pytest from google.api_core.exceptions import NotFound, PreconditionFailed from policyengine_simulation_executor import simulation_runtime -from policyengine_simulation_executor.stage12_cache import precompute +from policyengine_simulation_executor.stage12_cache import ( + precompute, + precompute_baselines, + precompute_datasets, + precompute_planning, +) from policyengine_simulation_executor.stage12_cache.models import ( BaselinePlanEntry, CacheBundleIdentity, @@ -153,11 +158,18 @@ def bucket(self, _name: str): def _use_store(monkeypatch, client: _FakeClient) -> None: - monkeypatch.setattr( + # Every precompute module that opens the cache store. + for module in ( precompute, - "Stage12CacheStore", - lambda bucket: Stage12CacheStore(bucket, client=client), - ) + precompute_planning, + precompute_datasets, + precompute_baselines, + ): + monkeypatch.setattr( + module, + "Stage12CacheStore", + lambda bucket: Stage12CacheStore(bucket, client=client), + ) def _store_every_planned_file(client: _FakeClient, plan: PrecomputePlan) -> None: @@ -228,10 +240,10 @@ def stubbed_planning(monkeypatch): """Stub everything plan_artifacts_impl reads except the state targets.""" client = _FakeClient() _use_store(monkeypatch, client) - monkeypatch.setattr(precompute, "load_stage12_bundle", lambda: "resolved") - monkeypatch.setattr(precompute, "_bundle_identity", _bundle) + monkeypatch.setattr(precompute_planning, "load_stage12_bundle", lambda: "resolved") + monkeypatch.setattr(precompute_planning, "_bundle_identity", _bundle) monkeypatch.setattr( - precompute, + precompute_planning, "collect_dataset_identity", lambda _resolved, *, country, year: SimpleNamespace( digest=f"dataset-{year}", @@ -239,10 +251,10 @@ def stubbed_planning(monkeypatch): filename=f"populace_year_{year}.h5", ), ) - monkeypatch.setattr(precompute, "US_REGION_GROUPS", (("state/ca",),)) + monkeypatch.setattr(precompute_planning, "US_REGION_GROUPS", (("state/ca",),)) monkeypatch.setattr( - precompute, - "_baseline_identity", + precompute_planning, + "baseline_identity", lambda year, _group: ( SimpleNamespace(output_plan=None), SimpleNamespace( @@ -253,7 +265,9 @@ def stubbed_planning(monkeypatch): ), ), ) - monkeypatch.setattr(precompute, "stage12_output_plan_sha256", lambda _: "b" * 64) + monkeypatch.setattr( + precompute_planning, "stage12_output_plan_sha256", lambda _: "b" * 64 + ) def fake_state_identity(resolved, *, country, year, dataset, **state_filter): assert resolved == "resolved" @@ -268,7 +282,7 @@ def fake_state_identity(resolved, *, country, year, dataset, **state_filter): ) monkeypatch.setattr( - precompute, "collect_state_dataset_identity", fake_state_identity + precompute_planning, "collect_state_dataset_identity", fake_state_identity ) targets = [ precompute.StateDatasetTarget( @@ -286,7 +300,9 @@ def fake_state_identity(resolved, *, country, year, dataset, **state_filter): dataset=REGIONAL_DATASET, ), ] - monkeypatch.setattr(precompute, "state_dataset_targets", lambda: tuple(targets)) + monkeypatch.setattr( + precompute_planning, "state_dataset_targets", lambda: tuple(targets) + ) return SimpleNamespace(client=client, targets=targets) @@ -326,7 +342,7 @@ def test_plan_lists_every_state_dataset_year_major_with_presence( def test_states_on_the_default_dataset_add_no_state_datasets( stubbed_planning, monkeypatch ) -> None: - monkeypatch.setattr(precompute, "state_dataset_targets", tuple) + monkeypatch.setattr(precompute_planning, "state_dataset_targets", tuple) assert precompute.plan_artifacts_impl("cache-bucket").state_datasets == () @@ -366,7 +382,9 @@ def region(state_code, fips, dataset_path): ), ) ) - monkeypatch.setattr(precompute, "_country_module", lambda _country: country_module) + monkeypatch.setattr( + precompute_planning, "_country_module", lambda _country: country_module + ) targets = precompute.state_dataset_targets() @@ -850,10 +868,10 @@ def test_determinism_check_records_only_a_pass_that_loaded_the_stored_file( probe = plan.baselines[0] (tmp_path / probe.filename).write_bytes(b"h5") monkeypatch.setattr(simulation_runtime, "resolve_data_folder", lambda: tmp_path) - monkeypatch.setattr(precompute, "_ensure_dataset_file", lambda *_: None) + monkeypatch.setattr(precompute_baselines, "_ensure_dataset_file", lambda *_: None) monkeypatch.setattr( - precompute, - "_scoped_planned_simulation", + precompute_baselines, + "scoped_planned_simulation", lambda _year, _group, *, reform=False: "reform" if reform else "baseline", ) calculations = { @@ -880,10 +898,10 @@ def test_determinism_check_records_only_a_pass_that_loaded_the_stored_file( def _use_determinism_calculations(monkeypatch, tmp_path, cached, fresh) -> None: monkeypatch.setattr(simulation_runtime, "resolve_data_folder", lambda: tmp_path) - monkeypatch.setattr(precompute, "_ensure_dataset_file", lambda *_: None) + monkeypatch.setattr(precompute_baselines, "_ensure_dataset_file", lambda *_: None) monkeypatch.setattr( - precompute, - "_scoped_planned_simulation", + precompute_baselines, + "scoped_planned_simulation", lambda _year, _group, *, reform=False: "reform" if reform else "baseline", ) calculations = { @@ -954,7 +972,7 @@ def download_file(self, path, local_path): def write_json(self, *_): pass - monkeypatch.setattr(precompute, "Stage12CacheStore", RecordingStore) + monkeypatch.setattr(precompute_baselines, "Stage12CacheStore", RecordingStore) plan = _plan(exists=True, verified=False) probe = plan.baselines[0] if already_local: @@ -1006,7 +1024,9 @@ def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( _use_store(monkeypatch, client) plan = precompute.plan_artifacts_impl("cache-bucket") entry = plan.baselines[0] - planned = precompute._scoped_planned_simulation(entry.year, entry.region_codes) + planned = precompute_planning.scoped_planned_simulation( + entry.year, entry.region_codes + ) params = simulation_module.simulation_runtime_params(planned) country_module = simulation_runtime._country_module("us") region = simulation_runtime._resolve_region( @@ -1078,9 +1098,11 @@ def target(state_code, fips): target("WY", sources.WYOMING_FIPS), ] monkeypatch.setattr( - precompute, "state_dataset_targets", lambda: tuple(state.targets) + precompute_datasets, "state_dataset_targets", lambda: tuple(state.targets) + ) + monkeypatch.setattr( + precompute_datasets, "load_stage12_bundle", lambda: "resolved" ) - monkeypatch.setattr(precompute, "load_stage12_bundle", lambda: "resolved") def fake_state_identity(resolved, *, country, year, dataset, **state_filter): digest = f"state-{year}-{state_filter['state_code']}" @@ -1091,7 +1113,7 @@ def fake_state_identity(resolved, *, country, year, dataset, **state_filter): ) monkeypatch.setattr( - precompute, "collect_state_dataset_identity", fake_state_identity + precompute_datasets, "collect_state_dataset_identity", fake_state_identity ) def fake_ensure_datasets(*, datasets, years, data_folder): @@ -1116,7 +1138,7 @@ def fake_ensure_datasets(*, datasets, years, data_folder): for geoid in (4901, 4902, 4903, 4904, 5600) ] monkeypatch.setattr( - precompute, + precompute_datasets, "_country_module", lambda _country: SimpleNamespace( ensure_datasets=fake_ensure_datasets, diff --git a/rules/baselines/python-file-length.txt b/rules/baselines/python-file-length.txt index dd28cdefd..dc30e2766 100644 --- a/rules/baselines/python-file-length.txt +++ b/rules/baselines/python-file-length.txt @@ -6,8 +6,5 @@ # baseline count, in the pull request that splits the file. projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py projects/policyengine-simulation-entry/tests/test_deployment_assets.py -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py -projects/policyengine-simulation-executor/tests/test_precompute.py projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py diff --git a/rules/registry.toml b/rules/registry.toml index 858f1ef21..c3fced421 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1190,7 +1190,7 @@ job = "quality-guards" invocation = "run_quality_guards.py" [rule.baseline] file = "rules/baselines/python-file-length.txt" -count = 7 +count = 4 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.python-file-length/bad_files.toml" good_fixture = "rules/fixtures/sim.quality.python-file-length/good_files.toml"