From 668bfbbb1c6337c3922771f6cd57d67a9c4e5ab2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:59:22 +0400 Subject: [PATCH 01/66] Import test modules with from-imports instead of aliases Co-Authored-By: Claude Opus 5.5 --- .../tests/test_hf_dataset.py | 2 +- .../tests/test_canonical_spm.py | 11 +++++------ .../tests/test_modal_bundle_image.py | 2 +- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py index 1b6a6efa4..90322da7a 100644 --- a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py +++ b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py @@ -4,8 +4,8 @@ import json -import policyengine_simulation_contract.hf_dataset as hf_dataset import pytest +from policyengine_simulation_contract import hf_dataset from policyengine_simulation_contract.hf_dataset import ( HuggingFaceDatasetReferenceError, parse_hf_dataset_uri, diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index 662ff425d..c54a5ec19 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -272,7 +272,7 @@ def test_result_transport_rejects_legacy_sibling_config(): def test_sync_compatibility_endpoint_preserves_explicit_null( monkeypatch, observability_runtime ): - import policyengine_simulation_executor.simulation as simulation + from policyengine_simulation_executor import simulation captured = [] @@ -350,7 +350,7 @@ def _stub_bundle(monkeypatch): @staticmethod def _stub_provider(monkeypatch, capability=CAPABILITY): """Count the two pieces of work a resolution repeats.""" - import policyengine_simulation_executor.spm as spm + from policyengine_simulation_executor import spm year_metadata = Mock(return_value={}) constructions = Mock( @@ -484,7 +484,7 @@ def test_collecting_a_baseline_identity_prevalidates_once(self, monkeypatch): def test_the_installed_capability_is_read_once(self, monkeypatch): import policyengine.bundle - import policyengine_simulation_executor.spm as spm + from policyengine_simulation_executor import spm reads = Mock(return_value={"measurements": {}}) monkeypatch.setattr(policyengine.bundle, "get_current_bundle", reads) @@ -498,7 +498,7 @@ def test_the_installed_capability_is_read_once(self, monkeypatch): def test_an_unavailable_capability_is_rederived_every_time(self, monkeypatch): import policyengine.bundle - import policyengine_simulation_executor.spm as spm + from policyengine_simulation_executor import spm reads = Mock(side_effect=ValueError("bundle configuration unavailable")) monkeypatch.setattr(policyengine.bundle, "get_current_bundle", reads) @@ -608,8 +608,7 @@ def test_fanout_and_budget_window_preserve_settings_and_receipts(): def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): import policyengine.core - import policyengine_simulation_executor.spm as spm - from policyengine_simulation_executor import baseline_artifacts + from policyengine_simulation_executor import baseline_artifacts, spm calls = [] monkeypatch.setattr(spm, "normalize_runtime_spm", lambda _params: SELECTION) diff --git a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py index 9220a8b07..e48359aaf 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py @@ -175,7 +175,7 @@ def read_manifest(self, digest): # The app imports ArtifactStore lazily inside the digest-set branch, # so patching the module attribute is enough — no GCS involved. - import policyengine_simulation_executor.artifact_store as artifact_store + from policyengine_simulation_executor import artifact_store monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) sys.modules.pop("src.modal.app", None) From dc32783d7176935a0db711eb54533ba01722ac28 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:59:22 +0400 Subject: [PATCH 02/66] Group test helper parameters into composite fixtures and option objects Split the long deployment-workflow test into focused tests and hoist closure-free test doubles to module level. Co-Authored-By: Claude Opus 5.5 --- .../tests/simulation/conftest.py | 24 +++ .../tests/simulation/test_budget_window.py | 9 +- .../tests/test_deploy_workflow_scripts.py | 20 ++- .../tests/test_deployment_assets.py | 54 ++++-- ...est_policyengine_package_update_scripts.py | 35 ++-- .../tests/test_budget_window_batch.py | 27 +-- .../tests/test_budget_window_scheduler.py | 53 +++--- ...est_policyengine_package_update_scripts.py | 7 +- .../tests/test_precompute.py | 163 ++++++++---------- .../tests/test_simulation_output_builder.py | 35 ++-- .../tests/test_stage12_runtime.py | 45 +++-- .../tests/test_spm_routes.py | 15 +- 12 files changed, 280 insertions(+), 207 deletions(-) diff --git a/projects/policyengine-apis-integ/tests/simulation/conftest.py b/projects/policyengine-apis-integ/tests/simulation/conftest.py index f5b6d5183..4903df629 100644 --- a/projects/policyengine-apis-integ/tests/simulation/conftest.py +++ b/projects/policyengine-apis-integ/tests/simulation/conftest.py @@ -1,5 +1,6 @@ import json import time +from dataclasses import dataclass from http import HTTPStatus import httpx @@ -169,6 +170,29 @@ def budget_window_request(us_model_version: str) -> BudgetWindowBatchRequest: ) +@dataclass(frozen=True) +class BudgetWindowSmokeCase: + """The staging budget-window smoke request and what its result must show.""" + + request: BudgetWindowBatchRequest + years: list[str] + model_version: str + + +@pytest.fixture +def budget_window_smoke_case( + budget_window_request: BudgetWindowBatchRequest, + budget_window_years: list[str], + us_model_version: str, +) -> BudgetWindowSmokeCase: + """Group the smoke request with its expected years and model version.""" + return BudgetWindowSmokeCase( + request=budget_window_request, + years=budget_window_years, + model_version=us_model_version, + ) + + @pytest.fixture def decode_response_content(): """Return a compact formatter for non-OK HTTP response payloads.""" diff --git a/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py b/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py index 19c4ae701..e6ef437d1 100644 --- a/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py +++ b/projects/policyengine-apis-integ/tests/simulation/test_budget_window.py @@ -18,12 +18,10 @@ @pytest.mark.beta_only def test_budget_window_multi_year_batch_completes( - budget_window_request, - budget_window_years, + budget_window_smoke_case, decode_response_content, submit_budget_window_batch, poll_budget_window_batch, - us_model_version: str, ): """A two-year US budget-window batch returns both annual impacts and totals. @@ -31,7 +29,8 @@ def test_budget_window_multi_year_batch_completes( When the batch is submitted and polled to completion Then the response contains 2026 and 2027 annual impacts plus totals. """ - submit_response = submit_budget_window_batch(budget_window_request) + budget_window_years = budget_window_smoke_case.years + submit_response = submit_budget_window_batch(budget_window_smoke_case.request) assert submit_response.status_code == HTTPStatus.OK, ( "Unexpected submit status " @@ -42,7 +41,7 @@ def test_budget_window_multi_year_batch_completes( f"Unexpected response type: {type(submit_response.parsed)}" ) assert submit_response.parsed.status == "submitted" - assert submit_response.parsed.version == us_model_version + assert submit_response.parsed.version == budget_window_smoke_case.model_version batch_job_id = submit_response.parsed.batch_job_id assert submit_response.parsed.poll_url == f"/budget-window-jobs/{batch_job_id}" diff --git a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py index 3ad8ab6ee..f861fad15 100644 --- a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py @@ -427,20 +427,22 @@ def test_deploy_workflow_checks_stage12_results_with_the_tested_script(): @pytest.mark.parametrize( - ("requested", "configuration", "deployment", "publication", "passes"), + ("requested", "results", "passes"), [ - ("true", "success", "success", "success", True), - ("true", "success", "success", "failure", False), - ("true", "success", "skipped", "skipped", False), - ("true", "skipped", "skipped", "skipped", False), - ("false", "skipped", "skipped", "skipped", True), - ("false", "success", "skipped", "skipped", False), - ("false", "skipped", "skipped", "failure", False), + ("true", ("success", "success", "success"), True), + ("true", ("success", "success", "failure"), False), + ("true", ("success", "skipped", "skipped"), False), + ("true", ("skipped", "skipped", "skipped"), False), + ("false", ("skipped", "skipped", "skipped"), True), + ("false", ("success", "skipped", "skipped"), False), + ("false", ("skipped", "skipped", "failure"), False), ], ) def test_stage12_ready_requires_all_or_none_of_the_v2_jobs( - fakes, requested, configuration, deployment, publication, passes + fakes, requested, results, passes ): + configuration, deployment, publication = results + result = run_script( STAGE12_READY, env=fakes.env( diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index 41b22d979..370fcfddf 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -310,12 +310,9 @@ def test_stage12_runtime_uses_typed_dml_without_raw_sql_or_schema_changes() -> N assert "create_all(" not in sources -def test_deployment_uses_gcloud_workflow_without_terraform(): - deploy_workflow = DEPLOY_WORKFLOW.read_text(encoding="utf-8") - reusable_workflow = REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8") - stage12_workflow = STAGE12_DEPLOY_WORKFLOW.read_text(encoding="utf-8") - cloud_run_deploy_script = CLOUD_RUN_DEPLOY_SCRIPT.read_text(encoding="utf-8") - +def _assert_gcloud_deployment_jobs( + reusable_workflow: str, cloud_run_deploy_script: str +) -> None: assert "deploy-cloud-run-simulation-entry.sh" in reusable_workflow assert "gcloud run deploy" in cloud_run_deploy_script assert "set-cloud-run-simulation-entry-revision.sh" in reusable_workflow @@ -332,6 +329,11 @@ def test_deployment_uses_gcloud_workflow_without_terraform(): "needs: [prepare, deploy_entrypoint, deploy_gateway, deploy_executor]" in reusable_workflow ) + + +def _assert_beta_then_prod_release( + deploy_workflow: str, reusable_workflow: str +) -> None: assert "needs: beta" in deploy_workflow assert deploy_workflow.index("beta:") < deploy_workflow.index("prod:") assert "release_environment: beta" in deploy_workflow @@ -347,6 +349,9 @@ def test_deployment_uses_gcloud_workflow_without_terraform(): assert "entrypoint_public_url" not in reusable_workflow assert "SIMULATION_ENTRYPOINT_PUBLIC_URL" not in reusable_workflow assert "staging.simulation.api.policyengine.org" not in deploy_workflow + + +def _assert_deployment_environment_settings(reusable_workflow: str) -> None: assert ( # Sentinel, restore and final routing acknowledgement use secrets too. reusable_workflow.count("environment: ${{ inputs.release_environment }}") == 15 @@ -367,6 +372,31 @@ def test_deployment_uses_gcloud_workflow_without_terraform(): reusable_workflow.count("vars.OLD_GATEWAY_AUTH_CLIENT_SECRET_SECRET_NAME") == 1 ) assert "simulation-entry-old-gateway-client-secret" not in reusable_workflow + + +def _assert_stage12_deploy_targets(stage12_workflow: str) -> None: + assert "deploy-stage12-staging:" in stage12_workflow + assert "deploy-stage12-production:" in stage12_workflow + assert stage12_workflow.count("deploy_existing_stack: false") == 2 + assert stage12_workflow.count("deploy_stage12_v2: true") == 2 + assert "release_environment: beta" in stage12_workflow + assert "deployment_environment: staging" in stage12_workflow + assert "modal_environment: staging" in stage12_workflow + assert "release_environment: prod" in stage12_workflow + assert "deployment_environment: production" in stage12_workflow + assert "modal_environment: main" in stage12_workflow + assert "DEPLOY_STAGE12_PRODUCTION" in stage12_workflow + + +def test_deployment_uses_gcloud_workflow_without_terraform(): + deploy_workflow = DEPLOY_WORKFLOW.read_text(encoding="utf-8") + reusable_workflow = REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8") + stage12_workflow = STAGE12_DEPLOY_WORKFLOW.read_text(encoding="utf-8") + cloud_run_deploy_script = CLOUD_RUN_DEPLOY_SCRIPT.read_text(encoding="utf-8") + + _assert_gcloud_deployment_jobs(reusable_workflow, cloud_run_deploy_script) + _assert_beta_then_prod_release(deploy_workflow, reusable_workflow) + _assert_deployment_environment_settings(reusable_workflow) assert AUTHENTICATED_TESTS.joinpath("test_deployed_auth.py").is_file() assert not list( (REPOSITORY_ROOT / "projects" / "policyengine-simulation-entry" / "infra").glob( @@ -379,17 +409,7 @@ def test_deployment_uses_gcloud_workflow_without_terraform(): / "scripts" / "configure-cloud-run-simulation-entry-domains.sh" ).exists() - assert "deploy-stage12-staging:" in stage12_workflow - assert "deploy-stage12-production:" in stage12_workflow - assert stage12_workflow.count("deploy_existing_stack: false") == 2 - assert stage12_workflow.count("deploy_stage12_v2: true") == 2 - assert "release_environment: beta" in stage12_workflow - assert "deployment_environment: staging" in stage12_workflow - assert "modal_environment: staging" in stage12_workflow - assert "release_environment: prod" in stage12_workflow - assert "deployment_environment: production" in stage12_workflow - assert "modal_environment: main" in stage12_workflow - assert "DEPLOY_STAGE12_PRODUCTION" in stage12_workflow + _assert_stage12_deploy_targets(stage12_workflow) def test_cloud_run_deployment_does_not_interpolate_github_values_in_shell(): diff --git a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py index 3aa97ee0b..1043c9956 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py @@ -4,6 +4,7 @@ import os import subprocess +from dataclasses import dataclass from typing import TYPE_CHECKING import pytest @@ -191,22 +192,30 @@ def install_fake_gh( ) +@dataclass(frozen=True) +class BundledVersions: + """Package versions and UK dataset digest the fake ``uv`` reports for 4.1.0.""" + + policyengine_version: str = "4.1.0" + core_version: str = "999.999.999" + spm_version: str = "1.0.0" + us_version: str = "1.1.0" + uk_version: str = "2.1.0" + uk_dataset_sha256: str = "a" * 64 + + def install_fake_uv( fake_bin: Path, *, log: Path, - bundled_core_version: str = "999.999.999", - bundled_us_version: str = "1.1.0", - bundled_uk_version: str = "2.1.0", - bundled_policyengine_version: str = "4.1.0", - bundled_spm_version: str = "1.0.0", - bundled_uk_dataset_sha256: str = "a" * 64, + bundled: BundledVersions | None = None, ) -> None: """Install a fake ``uv`` that appends each call to ``log``. Only the updater's isolated ``bundle-versions`` query for policyengine - 4.1.0 prints anything: the ``bundled_*`` versions as ``key=value`` lines. + 4.1.0 prints anything: the ``bundled`` versions as ``key=value`` lines. """ + bundled = bundled or BundledVersions() write_executable( fake_bin / "uv", f"""#!/usr/bin/env bash @@ -216,16 +225,16 @@ def install_fake_uv( if [[ "$1" == "run" \ && "$*" == *"--isolated --no-project --with policyengine==4.1.0"* \ && "$*" == *"update_policyengine_package.py bundle-versions"* ]]; then - echo "policyengine_version={bundled_policyengine_version}" - echo "policyengine_core_version={bundled_core_version}" - echo "spm_calculator_version={bundled_spm_version}" - echo "us_version={bundled_us_version}" + echo "policyengine_version={bundled.policyengine_version}" + echo "policyengine_core_version={bundled.core_version}" + echo "spm_calculator_version={bundled.spm_version}" + echo "us_version={bundled.us_version}" echo "us_data_version=1.10.0" - echo "uk_version={bundled_uk_version}" + echo "uk_version={bundled.uk_version}" echo "uk_data_version=1.20.0" echo "uk_data_package_version=1.20.0" echo "us_dataset_sha256=" - echo "uk_dataset_sha256={bundled_uk_dataset_sha256}" + echo "uk_dataset_sha256={bundled.uk_dataset_sha256}" echo "uk_certificate_sha256=" exit 0 fi diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py index af013dd02..115493e4f 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py @@ -88,6 +88,20 @@ def spawn( return call +class MockDict: + def __init__(self, data: dict): + self._data = data + + def __getitem__(self, key: str): + return self._data[key] + + def __setitem__(self, key: str, value): + self._data[key] = value + + def get(self, key: str, default=None): + return self._data.get(key, default) + + @pytest.fixture def mock_batch_modal(monkeypatch): dicts: dict[str, dict] = {} @@ -95,19 +109,6 @@ def mock_batch_modal(monkeypatch): functions: dict[tuple[str, str], object] = {} parent_call_id = "parent-123" - class MockDict: - def __init__(self, data: dict): - self._data = data - - def __getitem__(self, key: str): - return self._data[key] - - def __setitem__(self, key: str, value): - self._data[key] = value - - def get(self, key: str, default=None): - return self._data.get(key, default) - class MockModalDict: @staticmethod def from_name(name: str, *, create_if_missing: bool = False): diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py index 1e25f0f9a..972619109 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py @@ -217,39 +217,40 @@ def spawn(self, payload: dict, *, observability_context=None): return call -@pytest.fixture -def budget_window_semi_integration_client( - monkeypatch, -) -> tuple[TestClient, SemiIntegrationRuntime]: - runtime = SemiIntegrationRuntime() +class PropagatingRuntime: + def __init__(self): + self.context: dict = {} - class PropagatingRuntime: - def __init__(self): - self.context: dict = {} + def span(self, *_args, **_kwargs): + return nullcontext() - def span(self, *_args, **_kwargs): - return nullcontext() + @contextmanager + def operation(self, *_args, remote_context=None, **_kwargs): + previous = self.context + self.context = dict(remote_context or {}) + try: + yield + finally: + self.context = previous - @contextmanager - def operation(self, *_args, remote_context=None, **_kwargs): - previous = self.context - self.context = dict(remote_context or {}) - try: - yield - finally: - self.context = previous + def capture_context(self): + return dict(self.context) - def capture_context(self): - return dict(self.context) + def set_context(self, **attributes): + self.context.update(attributes) - def set_context(self, **attributes): - self.context.update(attributes) + def event(self, *_args, **_kwargs): + return None - def event(self, *_args, **_kwargs): - return None + def record_exception(self, *_args, **_kwargs): + return None - def record_exception(self, *_args, **_kwargs): - return None + +@pytest.fixture +def budget_window_semi_integration_client( + monkeypatch, +) -> tuple[TestClient, SemiIntegrationRuntime]: + runtime = SemiIntegrationRuntime() runtime.observability = PropagatingRuntime() runtime.dicts["simulation-api-us-versions"] = { diff --git a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py index 835239f39..0594c6ee3 100644 --- a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py @@ -9,6 +9,7 @@ import pytest from fixtures.test_policyengine_package_update_scripts import ( SCRIPT, + BundledVersions, install_fake_gh, install_fake_git, install_fake_uv, @@ -185,7 +186,7 @@ def test_update_policyengine_package_opens_pr_for_existing_branch_without_open_p ("uv_kwargs", "git_kwargs", "payload", "message"), [ ( - {"bundled_uk_version": "2.0.9"}, + {"bundled": BundledVersions(uk_version="2.0.9")}, {}, '{"version":"4.1.0","bundle":{"uk":{"model_version":"2.1.0"}}}', "does not match the release dispatch", @@ -198,7 +199,7 @@ def test_update_policyengine_package_opens_pr_for_existing_branch_without_open_p "does not match the release dispatch", ), ( - {"bundled_policyengine_version": "4.0.9"}, + {"bundled": BundledVersions(policyengine_version="4.0.9")}, {}, "", "reports bundle 4.0.9, expected 4.1.0", @@ -313,7 +314,7 @@ def test_update_policyengine_package_stops_when_loaded_bundle_is_not_target( install_fake_uv( fake_bin, log=uv_log, - bundled_policyengine_version="4.0.0", + bundled=BundledVersions(policyengine_version="4.0.0"), ) result = run_updater( diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 8c183fa80..694f954fa 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -949,133 +949,120 @@ class TestVerifyDeterminismImpl: It is the pipeline's only correctness check. """ - def _run( - self, - monkeypatch, - tmp_path, - 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) - ) + @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 - 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, monkeypatch, tmp_path): - verdict, downloads, _ = self._run( - monkeypatch, tmp_path, self._frames(), self._frames() - ) + 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, monkeypatch, tmp_path): - verdict, _, _ = self._run( - monkeypatch, tmp_path, self._frames(), self._frames(age=(30.0, 41.0)) - ) + 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, monkeypatch, tmp_path): + 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, _, _ = self._run(monkeypatch, tmp_path, left, right) + verdict, _, _ = run_verification(left, right) assert verdict.equal is False assert verdict.differences == ["person.age: dtype differs"] - def test_column_asymmetry_fails(self, monkeypatch, tmp_path): + def test_column_asymmetry_fails(self, run_verification): right = {"person": pd.DataFrame({"age": [30.0, 40.0], "income": [1.0, 2.0]})} - verdict, _, _ = self._run(monkeypatch, tmp_path, self._frames(), right) + 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, monkeypatch, tmp_path): - verdict, _, _ = self._run(monkeypatch, tmp_path, self._frames(), {}) + 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, monkeypatch, tmp_path - ): + 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, _, _ = self._run(monkeypatch, tmp_path, left, right) + verdict, _, _ = run_verification(left, right) assert verdict.equal is False assert len(verdict.differences) == 50 - def test_unloadable_artifact_fails_loudly(self, monkeypatch, tmp_path): + def test_unloadable_artifact_fails_loudly(self, run_verification): with pytest.raises(RuntimeError, match="could not load the artifact"): - self._run( - monkeypatch, - tmp_path, + run_verification( self._frames(), self._frames(), outcome="incomplete", ) - def test_absent_artifact_is_downloaded_first(self, monkeypatch, tmp_path): - verdict, downloads, _ = self._run( - monkeypatch, tmp_path, self._frames(), self._frames(), on_disk=False + 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, monkeypatch, tmp_path): - _, _, captured = self._run( - monkeypatch, tmp_path, self._frames(), self._frames() - ) + def test_fresh_run_reuses_the_baseline_inputs(self, run_verification): + _, _, captured = run_verification(self._frames(), self._frames()) assert captured == [ { "dataset": "dataset", diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index fe09f7051..5791fca7f 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -5,6 +5,7 @@ import json from collections import Counter from contextlib import AbstractContextManager +from dataclasses import dataclass from types import SimpleNamespace import pandas as pd @@ -198,17 +199,24 @@ def _macro_baseline_reform(): return baseline, reform +@dataclass(frozen=True) +class BuilderOptions: + """Optional inputs for a SimulationOutputBuilder built from fakes.""" + + analysis: object = None + include_cliffs: bool | None = None + resolved_region_code: str | None = "us" + uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + + def _simulation_output_builder( country: str, baseline, reform, - analysis=None, - *, - include_cliffs: bool | None = None, - resolved_region_code: str | None = "us", - uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, + options: BuilderOptions | None = None, ) -> SimulationOutputBuilder: - analysis = analysis or fake_analysis() + options = options or BuilderOptions() + analysis = options.analysis or fake_analysis() def economic_impact_analysis( _baseline_simulation, @@ -222,8 +230,8 @@ def economic_impact_analysis( "country": country, "data_version": "1.115.5" if country == "us" else "1.55.10", } - if include_cliffs is not None: - simulation_params["include_cliffs"] = include_cliffs + if options.include_cliffs is not None: + simulation_params["include_cliffs"] = options.include_cliffs country_module = SimpleNamespace( model=SimpleNamespace(version="1.715.2" if country == "us" else "2.88.20"), @@ -236,8 +244,8 @@ def economic_impact_analysis( dataset=SimpleNamespace(metadata={}), baseline=baseline, reform=reform, - resolved_region_code=resolved_region_code, - uk_local_authority_metadata=uk_local_authority_metadata, + resolved_region_code=options.resolved_region_code, + uk_local_authority_metadata=options.uk_local_authority_metadata, ) @@ -499,7 +507,7 @@ def fake_build_congressional_district_impact( "us", baseline, reform, - resolved_region_code="state/ut", + BuilderOptions(resolved_region_code="state/ut"), ) assert builder._build_congressional_district_impact() == ( @@ -1030,8 +1038,7 @@ def test_builder_serializes_cliff_impact_when_requested(monkeypatch): "us", baseline, reform, - analysis=analysis, - include_cliffs=True, + BuilderOptions(analysis=analysis, include_cliffs=True), ).build() assert output.cliff_impact is not None @@ -1828,7 +1835,7 @@ def build(country, baseline_simulation, reform_simulation, **kwargs): "uk", baseline, reform, - uk_local_authority_metadata=metadata, + BuilderOptions(uk_local_authority_metadata=metadata), )._build_uk_local_authority_impact() assert result is not None assert result.model_dump(mode="json") == expected diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index db2a0f3e2..7da9ac18f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -7,6 +7,7 @@ import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager, nullcontext +from dataclasses import dataclass from datetime import UTC, datetime, timedelta from hashlib import sha256 from threading import Lock @@ -1204,22 +1205,30 @@ def get(self, *, timeout=None): return self.result +@dataclass(frozen=True) +class InjectedChildFaults: + """Faults a ConcurrentInvoker injects into the child simulations it runs.""" + + fail_role: str | None = None + incompatible: bool = False + wrong_plan_digest: bool = False + missing_column: bool = False + + class ConcurrentInvoker: def __init__( self, artifacts, *, - fail_role=None, - incompatible=False, - wrong_plan_digest=False, - missing_column=False, + faults: InjectedChildFaults | None = None, production_result=None, ): + faults = faults or InjectedChildFaults() self.artifacts = artifacts - self.fail_role = fail_role - self.incompatible = incompatible - self.wrong_plan_digest = wrong_plan_digest - self.missing_column = missing_column + self.fail_role = faults.fail_role + self.incompatible = faults.incompatible + self.wrong_plan_digest = faults.wrong_plan_digest + self.missing_column = faults.missing_column self.executor = ThreadPoolExecutor(max_workers=2) self.intervals = {} self.events = [] @@ -1654,7 +1663,9 @@ def test_repeated_submission_does_not_retry_failed_comparison() -> None: def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None: store = FakeStore() artifacts = FakeArtifacts() - invoker = ConcurrentInvoker(artifacts, fail_role="reform") + invoker = ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) with pytest.raises( RuntimeError, match="Stage 12 report coordination failed" @@ -1688,8 +1699,12 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None @pytest.mark.parametrize( "invoker", [ - lambda artifacts: ConcurrentInvoker(artifacts, wrong_plan_digest=True), - lambda artifacts: ConcurrentInvoker(artifacts, missing_column=True), + lambda artifacts: ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(wrong_plan_digest=True) + ), + lambda artifacts: ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(missing_column=True) + ), ], ) def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( @@ -2169,7 +2184,9 @@ def get_simulation(self, simulation_execution_id): coordinator_invocation_id="coordinator-1", store=store, artifacts=artifacts, - invoker=ConcurrentInvoker(artifacts, fail_role="baseline"), + invoker=ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(fail_role="baseline") + ), aggregator=lambda **_: _aggregate_payload({"must": "not run"}), runtime=TelemetryRecordingRuntime(ordered_events), ) @@ -2238,7 +2255,9 @@ def test_coordinator_reports_a_child_dispatch_failure_once() -> None: def test_coordinator_rejects_incompatible_row_identity() -> None: store = FakeStore() artifacts = FakeArtifacts() - invoker = ConcurrentInvoker(artifacts, incompatible=True) + invoker = ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(incompatible=True) + ) with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): coordinate_report( diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py index f16ec7bbf..1dfbfea23 100644 --- a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -52,12 +52,13 @@ def legacy_route(mock_modal, source, model_version): } -@pytest.mark.parametrize(("endpoint", "extra"), ENDPOINTS) +@pytest.mark.parametrize("endpoint_case", ENDPOINTS) @pytest.mark.parametrize("source", LEGACY_SOURCES) @pytest.mark.parametrize("version", [None, "1.715.2"]) def test_historical_no_wrapper_route_submits_without_spm( - mock_modal, client, endpoint, extra, source, version + mock_modal, client, endpoint_case, source, version ): + endpoint, extra = endpoint_case legacy_route(mock_modal, source, "1.715.2") response = client.post( endpoint, json={"country": "us", "version": version, **extra} @@ -68,12 +69,13 @@ def test_historical_no_wrapper_route_submits_without_spm( assert "policyengine_version" not in response.json()["policyengine_bundle"] -@pytest.mark.parametrize(("endpoint", "extra"), ENDPOINTS) +@pytest.mark.parametrize("endpoint_case", ENDPOINTS) @pytest.mark.parametrize("source", LEGACY_SOURCES) @pytest.mark.parametrize("selection", [{}, {"geography_kind": "national"}]) def test_historical_route_rejects_any_explicit_spm( - mock_modal, client, endpoint, extra, source, selection + mock_modal, client, endpoint_case, source, selection ): + endpoint, extra = endpoint_case legacy_route(mock_modal, source, "1.715.2") response = client.post(endpoint, json={"country": "us", "spm": selection, **extra}) assert response.status_code == 400 @@ -81,7 +83,7 @@ def test_historical_route_rejects_any_explicit_spm( assert mock_modal["func"].calls == [] -@pytest.mark.parametrize(("endpoint", "extra"), ENDPOINTS) +@pytest.mark.parametrize("endpoint_case", ENDPOINTS) @pytest.mark.parametrize( ("source", "model_version"), [ @@ -93,8 +95,9 @@ def test_historical_route_rejects_any_explicit_spm( ], ) def test_missing_wrapper_does_not_certify_unknown_routes( - mock_modal, client, endpoint, extra, source, model_version + mock_modal, client, endpoint_case, source, model_version ): + endpoint, extra = endpoint_case legacy_route(mock_modal, source, model_version) response = client.post(endpoint, json={"country": "us", **extra}) assert response.status_code == 400 From d257ff2c81207aad4d28f5fd471f988ee7df6730 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:02:35 +0400 Subject: [PATCH 03/66] Pin contract and persistence validation behaviour before refactoring Characterize every branch of the comparison report record's state validation, the live column type checks, Hugging Face URI parsing, the SPM non-US and historical-wrapper rules, and the v2 application name limits. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_hf_dataset.py | 29 ++ .../tests/test_spm_selection.py | 57 ++++ .../tests/test_stage12_manifest.py | 21 ++ .../tests/test_stage12_report_record.py | 263 ++++++++++++++++++ .../tests/test_column_type_validation.py | 105 +++++++ 5 files changed, 475 insertions(+) create mode 100644 libs/policyengine-simulation-contract/tests/test_stage12_report_record.py create mode 100644 libs/policyengine-stage12-persistence/tests/test_column_type_validation.py diff --git a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py index 90322da7a..f575423fc 100644 --- a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py +++ b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py @@ -30,6 +30,35 @@ def read(self): return json.dumps(self.payload).encode("utf-8") +def test_parse_hf_dataset_uri_ignores_other_schemes(): + assert parse_hf_dataset_uri("gs://bucket/populace_us_2024.h5") is None + + +@pytest.mark.parametrize( + "dataset_uri", + [ + "hf://policyengine/populace-us", + "hf://policyengine//populace_us_2024.h5", + "hf:///populace-us/populace_us_2024.h5@v1", + "hf://policyengine/populace-us/@v1", + ], +) +def test_parse_hf_dataset_uri_rejects_uris_without_org_repo_and_path(dataset_uri): + with pytest.raises(HuggingFaceDatasetReferenceError) as error: + parse_hf_dataset_uri(dataset_uri) + + assert str(error.value) == f"Invalid Hugging Face dataset URI: {dataset_uri!r}" + + +def test_parse_hf_dataset_uri_keeps_nested_paths_and_the_last_revision(): + parsed = parse_hf_dataset_uri("hf://org/repo/nested/file@name.h5@rev") + + assert parsed is not None + assert parsed.repo_id == "org/repo" + assert parsed.path == "nested/file@name.h5" + assert parsed.revision == "rev" + + def test_parse_hf_dataset_uri_extracts_repo_path_and_revision(): parsed = parse_hf_dataset_uri( "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" diff --git a/libs/policyengine-simulation-contract/tests/test_spm_selection.py b/libs/policyengine-simulation-contract/tests/test_spm_selection.py index 17b701aa0..395803445 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_selection.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_selection.py @@ -202,3 +202,60 @@ def test_ordinary_serialization_is_unchanged(self): assert SPMSelection(geography_kind="national").model_dump(mode="json") == { "geography_kind": "national" } + + +@pytest.mark.parametrize("country", ["uk", "UK"]) +def test_non_us_countries_take_no_spm_settings(country): + assert ( + resolve_spm_selection( + country, + None, + capability=None, + policyengine_version="5.3.0", + model_version="1.824.7", + ) + is None + ) + with pytest.raises(SPMInputError) as error: + resolve_spm_selection( + country, + {}, + capability=None, + policyengine_version="5.3.0", + model_version="1.824.7", + ) + assert error.value.code == "SPM_SETTINGS_INVALID" + assert str(error.value) == "SPM settings are supported only for the US" + + +@pytest.mark.parametrize( + ("wrapper", "historical"), + [ + ("5.1.9", True), + ("4.10.0", True), + ("5.2.0", False), + ("5.2.1", False), + ("5.1", False), + ("5.1.9.0", False), + ("5.1.x", False), + ("v5.1.9", False), + ], +) +def test_wrappers_before_5_2_0_are_historical(wrapper, historical): + """Three numeric parts below 5.2.0 are historical whatever the model says.""" + + def resolve_unpinned(): + return resolve_spm_selection( + "us", + None, + capability=None, + policyengine_version=wrapper, + model_version="9.9.9", + ) + + if historical: + assert resolve_unpinned() is None + else: + with pytest.raises(SPMInputError) as error: + resolve_unpinned() + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 019bac88b..4d8b60547 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -120,6 +120,27 @@ def test_v2_application_name_cannot_collide_with_v1() -> None: assert not v2_application_name("5.2.0").startswith("policyengine-simulation-py") +def test_v2_application_name_allows_exactly_64_characters() -> None: + version = "5." + "1" * 33 + + assert len(v2_application_name(version)) == 64 + with pytest.raises( + ValueError, match=r"^v2 Modal application name exceeds 64 characters$" + ): + v2_application_name(version + "1") + + +@pytest.mark.parametrize("version", ["", ".5.2.0", "5.2.0 ", "5/2"]) +def test_v2_application_name_rejects_versions_a_modal_name_cannot_hold( + version: str, +) -> None: + with pytest.raises( + ValueError, + match=r"^PolicyEngine\.py version cannot form a Modal application name$", + ): + v2_application_name(version) + + def test_publish_writes_one_complete_v2_document() -> None: store = FakeStore() diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_report_record.py b/libs/policyengine-simulation-contract/tests/test_stage12_report_record.py new file mode 100644 index 000000000..c3061315d --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_stage12_report_record.py @@ -0,0 +1,263 @@ +"""Characterization tests for the comparison report record's state validation.""" + +from __future__ import annotations + +import re +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from policyengine_simulation_contract.stage12_execution import ComparisonReportRecord +from pydantic import ValidationError + +NOW = datetime(2026, 9, 22, tzinfo=UTC) +AGGREGATE_OUTPUT = { + "aggregate_output_uri": "gs://bucket/aggregate.json", + "aggregate_output_sha256": "c" * 64, + "aggregate_schema_version": 1, +} +COMPARISON_OUTPUT = { + "comparison_output_uri": "gs://bucket/comparison.json", + "comparison_output_sha256": "d" * 64, + "comparison_schema_version": 1, +} +SUCCEEDED_REPORT = { + "status": "succeeded", + "aggregation_status": "succeeded", + **AGGREGATE_OUTPUT, + "completed_at": NOW, +} + + +def _record_fields(**overrides: Any) -> dict[str, Any]: + return { + "evaluation_id": "00000000-0000-0000-0000-000000000001", + "status": "pending", + "aggregation_status": "not_started", + "environment": "staging", + "calculation_flow": "economy", + "originating_request_id": "request-1", + "production_identity": "job-1", + "worker_version": "5.2.0", + "modal_application": "policyengine-simulation-v2-py5-2-0", + "report_coordinator_callable": "coordinate_report", + "version_manifest_sha256": "a" * 64, + "policyengine_version": "5.2.0", + "country_package_name": "policyengine-us", + "country_package_version": "1.764.6", + "country": "us", + "dataset_identity": "populace_us_2024", + "dataset_uri": "hf://policyengine/data/populace_us_2024.h5@revision", + "data_package_name": "populace-data", + "data_package_version": "0.1.0", + "data_artifact_revision": "revision", + "created_at": NOW, + "updated_at": NOW, + "retention_expires_at": NOW + timedelta(days=30), + **overrides, + } + + +@pytest.mark.parametrize( + "overrides", + [ + pytest.param({}, id="pending"), + pytest.param(SUCCEEDED_REPORT, id="succeeded"), + pytest.param({"status": "failed", "error_code": "E"}, id="failed"), + pytest.param({"status": "skipped", "error_code": "E"}, id="skipped"), + pytest.param({"status": "running", **AGGREGATE_OUTPUT}, id="running-output"), + pytest.param({"comparison_status": "pending"}, id="comparison-pending"), + pytest.param({"comparison_status": "running"}, id="comparison-running"), + pytest.param( + { + "comparison_status": "matched", + **COMPARISON_OUTPUT, + "comparison_completed_at": NOW, + }, + id="comparison-matched", + ), + pytest.param( + { + "comparison_status": "different", + **COMPARISON_OUTPUT, + "comparison_completed_at": NOW, + }, + id="comparison-different", + ), + pytest.param( + { + "comparison_status": "failed", + "comparison_completed_at": NOW, + "comparison_error_code": "E", + "comparison_error_summary": "summary", + }, + id="comparison-failed", + ), + ], +) +def test_consistent_report_states_validate(overrides): + fields = _record_fields(**overrides) + + record = ComparisonReportRecord.model_validate(fields) + + assert record.status == fields["status"] + assert record.comparison_status == fields.get("comparison_status", "not_requested") + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + pytest.param( + {"created_at": NOW.replace(tzinfo=None)}, + "comparison retention timestamps must include a timezone", + id="naive-created-at", + ), + pytest.param( + {"retention_expires_at": NOW}, + "comparison retention must be greater than zero", + id="zero-retention", + ), + pytest.param( + {"retention_expires_at": NOW + timedelta(days=30, seconds=1)}, + "comparison retention must not exceed 30 days", + id="long-retention", + ), + pytest.param( + {"aggregate_output_uri": "gs://bucket/aggregate.json"}, + "aggregate output fields must be supplied together", + id="partial-aggregate-output", + ), + pytest.param( + {**SUCCEEDED_REPORT, "aggregation_status": "running"}, + "a successful report requires successful aggregation", + id="succeeded-without-aggregation", + ), + pytest.param( + { + **SUCCEEDED_REPORT, + "aggregate_output_uri": None, + "aggregate_output_sha256": None, + "aggregate_schema_version": None, + }, + "a successful report requires an aggregate output", + id="succeeded-without-output", + ), + pytest.param( + {**SUCCEEDED_REPORT, "completed_at": None}, + "a successful report requires completed_at", + id="succeeded-without-completed-at", + ), + pytest.param( + {"status": "failed"}, + "a failed or skipped report requires an error_code", + id="failed-without-error-code", + ), + pytest.param( + {"status": "skipped"}, + "a failed or skipped report requires an error_code", + id="skipped-without-error-code", + ), + pytest.param( + {"comparison_output_uri": "gs://bucket/comparison.json"}, + "an incomplete comparison cannot reference an output artifact", + id="not-requested-with-output", + ), + pytest.param( + {"comparison_status": "running", "comparison_completed_at": NOW}, + "an incomplete comparison cannot have completed_at", + id="running-with-completed-at", + ), + pytest.param( + {"comparison_status": "pending", "comparison_error_code": "E"}, + "an incomplete comparison cannot have an error_code", + id="pending-with-error-code", + ), + pytest.param( + {"comparison_error_summary": "summary"}, + "an incomplete comparison cannot have an error_summary", + id="not-requested-with-error-summary", + ), + pytest.param( + {"comparison_status": "matched", "comparison_completed_at": NOW}, + "a completed comparison requires an output artifact", + id="matched-without-output", + ), + pytest.param( + {"comparison_status": "different", **COMPARISON_OUTPUT}, + "a completed comparison requires completed_at", + id="different-without-completed-at", + ), + pytest.param( + { + "comparison_status": "matched", + **COMPARISON_OUTPUT, + "comparison_completed_at": NOW, + "comparison_error_code": "E", + }, + "a completed comparison cannot have an error_code", + id="matched-with-error-code", + ), + pytest.param( + { + "comparison_status": "matched", + **COMPARISON_OUTPUT, + "comparison_completed_at": NOW, + "comparison_error_summary": "summary", + }, + "a completed comparison cannot have an error_summary", + id="matched-with-error-summary", + ), + pytest.param( + { + "comparison_status": "failed", + "comparison_output_sha256": "d" * 64, + "comparison_completed_at": NOW, + "comparison_error_code": "E", + }, + "a failed comparison cannot reference an output artifact", + id="failed-comparison-with-output", + ), + pytest.param( + {"comparison_status": "failed", "comparison_error_code": "E"}, + "a failed comparison requires completed_at", + id="failed-comparison-without-completed-at", + ), + pytest.param( + {"comparison_status": "failed", "comparison_completed_at": NOW}, + "a failed comparison requires an error_code", + id="failed-comparison-without-error-code", + ), + pytest.param( + { + "status": "failed", + "aggregate_schema_version": 1, + "comparison_status": "matched", + }, + "aggregate output fields must be supplied together", + id="aggregate-checked-before-status", + ), + pytest.param( + {"status": "skipped", "comparison_status": "failed"}, + "a failed or skipped report requires an error_code", + id="report-checked-before-comparison", + ), + pytest.param( + { + "comparison_status": "pending", + **COMPARISON_OUTPUT, + "comparison_completed_at": NOW, + "comparison_error_code": "E", + }, + "an incomplete comparison cannot reference an output artifact", + id="comparison-output-checked-first", + ), + ], +) +def test_inconsistent_report_states_are_rejected(overrides, message): + with pytest.raises(ValidationError) as error: + ComparisonReportRecord.model_validate(_record_fields(**overrides)) + + assert error.value.error_count() == 1 + assert re.fullmatch( + f"Value error, {re.escape(message)}", error.value.errors()[0]["msg"] + ) diff --git a/libs/policyengine-stage12-persistence/tests/test_column_type_validation.py b/libs/policyengine-stage12-persistence/tests/test_column_type_validation.py new file mode 100644 index 000000000..febef3684 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_column_type_validation.py @@ -0,0 +1,105 @@ +"""Characterization tests for live column type checks in schema qualification.""" + +from __future__ import annotations + +import pytest +from policyengine_stage12_persistence.validation import ( + Stage12RuntimeAccessError, + _require_column_type, +) +from sqlalchemy import JSON, BigInteger, Boolean, DateTime, Integer, String, Text +from sqlalchemy.dialects.postgresql import ENUM, JSONB, UUID + +STATUS_ENUM = ENUM("pending", "running", name="comparison_status") + + +@pytest.mark.parametrize( + ("expected", "actual"), + [ + pytest.param( + STATUS_ENUM, + ENUM("pending", "running", name="comparison_status"), + id="same-enum", + ), + pytest.param(UUID(), UUID(), id="uuid"), + pytest.param(JSON(), JSON(), id="json"), + pytest.param(JSON(), JSONB(), id="jsonb-for-json"), + pytest.param( + DateTime(timezone=True), DateTime(timezone=True), id="aware-datetime" + ), + pytest.param(String(64), String(64), id="same-length-string"), + pytest.param(String(), Text(), id="unbounded-text-for-string"), + pytest.param(Integer(), Integer(), id="integer"), + pytest.param(Integer(), BigInteger(), id="biginteger-for-integer"), + pytest.param(Boolean(), String(), id="unchecked-type"), + ], +) +def test_matching_column_types_pass(expected, actual): + _require_column_type(expected, actual, column="table.column") + + +@pytest.mark.parametrize( + ("expected", "actual", "message"), + [ + pytest.param( + STATUS_ENUM, + String(), + "table.column has an unexpected type", + id="enum-as-string", + ), + pytest.param( + STATUS_ENUM, + ENUM("pending", "running", name="other_status"), + "table.column has an unexpected enum", + id="enum-name", + ), + pytest.param( + STATUS_ENUM, + ENUM("running", "pending", name="comparison_status"), + "table.column has an unexpected enum", + id="enum-order", + ), + pytest.param( + STATUS_ENUM, + ENUM("pending", name="comparison_status"), + "table.column has an unexpected enum", + id="enum-values", + ), + pytest.param( + UUID(), String(), "table.column has an unexpected type", id="uuid" + ), + pytest.param(JSON(), Text(), "table.column has an unexpected type", id="json"), + pytest.param( + DateTime(timezone=True), + DateTime(), + "table.column has an unexpected type", + id="naive-datetime", + ), + pytest.param( + DateTime(timezone=True), + String(), + "table.column has an unexpected type", + id="datetime", + ), + pytest.param( + String(64), + String(32), + "table.column has an unexpected type", + id="string-length", + ), + pytest.param( + String(64), + Integer(), + "table.column has an unexpected type", + id="string", + ), + pytest.param( + Integer(), String(), "table.column has an unexpected type", id="integer" + ), + ], +) +def test_mismatched_column_types_are_rejected(expected, actual, message): + with pytest.raises(Stage12RuntimeAccessError) as error: + _require_column_type(expected, actual, column="table.column") + + assert str(error.value) == message From 5420a652f8a1337956c56f573d1ee1fd57d2224f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:03:33 +0400 Subject: [PATCH 04/66] Split the comparison report record's state validation into report and comparison phases Co-Authored-By: Claude Opus 5.5 --- .../stage12_execution.py | 101 +++++++++++------- 1 file changed, 62 insertions(+), 39 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py index 7ef082277..cc41fb2a4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py @@ -726,6 +726,12 @@ def validate_state(self) -> ComparisonReportRecord: ``error_code``, and the retention window is at most 30 days. """ _validate_retention(self.created_at, self.retention_expires_at) + self._validate_report_state() + self._validate_comparison_state() + return self + + def _validate_report_state(self) -> None: + """Check the aggregate output fields against the report's status.""" output_fields = ( self.aggregate_output_uri, self.aggregate_output_sha256, @@ -756,55 +762,72 @@ def validate_state(self) -> ComparisonReportRecord: ): msg = "a failed or skipped report requires an error_code" raise ValueError(msg) - comparison_output_fields = ( - self.comparison_output_uri, - self.comparison_output_sha256, - self.comparison_schema_version, - ) + + def _validate_comparison_state(self) -> None: + """Check the result comparison fields against the comparison's status.""" if self.comparison_status in { ResultComparisonStatus.NOT_REQUESTED, ResultComparisonStatus.PENDING, ResultComparisonStatus.RUNNING, }: - if any(value is not None for value in comparison_output_fields): - msg = "an incomplete comparison cannot reference an output artifact" - raise ValueError(msg) - if self.comparison_completed_at is not None: - msg = "an incomplete comparison cannot have completed_at" - raise ValueError(msg) - if self.comparison_error_code is not None: - msg = "an incomplete comparison cannot have an error_code" - raise ValueError(msg) - if self.comparison_error_summary is not None: - msg = "an incomplete comparison cannot have an error_summary" - raise ValueError(msg) + self._validate_incomplete_comparison() elif self.comparison_status in { ResultComparisonStatus.MATCHED, ResultComparisonStatus.DIFFERENT, }: - if not all(value is not None for value in comparison_output_fields): - msg = "a completed comparison requires an output artifact" - raise ValueError(msg) - if self.comparison_completed_at is None: - msg = "a completed comparison requires completed_at" - raise ValueError(msg) - if self.comparison_error_code is not None: - msg = "a completed comparison cannot have an error_code" - raise ValueError(msg) - if self.comparison_error_summary is not None: - msg = "a completed comparison cannot have an error_summary" - raise ValueError(msg) + self._validate_completed_comparison() else: - if any(value is not None for value in comparison_output_fields): - msg = "a failed comparison cannot reference an output artifact" - raise ValueError(msg) - if self.comparison_completed_at is None: - msg = "a failed comparison requires completed_at" - raise ValueError(msg) - if self.comparison_error_code is None: - msg = "a failed comparison requires an error_code" - raise ValueError(msg) - return self + self._validate_failed_comparison() + + @property + def _comparison_output_fields(self) -> tuple[object, ...]: + return ( + self.comparison_output_uri, + self.comparison_output_sha256, + self.comparison_schema_version, + ) + + def _validate_incomplete_comparison(self) -> None: + """Reject outputs, completion and errors on a comparison still to finish.""" + if any(value is not None for value in self._comparison_output_fields): + msg = "an incomplete comparison cannot reference an output artifact" + raise ValueError(msg) + if self.comparison_completed_at is not None: + msg = "an incomplete comparison cannot have completed_at" + raise ValueError(msg) + if self.comparison_error_code is not None: + msg = "an incomplete comparison cannot have an error_code" + raise ValueError(msg) + if self.comparison_error_summary is not None: + msg = "an incomplete comparison cannot have an error_summary" + raise ValueError(msg) + + def _validate_completed_comparison(self) -> None: + """Require an output and completion, and no error, on a finished comparison.""" + if not all(value is not None for value in self._comparison_output_fields): + msg = "a completed comparison requires an output artifact" + raise ValueError(msg) + if self.comparison_completed_at is None: + msg = "a completed comparison requires completed_at" + raise ValueError(msg) + if self.comparison_error_code is not None: + msg = "a completed comparison cannot have an error_code" + raise ValueError(msg) + if self.comparison_error_summary is not None: + msg = "a completed comparison cannot have an error_summary" + raise ValueError(msg) + + def _validate_failed_comparison(self) -> None: + """Require completion and an error code, but no output, on a failure.""" + if any(value is not None for value in self._comparison_output_fields): + msg = "a failed comparison cannot reference an output artifact" + raise ValueError(msg) + if self.comparison_completed_at is None: + msg = "a failed comparison requires completed_at" + raise ValueError(msg) + if self.comparison_error_code is None: + msg = "a failed comparison requires an error_code" + raise ValueError(msg) class ComparisonSimulationRecord(StrictContractModel): From 18867e4758a42bd1f2357828718463940cae43d0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:03:33 +0400 Subject: [PATCH 05/66] Name the contract's URI, version and name-length limits Parse SPM wrapper and model release versions through one helper, and keep the keyword signatures of the public contract builders that the gateway and executor call. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_state.py | 2 +- .../hf_dataset.py | 4 ++- .../policyengine_simulation_contract/spm.py | 31 ++++++++++++------- .../stage12_manifest.py | 8 +++-- 4 files changed, 30 insertions(+), 15 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py index 5d0e0b289..26728f827 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py @@ -49,7 +49,7 @@ def _touch(state: BudgetWindowBatchState) -> BudgetWindowBatchState: return state -def create_initial_batch_state( +def create_initial_batch_state( # noqa: PLR0913 # public contract API: the gateway and the executor's batch scheduler call it with these keywords *, batch_job_id: str, request: BudgetWindowBatchRequest, diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py index 49b053282..ecb4a0723 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/hf_dataset.py @@ -20,6 +20,8 @@ HF_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/") HF_REQUEST_TIMEOUT_SECONDS = 30 +# An hf:// URI names the organization, the repo and the in-repo path. +_URI_ORG_REPO_PATH_PART_COUNT = 3 class HuggingFaceDatasetReferenceError(ValueError): @@ -59,7 +61,7 @@ def parse_hf_dataset_uri(dataset_uri: str) -> HFDatasetReference | None: else (without_scheme, None) ) parts = path_with_repo.split("/", maxsplit=2) - if len(parts) != 3 or not all(parts): + if len(parts) != _URI_ORG_REPO_PATH_PART_COUNT or not all(parts): msg = f"Invalid Hugging Face dataset URI: {dataset_uri!r}" raise HuggingFaceDatasetReferenceError(msg) return HFDatasetReference( diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 09b1462bc..4fffe6b83 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -494,7 +494,21 @@ def _certified_spm_selection( return SPMSelection.model_validate(values).model_dump(mode="json") -def resolve_spm_selection( +# A release version this module compares numerically is major.minor.patch. +_RELEASE_VERSION_PART_COUNT = 3 + + +def _numeric_release_version(version: str | None) -> tuple[int, ...] | None: + """Return ``major.minor.patch`` as integers, or None for any other shape.""" + parts = str(version or "").split(".") + if len(parts) != _RELEASE_VERSION_PART_COUNT or not all( + re.fullmatch(r"[0-9]+", part) for part in parts + ): + return None + return tuple(map(int, parts)) + + +def resolve_spm_selection( # noqa: PLR0913 # public contract API: the gateway and executor call it with these keywords country: str, selection: SPMSelection | Mapping[str, object] | None, *, @@ -514,12 +528,8 @@ def resolve_spm_selection( if capability is None: # Existing versioned worker bundles predate canonical SPM. This does # not certify a future bundle or change the routing registry default. - version_parts = str(policyengine_version or "").split(".") - historical = ( - len(version_parts) == 3 - and all(re.fullmatch(r"[0-9]+", p) for p in version_parts) - and tuple(map(int, version_parts)) < (5, 2, 0) - ) + wrapper_version = _numeric_release_version(policyengine_version) + historical = wrapper_version is not None and wrapper_version < (5, 2, 0) # The last two pre-canonical wrappers pin US 1.764.6. A route the # registry carries no manifest for states no model version at all, # and states nothing that contradicts the wrapper pin; a route that @@ -533,13 +543,12 @@ def resolve_spm_selection( # shape, not from the registry's rewritable generation marker -- and # a pre-canonical US model; absence of capability alone proves # nothing. - model_parts = str(model_version or "").split(".") + model_release = _numeric_release_version(model_version) historical = historical or ( policyengine_version is None and route_provenance in {"legacy-country-dict", "legacy-country-route"} - and len(model_parts) == 3 - and all(re.fullmatch(r"[0-9]+", p) for p in model_parts) - and tuple(map(int, model_parts)) <= (1, 764, 6) + and model_release is not None + and model_release <= (1, 764, 6) ) if selection is None and historical: return None diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 009c29112..488dcccd1 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -24,6 +24,7 @@ ACTIVE_MANIFEST_KEY = "active" V2_APPLICATION_PREFIX = "policyengine-simulation-v2-py" V1_APPLICATION_PREFIX = "policyengine-simulation-py" +_MAX_MODAL_APPLICATION_NAME_LENGTH = 64 _CALLABLE_PATTERN = re.compile(r"^[a-z][a-z0-9_]{0,127}$") ManifestText = Annotated[str, Field(min_length=1, max_length=255)] @@ -165,8 +166,11 @@ def v2_application_name(policyengine_version: str) -> str: name = V2_APPLICATION_PREFIX + policyengine_version.replace(".", "-").replace( "+", "-" ) - if len(name) > 64: - msg = "v2 Modal application name exceeds 64 characters" + if len(name) > _MAX_MODAL_APPLICATION_NAME_LENGTH: + msg = ( + "v2 Modal application name exceeds " + f"{_MAX_MODAL_APPLICATION_NAME_LENGTH} characters" + ) raise ValueError(msg) return name From 96ec811e759927773c07b2fcdb7d985d1b2e22d5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:03:49 +0400 Subject: [PATCH 06/66] Split the live column type check into enum and plain type checks Co-Authored-By: Claude Opus 5.5 --- .../validation.py | 53 ++++++++++--------- 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py index d1f0be1a8..8c96658a1 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/validation.py @@ -44,36 +44,39 @@ class Stage12RuntimeAccessError(RuntimeError): def _require_column_type(expected: object, actual: object, *, column: str) -> None: if isinstance(expected, ENUM): - if not isinstance(actual, ENUM): - msg = f"{column} has an unexpected type" - raise Stage12RuntimeAccessError(msg) - if expected.name != actual.name or tuple(expected.enums) != tuple(actual.enums): - msg = f"{column} has an unexpected enum" - raise Stage12RuntimeAccessError(msg) + _require_enum_type(expected, actual, column=column) return + if not _column_type_matches(expected, actual): + msg = f"{column} has an unexpected type" + raise Stage12RuntimeAccessError(msg) + + +def _require_enum_type(expected: ENUM, actual: object, *, column: str) -> None: + """Require a live enum with the declared name and values in order.""" + if not isinstance(actual, ENUM): + msg = f"{column} has an unexpected type" + raise Stage12RuntimeAccessError(msg) + if expected.name != actual.name or tuple(expected.enums) != tuple(actual.enums): + msg = f"{column} has an unexpected enum" + raise Stage12RuntimeAccessError(msg) + + +def _column_type_matches(expected: object, actual: object) -> bool: + """Return whether a live non-enum column type satisfies the declared one. + + Declared types this module does not check always match. + """ if isinstance(expected, UUID): - if not isinstance(actual, UUID): - msg = f"{column} has an unexpected type" - raise Stage12RuntimeAccessError(msg) - return + return isinstance(actual, UUID) if isinstance(expected, JSON): - if not isinstance(actual, JSON): - msg = f"{column} has an unexpected type" - raise Stage12RuntimeAccessError(msg) - return + return isinstance(actual, JSON) if isinstance(expected, DateTime): - if not isinstance(actual, DateTime) or not actual.timezone: - msg = f"{column} has an unexpected type" - raise Stage12RuntimeAccessError(msg) - return + return isinstance(actual, DateTime) and bool(actual.timezone) if isinstance(expected, String): - if not isinstance(actual, String) or expected.length != actual.length: - msg = f"{column} has an unexpected type" - raise Stage12RuntimeAccessError(msg) - return - if isinstance(expected, Integer) and not isinstance(actual, Integer): - msg = f"{column} has an unexpected type" - raise Stage12RuntimeAccessError(msg) + return isinstance(actual, String) and expected.length == actual.length + if isinstance(expected, Integer): + return isinstance(actual, Integer) + return True def _validate_table_schema( From 7a0435154d998ec9b9dcd9ce605cbb73b76bb266 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:03:52 +0400 Subject: [PATCH 07/66] Keep the observability initializer's keyword-only identity arguments Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_observability/observability.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py index 37360d368..d07de8202 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/observability.py @@ -120,7 +120,7 @@ def build_runtime( return configure(config) -def init_simulation_observability( +def init_simulation_observability( # noqa: PLR0913 # public observability API: the entry, gateway and executor apps pass these keyword-only identity arguments app: FastAPI, *, service_name: str, From c5b5435b1a4154e3f24503d1542214177e763162 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:05:30 +0400 Subject: [PATCH 08/66] Pin the order and messages of entry settings validation Co-Authored-By: Claude Opus 5.5 --- .../tests/test_config.py | 135 ++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/projects/policyengine-simulation-entry/tests/test_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index dc4252119..73149ce33 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -132,3 +132,138 @@ def test_stage12_database_url_requires_the_canonical_scheme(database_url): def test_stage12_retention_is_fixed_at_thirty_days(): with pytest.raises(ConfigurationError, match="exactly 30"): make_settings(stage12_retention_days=31).validate() + + +STAGE12_DATABASE_URL = ( + "postgresql://policyengine_v2_runtime:secret@" + "test.pooler.supabase.com:5432/postgres?sslmode=require" +) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + pytest.param( + { + "environment": "prod", + "auth_required": False, + "auth_audience": "", + "old_gateway_url": "", + }, + "Caller authentication cannot be disabled in production.", + id="production-auth-first", + ), + pytest.param( + {"auth_issuer": "", "old_gateway_url": ""}, + "Set both SIMULATION_ENTRYPOINT_AUTH_ISSUER and " + "SIMULATION_ENTRYPOINT_AUTH_AUDIENCE.", + id="issuer-without-audience", + ), + pytest.param( + {"auth_issuer": "", "auth_audience": "", "old_gateway_url": ""}, + "Caller authentication is required but issuer/audience are missing.", + id="required-auth-without-issuer", + ), + pytest.param( + {"auth_issuer": "https://issuer.example/?x=1", "old_gateway_url": ""}, + "SIMULATION_ENTRYPOINT_AUTH_ISSUER must not contain parameters, a " + "query string, or a fragment.", + id="caller-issuer-before-backend", + ), + pytest.param( + { + "old_gateway_url": "", + "old_gateway_auth_client_id": "", + "old_gateway_auth_client_secret": "", + }, + "Missing old-gateway configuration: OLD_GATEWAY_URL, " + "OLD_GATEWAY_AUTH_CLIENT_ID, OLD_GATEWAY_AUTH_CLIENT_SECRET.", + id="missing-backend", + ), + pytest.param( + { + "old_gateway_auth_issuer": "http://issuer.example/", + "connect_timeout_seconds": 0, + }, + "OLD_GATEWAY_AUTH_ISSUER must be an absolute HTTPS URL.", + id="backend-issuer-before-timeouts", + ), + pytest.param( + {"connect_timeout_seconds": 0, "old_gateway_url": "http://x.modal.run"}, + "Old-gateway timeouts must be positive.", + id="zero-connect-timeout", + ), + pytest.param( + {"request_timeout_seconds": -1.0}, + "Old-gateway timeouts must be positive.", + id="negative-request-timeout", + ), + pytest.param( + { + "old_gateway_url": "https://gateway.example", + "stage12_v2_manifest_name": "simulation-api-routing-state", + }, + "OLD_GATEWAY_URL must point to the existing Modal gateway.", + id="gateway-host-before-stage12", + ), + pytest.param( + { + "stage12_retention_days": 29, + "stage12_v2_manifest_environment": "staging", + }, + "STAGE12_RETENTION_DAYS must be exactly 30.", + id="retention-before-resources", + ), + pytest.param( + {"stage12_artifact_bucket": "bucket", "stage12_enabled": True}, + "Stage 12 resources require: STAGE12_V2_MANIFEST_ENVIRONMENT, " + "STAGE12_DATABASE_URL.", + id="partial-resources-before-enabled", + ), + pytest.param( + { + "stage12_v2_manifest_environment": "staging", + "stage12_database_url": "mysql://db.example/stage12", + "stage12_artifact_bucket": "bucket", + }, + "STAGE12_DATABASE_URL must use the canonical postgresql:// scheme.", + id="database-scheme", + ), + ], +) +def test_validation_reports_the_first_failed_check(overrides, message): + with pytest.raises(ConfigurationError) as error: + make_settings(**overrides).validate() + + assert str(error.value) == message + + +@pytest.mark.parametrize( + "overrides", + [ + pytest.param({}, id="complete"), + pytest.param( + {"auth_required": False, "auth_issuer": "", "auth_audience": ""}, + id="local-without-caller-auth", + ), + pytest.param( + { + "stage12_enabled": True, + "stage12_v2_manifest_environment": "staging", + "stage12_database_url": STAGE12_DATABASE_URL, + "stage12_artifact_bucket": "bucket", + }, + id="stage12-enabled", + ), + pytest.param( + { + "stage12_v2_manifest_environment": "staging", + "stage12_database_url": STAGE12_DATABASE_URL, + "stage12_artifact_bucket": "bucket", + }, + id="stage12-configured-but-disabled", + ), + ], +) +def test_consistent_configuration_validates(overrides): + make_settings(**overrides).validate() From 9248a229d4fe58cd6a86dd4634e9de81c908a95a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:05:30 +0400 Subject: [PATCH 09/66] Split entry settings validation into caller auth, old gateway and Stage 12 checks Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_entry/config.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index 399684706..a2c856489 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -14,6 +14,8 @@ PRODUCTION_ENVIRONMENTS = frozenset({"main", "prod", "production"}) MODAL_GATEWAY_HOST_SUFFIX = ".modal.run" +# Stage 12 comparison artifacts and rows are kept for exactly this many days. +REQUIRED_STAGE12_RETENTION_DAYS = 30 class ConfigurationError(RuntimeError): @@ -85,7 +87,7 @@ class Settings: stage12_v2_worker_version: str | None = None stage12_database_url: str = "" stage12_artifact_bucket: str = "" - stage12_retention_days: int = 30 + stage12_retention_days: int = REQUIRED_STAGE12_RETENTION_DAYS @classmethod def from_env(cls) -> Settings: @@ -158,6 +160,12 @@ def validate(self) -> None: startup rather than on its first request. Stage 12 resources must be configured all together or not at all. """ + self._validate_caller_auth() + self._validate_old_gateway() + self._validate_stage12() + + def _validate_caller_auth(self) -> None: + """Require caller authentication in production, with an HTTPS issuer.""" if self.production and not self.auth_required: msg = "Caller authentication cannot be disabled in production." raise ConfigurationError(msg) @@ -177,6 +185,8 @@ def validate(self) -> None: self.auth_issuer, ) + def _validate_old_gateway(self) -> None: + """Require complete old-gateway credentials and the HTTPS Modal gateway.""" required_backend = { "OLD_GATEWAY_URL": self.old_gateway_url, "OLD_GATEWAY_AUTH_ISSUER": self.old_gateway_auth_issuer, @@ -211,6 +221,8 @@ def validate(self) -> None: msg = "OLD_GATEWAY_URL must point to the existing Modal gateway." raise ConfigurationError(msg) + def _validate_stage12(self) -> None: + """Require separate manifest storage and all-or-nothing Stage 12 resources.""" try: assert_separate_manifest_names( v1_name=V1_ROUTING_STATE_NAME, @@ -218,8 +230,11 @@ def validate(self) -> None: ) except ValueError as error: raise ConfigurationError(str(error)) from error - if self.stage12_retention_days != 30: - msg = "STAGE12_RETENTION_DAYS must be exactly 30." + if self.stage12_retention_days != REQUIRED_STAGE12_RETENTION_DAYS: + msg = ( + "STAGE12_RETENTION_DAYS must be exactly " + f"{REQUIRED_STAGE12_RETENTION_DAYS}." + ) raise ConfigurationError(msg) required_stage12 = { "STAGE12_V2_MANIFEST_ENVIRONMENT": self.stage12_v2_manifest_environment, From 1d917e568e753e3d881d02a9d5d6d769b96bf7d1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:06:30 +0400 Subject: [PATCH 10/66] Pin the Stage 12 adapter's complete report input and check order Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_adapter.py | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index c1b8ef4a9..cb1631de8 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -138,3 +138,157 @@ def test_adapter_returns_bounded_skip_reason(update, reason) -> None: assert result.report is None assert result.skip_reason is reason + + +def _expected_simulation(role: str, execution_id: str, policy: dict) -> dict: + dataset_uri = "hf://policyengine/data/populace_us_2024.h5@populace_us_2024-revision" + return { + "contract_version": 1, + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": execution_id, + "role": role, + "policy": policy, + "population": { + "kind": "dataset", + "artifact": { + "uri": dataset_uri, + "media_type": "application/x-hdf5", + "content_sha256": "a" * 64, + "size_bytes": None, + }, + }, + "year": 2026, + "geography": { + "country": "us", + "region": "us", + "filter_field": None, + "filter_value": None, + "filter_strategy": None, + }, + "options": {"spm": {"geography_kind": "national"}, "include_cliffs": True}, + "bundle": { + "policyengine_version": "5.2.0", + "country_package_name": "policyengine-us", + "country_package_version": "1.764.6", + "dataset": { + "identity": "populace_us_2024", + "uri": dataset_uri, + "artifact_revision": "populace_us_2024-revision", + "data_package_name": "populace-data", + "data_package_version": "0.1.0", + }, + "bundle_manifest_sha256": "b" * 64, + }, + } + + +@pytest.mark.parametrize( + "data", + [ + "default", + "populace_us_2024", + "hf://policyengine/data/populace_us_2024.h5@populace_us_2024-revision", + "populace_us_2024@populace_us_2024-revision", + ], +) +def test_adapter_builds_the_complete_report_input(data) -> None: + payload = { + **eligible_payload(), + "include_cliffs": True, + "spm": {"geography_kind": "national"}, + "data": data, + } + + result = adapt_annual_comparison( + payload, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.skip_reason is None + assert result.report is not None + assert result.report.model_dump(mode="json") == { + "contract_version": 1, + "evaluation_id": str(EVALUATION_ID), + "baseline": _expected_simulation( + "baseline", "20eaf597-1c15-549e-a96e-3661052c3f31", {} + ), + "reform": _expected_simulation( + "reform", + "5a17bbb2-20a6-5789-a508-9d26900587e4", + {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3000}}, + ), + "requested_aggregates": [ + "budget", + "poverty", + "inequality", + "distributional", + "winners_and_losers", + "geographic", + "program_statistics", + ], + } + + +@pytest.mark.parametrize( + ("update", "reason"), + [ + ({"region_group": None}, None), + ({"spm": None, "include_cliffs": None, "segmented": False}, None), + ({"scope": None}, ComparisonSkipReason.UNSUPPORTED_SCOPE), + ({"baseline": None}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), + ({"reform": []}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), + ({"time_period": 2026}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), + ({"time_period": "20a6"}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), + ({"time_period": "02026"}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), + ({"region": ""}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), + ({"country": None}, ComparisonSkipReason.UNSUPPORTED_COUNTRY), + ( + {"region_group": ["state/ca"], "scope": "household"}, + ComparisonSkipReason.UNSUPPORTED_OPTIONS, + ), + ( + {"scope": "household", "include_cliffs": "true"}, + ComparisonSkipReason.UNSUPPORTED_SCOPE, + ), + ( + {"include_cliffs": "true", "segmented": True}, + ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE, + ), + ( + {"segmented": True, "country": "ca"}, + ComparisonSkipReason.UNSUPPORTED_OPTIONS, + ), + ( + {"country": "ca", "policyengine_version": "another"}, + ComparisonSkipReason.UNSUPPORTED_COUNTRY, + ), + ( + {"policyengine_version": "another", "baseline": None}, + ComparisonSkipReason.MISSING_BUNDLE_PROVENANCE, + ), + ( + {"version": "another", "time_period": "2026-01"}, + ComparisonSkipReason.MISSING_BUNDLE_PROVENANCE, + ), + ( + {"baseline": None, "time_period": "2026-01"}, + ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE, + ), + ( + {"region": None, "data": "another"}, + ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE, + ), + ], +) +def test_adapter_checks_eligibility_in_order(update, reason) -> None: + payload = {**eligible_payload(), **update} + + result = adapt_annual_comparison( + payload, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.skip_reason is reason + assert result.eligible is (reason is None) From dca5ea2a1516abfdbbad7fa30b994ecb10170849 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:07:04 +0400 Subject: [PATCH 11/66] Split the Stage 12 annual-comparison adapter into eligibility checks and a report builder Co-Authored-By: Claude Opus 5.5 --- .../stage12_adapter.py | 148 ++++++++++++------ 1 file changed, 101 insertions(+), 47 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index ba598f9f2..57baddc5b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -21,7 +21,11 @@ ) if TYPE_CHECKING: - from policyengine_simulation_contract.stage12_bundle import CountryId + from policyengine_simulation_contract.stage12_bundle import ( + CountryId, + Stage12CountryBundle, + Stage12Dataset, + ) from policyengine_simulation_contract.stage12_manifest import ( V2CountryWorker, V2WorkerVersion, @@ -81,6 +85,7 @@ def eligible(self) -> bool: } ) _AGGREGATES = tuple(ReportAggregate) +_YEAR_DIGIT_COUNT = 4 def _skip(reason: ComparisonSkipReason) -> AdaptedComparisonRun: @@ -94,75 +99,126 @@ def _country_worker(worker: V2WorkerVersion, country: str) -> V2CountryWorker | ) -def adapt_annual_comparison( - payload: dict[str, Any], - *, - evaluation_id: UUID, - worker: V2WorkerVersion, -) -> AdaptedComparisonRun: - """Convert only the reviewed production request shape without guessing.""" +def _options_skip_reason(payload: dict[str, Any]) -> ComparisonSkipReason | None: + """Return why the request's fields or options are unsupported, if they are.""" if any( value is not None for key, value in payload.items() if key not in _ALLOWED_FIELDS ): - return _skip(ComparisonSkipReason.UNSUPPORTED_OPTIONS) + return ComparisonSkipReason.UNSUPPORTED_OPTIONS if payload.get("scope") != "macro": - return _skip(ComparisonSkipReason.UNSUPPORTED_SCOPE) + return ComparisonSkipReason.UNSUPPORTED_SCOPE include_cliffs = payload.get("include_cliffs") if include_cliffs is not None and not isinstance(include_cliffs, bool): - return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) + return ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE if payload.get("segmented") not in {None, False}: - return _skip(ComparisonSkipReason.UNSUPPORTED_OPTIONS) - country = payload.get("country") - if not isinstance(country, str) or _country_worker(worker, country) is None: - return _skip(ComparisonSkipReason.UNSUPPORTED_COUNTRY) - bundle_country = next( - item for item in worker.bundle.countries if item.country == country - ) + return ComparisonSkipReason.UNSUPPORTED_OPTIONS + return None + + +def _bundle_version_skip_reason( + payload: dict[str, Any], + worker: V2WorkerVersion, + bundle_country: Stage12CountryBundle, +) -> ComparisonSkipReason | None: + """Return a skip reason if a requested package version is not the bundle's.""" requested_policyengine_version = payload.get("policyengine_version") if ( requested_policyengine_version is not None and requested_policyengine_version != worker.bundle.policyengine_version ): - return _skip(ComparisonSkipReason.MISSING_BUNDLE_PROVENANCE) + return ComparisonSkipReason.MISSING_BUNDLE_PROVENANCE requested_country_package_version = payload.get("version") if ( requested_country_package_version is not None and requested_country_package_version != bundle_country.country_package_version ): - return _skip(ComparisonSkipReason.MISSING_BUNDLE_PROVENANCE) - baseline = payload.get("baseline") - reform = payload.get("reform") - if not isinstance(baseline, dict) or not isinstance(reform, dict): - return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) + return ComparisonSkipReason.MISSING_BUNDLE_PROVENANCE + return None + + +def _request_shape_skip_reason(payload: dict[str, Any]) -> ComparisonSkipReason | None: + """Return a skip reason if the policies, year or region are malformed.""" + if not isinstance(payload.get("baseline"), dict) or not isinstance( + payload.get("reform"), dict + ): + return ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE year_text = payload.get("time_period") if ( not isinstance(year_text, str) - or len(year_text) != 4 + or len(year_text) != _YEAR_DIGIT_COUNT or not year_text.isdecimal() ): - return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) + return ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE region = payload.get("region") if not isinstance(region, str) or not region: - return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) + return ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE + return None - selected_dataset = next( - item - for item in bundle_country.datasets - if item.identity == bundle_country.default_dataset - ) - requested_dataset = payload.get("data") - accepted_dataset_values = { + +def _accepted_dataset_values(selected_dataset: Stage12Dataset) -> set[str | None]: + """Return the ``data`` values that name the bundle's default dataset.""" + return { None, "default", selected_dataset.identity, selected_dataset.uri, f"{selected_dataset.identity}@{selected_dataset.artifact_revision}", } - if requested_dataset not in accepted_dataset_values: + + +def adapt_annual_comparison( + payload: dict[str, Any], + *, + evaluation_id: UUID, + worker: V2WorkerVersion, +) -> AdaptedComparisonRun: + """Convert only the reviewed production request shape without guessing.""" + skip_reason = _options_skip_reason(payload) + if skip_reason is not None: + return _skip(skip_reason) + country = payload.get("country") + if not isinstance(country, str) or _country_worker(worker, country) is None: + return _skip(ComparisonSkipReason.UNSUPPORTED_COUNTRY) + bundle_country = next( + item for item in worker.bundle.countries if item.country == country + ) + skip_reason = _bundle_version_skip_reason(payload, worker, bundle_country) + if skip_reason is None: + skip_reason = _request_shape_skip_reason(payload) + if skip_reason is not None: + return _skip(skip_reason) + + selected_dataset = next( + item + for item in bundle_country.datasets + if item.identity == bundle_country.default_dataset + ) + if payload.get("data") not in _accepted_dataset_values(selected_dataset): return _skip(ComparisonSkipReason.UNSUPPORTED_DATASET) + return AdaptedComparisonRun( + report=_report_input( + payload, + evaluation_id=evaluation_id, + worker=worker, + bundle_country=bundle_country, + selected_dataset=selected_dataset, + ), + skip_reason=None, + ) + + +def _report_input( + payload: dict[str, Any], + *, + evaluation_id: UUID, + worker: V2WorkerVersion, + bundle_country: Stage12CountryBundle, + selected_dataset: Stage12Dataset, +) -> ReportExecutionInput: + """Build the baseline and reform inputs for an eligible request.""" dataset = DatasetProvenance( identity=selected_dataset.identity, uri=selected_dataset.uri, @@ -185,12 +241,13 @@ def adapt_annual_comparison( ) ) geography = GeographySelection( - country=cast("CountryId", country), - region=region, + country=cast("CountryId", payload["country"]), + region=payload["region"], ) options = {key: payload[key] for key in ("spm",) if payload.get(key) is not None} - if include_cliffs is True: + if payload.get("include_cliffs") is True: options["include_cliffs"] = True + year = int(payload["time_period"]) def simulation(role: SimulationRole, policy: dict) -> SimulationExecutionInput: return SimulationExecutionInput( @@ -199,18 +256,15 @@ def simulation(role: SimulationRole, policy: dict) -> SimulationExecutionInput: role=role, policy=policy, population=population, - year=int(year_text), + year=year, geography=geography, options=options, bundle=provenance, ) - return AdaptedComparisonRun( - report=ReportExecutionInput( - evaluation_id=evaluation_id, - baseline=simulation(SimulationRole.BASELINE, baseline), - reform=simulation(SimulationRole.REFORM, reform), - requested_aggregates=_AGGREGATES, - ), - skip_reason=None, + return ReportExecutionInput( + evaluation_id=evaluation_id, + baseline=simulation(SimulationRole.BASELINE, payload["baseline"]), + reform=simulation(SimulationRole.REFORM, payload["reform"]), + requested_aggregates=_AGGREGATES, ) From 4cee9d172e360987ba16376332f8f213f902f206 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:09:45 +0400 Subject: [PATCH 12/66] Pin the Stage 12 dispatch calls, acknowledged parent and skip event Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_backend.py | 233 ++++++++++++++++++ 1 file changed, 233 insertions(+) diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py index 0145e54dc..ec75a064b 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py @@ -4,8 +4,10 @@ import asyncio import json +from datetime import UTC, datetime from types import SimpleNamespace from unittest.mock import MagicMock +from uuid import UUID, uuid5 import pytest from conftest import make_settings @@ -14,6 +16,7 @@ ComparisonRunLifecycleStatus, ) from policyengine_simulation_entry import stage12_backend as backend_module +from policyengine_simulation_entry.stage12_adapter import adapt_annual_comparison from policyengine_simulation_entry.stage12_backend import ( ModalReportInvoker, Stage12ComparisonBackend, @@ -418,3 +421,233 @@ async def fake_to_thread(function, *args, **kwargs): ) ] assert report.status is ComparisonRunLifecycleStatus.RUNNING + + +FIXED_NOW = datetime(2026, 9, 22, 12, 30, tzinfo=UTC) +TEMPORARY_EVALUATION_ID = UUID("00000000-0000-4000-8000-0000000000aa") + + +class FixedClock(datetime): + @classmethod + def now(cls, tz=None): + assert tz is UTC + return FIXED_NOW + + +def _expected_parent(evaluation_id, **overrides) -> dict: + return { + "contract_version": 1, + "evaluation_id": str(evaluation_id), + "status": "pending", + "aggregation_status": "not_started", + "environment": "staging", + "calculation_flow": "economy", + "originating_request_id": "request-1", + "observability_id": OBSERVABILITY_ID, + "production_identity": "job-1", + "incumbent_execution_id": "job-1", + "worker_version": "5.2.0", + "modal_application": "policyengine-simulation-v2-py5-2-0", + "report_coordinator_callable": "coordinate_report", + "version_manifest_sha256": "c" * 64, + "policyengine_version": "5.2.0", + "country_package_name": "policyengine-us", + "country_package_version": "1.764.6", + "country": "us", + "dataset_identity": "populace_us_2024", + "dataset_uri": ( + "hf://policyengine/data/populace_us_2024.h5@populace_us_2024-revision" + ), + "data_package_name": "populace-data", + "data_package_version": "0.1.0", + "data_artifact_revision": "populace_us_2024-revision", + "coordinator_invocation_id": None, + "error_code": None, + "error_summary": None, + "aggregate_output_uri": None, + "aggregate_output_sha256": None, + "aggregate_schema_version": None, + "comparison_status": "pending", + "comparison_output_uri": None, + "comparison_output_sha256": None, + "comparison_schema_version": None, + "comparison_completed_at": None, + "comparison_error_code": None, + "comparison_error_summary": None, + "created_at": "2026-09-22T12:30:00Z", + "updated_at": "2026-09-22T12:30:00Z", + "started_at": None, + "completed_at": None, + "retention_expires_at": "2026-10-22T12:30:00Z", + **overrides, + } + + +def _expected_context(evaluation_id, **overrides) -> dict: + return { + "request_id": "request-1", + "environment": "staging", + "modal_environment": "staging", + "worker_version": "5.2.0", + "modal_application": "policyengine-simulation-v2-py5-2-0", + "simulation_callable": "run_single_simulation_us", + "version_manifest_sha256": "c" * 64, + "bundle_manifest_sha256": "b" * 64, + "artifact_prefix": f"stage-12-runs/staging/2026/09/{evaluation_id}", + "production_function_call_id": "job-1", + "created_at": "2026-09-22T12:30:00+00:00", + "retention_expires_at": "2026-10-22T12:30:00+00:00", + **overrides, + } + + +@pytest.fixture +def fixed_clock(monkeypatch): + monkeypatch.setattr(backend_module, "datetime", FixedClock) + monkeypatch.setattr(backend_module, "uuid4", lambda: TEMPORARY_EVALUATION_ID) + + +@pytest.mark.usefixtures("fixed_clock") +def test_automatic_dispatch_sends_the_complete_coordinator_call() -> None: + invoker = FakeInvoker() + backend = _backend(FakeStore(), invoker) + + asyncio.run( + backend.dispatch_after_production( + request_payload=eligible_payload(), + production_response=_response(), + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + [call] = invoker.calls + evaluation_id = call["report_payload"]["evaluation_id"] + assert evaluation_id == str( + uuid5( + backend_module.AUTOMATIC_REPORT_NAMESPACE, + f"staging|economy|job-1|5.2.0|{'c' * 64}", + ) + ) + assert call["worker"] == backend._manifest_loader.worker + assert call["report_payload"] == ( + adapt_annual_comparison( + eligible_payload(), + evaluation_id=UUID(evaluation_id), + worker=worker(), + ).report.model_dump(mode="json") + ) + assert call["parent_payload"] == _expected_parent(evaluation_id) + assert call["context_payload"] == _expected_context(evaluation_id) + assert call["observability_context"] == { + "traceparent": "00-11111111111111111111111111111111-2222222222222222-01", + "observability_id": OBSERVABILITY_ID, + } + assert backend._manifest_loader.resolved_versions == [None] + backend._runtime.span.assert_called_once_with( + "stage12_entry_dispatch", + attributes={ + "evaluation_id": evaluation_id, + "execution_mode": "shadow", + "runner_name": "stage12", + }, + ) + backend._runtime.event.assert_called_once_with( + "stage12_comparison_run_dispatched", + attributes={ + "request_id": "request-1", + "evaluation_id": evaluation_id, + "production_identity": "job-1", + "worker_version": "5.2.0", + "modal_application": "policyengine-simulation-v2-py5-2-0", + "coordinator_invocation_id": "modal-call-1", + }, + ) + + +@pytest.mark.usefixtures("fixed_clock") +def test_temporary_submission_returns_the_complete_acknowledged_parent() -> None: + invoker = FakeInvoker() + backend = _backend(FakeStore(), invoker) + + report = asyncio.run( + backend.submit_temporary_report( + request_payload=eligible_payload(), + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + direct_identity = f"direct:{TEMPORARY_EVALUATION_ID}" + [call] = invoker.calls + assert call["parent_payload"] == _expected_parent( + TEMPORARY_EVALUATION_ID, + production_identity=direct_identity, + incumbent_execution_id=None, + comparison_status="not_requested", + ) + assert call["context_payload"] == _expected_context( + TEMPORARY_EVALUATION_ID, production_function_call_id=None + ) + assert report.model_dump(mode="json") == _expected_parent( + TEMPORARY_EVALUATION_ID, + status="running", + production_identity=direct_identity, + incumbent_execution_id=None, + comparison_status="not_requested", + coordinator_invocation_id="modal-call-1", + started_at="2026-09-22T12:30:00Z", + ) + backend._runtime.span.assert_called_once_with( + "stage12_entry_dispatch", + attributes={ + "evaluation_id": str(TEMPORARY_EVALUATION_ID), + "execution_mode": "authoritative", + "runner_name": "stage12", + }, + ) + + +def test_automatic_dispatch_skips_unsupported_requests_with_an_event() -> None: + invoker = FakeInvoker() + backend = _backend(FakeStore(), invoker) + + asyncio.run( + backend.dispatch_after_production( + request_payload={**eligible_payload(), "scope": "household"}, + production_response=_response(), + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert invoker.calls == [] + backend._runtime.event.assert_called_once_with( + "stage12_comparison_run_skipped", + attributes={ + "request_id": "request-1", + "production_identity": "job-1", + "reason": "unsupported_scope", + }, + ) + + +@pytest.mark.parametrize("response", [b"[]", b'{"job_id": ""}', b'{"job_id": 7}']) +def test_automatic_dispatch_requires_a_production_job_identifier(response) -> None: + invoker = FakeInvoker() + backend = _backend(FakeStore(), invoker) + + with pytest.raises( + ValueError, match=r"^accepted production response has no job identifier$" + ): + asyncio.run( + backend.dispatch_after_production( + request_payload=eligible_payload(), + production_response=response, + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert invoker.calls == [] + assert backend._manifest_loader.resolved_versions == [] From fa1ae11e727e2b627ee7b4197a934de3cb916513 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:10:04 +0400 Subject: [PATCH 13/66] Group Stage 12 dispatch inputs into a prepared report, origin and worker version Co-Authored-By: Claude Opus 5.5 --- .../stage12_backend.py | 163 +++++++++--------- 1 file changed, 81 insertions(+), 82 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py index 91bfd9991..29dad7567 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py @@ -5,8 +5,9 @@ import asyncio import json import logging +from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, NamedTuple, Protocol from uuid import NAMESPACE_URL, UUID, uuid4, uuid5 import modal @@ -78,14 +79,32 @@ async def spawn( ... -type PreparedReport = tuple[ - ComparisonReportRecord, - ReportExecutionInput, - V2WorkerVersion, - V2CountryWorker, - str, - str, -] +class ResolvedWorkerVersion(NamedTuple): + """A v2 worker version as the manifest loader resolves it.""" + + version: str + worker: V2WorkerVersion + manifest_sha256: str + + +@dataclass(frozen=True) +class ReportOrigin: + """The request a comparison run answers, and the production job it shadows.""" + + request_id: str + production_identity: str + incumbent_execution_id: str | None + comparison_status: ResultComparisonStatus + + +@dataclass(frozen=True) +class PreparedReport: + """A comparison run ready to dispatch: its parent row, inputs and worker.""" + + parent: ComparisonReportRecord + report: ReportExecutionInput + worker_version: ResolvedWorkerVersion + country_worker: V2CountryWorker class TemporaryStage12UnsupportedRequestError(ValueError): @@ -217,16 +236,10 @@ async def dispatch_after_production( ) if prepared is None: return - parent, report, worker, country_worker, version, manifest_sha256 = prepared await self._spawn_report( - parent=parent, - report=report, - worker=worker, - country_worker=country_worker, - version=version, - manifest_sha256=manifest_sha256, + prepared, request_id=request_id, - production_function_call_id=parent.production_identity, + production_function_call_id=prepared.parent.production_identity, observability_id=observability_id, ) @@ -245,14 +258,8 @@ async def submit_temporary_report( request_payload=request_payload, request_id=request_id, ) - parent, report, worker, country_worker, version, manifest_sha256 = prepared return await self._spawn_report( - parent=parent, - report=report, - worker=worker, - country_worker=country_worker, - version=version, - manifest_sha256=manifest_sha256, + prepared, request_id=request_id, production_function_call_id=None, observability_id=observability_id, @@ -282,15 +289,14 @@ def _resolve_report( evaluation_id: UUID, ) -> tuple[ ReportExecutionInput | None, - V2WorkerVersion, + ResolvedWorkerVersion, V2CountryWorker, - str, - str, str | None, ]: - version, worker, manifest_sha256 = self._manifest_loader.resolve( - self._settings.stage12_v2_worker_version + worker_version = ResolvedWorkerVersion( + *self._manifest_loader.resolve(self._settings.stage12_v2_worker_version) ) + worker = worker_version.worker adapted = adapt_annual_comparison( request_payload, evaluation_id=evaluation_id, @@ -307,10 +313,8 @@ def _resolve_report( skip_reason = adapted.skip_reason.value if adapted.skip_reason else None return ( adapted.report, - worker, + worker_version, country_worker, - version, - manifest_sha256, skip_reason, ) @@ -318,31 +322,27 @@ def _parent_record( self, *, evaluation_id: UUID, - request_id: str, - production_identity: str, - incumbent_execution_id: str | None, - worker: V2WorkerVersion, + origin: ReportOrigin, + worker_version: ResolvedWorkerVersion, report: ReportExecutionInput, - version: str, - manifest_sha256: str, - comparison_status: ResultComparisonStatus, ) -> ComparisonReportRecord: bundle = report.baseline.bundle + worker = worker_version.worker now = datetime.now(UTC) return ComparisonReportRecord( evaluation_id=evaluation_id, status=ComparisonRunLifecycleStatus.PENDING, aggregation_status=ComparisonRunAggregationStatus.NOT_STARTED, - comparison_status=comparison_status, + comparison_status=origin.comparison_status, environment=self._settings.environment, calculation_flow="economy", - originating_request_id=request_id, - production_identity=production_identity, - incumbent_execution_id=incumbent_execution_id, - worker_version=version, + originating_request_id=origin.request_id, + production_identity=origin.production_identity, + incumbent_execution_id=origin.incumbent_execution_id, + worker_version=worker_version.version, modal_application=worker.application_name, report_coordinator_callable=worker.report_coordinator_callable, - version_manifest_sha256=manifest_sha256, + version_manifest_sha256=worker_version.manifest_sha256, policyengine_version=bundle.policyengine_version, country_package_name=bundle.country_package_name, country_package_version=bundle.country_package_version, @@ -364,11 +364,9 @@ def _prepare_temporary_report( request_id: str, ) -> PreparedReport: evaluation_id = uuid4() - report, worker, country_worker, version, manifest_sha256, skip_reason = ( - self._resolve_report( - request_payload=request_payload, - evaluation_id=evaluation_id, - ) + report, worker_version, country_worker, skip_reason = self._resolve_report( + request_payload=request_payload, + evaluation_id=evaluation_id, ) if report is None: raise TemporaryStage12UnsupportedRequestError( @@ -376,16 +374,21 @@ def _prepare_temporary_report( ) parent = self._parent_record( evaluation_id=evaluation_id, - request_id=request_id, - production_identity=f"direct:{evaluation_id}", - incumbent_execution_id=None, - worker=worker, + origin=ReportOrigin( + request_id=request_id, + production_identity=f"direct:{evaluation_id}", + incumbent_execution_id=None, + comparison_status=ResultComparisonStatus.NOT_REQUESTED, + ), + worker_version=worker_version, report=report, - version=version, - manifest_sha256=manifest_sha256, - comparison_status=ResultComparisonStatus.NOT_REQUESTED, ) - return parent, report, worker, country_worker, version, manifest_sha256 + return PreparedReport( + parent=parent, + report=report, + worker_version=worker_version, + country_worker=country_worker, + ) def _prepare_automatic_report( self, @@ -402,9 +405,10 @@ def _prepare_automatic_report( msg = "accepted production response has no job identifier" raise ValueError(msg) - version, worker, manifest_sha256 = self._manifest_loader.resolve( - self._settings.stage12_v2_worker_version + worker_version = ResolvedWorkerVersion( + *self._manifest_loader.resolve(self._settings.stage12_v2_worker_version) ) + version, worker, manifest_sha256 = worker_version evaluation_id = uuid5( AUTOMATIC_REPORT_NAMESPACE, f"{self._settings.environment}|economy|{production_identity}|{version}|{manifest_sha256}", @@ -438,38 +442,33 @@ def _prepare_automatic_report( ) parent = self._parent_record( evaluation_id=evaluation_id, - request_id=request_id, - production_identity=production_identity, - incumbent_execution_id=production_identity, - worker=worker, + origin=ReportOrigin( + request_id=request_id, + production_identity=production_identity, + incumbent_execution_id=production_identity, + comparison_status=ResultComparisonStatus.PENDING, + ), + worker_version=worker_version, report=adapted.report, - version=version, - manifest_sha256=manifest_sha256, - comparison_status=ResultComparisonStatus.PENDING, ) - return ( - parent, - adapted.report, - worker, - country_worker, - version, - manifest_sha256, + return PreparedReport( + parent=parent, + report=adapted.report, + worker_version=worker_version, + country_worker=country_worker, ) async def _spawn_report( self, + prepared: PreparedReport, *, - parent: ComparisonReportRecord, - report: ReportExecutionInput, - worker: V2WorkerVersion, - country_worker: V2CountryWorker, - version: str, - manifest_sha256: str, request_id: str, production_function_call_id: str | None, observability_id: str, ) -> ComparisonReportRecord: """Return after Modal accepts the coordinator; perform no database writes.""" + parent = prepared.parent + version, worker, manifest_sha256 = prepared.worker_version artifact_prefix = ( f"stage-12-runs/{self._settings.environment}/" f"{parent.created_at:%Y}/{parent.created_at:%m}/" @@ -507,7 +506,7 @@ async def _spawn_report( ): invocation_id = await self._invoker.spawn( worker=worker, - report_payload=report.model_dump(mode="json"), + report_payload=prepared.report.model_dump(mode="json"), parent_payload=parent.model_dump(mode="json"), observability_context=captured_context or None, context_payload={ @@ -519,7 +518,7 @@ async def _spawn_report( "worker_version": version, "modal_application": worker.application_name, "simulation_callable": ( - country_worker.single_simulation_callable + prepared.country_worker.single_simulation_callable ), "version_manifest_sha256": manifest_sha256, "bundle_manifest_sha256": worker.bundle_manifest_sha256, From 6be69509ed478f4aedeb627543308cee541adcff Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:13:33 +0400 Subject: [PATCH 14/66] Pin entry app telemetry, Stage 12 route failures and backend readiness Co-Authored-By: Claude Opus 5.5 --- .../tests/test_app_characterization.py | 442 ++++++++++++++++++ .../tests/test_backend.py | 58 +++ 2 files changed, 500 insertions(+) create mode 100644 projects/policyengine-simulation-entry/tests/test_app_characterization.py diff --git a/projects/policyengine-simulation-entry/tests/test_app_characterization.py b/projects/policyengine-simulation-entry/tests/test_app_characterization.py new file mode 100644 index 000000000..3e32c1750 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_app_characterization.py @@ -0,0 +1,442 @@ +"""Characterization tests for entry app branches the route tests leave uncovered.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import override +from unittest.mock import Mock + +import pytest +from conftest import FakeBackend, make_settings +from fastapi.testclient import TestClient +from policyengine_observability import REQUEST_ID_HEADER +from policyengine_simulation_entry import app as app_module +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.backend import BackendResponse +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER +from stage12_fixtures import EVALUATION_ID, eligible_payload + +STAGE12_SETTINGS = { + "stage12_enabled": True, + "stage12_v2_manifest_environment": "staging", + "stage12_database_url": ( + "postgresql://policyengine_v2_runtime:secret@" + "test.pooler.supabase.com:5432/postgres?sslmode=require" + ), + "stage12_artifact_bucket": "policyengine-stage12-staging", +} +ENTRY_OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000001" +BACKEND_OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000099" +BUDGET_WINDOW_BODY = { + "country": "us", + "region": "us", + "scope": "macro", + "reform": {}, + "start_year": "2026", + "window_size": 1, +} + + +def _backend_response( + content: bytes, headers: dict[str, str] | None = None +) -> BackendResponse: + return BackendResponse( + status_code=200, + content=content, + headers={"content-type": "application/json", **(headers or {})}, + ) + + +def _event_attributes(event: Mock, name: str) -> list[dict]: + return [ + call.kwargs["attributes"] + for call in event.call_args_list + if call.args == (name,) + ] + + +@pytest.fixture +def recorded_events(client, monkeypatch) -> Mock: + runtime = client.app.state.policyengine_observability + event = Mock(wraps=runtime.event) + monkeypatch.setattr(runtime, "event", event) + return event + + +@dataclass(frozen=True) +class ForwardedCase: + method: str + path: str + route: str + content: bytes + expected: dict[str, object] + body: dict | None = None + + +@pytest.mark.parametrize( + "case", + [ + pytest.param( + ForwardedCase( + "POST", + "/simulate/economy/comparison", + "/simulate/economy/comparison", + b'{"status": "submitted", "job_id": "fc-1"}', + {"job_state": "submitted", "job_id": "fc-1"}, + body={"country": "us", "scope": "macro", "reform": {}}, + ), + id="submission-records-job-id", + ), + pytest.param( + ForwardedCase( + "POST", + "/simulate/economy/budget-window", + "/simulate/economy/budget-window", + b'{"status": "submitted", "batch_job_id": "batch-1"}', + {"job_state": "submitted", "batch_job_id": "batch-1"}, + body=BUDGET_WINDOW_BODY, + ), + id="batch-records-batch-id", + ), + pytest.param( + ForwardedCase( + "POST", + "/simulate/economy/comparison", + "/simulate/economy/comparison", + b'{"job_id": "fc-1"}', + {"job_id": "fc-1"}, + body={"country": "us", "scope": "macro", "reform": {}}, + ), + id="submission-without-status", + ), + pytest.param( + ForwardedCase("GET", "/jobs/fc-1", "/jobs/{job_id}", b"not json", {}), + id="unparsed-poll", + ), + pytest.param( + ForwardedCase( + "GET", + "/budget-window-jobs/batch-1", + "/budget-window-jobs/{batch_job_id}", + b'{"status": "running", "batch_job_id": "ignored"}', + {"job_state": "running"}, + ), + id="batch-poll-records-path-id", + ), + pytest.param( + ForwardedCase("GET", "/versions", "/versions", b"[]", {}), + id="unparsed-versions", + ), + ], +) +def test_backend_response_event_records_every_attribute( + recorded_events, client, backend, case +): + backend.responses[(case.method, case.path)] = _backend_response(case.content) + + result = client.request( + case.method, + case.path, + json=case.body, + headers={REQUEST_ID_HEADER: "request-1"}, + ) + + assert result.status_code == 200 + [attributes] = _event_attributes( + recorded_events, "simulation_entry_backend_response" + ) + elapsed_ms = attributes.pop("elapsed_ms") + assert isinstance(elapsed_ms, float) + assert elapsed_ms >= 0 + identifiers = { + key: case.path.rsplit("/", 1)[1] + for key in ("job_id", "batch_job_id") + if f"{{{key}}}" in case.route + } + assert list(attributes.items()) == list( + { + "request_id": "request-1", + "route": case.route, + "method": case.method, + "status_code": 200, + **identifiers, + **case.expected, + }.items() + ) + + +def test_invalid_backend_observability_header_is_dropped(client, backend): + backend.responses[("GET", "/versions")] = _backend_response( + b"{}", {OBSERVABILITY_ID_HEADER: "not-an-identifier"} + ) + + result = client.get("/versions") + + assert result.status_code == 200 + assert OBSERVABILITY_ID_HEADER not in result.headers + + +def test_non_workflow_route_adopts_the_backend_observability_id(client, backend): + backend.responses[("GET", "/versions")] = _backend_response( + b"{}", {OBSERVABILITY_ID_HEADER: BACKEND_OBSERVABILITY_ID} + ) + + result = client.get("/versions") + + assert result.headers[OBSERVABILITY_ID_HEADER] == BACKEND_OBSERVABILITY_ID + + +def test_telemetry_failures_never_change_responses(client, backend, monkeypatch): + runtime = client.app.state.policyengine_observability + + def failing_set_context(**_attributes): + msg = "context unavailable" + raise RuntimeError(msg) + + def failing_mismatch_event(name, **_kwargs): + if name == "simulation_entry_observability_id_mismatch": + msg = "event unavailable" + raise RuntimeError(msg) + + monkeypatch.setattr(runtime, "set_context", failing_set_context) + monkeypatch.setattr(runtime, "event", failing_mismatch_event) + backend.responses[("POST", "/simulate/economy/comparison")] = _backend_response( + b'{"status": "submitted", "job_id": "fc-1"}', + {OBSERVABILITY_ID_HEADER: BACKEND_OBSERVABILITY_ID}, + ) + + result = client.post( + "/simulate/economy/comparison", + json={"country": "us", "scope": "macro", "reform": {}}, + headers={OBSERVABILITY_ID_HEADER: ENTRY_OBSERVABILITY_ID}, + ) + + assert result.status_code == 200 + assert result.json() == {"status": "submitted", "job_id": "fc-1"} + assert result.headers[OBSERVABILITY_ID_HEADER] == ENTRY_OBSERVABILITY_ID + + +def test_observability_id_mismatch_is_reported_without_identifiers( + recorded_events, client, backend +): + backend.responses[("POST", "/simulate/economy/comparison")] = _backend_response( + b"{}", {OBSERVABILITY_ID_HEADER: BACKEND_OBSERVABILITY_ID} + ) + + client.post( + "/simulate/economy/comparison", + json={"country": "us", "scope": "macro", "reform": {}}, + headers={OBSERVABILITY_ID_HEADER: ENTRY_OBSERVABILITY_ID}, + ) + + assert _event_attributes( + recorded_events, "simulation_entry_observability_id_mismatch" + ) == [{"backend": "old_gateway", "reason": "downstream_identifier_differed"}] + + +def test_unhandled_failures_log_the_batch_identifier(caplog): + class FailingBackend(FakeBackend): + @override + async def request(self, *args, **kwargs): + msg = "unexpected backend failure" + raise RuntimeError(msg) + + app = create_app( + settings=make_settings(), + backend=FailingBackend(), + auth_dependency=lambda: None, + ) + + with caplog.at_level(logging.INFO), TestClient(app) as client: + result = client.get( + "/budget-window-jobs/batch-1", headers={REQUEST_ID_HEADER: "request-500"} + ) + + assert result.status_code == 500 + record = next( + item + for item in caplog.records + if item.getMessage() == "simulation_entry_unhandled_request" + ) + assert record.request_id == "request-500" + assert record.method == "GET" + assert record.path == "/budget-window-jobs/{batch_job_id}" + assert record.batch_job_id == "batch-1" + assert not hasattr(record, "job_id") + + +def test_unhandled_status_failures_log_the_evaluation_identifier(caplog): + class MalformedComparison: + async def get_temporary_report(self, _evaluation_id): + return object(), () + + with ( + caplog.at_level(logging.INFO), + _stage12_client(MalformedComparison()) as client, + ): + result = client.get( + f"/internal/stage12/reports/{EVALUATION_ID}", + headers={REQUEST_ID_HEADER: "request-500"}, + ) + + assert result.status_code == 500 + assert result.text == "Internal Server Error" + record = next( + item + for item in caplog.records + if item.getMessage() == "simulation_entry_unhandled_request" + ) + assert record.path == "/internal/stage12/reports/{evaluation_id}" + assert record.evaluation_id == str(EVALUATION_ID) + + +def _stage12_client(comparison, **settings_overrides) -> TestClient: + return TestClient( + create_app( + settings=make_settings(**{**STAGE12_SETTINGS, **settings_overrides}), + backend=FakeBackend(), + auth_dependency=lambda: None, + comparison_backend=comparison, + ) + ) + + +def test_configured_stage12_backend_is_built_from_settings(monkeypatch): + built = [] + + class UnavailableComparisonBackend: + @classmethod + def from_settings(cls, settings, *, runtime): + built.append((settings.stage12_artifact_bucket, runtime)) + + monkeypatch.setattr( + app_module, "Stage12ComparisonBackend", UnavailableComparisonBackend + ) + + with _stage12_client(None) as client: + submission = client.post("/internal/stage12/reports", json=eligible_payload()) + status = client.get(f"/internal/stage12/reports/{EVALUATION_ID}") + + assert [bucket for bucket, _runtime in built] == ["policyengine-stage12-staging"] + assert built[0][1] is client.app.state.policyengine_observability + for result in (submission, status): + assert result.status_code == 503 + assert result.json() == {"detail": "Stage 12 direct execution is unavailable."} + assert result.headers["retry-after"] == "10" + + +def test_unconfigured_stage12_status_is_unavailable(): + app = create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + ) + + with TestClient(app) as client: + result = client.get(f"/internal/stage12/reports/{EVALUATION_ID}") + + assert result.status_code == 503 + assert result.json() == {"detail": "Stage 12 direct execution is unavailable."} + assert result.headers["retry-after"] == "10" + + +class SecretBearingError(RuntimeError): + pass + + +class FailingComparisonBackend: + async def submit_temporary_report(self, **_submission): + msg = "secret request details" + raise SecretBearingError(msg) + + async def get_temporary_report(self, _evaluation_id): + msg = "secret report details" + raise SecretBearingError(msg) + + +def test_unexpected_stage12_submission_failure_logs_only_the_error_type(caplog): + with ( + caplog.at_level(logging.INFO), + _stage12_client(FailingComparisonBackend()) as client, + ): + result = client.post( + "/internal/stage12/reports", + json=eligible_payload(), + headers={REQUEST_ID_HEADER: "request-1"}, + ) + + assert result.status_code == 503 + assert result.json() == {"detail": "Stage 12 direct execution is unavailable."} + assert result.headers["retry-after"] == "10" + [record] = [ + item + for item in caplog.records + if item.getMessage() == "stage12_direct_submission_failed" + ] + assert record.name == "policyengine_simulation_entry.app" + assert record.levelno == logging.ERROR + assert record.exc_info is None + assert record.request_id == "request-1" + assert record.error_type == "SecretBearingError" + assert "secret" not in caplog.text + + +def test_unexpected_stage12_status_failure_logs_only_the_error_type(caplog): + with ( + caplog.at_level(logging.INFO), + _stage12_client(FailingComparisonBackend()) as client, + ): + result = client.get( + f"/internal/stage12/reports/{EVALUATION_ID}", + headers={REQUEST_ID_HEADER: "request-1"}, + ) + + assert result.status_code == 503 + assert result.json() == {"detail": "Stage 12 report status is unavailable."} + assert result.headers["retry-after"] == "10" + [record] = [ + item + for item in caplog.records + if item.getMessage() == "stage12_direct_status_failed" + ] + assert record.name == "policyengine_simulation_entry.app" + assert record.levelno == logging.ERROR + assert record.exc_info is None + assert record.request_id == "request-1" + assert record.evaluation_id == str(EVALUATION_ID) + assert record.error_type == "SecretBearingError" + assert "secret" not in caplog.text + + +@pytest.mark.parametrize("status_code", [200, 202, 400, 500]) +def test_automatic_dispatch_runs_only_for_accepted_submissions(status_code): + class RecordingComparison: + def __init__(self): + self.calls = [] + + async def dispatch_after_production(self, **call): + self.calls.append(call) + + comparison = RecordingComparison() + backend = FakeBackend() + backend.responses[("POST", "/simulate/economy/comparison")] = BackendResponse( + status_code=status_code, + content=b'{"job_id": "fc-1"}', + headers={"content-type": "application/json"}, + ) + app = create_app( + settings=make_settings(**STAGE12_SETTINGS), + backend=backend, + auth_dependency=lambda: None, + comparison_backend=comparison, + ) + + with TestClient(app) as client: + result = client.post( + "/simulate/economy/comparison", + json={"country": "us", "scope": "macro", "reform": {}}, + ) + + assert result.status_code == status_code + assert len(comparison.calls) == (1 if status_code in {200, 202} else 0) diff --git a/projects/policyengine-simulation-entry/tests/test_backend.py b/projects/policyengine-simulation-entry/tests/test_backend.py index 502542868..cee712338 100644 --- a/projects/policyengine-simulation-entry/tests/test_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_backend.py @@ -249,3 +249,61 @@ def handler(request: httpx.Request) -> httpx.Response: assert result.status_code == 200 assert gateway_attempts == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("health_status", "expected"), + [(200, True), (204, False), (503, False)], +) +async def test_readiness_requires_a_200_health_response(health_status, expected): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/oauth/token": + return httpx.Response( + 200, + json={"access_token": "service-token", "expires_in": 3600}, + ) + assert request.method == "GET" + assert request.url.path == "/health" + return httpx.Response(health_status, json={"status": "healthy"}) + + backend = OldGatewayBackend( + make_settings(), + transport=httpx.MockTransport(handler), + ) + await backend.start() + try: + assert await backend.ready() is expected + finally: + await backend.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", + [httpx.ConnectError("refused"), httpx.ReadTimeout("slow")], + ids=["unreachable", "timed-out"], +) +async def test_readiness_is_false_when_the_gateway_cannot_answer(failure): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/oauth/token": + return httpx.Response( + 200, + json={"access_token": "service-token", "expires_in": 3600}, + ) + raise failure + + backend = OldGatewayBackend( + make_settings(), + transport=httpx.MockTransport(handler), + ) + await backend.start() + try: + assert await backend.ready() is False + finally: + await backend.close() + + +@pytest.mark.asyncio +async def test_readiness_is_false_before_the_backend_starts(): + assert await OldGatewayBackend(make_settings()).ready() is False From 18d49cdd8b7a0c1c1704c68adb90bec924672f73 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:16:03 +0400 Subject: [PATCH 15/66] Split the entry app factory into a shared runtime and route registration steps The routes keep their names, decorators, parameters and registration order, so the published OpenAPI document is unchanged. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/app.py | 651 +++++++++++------- 1 file changed, 389 insertions(+), 262 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 3e998db5f..0ba68c2ed 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -8,7 +8,8 @@ import uuid from collections.abc import Callable from contextlib import asynccontextmanager -from typing import TYPE_CHECKING, Any, Literal, Protocol +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Protocol from uuid import UUID from fastapi import Depends, FastAPI, Request @@ -72,6 +73,7 @@ if TYPE_CHECKING: from collections.abc import AsyncIterator + from policyengine_observability import ObservabilityRuntime from starlette.middleware.base import RequestResponseEndpoint logger = logging.getLogger(__name__) @@ -165,77 +167,146 @@ def _request_identifiers(request: Request) -> RequestIdentifiers: return identifiers -def create_app( - *, - settings: Settings | None = None, - backend: SimulationBackend | None = None, - auth_dependency: AuthenticationDependency | None = None, - comparison_backend: ComparisonBackend | None = None, -) -> FastAPI: - """Build the app with injectable auth/backend seams for hermetic tests.""" - runtime_settings = settings or Settings.from_env() +@dataclass(frozen=True) +class _BackendCall: + """One request a route forwards to the old gateway, and what its event records. - @asynccontextmanager - async def lifespan(_: FastAPI) -> AsyncIterator[None]: - runtime_settings.validate() - await runtime_backend.start() - try: - yield - finally: - await runtime_backend.close() - runtime.shutdown() + ``identifiers`` come from the route's path; ``response_identifier_key`` + names the response field to copy into the backend response event. + """ - app = FastAPI( - title="PolicyEngine Simulation Entrypoint", - description=("Authenticated simulation submission and polling control plane."), - version="1.0.0", - lifespan=lifespan, + method: str + path: str + body: JsonObject | None = None + identifiers: RequestIdentifiers | None = None + response_identifier_key: ResponseIdentifier | None = None + + +class _BackendFailure(NamedTuple): + """The event and sanitized response for an old-gateway call that failed.""" + + event_name: str + status_code: int + detail: str + + +_BACKEND_TIMEOUT = _BackendFailure( + "simulation_entry_backend_timeout", 504, "Simulation backend timed out." +) +_BACKEND_UNAVAILABLE = _BackendFailure( + "simulation_entry_backend_unavailable", 503, "Simulation backend is unavailable." +) + + +def _backend_observability_id(headers: dict[str, str]) -> str | None: + """Return the backend's normalized observability ID header, if it sent one.""" + return next( + ( + normalize_observability_id(value) + for key, value in headers.items() + if key.lower() == OBSERVABILITY_ID_HEADER.lower() + ), + None, ) - runtime = init_simulation_observability( - app, - service_name="policyengine-simulation-entry", - service_role="simulation_entry", - platform="google_cloud_run", - environment=runtime_settings.environment, + +def _set_observability_header(response: Response, observability_id: str | None) -> None: + """Return ``observability_id`` on the response, or drop any stale header.""" + if observability_id is not None: + response.headers[OBSERVABILITY_ID_HEADER] = observability_id + elif OBSERVABILITY_ID_HEADER in response.headers: + del response.headers[OBSERVABILITY_ID_HEADER] + + +def _add_response_telemetry( + attributes: BackendTelemetryAttributes, + content: bytes, + response_identifier_key: ResponseIdentifier | None, +) -> None: + """Add the job state and requested identifier from a JSON backend response.""" + try: + response_payload = BackendTelemetryPayload.model_validate_json(content) + except ValidationError: + response_payload = None + job_state = response_payload.status if response_payload is not None else None + if job_state is not None: + attributes["job_state"] = job_state + response_identifier = ( + getattr(response_payload, response_identifier_key) + if response_payload is not None and response_identifier_key + else None + ) + if response_identifier_key and response_identifier is not None: + attributes[response_identifier_key] = response_identifier + + +def _temporary_submission_payload( + report: ComparisonReportRecord, +) -> TemporaryStage12SubmissionResponse: + """Describe a submitted temporary report and where to poll it.""" + return TemporaryStage12SubmissionResponse( + evaluation_id=report.evaluation_id, + status=report.status, + poll_url=f"/internal/stage12/reports/{report.evaluation_id}", ) - def bind_observability_id(request: Request, value: object) -> str | None: + +def _stage12_unavailable_response(detail: str) -> JSONResponse: + """Return the 503 the temporary Stage 12 routes send when they cannot run.""" + return JSONResponse( + status_code=503, + content={"detail": detail}, + headers={"Retry-After": "10"}, + ) + + +class _EntryRuntime: + """The settings, telemetry runtime and backends the entry routes share.""" + + def __init__( + self, + *, + settings: Settings, + runtime: ObservabilityRuntime, + backend: SimulationBackend, + comparison: ComparisonBackend | None, + ) -> None: + """Keep the dependencies that ``create_app`` resolved.""" + self.settings = settings + self.runtime = runtime + self.backend = backend + self.comparison = comparison + + def bind_observability_id(self, request: Request, value: object) -> str | None: + """Bind a normalized observability ID to the request and telemetry context.""" observability_id = normalize_observability_id(value) request.state.observability_id = observability_id if observability_id is not None: try: # noqa: SIM105 # try/except keeps the swallow's reason on its except line - runtime.set_context(observability_id=observability_id) + self.runtime.set_context(observability_id=observability_id) except Exception: # noqa: BLE001, S110 # telemetry is non-fatal pass return observability_id - def start_workflow_observability(request: Request) -> str: + def start_workflow_observability(self, request: Request) -> str: + """Bind the caller's observability ID, or a new one, to a workflow request.""" observability_id = resolve_observability_id( getattr(request.state, "incoming_observability_id", None) ) - bind_observability_id(request, observability_id) + self.bind_observability_id(request, observability_id) return observability_id - def backend_observability_id(headers: dict[str, str]) -> str | None: - return next( - ( - normalize_observability_id(value) - for key, value in headers.items() - if key.lower() == OBSERVABILITY_ID_HEADER.lower() - ), - None, - ) - def adopt_backend_observability_id( + self, request: Request, headers: dict[str, str], ) -> str | None: + """Keep the entry's observability ID, else adopt the backend's.""" current = normalize_observability_id(request.state.observability_id) - downstream = backend_observability_id(headers) + downstream = _backend_observability_id(headers) if current is not None and downstream is not None and current != downstream: try: # noqa: SIM105 # try/except keeps the swallow's reason on its except line - runtime.event( + self.runtime.event( "simulation_entry_observability_id_mismatch", attributes={ "backend": "old_gateway", @@ -245,93 +316,37 @@ def adopt_backend_observability_id( except Exception: # noqa: BLE001, S110 # telemetry is non-fatal pass return current - return bind_observability_id(request, current or downstream) + return self.bind_observability_id(request, current or downstream) - runtime_backend = backend or OldGatewayBackend(runtime_settings, runtime=runtime) - authenticate = auth_dependency or CallerAuthenticator(runtime_settings, runtime) - runtime_comparison = comparison_backend - if runtime_comparison is None and runtime_settings.stage12_resources_configured: - runtime_comparison = Stage12ComparisonBackend.from_settings( - runtime_settings, - runtime=runtime, - ) + async def forward(self, request: Request, call: _BackendCall) -> Response: + """Forward ``call`` to the old gateway and relay its response. - @app.middleware("http") - async def request_context( - request: Request, call_next: RequestResponseEndpoint - ) -> Response: - headers = MutableHeaders(scope=request.scope) - request_id = headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4()) - incoming_observability_id = normalize_observability_id( - headers.get(OBSERVABILITY_ID_HEADER) - ) - headers[REQUEST_ID_HEADER] = request_id - request.state.request_id = request_id - request.state.incoming_observability_id = incoming_observability_id - request.state.observability_id = None - try: - response = await call_next(request) - except Exception: - logger.exception( - "simulation_entry_unhandled_request", - extra={ - "request_id": request_id, - "method": request.method, - "path": _route_template(request), - **_request_identifiers(request), - }, - ) - response = Response( - content="Internal Server Error", - status_code=500, - media_type="text/plain", - headers={REQUEST_ID_HEADER: request_id}, - ) - bound_observability_id = normalize_observability_id( - request.state.observability_id - ) - if bound_observability_id is not None: - response.headers[OBSERVABILITY_ID_HEADER] = bound_observability_id - elif OBSERVABILITY_ID_HEADER in response.headers: - del response.headers[OBSERVABILITY_ID_HEADER] - if runtime_settings.revision: - response.headers["X-PolicyEngine-Simulation-Revision"] = ( - runtime_settings.revision - ) - return response - - async def forward( - request: Request, - method: str, - path: str, - body: JsonObject | None = None, - *, - identifiers: RequestIdentifiers | None = None, - response_identifier_key: ResponseIdentifier | None = None, - ) -> Response: + Timeouts and an unavailable gateway become sanitized 504 and 503 + responses; every outcome is reported as a telemetry event. + """ backend_started = time.monotonic() route = _route_template(request) - event_identifiers: RequestIdentifiers = identifiers or {} + event_identifiers: RequestIdentifiers = call.identifiers or {} try: # noqa: SIM105 # try/except keeps the swallow's reason on its except line - runtime.set_context(backend="old_gateway", **event_identifiers) + self.runtime.set_context(backend="old_gateway", **event_identifiers) except Exception: # noqa: BLE001, S110 # telemetry is non-fatal pass try: - result = await runtime_backend.request( - method, - path, - json_body=body, + result = await self.backend.request( + call.method, + call.path, + json_body=call.body, request_id=request.state.request_id, observability_id=request.state.observability_id, ) - observability_id = adopt_backend_observability_id( + observability_id = self.adopt_backend_observability_id( request, result.headers, ) attributes: BackendTelemetryAttributes = { "request_id": request.state.request_id, "route": route, - "method": method, + "method": call.method, "status_code": result.status_code, "elapsed_ms": round( (time.monotonic() - backend_started) * 1000, @@ -339,80 +354,153 @@ async def forward( ), **event_identifiers, } - try: - response_payload = BackendTelemetryPayload.model_validate_json( - result.content - ) - except ValidationError: - response_payload = None - job_state = ( - response_payload.status if response_payload is not None else None - ) - if job_state is not None: - attributes["job_state"] = job_state - response_identifier = ( - getattr(response_payload, response_identifier_key) - if response_payload is not None and response_identifier_key - else None + _add_response_telemetry( + attributes, result.content, call.response_identifier_key ) - if response_identifier_key and response_identifier is not None: - attributes[response_identifier_key] = response_identifier - runtime.event( + self.runtime.event( "simulation_entry_backend_response", attributes=attributes, ) response = _response(result) - if observability_id is not None: - response.headers[OBSERVABILITY_ID_HEADER] = observability_id - elif OBSERVABILITY_ID_HEADER in response.headers: - del response.headers[OBSERVABILITY_ID_HEADER] + _set_observability_header(response, observability_id) except BackendTimeoutError: - runtime.event( - "simulation_entry_backend_timeout", - attributes={ - "request_id": request.state.request_id, - "route": route, - **event_identifiers, - }, - ) - return JSONResponse( - status_code=504, - content={"detail": "Simulation backend timed out."}, - headers={ - "Retry-After": "10", - **BACKEND_RESPONSE_HEADER, - }, + return self._backend_failure_response( + request, _BACKEND_TIMEOUT, route=route, identifiers=event_identifiers ) except BackendUnavailableError: - runtime.event( - "simulation_entry_backend_unavailable", - attributes={ + return self._backend_failure_response( + request, + _BACKEND_UNAVAILABLE, + route=route, + identifiers=event_identifiers, + ) + else: + return response + + def _backend_failure_response( + self, + request: Request, + failure: _BackendFailure, + *, + route: str, + identifiers: RequestIdentifiers, + ) -> JSONResponse: + """Report a failed old-gateway call and return its sanitized response.""" + self.runtime.event( + failure.event_name, + attributes={ + "request_id": request.state.request_id, + "route": route, + **identifiers, + }, + ) + return JSONResponse( + status_code=failure.status_code, + content={"detail": failure.detail}, + headers={ + "Retry-After": "10", + **BACKEND_RESPONSE_HEADER, + }, + ) + + async def dispatch_comparison_run( + self, + request: Request, + request_payload: JsonObject, + response: Response, + observability_id: str, + ) -> None: + """Start a Stage 12 comparison of an accepted submission, if enabled. + + Waits only for Modal to acknowledge the already-deployed report + coordinator; a timeout or failure is logged with no request data and + never changes the production response. + """ + comparison = self.comparison + if not ( + self.settings.stage12_enabled + and comparison is not None + and response.status_code in {200, 202} + ): + return + try: + await asyncio.wait_for( + comparison.dispatch_after_production( + request_payload=request_payload, + production_response=bytes(response.body), + request_id=request.state.request_id, + observability_id=observability_id, + ), + timeout=STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS, + ) + except TimeoutError: + logger.error( # noqa: TRY400 # an expected timeout, fully described by timeout_seconds; the traceback would show only asyncio cancellation + "stage12_comparison_dispatch_timed_out", + extra={ "request_id": request.state.request_id, - "route": route, - **event_identifiers, + "timeout_seconds": STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS, }, ) - return JSONResponse( - status_code=503, - content={"detail": "Simulation backend is unavailable."}, - headers={ - "Retry-After": "10", - **BACKEND_RESPONSE_HEADER, + # Automatic Stage 12 work must never change the production v1 + # response, including for unexpected integration failures. + except Exception as error: # noqa: BLE001 # Stage 12 dispatch must never change the v1 response + logger.error( # noqa: TRY400 # Stage 12 failure logs carry only the error type; a traceback could expose request or calculation data + "stage12_comparison_dispatch_failed", + extra={ + "request_id": request.state.request_id, + "error_type": type(error).__name__, }, ) - else: - return response - protected = [Depends(authenticate)] - def temporary_submission_payload( - report: ComparisonReportRecord, - ) -> TemporaryStage12SubmissionResponse: - return TemporaryStage12SubmissionResponse( - evaluation_id=report.evaluation_id, - status=report.status, - poll_url=f"/internal/stage12/reports/{report.evaluation_id}", +def _install_request_context(app: FastAPI, settings: Settings) -> None: + """Add the middleware that assigns request IDs and correlation headers.""" + + @app.middleware("http") + async def request_context( + request: Request, call_next: RequestResponseEndpoint + ) -> Response: + headers = MutableHeaders(scope=request.scope) + request_id = headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4()) + incoming_observability_id = normalize_observability_id( + headers.get(OBSERVABILITY_ID_HEADER) ) + headers[REQUEST_ID_HEADER] = request_id + request.state.request_id = request_id + request.state.incoming_observability_id = incoming_observability_id + request.state.observability_id = None + try: + response = await call_next(request) + except Exception: + logger.exception( + "simulation_entry_unhandled_request", + extra={ + "request_id": request_id, + "method": request.method, + "path": _route_template(request), + **_request_identifiers(request), + }, + ) + response = Response( + content="Internal Server Error", + status_code=500, + media_type="text/plain", + headers={REQUEST_ID_HEADER: request_id}, + ) + _set_observability_header( + response, normalize_observability_id(request.state.observability_id) + ) + if settings.revision: + response.headers["X-PolicyEngine-Simulation-Revision"] = settings.revision + return response + + +def _register_temporary_stage12_submission( + app: FastAPI, + entry: _EntryRuntime, + protected: list[Any], +) -> None: + """Add the operator-only route that submits a temporary Stage 12 report.""" # TEMPORARY(Stage 12): This operator-only route is deliberately excluded # from OpenAPI. Remove it when Stage 14 provides authoritative v2 report @@ -426,20 +514,14 @@ async def submit_temporary_stage12_report( body: SimulationRequest, request: Request, ) -> Response: - if not runtime_settings.stage12_enabled: - return JSONResponse( - status_code=503, - content={"detail": "Stage 12 execution is disabled."}, - headers={"Retry-After": "10"}, - ) - comparison = runtime_comparison - if comparison is None: - return JSONResponse( - status_code=503, - content={"detail": "Stage 12 direct execution is unavailable."}, - headers={"Retry-After": "10"}, + comparison = entry.comparison + if not entry.settings.stage12_enabled or comparison is None: + return _stage12_unavailable_response( + "Stage 12 direct execution is unavailable." + if entry.settings.stage12_enabled + else "Stage 12 execution is disabled." ) - observability_id = start_workflow_observability(request) + observability_id = entry.start_workflow_observability(request) try: report = await asyncio.wait_for( comparison.submit_temporary_report( @@ -466,7 +548,7 @@ async def submit_temporary_stage12_report( except TemporaryStage12DispatchFailed as error: return JSONResponse( status_code=502, - content=temporary_submission_payload(error.report).model_dump( + content=_temporary_submission_payload(error.report).model_dump( mode="json" ), ) @@ -480,19 +562,25 @@ async def submit_temporary_stage12_report( "error_type": type(error).__name__, }, ) - return JSONResponse( - status_code=503, - content={"detail": "Stage 12 direct execution is unavailable."}, - headers={"Retry-After": "10"}, + return _stage12_unavailable_response( + "Stage 12 direct execution is unavailable." ) return JSONResponse( status_code=202, - content=temporary_submission_payload(report).model_dump(mode="json"), + content=_temporary_submission_payload(report).model_dump(mode="json"), # The Modal coordinator, rather than Cloud Run, creates the durable # parent row. Give it a brief head start before the first status read. headers={"Retry-After": "1"}, ) + +def _register_temporary_stage12_status( + app: FastAPI, + entry: _EntryRuntime, + protected: list[Any], +) -> None: + """Add the operator-only route that reads a temporary Stage 12 report.""" + # TEMPORARY(Stage 12): Read temporary comparison state through the # API-owned persistence service; never wait on or poll a Modal FunctionCall # from this Cloud Run request. @@ -505,12 +593,10 @@ async def get_temporary_stage12_report( evaluation_id: UUID, request: Request, ) -> Response: - comparison = runtime_comparison + comparison = entry.comparison if comparison is None: - return JSONResponse( - status_code=503, - content={"detail": "Stage 12 direct execution is unavailable."}, - headers={"Retry-After": "10"}, + return _stage12_unavailable_response( + "Stage 12 direct execution is unavailable." ) try: report, simulations = await comparison.get_temporary_report(evaluation_id) @@ -534,10 +620,8 @@ async def get_temporary_stage12_report( "error_type": type(error).__name__, }, ) - return JSONResponse( - status_code=503, - content={"detail": "Stage 12 report status is unavailable."}, - headers={"Retry-After": "10"}, + return _stage12_unavailable_response( + "Stage 12 report status is unavailable." ) payload = TemporaryStage12ReportResponse( report=report, @@ -548,7 +632,7 @@ async def get_temporary_stage12_report( ComparisonRunLifecycleStatus.RUNNING, } response_headers = {"Retry-After": "5"} if running else {} - report_observability_id = bind_observability_id( + report_observability_id = entry.bind_observability_id( request, report.observability_id, ) @@ -560,6 +644,14 @@ async def get_temporary_stage12_report( headers=response_headers or None, ) + +def _register_simulation_routes( + app: FastAPI, + entry: _EntryRuntime, + protected: list[Any], +) -> None: + """Add the authenticated submission and polling routes the gateway serves.""" + @app.post( "/simulate/economy/comparison", summary="Submit Simulation", @@ -584,50 +676,20 @@ async def submit_comparison( body: SimulationRequest, request: Request, ) -> Response: - observability_id = start_workflow_observability(request) + observability_id = entry.start_workflow_observability(request) request_payload = _model_json(body) - response = await forward( + response = await entry.forward( request, - "POST", - "/simulate/economy/comparison", - request_payload, - response_identifier_key="job_id", + _BackendCall( + "POST", + "/simulate/economy/comparison", + request_payload, + response_identifier_key="job_id", + ), + ) + await entry.dispatch_comparison_run( + request, request_payload, response, observability_id ) - if ( - runtime_settings.stage12_enabled - and runtime_comparison is not None - and response.status_code in {200, 202} - ): - # Wait only for Modal to acknowledge the already-deployed report - # coordinator invocation. Calculation and comparison run later. - try: - await asyncio.wait_for( - runtime_comparison.dispatch_after_production( - request_payload=request_payload, - production_response=bytes(response.body), - request_id=request.state.request_id, - observability_id=observability_id, - ), - timeout=STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS, - ) - except TimeoutError: - logger.error( # noqa: TRY400 # an expected timeout, fully described by timeout_seconds; the traceback would show only asyncio cancellation - "stage12_comparison_dispatch_timed_out", - extra={ - "request_id": request.state.request_id, - "timeout_seconds": STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS, - }, - ) - # Automatic Stage 12 work must never change the production v1 - # response, including for unexpected integration failures. - except Exception as error: # noqa: BLE001 # Stage 12 dispatch must never change the v1 response - logger.error( # noqa: TRY400 # Stage 12 failure logs carry only the error type; a traceback could expose request or calculation data - "stage12_comparison_dispatch_failed", - extra={ - "request_id": request.state.request_id, - "error_type": type(error).__name__, - }, - ) return response @app.post( @@ -653,13 +715,15 @@ async def submit_budget_window( body: BudgetWindowBatchRequest, request: Request, ) -> Response: - start_workflow_observability(request) - return await forward( + entry.start_workflow_observability(request) + return await entry.forward( request, - "POST", - "/simulate/economy/budget-window", - _model_json(body), - response_identifier_key="batch_job_id", + _BackendCall( + "POST", + "/simulate/economy/budget-window", + _model_json(body), + response_identifier_key="batch_job_id", + ), ) @app.get( @@ -689,11 +753,9 @@ async def submit_budget_window( dependencies=protected, ) async def get_job(job_id: str, request: Request) -> Response: - return await forward( + return await entry.forward( request, - "GET", - f"/jobs/{job_id}", - identifiers={"job_id": job_id}, + _BackendCall("GET", f"/jobs/{job_id}", identifiers={"job_id": job_id}), ) @app.get( @@ -724,13 +786,19 @@ async def get_budget_window_job( batch_job_id: str, request: Request, ) -> Response: - return await forward( + return await entry.forward( request, - "GET", - f"/budget-window-jobs/{batch_job_id}", - identifiers={"batch_job_id": batch_job_id}, + _BackendCall( + "GET", + f"/budget-window-jobs/{batch_job_id}", + identifiers={"batch_job_id": batch_job_id}, + ), ) + +def _register_service_routes(app: FastAPI, entry: _EntryRuntime) -> None: + """Add the public version, health, readiness and ping routes.""" + @app.get( "/versions", summary="List Versions", @@ -739,7 +807,7 @@ async def get_budget_window_job( response_model=VersionsResponse, ) async def list_versions(request: Request) -> Response: - return await forward(request, "GET", "/versions") + return await entry.forward(request, _BackendCall("GET", "/versions")) @app.get( "/versions/{kind}", @@ -749,7 +817,7 @@ async def list_versions(request: Request) -> Response: response_model=VersionMap, ) async def get_country_versions(kind: str, request: Request) -> Response: - return await forward(request, "GET", f"/versions/{kind}") + return await entry.forward(request, _BackendCall("GET", f"/versions/{kind}")) @app.get( "/health", @@ -762,7 +830,7 @@ async def health() -> HealthResponse: @app.get("/ready", response_model=ReadinessResponse) async def ready() -> Response: - if await runtime_backend.ready(): + if await entry.backend.ready(): return JSONResponse(ReadinessResponse(status="ready").model_dump()) return JSONResponse( ReadinessResponse(status="not_ready").model_dump(), @@ -776,8 +844,67 @@ async def ready() -> Response: response_model=PingResponse, ) async def ping(body: PingRequest, request: Request) -> Response: - return await forward(request, "POST", "/ping", _model_json(body)) + return await entry.forward( + request, _BackendCall("POST", "/ping", _model_json(body)) + ) + + +def create_app( + *, + settings: Settings | None = None, + backend: SimulationBackend | None = None, + auth_dependency: AuthenticationDependency | None = None, + comparison_backend: ComparisonBackend | None = None, +) -> FastAPI: + """Build the app with injectable auth/backend seams for hermetic tests.""" + runtime_settings = settings or Settings.from_env() + @asynccontextmanager + async def lifespan(_: FastAPI) -> AsyncIterator[None]: + runtime_settings.validate() + await runtime_backend.start() + try: + yield + finally: + await runtime_backend.close() + runtime.shutdown() + + app = FastAPI( + title="PolicyEngine Simulation Entrypoint", + description=("Authenticated simulation submission and polling control plane."), + version="1.0.0", + lifespan=lifespan, + ) + + runtime = init_simulation_observability( + app, + service_name="policyengine-simulation-entry", + service_role="simulation_entry", + platform="google_cloud_run", + environment=runtime_settings.environment, + ) + + runtime_backend = backend or OldGatewayBackend(runtime_settings, runtime=runtime) + authenticate = auth_dependency or CallerAuthenticator(runtime_settings, runtime) + runtime_comparison = comparison_backend + if runtime_comparison is None and runtime_settings.stage12_resources_configured: + runtime_comparison = Stage12ComparisonBackend.from_settings( + runtime_settings, + runtime=runtime, + ) + entry = _EntryRuntime( + settings=runtime_settings, + runtime=runtime, + backend=runtime_backend, + comparison=runtime_comparison, + ) + + _install_request_context(app, runtime_settings) + protected = [Depends(authenticate)] + _register_temporary_stage12_submission(app, entry, protected) + _register_temporary_stage12_status(app, entry, protected) + _register_simulation_routes(app, entry, protected) + _register_service_routes(app, entry) return app From 79ec691f468670bfaa8bc26fe07125f5e1df8571 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:16:17 +0400 Subject: [PATCH 16/66] Compare the old gateway's health status with HTTPStatus.OK Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/backend.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py index 955dbc29b..19edbe7f5 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py @@ -5,6 +5,7 @@ import asyncio import time from dataclasses import dataclass +from http import HTTPStatus from typing import TYPE_CHECKING, Protocol import httpx @@ -205,7 +206,7 @@ async def ready(self) -> bool: response = await self.request("GET", "/health") except (BackendUnavailableError, BackendTimeoutError): return False - return response.status_code == 200 + return response.status_code == HTTPStatus.OK async def request( self, From 7b8eeb26b18dc3dfefc7c64f35f9a76277b7d3f5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:17:43 +0400 Subject: [PATCH 17/66] Pin gateway route resolution from the active routing state Co-Authored-By: Claude Opus 5.5 --- .../tests/test_active_state_routes.py | 198 ++++++++++++++++++ 1 file changed, 198 insertions(+) create mode 100644 projects/policyengine-simulation-gateway/tests/test_active_state_routes.py diff --git a/projects/policyengine-simulation-gateway/tests/test_active_state_routes.py b/projects/policyengine-simulation-gateway/tests/test_active_state_routes.py new file mode 100644 index 000000000..0e10306fe --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_active_state_routes.py @@ -0,0 +1,198 @@ +"""Characterization tests for route resolution from the active routing state.""" + +import re +from copy import deepcopy +from dataclasses import asdict + +import pytest +from fixtures.gateway_endpoints import TEST_APP_RELEASE_BUNDLE, TEST_ROUTING_STATE +from policyengine_simulation_gateway.endpoints import _resolve_from_active_state + +CURRENT_APP = "policyengine-simulation-py4-10-0" +PREVIOUS_APP = "policyengine-simulation-py3-9-0" +PREVIOUS_BUNDLE = TEST_ROUTING_STATE["bundles"]["3.9.0"] + + +def _resolve(state=None, *, country="us", version=None, policyengine_version=None): + return asdict( + _resolve_from_active_state( + state=deepcopy(TEST_ROUTING_STATE) if state is None else state, + country=country, + version=version, + policyengine_version=policyengine_version, + ) + ) + + +def _policyengine_route(version, app_name, bundle, *, response_version=None): + return { + "app_name": app_name, + "response_version": response_version or version, + "policyengine_version": version, + "bundle_manifest": bundle, + "route_provenance": None, + "country_model_version": None, + } + + +def _state_without_policyengine_latest(): + state = deepcopy(TEST_ROUTING_STATE) + del state["latest"]["policyengine"] + return state + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + pytest.param( + {"policyengine_version": "4.10.0"}, + _policyengine_route("4.10.0", CURRENT_APP, TEST_APP_RELEASE_BUNDLE), + id="pinned-wrapper", + ), + pytest.param( + {"policyengine_version": "3.9.0", "version": "1.459.0"}, + _policyengine_route("3.9.0", PREVIOUS_APP, PREVIOUS_BUNDLE), + id="pinned-wrapper-with-matching-model", + ), + pytest.param( + {}, + _policyengine_route("4.10.0", CURRENT_APP, TEST_APP_RELEASE_BUNDLE), + id="latest-wrapper", + ), + pytest.param( + {"state": _state_without_policyengine_latest()}, + { + "app_name": CURRENT_APP, + "response_version": "1.500.0", + "policyengine_version": "4.10.0", + "bundle_manifest": TEST_APP_RELEASE_BUNDLE, + "route_provenance": None, + "country_model_version": "1.500.0", + }, + id="latest-country-model", + ), + pytest.param( + {"version": "1.459.0"}, + { + "app_name": PREVIOUS_APP, + "response_version": "1.459.0", + "policyengine_version": "3.9.0", + "bundle_manifest": PREVIOUS_BUNDLE, + "route_provenance": None, + "country_model_version": "1.459.0", + }, + id="country-model", + ), + pytest.param( + {"version": "3.9.0"}, + _policyengine_route("3.9.0", PREVIOUS_APP, PREVIOUS_BUNDLE), + id="wrapper-version-as-version", + ), + ], +) +def test_active_state_resolves_routes(arguments, expected): + assert _resolve(**arguments) == expected + + +def test_version_naming_one_app_as_both_kinds_resolves_the_wrapper(): + state = deepcopy(TEST_ROUTING_STATE) + state["routes"]["us"]["4.10.0"] = CURRENT_APP + state["bundles"]["4.10.0"] = { + **TEST_APP_RELEASE_BUNDLE, + "us": {**TEST_APP_RELEASE_BUNDLE["us"], "model_version": "4.10.0"}, + } + + assert _resolve(state, version="4.10.0") == _policyengine_route( + "4.10.0", CURRENT_APP, state["bundles"]["4.10.0"] + ) + + +def test_legacy_country_route_without_a_bundle_is_marked_pre_canonical(): + state = deepcopy(TEST_ROUTING_STATE) + state["routes"]["us"]["1.400.0"] = "legacy-app" + + assert _resolve(state, version="1.400.0") == { + "app_name": "legacy-app", + "response_version": "1.400.0", + "policyengine_version": None, + "bundle_manifest": {}, + "route_provenance": "legacy-country-route", + "country_model_version": "1.400.0", + } + + +def _ambiguous_state(): + state = deepcopy(TEST_ROUTING_STATE) + state["routes"]["us"]["3.9.0"] = "legacy-app" + return state + + +def _stale_country_route_state(): + state = deepcopy(TEST_ROUTING_STATE) + state["routes"]["us"]["3.9.0"] = CURRENT_APP + return state + + +def _state_without_latest_routes(): + state = _state_without_policyengine_latest() + state["latest"]["us"] = "9.9.9" + return state + + +@pytest.mark.parametrize( + ("arguments", "message"), + [ + pytest.param( + {"policyengine_version": "9.9.9"}, + "Unknown policyengine.py version 9.9.9", + id="unknown-wrapper", + ), + pytest.param( + {"policyengine_version": "3.9.0", "version": "1.459.0", "country": "uk"}, + "Requested uk version 1.459.0 does not match policyengine.py bundle " + "3.9.0 (uk model version 2.66.0)", + id="pinned-wrapper-checks-the-requested-country", + ), + pytest.param( + {"policyengine_version": "4.10.0", "version": "1.459.0"}, + "Requested us version 1.459.0 does not match policyengine.py bundle " + "4.10.0 (us model version 1.500.0)", + id="pinned-wrapper-with-other-model", + ), + pytest.param( + {"state": _state_without_latest_routes()}, + "Routing state does not define a latest route", + id="latest-country-route-missing", + ), + pytest.param( + {"state": {"routes": {}}}, + "Routing state does not define a latest route", + id="no-latest", + ), + pytest.param( + {"state": _ambiguous_state(), "version": "3.9.0"}, + "Ambiguous version 3.9.0 for country us; pass policyengine_version " + "to select a bundle explicitly", + id="ambiguous-version", + ), + pytest.param( + {"state": _stale_country_route_state(), "version": "3.9.0"}, + "Requested us version 3.9.0 does not match policyengine.py bundle " + "4.10.0 (us model version 1.500.0)", + id="country-route-checked-before-ambiguity", + ), + pytest.param( + {"version": "0.0.1"}, + "Unknown version 0.0.1 for country us", + id="unknown-version", + ), + pytest.param( + {"version": "0.0.1", "country": "uk"}, + "Unknown version 0.0.1 for country uk", + id="unknown-uk-version", + ), + ], +) +def test_active_state_rejects_unresolvable_routes(arguments, message): + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + _resolve(**arguments) From 29d3cfc5a7650891ab41ea18869a6747c8e8376e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:17:49 +0400 Subject: [PATCH 18/66] Split gateway active-state route resolution into pinned, latest and requested-version steps Co-Authored-By: Claude Opus 5.5 --- .../endpoints.py | 81 +++++++++++++------ 1 file changed, 55 insertions(+), 26 deletions(-) diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py index 261a8d05f..046b63910 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -398,37 +398,66 @@ def _resolve_from_active_state( policyengine_version: str | None, ) -> RouteResolution: if policyengine_version is not None: - resolution = _resolve_policyengine_route( + return _resolve_pinned_policyengine_route( state, + country=country, + version=version, policyengine_version=policyengine_version, ) - if version is not None: - _validate_legacy_version_matches_bundle( - country=country, - requested_version=version, - manifest=resolution.bundle_manifest, - ) - return resolution - if version is None: - latest_policyengine = _routing_state_latest(state, "policyengine") - if latest_policyengine is not None: - return _resolve_policyengine_route( - state, - policyengine_version=latest_policyengine, - ) - latest_country_version = _routing_state_latest(state, country) - if latest_country_version is not None: - country_resolution = _resolve_country_route( - state, - country=country, - version=latest_country_version, - ) - if country_resolution is not None: - return country_resolution - msg = "Routing state does not define a latest route" - raise ValueError(msg) + return _resolve_latest_route(state, country=country) + return _resolve_requested_version_route(state, country=country, version=version) + + +def _resolve_pinned_policyengine_route( + state: dict, + *, + country: str, + version: str | None, + policyengine_version: str, +) -> RouteResolution: + """Route to an explicit wrapper version, checking any requested model version.""" + resolution = _resolve_policyengine_route( + state, + policyengine_version=policyengine_version, + ) + if version is not None: + _validate_legacy_version_matches_bundle( + country=country, + requested_version=version, + manifest=resolution.bundle_manifest, + ) + return resolution + + +def _resolve_latest_route(state: dict, *, country: str) -> RouteResolution: + """Route to the latest wrapper version, else the latest country model.""" + latest_policyengine = _routing_state_latest(state, "policyengine") + if latest_policyengine is not None: + return _resolve_policyengine_route( + state, + policyengine_version=latest_policyengine, + ) + latest_country_version = _routing_state_latest(state, country) + if latest_country_version is not None: + country_resolution = _resolve_country_route( + state, + country=country, + version=latest_country_version, + ) + if country_resolution is not None: + return country_resolution + msg = "Routing state does not define a latest route" + raise ValueError(msg) + +def _resolve_requested_version_route( + state: dict, + *, + country: str, + version: str, +) -> RouteResolution: + """Route ``version`` as a country model or wrapper version, refusing ambiguity.""" country_resolution = _resolve_country_route( state, country=country, From 133aca33e0e809068044b024009a6e5e38af8b1d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:20:38 +0400 Subject: [PATCH 19/66] Pass the resolved region to _build_simulation instead of its code and scoping strategy Every caller already held a RegionResolution, so the builder now takes it whole. Also name the simulation runtime's preview, date-range and state-code limits. Co-Authored-By: Claude Opus 5.5 --- .../precompute.py | 3 +-- .../simulation_runtime.py | 26 ++++++++++++------- .../stage12_runtime/simulation.py | 3 +-- .../tests/test_baseline_artifacts.py | 8 +++--- .../test_baseline_artifacts_spm_native.py | 3 ++- .../tests/test_canonical_spm.py | 2 +- .../tests/test_canonical_spm_native.py | 14 +++++----- .../tests/test_simulation_output_builder.py | 20 +++++++------- .../tests/test_simulation_runtime_policy.py | 5 ++-- 9 files changed, 44 insertions(+), 40 deletions(-) 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 c4aeddbbf..85b74a551 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -330,8 +330,7 @@ def _prepare_cohort_baseline( dataset=dataset, dataset_selection=dataset_selection, policy=None, - scoping_strategy=resolution.scoping_strategy, - region_code=resolution.code, + region=resolution, ) if baseline.id != expected.simulation_id: msg = ( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 51852c884..b0e2a03e3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -51,6 +51,12 @@ os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1") DEFAULT_YEAR = 2026 +# Credential-adjacent environment values are cut to this many characters. +_ENV_VALUE_PREVIEW_CHARS = 50 +# A policy period key is one date or a ``start.stop`` pair of dates. +_DATE_RANGE_PART_COUNT = 2 +# A bare two-letter US region is a state code, such as "ca". +_US_STATE_CODE_LENGTH = 2 @dataclass(frozen=True) @@ -119,7 +125,9 @@ def setup_gcp_credentials( with credential_span: # Log available GCP-related env vars for debugging. gcp_vars = { - k: v[:50] + "..." if len(v) > 50 else v + k: v[:_ENV_VALUE_PREVIEW_CHARS] + "..." + if len(v) > _ENV_VALUE_PREVIEW_CHARS + else v for k, v in os.environ.items() if "GOOGLE" in k or "GCP" in k or "CREDENTIAL" in k } @@ -211,7 +219,7 @@ def _parse_policy_period_dates(text: str) -> tuple[datetime, datetime | None]: # unpadded months and days ("2027-4-6"); accept what it accepted. dates = [datetime.strptime(part, "%Y-%m-%d") for part in parts] # noqa: DTZ007 # naive like policyengine.py's own period parsing; its ParameterValue dates are naive start = dates[0] - stop = dates[1] if len(dates) == 2 else None + stop = dates[1] if len(dates) == _DATE_RANGE_PART_COUNT else None if stop is not None and stop < start: raise ValueError return start, stop @@ -312,7 +320,7 @@ def _normalise_region_code(country: str, region: object) -> str: return raw.lower() if "/" not in raw: - if country == "us" and len(raw) == 2: + if country == "us" and len(raw) == _US_STATE_CODE_LENGTH: return f"state/{raw.lower()}" if country == "uk": return f"country/{raw.lower().replace(' ', '_')}" @@ -729,8 +737,7 @@ def _build_simulation( dataset: Dataset, dataset_selection: DatasetSelection, policy: dict[str, Any] | None, - scoping_strategy: ScopingStrategy | None = None, - region_code: str | None = None, + region: RegionResolution, ) -> Simulation: from policyengine.core import ( # noqa: PLC0415 # resolved per call so tests can patch policyengine.core.Simulation Simulation, @@ -752,12 +759,13 @@ def _build_simulation( policy = _compile_interval_policy( policy, country_module=country_module, year=_parse_year(params) ) + scoping_strategy = region.scoping_strategy simulation_id = deterministic_baseline_id( params, country=country, dataset_is_default=dataset_selection.is_default, policy=policy, - region_code=region_code, + region_code=region.code, scoping_strategy=scoping_strategy, year=_parse_year(params), ) @@ -859,8 +867,7 @@ def _run_simulation_impl_core( dataset=dataset, dataset_selection=dataset_selection, policy=baseline_policy, - scoping_strategy=region_resolution.scoping_strategy, - region_code=region_resolution.code, + region=region_resolution, ) with runtime.span( ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), @@ -871,8 +878,7 @@ def _run_simulation_impl_core( dataset=dataset, dataset_selection=dataset_selection, policy=reform_policy, - scoping_strategy=region_resolution.scoping_strategy, - region_code=region_resolution.code, + region=region_resolution, ) logger.info("Calculating economic impact") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 89c3fb4d1..28a721055 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -206,8 +206,7 @@ def calculate_simulation_frames( dataset=dataset, dataset_selection=dataset_selection, policy=policy, - scoping_strategy=region.scoping_strategy, - region_code=region.code, + region=region, ) with stage_span(Stage.STAGE12_CALCULATION): apply_output_plan(model, simulation.output_plan) diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index a9544cfe3..7e90411c1 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -419,8 +419,7 @@ def test_qualifying_baseline_uses_artifact_class(self, wired): dataset="dataset", dataset_selection=SimpleNamespace(is_default=True), policy=None, - scoping_strategy=None, - region_code="us", + region=sr.RegionResolution(code="us"), ) assert wired.plain_kwargs is None assert wired.artifact_kwargs["id"] == "bl1-deadbeefdeadbeef" @@ -455,8 +454,7 @@ def test_non_qualifying_uses_plain_simulation(self, wired): dataset="dataset", dataset_selection=SimpleNamespace(is_default=False), policy={"gov.x": 1}, - scoping_strategy=None, - region_code="us", + region=sr.RegionResolution(code="us"), ) assert wired.artifact_kwargs is None assert wired.plain_kwargs["policy"] == {"gov.x": 1} @@ -495,7 +493,7 @@ def collect_identity(*args, **kwargs): kwargs = { "dataset": "dataset", "dataset_selection": sr.DatasetSelection("default", "default", is_default=True), - "region_code": "uk", + "region": sr.RegionResolution(code="uk"), } params = {"country": "uk", "scope": "macro", "time_period": 2027} baseline = sr._build_simulation(params, policy=None, **kwargs) diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py index 6e3e9447b..15c152fd7 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py @@ -28,6 +28,7 @@ def _baseline(dataset, *, geography="national", tax_only=False): ) from policyengine_simulation_executor.simulation_runtime import ( DatasetSelection, + RegionResolution, _build_simulation, ) @@ -44,7 +45,7 @@ def _baseline(dataset, *, geography="national", tax_only=False): "bounded-native-test", "native", is_default=False ), policy=None, - region_code="us", + region=RegionResolution(code="us"), ) baseline = ArtifactBaselineSimulation.model_construct(**built.__dict__) baseline.id = "bl1-native-receipt" diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index c54a5ec19..e126a165b 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -632,7 +632,7 @@ def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): "tiny", "tiny", is_default=True ), policy=policy, - region_code="us", + region=simulation_runtime.RegionResolution(code="us"), ) assert len(calls) == 2 assert all(call["spm"] == SELECTION for call in calls) diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py index cc016fc2a..967ce39fe 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py @@ -31,6 +31,7 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( from policyengine.core.simulation import _cache from policyengine_simulation_executor.simulation_runtime import ( DatasetSelection, + RegionResolution, _build_simulation, ) from policyengine_simulation_executor.spm import simulation_spm_result @@ -50,14 +51,14 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( dataset=native_dataset, dataset_selection=dataset_selection, policy=None, - region_code="us", + region=RegionResolution(code="us"), ) reform = _build_simulation( params, dataset=native_dataset, dataset_selection=dataset_selection, policy={"gov.irs.credits.ctc.amount.base[0].amount": 3000}, - region_code="us", + region=RegionResolution(code="us"), ) for simulation in (baseline, reform): simulation.extra_variables = {"spm_unit": ["spm_unit_spm_threshold"]} @@ -84,7 +85,7 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( dataset=native_dataset, dataset_selection=dataset_selection, policy=None, - region_code="us", + region=RegionResolution(code="us"), ) replay.id = baseline.id replay.ensure() @@ -96,7 +97,7 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( dataset=native_dataset, dataset_selection=dataset_selection, policy=None, - region_code="us", + region=RegionResolution(code="us"), ) local.id = baseline.id assert local.storage_id != baseline.storage_id @@ -125,6 +126,7 @@ def test_native_state_only_worker_requires_geography_and_explicit_national_works from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_executor.simulation_runtime import ( DatasetSelection, + RegionResolution, _build_simulation, ) @@ -143,7 +145,7 @@ def test_native_state_only_worker_requires_geography_and_explicit_national_works dataset=native_dataset, dataset_selection=dataset_selection, policy=None, - region_code="us", + region=RegionResolution(code="us"), ) with pytest.raises(ValueError) as error: # noqa: PT011 # the typed code is asserted below; the message depends on which installed layer rejects the missing county simulation.run() @@ -153,7 +155,7 @@ def test_native_state_only_worker_requires_geography_and_explicit_national_works dataset=native_dataset, dataset_selection=dataset_selection, policy=None, - region_code="us", + region=RegionResolution(code="us"), ) national.run() assert national.spm_provenance()["geography_kind"] == "national" diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 5791fca7f..974354360 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -532,10 +532,11 @@ def fake_build_simulation( dataset, dataset_selection, # noqa: ARG001 # keyword-only, as the runtime passes it policy, - scoping_strategy=None, - region_code=None, + region, ): - build_calls.append((params, dataset, policy, scoping_strategy, region_code)) + build_calls.append( + (params, dataset, policy, region.scoping_strategy, region.code) + ) return baseline_simulation if len(build_calls) == 1 else reform_simulation class FakeSimulationOutputBuilder: @@ -666,8 +667,7 @@ def fake_build_simulation( dataset, # noqa: ARG001 # keyword-only, as the runtime passes it dataset_selection, # noqa: ARG001 # keyword-only, as the runtime passes it policy, # noqa: ARG001 # keyword-only, as the runtime passes it - scoping_strategy=None, # noqa: ARG001 # keyword-only, as the runtime passes it - region_code=None, # noqa: ARG001 # keyword-only, as the runtime passes it + region, # noqa: ARG001 # keyword-only, as the runtime passes it ): build_count[0] += 1 return simulations["baseline"] if build_count[0] == 1 else simulations["reform"] @@ -1073,10 +1073,9 @@ def fake_build_simulation( dataset, dataset_selection, # noqa: ARG001 # keyword-only, as the runtime passes it policy, - scoping_strategy=None, - region_code=None, # noqa: ARG001 # keyword-only, as the runtime passes it + region, ): - build_calls.append((params, dataset, policy, scoping_strategy)) + build_calls.append((params, dataset, policy, region.scoping_strategy)) return baseline_simulation if len(build_calls) == 1 else reform_simulation class FakeSimulationOutputBuilder: @@ -1166,11 +1165,10 @@ def fake_build_simulation( dataset, dataset_selection, policy, - scoping_strategy=None, - region_code=None, # noqa: ARG001 # keyword-only, as the runtime passes it + region, ): build_calls.append( - (params, dataset, policy, scoping_strategy, dataset_selection) + (params, dataset, policy, region.scoping_strategy, dataset_selection) ) return baseline_simulation if len(build_calls) == 1 else reform_simulation diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py b/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py index de1858686..65c589bbd 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py @@ -10,6 +10,7 @@ from policyengine_simulation_executor.simulation_runtime import ( DEFAULT_YEAR, DatasetSelection, + RegionResolution, _build_simulation, _compile_interval_policy, _country_module, @@ -167,7 +168,7 @@ def test_scalar_in_interval_policy_uses_request_year(period, year): dataset=dataset, dataset_selection=DatasetSelection("in-memory", "in-memory", is_default=False), policy=raw, - region_code="uk", + region=RegionResolution(code="uk"), ) assert _parse_year(params) == year @@ -229,7 +230,7 @@ def recording_compiler(reform, *, year, model_version): dataset=Dataset(name="in-memory", description="Compilation only", year=2027), dataset_selection=DatasetSelection("in-memory", "in-memory", is_default=False), policy=raw, - region_code="uk", + region=RegionResolution(code="uk"), ) assert len(calls) == 1 From e359319455bbc6891d7523d65ea9966a21989182 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:21:59 +0400 Subject: [PATCH 20/66] Pass artifact key inputs to the dataset and baseline key builders as dataclasses The payloads are built from the same fields in the same way; the golden digest tests still pass unchanged. Co-Authored-By: Claude Opus 5.5 --- .../artifact_keys.py | 138 ++++++++++-------- .../tests/test_artifact_keys.py | 18 ++- .../tests/test_canonical_spm.py | 12 +- 3 files changed, 93 insertions(+), 75 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index 06365584a..cb1eabcb8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -63,20 +63,9 @@ def canonical_digest(payload: Mapping[str, Any]) -> str: return hashlib.sha256(encoded.encode("utf-8")).hexdigest() -def dataset_key( - *, - country: str, - dataset: str, - year: int, - data_version: str, - data_artifact_revision: str, - source_sha256: str | None, - data_build_fingerprint: str | None, - model_version: str, - policyengine_version: str, - spm: dict | None = None, -) -> str: - """Digest identifying one single-year dataset artifact. +@dataclass(frozen=True) +class DatasetKeyInputs: + """Every input that can change one single-year dataset artifact's bytes. ``model_version`` and ``policyengine_version`` belong here because the single-year build runs a full Microsimulation pass — its output depends @@ -85,52 +74,71 @@ def dataset_key( closes the same-version re-release hole), ``data_build_fingerprint`` the certification fingerprint; both may be None when unavailable. """ + + country: str + dataset: str + year: int + data_version: str + data_artifact_revision: str + source_sha256: str | None + data_build_fingerprint: str | None + model_version: str + policyengine_version: str + spm: dict | None = None + + +def dataset_key(inputs: DatasetKeyInputs) -> str: + """Digest identifying one single-year dataset artifact.""" return canonical_digest( { - "schema": DATASET_KEY_SCHEMA if spm is None else "ds2-spm", - **({"spm": spm} if spm is not None else {}), - "country": country.lower(), - "dataset": dataset, - "year": int(year), - "data_version": data_version, - "data_artifact_revision": data_artifact_revision, - "source_sha256": source_sha256, - "data_build_fingerprint": data_build_fingerprint, - "model_version": model_version, - "policyengine_version": policyengine_version, + "schema": DATASET_KEY_SCHEMA if inputs.spm is None else "ds2-spm", + **({"spm": inputs.spm} if inputs.spm is not None else {}), + "country": inputs.country.lower(), + "dataset": inputs.dataset, + "year": int(inputs.year), + "data_version": inputs.data_version, + "data_artifact_revision": inputs.data_artifact_revision, + "source_sha256": inputs.source_sha256, + "data_build_fingerprint": inputs.data_build_fingerprint, + "model_version": inputs.model_version, + "policyengine_version": inputs.policyengine_version, } ) -def baseline_key( - *, - country: str, - region: str, - scope_key: str | None, - dataset_digest: str, - model_version: str, - policyengine_version: str, - spm: dict | None = None, - policy: str = CURRENT_LAW_POLICY, -) -> str: - """Digest identifying one precomputed baseline simulation artifact. +@dataclass(frozen=True) +class BaselineKeyInputs: + """Every input that can change one precomputed baseline artifact's bytes. ``scope_key`` is the ``ScopingStrategy.cache_key`` string (None for an unscoped national run); ``region`` is the human-readable label ("national", "state/ca", ...). The simulation year rides in via ``dataset_digest``. """ + + country: str + region: str + scope_key: str | None + dataset_digest: str + model_version: str + policyengine_version: str + spm: dict | None = None + policy: str = CURRENT_LAW_POLICY + + +def baseline_key(inputs: BaselineKeyInputs) -> str: + """Digest identifying one precomputed baseline simulation artifact.""" return canonical_digest( { - "schema": BASELINE_KEY_SCHEMA if spm is None else "bl2-spm", - **({"spm": spm} if spm is not None else {}), - "country": country.lower(), - "region": region, - "scope_key": scope_key, - "dataset_key": dataset_digest, - "policy": policy, - "model_version": model_version, - "policyengine_version": policyengine_version, + "schema": BASELINE_KEY_SCHEMA if inputs.spm is None else "bl2-spm", + **({"spm": inputs.spm} if inputs.spm is not None else {}), + "country": inputs.country.lower(), + "region": inputs.region, + "scope_key": inputs.scope_key, + "dataset_key": inputs.dataset_digest, + "policy": inputs.policy, + "model_version": inputs.model_version, + "policyengine_version": inputs.policyengine_version, } ) @@ -222,16 +230,18 @@ class DatasetArtifactIdentity: def digest(self) -> str: """Content key over every input that can change this dataset's bytes.""" return dataset_key( - country=self.country, - spm=self.spm, - dataset=self.dataset, - year=self.year, - data_version=self.data_version, - data_artifact_revision=self.data_artifact_revision, - source_sha256=self.source_sha256, - data_build_fingerprint=self.data_build_fingerprint, - model_version=self.model_version, - policyengine_version=self.policyengine_version, + DatasetKeyInputs( + country=self.country, + spm=self.spm, + dataset=self.dataset, + year=self.year, + data_version=self.data_version, + data_artifact_revision=self.data_artifact_revision, + source_sha256=self.source_sha256, + data_build_fingerprint=self.data_build_fingerprint, + model_version=self.model_version, + policyengine_version=self.policyengine_version, + ) ) @property @@ -259,13 +269,15 @@ class BaselineArtifactIdentity: def digest(self) -> str: """Content key; nests the dataset digest, so a dataset change rotates it.""" return baseline_key( - country=self.country, - region=self.region, - spm=self.spm, - scope_key=self.scope_key, - dataset_digest=self.dataset.digest, - model_version=self.dataset.model_version, - policyengine_version=self.dataset.policyengine_version, + BaselineKeyInputs( + country=self.country, + region=self.region, + spm=self.spm, + scope_key=self.scope_key, + dataset_digest=self.dataset.digest, + model_version=self.dataset.model_version, + policyengine_version=self.dataset.policyengine_version, + ) ) @property diff --git a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py index d8eaad533..0cfb6dc18 100644 --- a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py +++ b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py @@ -52,25 +52,30 @@ class TestDigests: def test_dataset_key_golden(self): - assert ak.dataset_key(**_DATASET_KWARGS) == _DATASET_GOLDEN + assert ak.dataset_key(ak.DatasetKeyInputs(**_DATASET_KWARGS)) == _DATASET_GOLDEN def test_none_field_is_a_distinct_identity(self): - digest = ak.dataset_key(**{**_DATASET_KWARGS, "source_sha256": None}) + digest = ak.dataset_key( + ak.DatasetKeyInputs(**{**_DATASET_KWARGS, "source_sha256": None}) + ) assert digest == _DATASET_NO_SHA_GOLDEN assert digest != _DATASET_GOLDEN def test_baseline_key_golden(self): - assert ak.baseline_key(**_BASELINE_KWARGS) == _BASELINE_GOLDEN + assert ( + ak.baseline_key(ak.BaselineKeyInputs(**_BASELINE_KWARGS)) + == _BASELINE_GOLDEN + ) @pytest.mark.parametrize("field", sorted(_DATASET_KWARGS)) def test_every_dataset_field_perturbs_digest(self, field): perturbed = {**_DATASET_KWARGS, field: 2027 if field == "year" else "other"} - assert ak.dataset_key(**perturbed) != _DATASET_GOLDEN + assert ak.dataset_key(ak.DatasetKeyInputs(**perturbed)) != _DATASET_GOLDEN @pytest.mark.parametrize("field", sorted(_BASELINE_KWARGS)) def test_every_baseline_field_perturbs_digest(self, field): perturbed = {**_BASELINE_KWARGS, field: "other"} - assert ak.baseline_key(**perturbed) != _BASELINE_GOLDEN + assert ak.baseline_key(ak.BaselineKeyInputs(**perturbed)) != _BASELINE_GOLDEN def test_canonical_digest_ignores_insertion_order(self): forward = {"a": 1, "b": {"c": 2, "d": None}} @@ -213,8 +218,7 @@ def entry(self, year, *, scenario, as_of): assert identity.spm == expected_selection assert identity.digest == ak.dataset_key( - **_DATASET_KWARGS, - spm=expected_selection, + ak.DatasetKeyInputs(**_DATASET_KWARGS, spm=expected_selection) ) assert identity.digest != _DATASET_GOLDEN diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index e126a165b..4679cfefe 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -560,8 +560,10 @@ def test_each_setting_rotates_artifact_identity(key, value): } changed = {**SELECTION, key: value} assert artifact_keys.baseline_key( - **common, spm=SELECTION - ) != artifact_keys.baseline_key(**common, spm=changed) + artifact_keys.BaselineKeyInputs(**common, spm=SELECTION) + ) != artifact_keys.baseline_key( + artifact_keys.BaselineKeyInputs(**common, spm=changed) + ) ds = { "country": "us", "dataset": "test", @@ -573,9 +575,9 @@ def test_each_setting_rotates_artifact_identity(key, value): "model_version": "test", "policyengine_version": "test", } - assert artifact_keys.dataset_key(**ds, spm=SELECTION) != artifact_keys.dataset_key( - **ds, spm=changed - ) + assert artifact_keys.dataset_key( + artifact_keys.DatasetKeyInputs(**ds, spm=SELECTION) + ) != artifact_keys.dataset_key(artifact_keys.DatasetKeyInputs(**ds, spm=changed)) def test_fanout_and_budget_window_preserve_settings_and_receipts(): From 6b18009bbf1a7aa8e64dc8b5258e7ae0762bb719 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:23:37 +0400 Subject: [PATCH 21/66] Pass baseline identity facts as one request and split out the artifact scope check Co-Authored-By: Claude Opus 5.5 --- .../baseline_artifacts.py | 92 ++++++++++--------- .../precompute.py | 17 ++-- .../simulation_runtime.py | 17 ++-- .../tests/test_baseline_artifacts.py | 47 ++++++---- .../tests/test_canonical_spm.py | 18 ++-- .../tests/test_precompute.py | 17 ++-- 6 files changed, 117 insertions(+), 91 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py index 1ffe75907..dd8e1d06f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py @@ -21,6 +21,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from policyengine.core import Simulation @@ -45,15 +46,44 @@ OUTCOME_MISS = "miss" +@dataclass(frozen=True) +class BaselineIdentityRequest: + """The request facts that decide whether a baseline has a deterministic id. + + ``region_code`` is the normalized region (None when unresolved) and + ``scoping_strategy`` narrows the simulation to it (None for a whole + dataset, as a national run is). + """ + + params: dict[str, Any] + country: str + dataset_is_default: bool + policy: object + region_code: str | None + scoping_strategy: ScopingStrategy | None + year: int + + +def _artifact_scope( + request: BaselineIdentityRequest, +) -> tuple[str, str | None] | None: + """Return the artifact's region label and scope key, or None if unsupported. + + Supported scopes are an unscoped national run and a region group. + """ + if request.region_code is None: + return None + if request.scoping_strategy is None: + if request.region_code != request.country.lower(): + return None + return artifact_keys.NATIONAL_REGION, None + if isinstance(request.scoping_strategy, RegionGroupStrategy): + return request.region_code, request.scoping_strategy.cache_key + return None + + def qualifying_baseline_identity( - params: dict[str, Any], - *, - country: str, - dataset_is_default: bool, - policy: object, - region_code: str | None, - scoping_strategy: ScopingStrategy | None, - year: int, + request: BaselineIdentityRequest, ) -> artifact_keys.BaselineArtifactIdentity | None: """Full artifact identity for a qualifying baseline, else None. @@ -73,44 +103,28 @@ def qualifying_baseline_identity( reader cannot disagree by construction. Identity-collection errors propagate: the writer must fail loudly on them. """ - if policy is not None: + params = request.params + if request.policy is not None: return None if str(params.get("scope") or "").lower() != "macro": return None - if not dataset_is_default: + if not request.dataset_is_default: return None - if region_code is None: - return None - - country = country.lower() - if scoping_strategy is None: - if region_code != country: - return None - region, scope_key = artifact_keys.NATIONAL_REGION, None - elif isinstance(scoping_strategy, RegionGroupStrategy): - region, scope_key = region_code, scoping_strategy.cache_key - else: + scope = _artifact_scope(request) + if scope is None: return None + region, scope_key = scope return artifact_keys.collect_baseline_identity( - country, - year, + request.country.lower(), + request.year, region=region, scope_key=scope_key, **({"spm": params["spm"]} if params.get("spm") is not None else {}), ) -def deterministic_baseline_id( - params: dict[str, Any], - *, - country: str, - dataset_is_default: bool, - policy: object, - region_code: str | None, - scoping_strategy: ScopingStrategy | None, - year: int, -) -> str | None: +def deterministic_baseline_id(request: BaselineIdentityRequest) -> str | None: """Return the deterministic simulation id for a qualifying baseline, else None. This is the runtime reader's wrapper around ``qualifying_baseline_identity``. @@ -120,22 +134,14 @@ def deterministic_baseline_id( errors still propagate, since they mean the request itself is invalid. """ try: - identity = qualifying_baseline_identity( - params, - country=country, - dataset_is_default=dataset_is_default, - policy=policy, - region_code=region_code, - scoping_strategy=scoping_strategy, - year=year, - ) + identity = qualifying_baseline_identity(request) except Exception as exc: if spm_error_detail(exc): raise logger.warning( "Could not collect baseline artifact identity for %s; " "using a random simulation id", - country, + request.country, exc_info=True, ) return None 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 85b74a551..cd360501e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -121,6 +121,7 @@ def _cohort_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 ( # noqa: PLC0415 # resolved per call so tests can patch baseline_artifacts + BaselineIdentityRequest, qualifying_baseline_identity, ) from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers @@ -130,13 +131,15 @@ def cohort_identity(year: int, group: list[str]) -> BaselineArtifactIdentity: params, _, resolution = _cohort_resolution(year, group) dataset_selection = _resolve_dataset_selection(params, region_resolution=resolution) identity = qualifying_baseline_identity( - params, - country=PRECOMPUTE_COUNTRY, - dataset_is_default=dataset_selection.is_default, - policy=None, - region_code=resolution.code, - scoping_strategy=resolution.scoping_strategy, - year=year, + 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 = ( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index b0e2a03e3..3b69702f0 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -745,6 +745,7 @@ def _build_simulation( from policyengine_simulation_executor.baseline_artifacts import ( # noqa: PLC0415 # resolved per call so tests can patch baseline_artifacts ArtifactBaselineSimulation, + BaselineIdentityRequest, deterministic_baseline_id, ) from policyengine_simulation_executor.spm import ( # noqa: PLC0415 # resolved per call so tests can patch spm.normalize_runtime_spm @@ -761,13 +762,15 @@ def _build_simulation( ) scoping_strategy = region.scoping_strategy simulation_id = deterministic_baseline_id( - params, - country=country, - dataset_is_default=dataset_selection.is_default, - policy=policy, - region_code=region.code, - scoping_strategy=scoping_strategy, - year=_parse_year(params), + BaselineIdentityRequest( + params, + country=country, + dataset_is_default=dataset_selection.is_default, + policy=policy, + region_code=region.code, + scoping_strategy=scoping_strategy, + year=_parse_year(params), + ) ) if simulation_id is not None: # Deterministic id: ensure() loads a precomputed baseline artifact diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 7e90411c1..bc568d457 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -62,7 +62,7 @@ def _id(self, _collected, **overrides): "year": 2026, } kwargs.update(overrides) - return ba.deterministic_baseline_id(kwargs.pop("params"), **kwargs) + return ba.deterministic_baseline_id(ba.BaselineIdentityRequest(**kwargs)) def test_national_baseline_qualifies(self, collected): assert self._id(collected) == "bl1-deadbeefdeadbeef" @@ -102,13 +102,15 @@ def test_reforms_keep_current_law_artifact_identity_separate(self, collected, ki assert self._id(collected, policy=policy) is None assert ( ba.qualifying_baseline_identity( - {"scope": "macro"}, - country="us", - dataset_is_default=True, - policy=policy, - region_code="us", - scoping_strategy=None, - year=2026, + ba.BaselineIdentityRequest( + {"scope": "macro"}, + country="us", + dataset_is_default=True, + policy=policy, + region_code="us", + scoping_strategy=None, + year=2026, + ) ) is None ) @@ -145,13 +147,15 @@ def boom(*_args, **_kwargs): monkeypatch.setattr(ak, "collect_baseline_identity", boom) assert ( ba.deterministic_baseline_id( - {"scope": "macro"}, - country="us", - dataset_is_default=True, - policy=None, - region_code="us", - scoping_strategy=None, - year=2026, + ba.BaselineIdentityRequest( + {"scope": "macro"}, + country="us", + dataset_is_default=True, + policy=None, + region_code="us", + scoping_strategy=None, + year=2026, + ) ) is None ) @@ -403,8 +407,8 @@ def __init__(self, **kwargs): lambda _country: SimpleNamespace(model=SimpleNamespace(id="us-model")), ) - def fake_deterministic_id(params, **kwargs): - recorded.id_kwargs = {"params": params, **kwargs} + def fake_deterministic_id(request): + recorded.id_kwargs = vars(request) return recorded.id monkeypatch.setattr(ba, "deterministic_baseline_id", fake_deterministic_id) @@ -538,7 +542,9 @@ def boom(*_args, **_kwargs): monkeypatch.setattr(ak, "collect_baseline_identity", boom) with pytest.raises(RuntimeError, match="receipt unreadable"): - ba.qualifying_baseline_identity(**self._kwargs()) + ba.qualifying_baseline_identity( + ba.BaselineIdentityRequest(**self._kwargs()) + ) def test_deterministic_id_swallows_collection_errors(self, monkeypatch): def boom(*_args, **_kwargs): @@ -546,7 +552,10 @@ def boom(*_args, **_kwargs): raise RuntimeError(msg) monkeypatch.setattr(ak, "collect_baseline_identity", boom) - assert ba.deterministic_baseline_id(**self._kwargs()) is None + assert ( + ba.deterministic_baseline_id(ba.BaselineIdentityRequest(**self._kwargs())) + is None + ) def _year_data(person_columns=None): diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index 4679cfefe..fa28cd15e 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -620,7 +620,7 @@ def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): monkeypatch.setattr( baseline_artifacts, "deterministic_baseline_id", - lambda *_args, **kwargs: "baseline-id" if kwargs["policy"] is None else None, + lambda request: "baseline-id" if request.policy is None else None, ) monkeypatch.setattr( baseline_artifacts, "ArtifactBaselineSimulation", lambda **k: calls.append(k) @@ -651,13 +651,15 @@ def test_identity_errors_do_not_degrade_to_an_unidentified_baseline(monkeypatch) ) with pytest.raises(SPMInputError): baseline_artifacts.deterministic_baseline_id( - {}, - country="us", - dataset_is_default=True, - policy=None, - region_code="us", - scoping_strategy=None, - year=2026, + baseline_artifacts.BaselineIdentityRequest( + {}, + country="us", + dataset_is_default=True, + policy=None, + region_code="us", + scoping_strategy=None, + year=2026, + ) ) diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 694f954fa..501697b33 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -384,6 +384,7 @@ def get_region(code): 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, ) @@ -398,13 +399,15 @@ def test_precompute_identity_equals_runtime_id(self): params, region_resolution=resolution ) reader_id = deterministic_baseline_id( - params, - country="us", - dataset_is_default=dataset_selection.is_default, - policy=None, - region_code=resolution.code, - scoping_strategy=resolution.scoping_strategy, - year=2026, + 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") From f78e9397cf09b4bb5a1fe792297602b0dc7570fb Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:24:47 +0400 Subject: [PATCH 22/66] Pass the precompute's Modal function handles to run_precompute as one group The precompute local entrypoint keeps its --force CLI signature. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/precompute_app.py | 13 +++++--- .../precompute.py | 32 ++++++++++++------- .../tests/test_precompute.py | 5 ++- .../tests/test_precompute_app.py | 4 +-- 4 files changed, 35 insertions(+), 19 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index bd2365642..9d4f79e16 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -154,15 +154,18 @@ def main(force: bool = False) -> None: # noqa: FBT001, FBT002 # Modal local en resolve_bucket_name, ) from policyengine_simulation_executor.precompute import ( # noqa: PLC0415 # local entrypoint: loaded only when the precompute runs + PrecomputeFunctions, run_precompute, ) run_precompute( resolve_bucket_name(), force=force, - plan_artifacts=plan_artifacts, - build_dataset=build_dataset, - compute_baseline=compute_baseline, - verify_determinism=verify_determinism, - publish_manifest=publish_manifest, + functions=PrecomputeFunctions( + plan_artifacts=plan_artifacts, + build_dataset=build_dataset, + compute_baseline=compute_baseline, + verify_determinism=verify_determinism, + publish_manifest=publish_manifest, + ), ) 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 cd360501e..ace37643f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -34,6 +34,7 @@ import logging import time +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -477,20 +478,27 @@ def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: return ArtifactStore(bucket).write_manifest(manifest.canonical_payload()) +@dataclass(frozen=True) +class PrecomputeFunctions: + """The five Modal function handles the precompute orchestration calls.""" + + plan_artifacts: RemoteFunction + build_dataset: RemoteFunction + compute_baseline: RemoteFunction + verify_determinism: RemoteFunction + publish_manifest: RemoteFunction + + def run_precompute( bucket: str, *, force: bool, - plan_artifacts: RemoteFunction, - build_dataset: RemoteFunction, - compute_baseline: RemoteFunction, - verify_determinism: RemoteFunction, - publish_manifest: RemoteFunction, + functions: PrecomputeFunctions, echo: Callable[[str], None] = print, ) -> str: """Orchestrate the local entrypoint: plan, spawn misses, gate, publish. - Takes the five Modal function handles as parameters (each offering + Takes the five Modal function handles as ``functions`` (each offering ``.remote``/``.spawn``) so the flow is testable with plain fakes. Everything crossing a handle is a plain dict (Modal's serialization boundary); it is validated back into models immediately on return. @@ -500,7 +508,7 @@ def run_precompute( parses. """ echo(f"Planning against gs://{bucket} (force={force})") - plan = PrecomputePlan.model_validate(plan_artifacts.remote(bucket)) + plan = PrecomputePlan.model_validate(functions.plan_artifacts.remote(bucket)) work = select_work(plan, force=force) echo( f"Plan: {len(plan.datasets)} datasets ({len(work.datasets)} to " @@ -509,7 +517,8 @@ def run_precompute( ) dataset_handles = [ - build_dataset.spawn(bucket, entry.model_dump()) for entry in work.datasets + functions.build_dataset.spawn(bucket, entry.model_dump()) + for entry in work.datasets ] for handle in dataset_handles: result = DatasetBuildResult.model_validate(handle.get()) @@ -520,7 +529,8 @@ def run_precompute( ) baseline_handles = [ - compute_baseline.spawn(bucket, entry.model_dump()) for entry in work.baselines + functions.compute_baseline.spawn(bucket, entry.model_dump()) + for entry in work.baselines ] for handle in baseline_handles: result = BaselineComputeResult.model_validate(handle.get()) @@ -534,7 +544,7 @@ def run_precompute( probe = work.baselines[0] echo(f"Determinism gate: recomputing year={probe.year} group={probe.group}") verdict = DeterminismVerdict.model_validate( - verify_determinism.remote(bucket, probe.model_dump()) + functions.verify_determinism.remote(bucket, probe.model_dump()) ) if not verdict.equal: raise SystemExit( @@ -542,6 +552,6 @@ def run_precompute( ) echo("Determinism gate passed") - manifest_digest = publish_manifest.remote(bucket, plan.model_dump()) + manifest_digest = functions.publish_manifest.remote(bucket, plan.model_dump()) echo(manifest_digest_line(manifest_digest)) return manifest_digest diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 501697b33..b5606da93 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -282,7 +282,10 @@ def _functions(self, verify_result=None): def _run(self, functions, *, force=False): lines = [] digest = precompute.run_precompute( - "bucket-x", force=force, echo=lines.append, **functions + "bucket-x", + force=force, + functions=precompute.PrecomputeFunctions(**functions), + echo=lines.append, ) return digest, lines diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index 78409b902..16d710103 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_app.py @@ -225,10 +225,10 @@ def test_local_entrypoint_wires_bucket_force_and_handles( captured = {} - def fake_run_precompute(bucket, *, force, **handles): + def fake_run_precompute(bucket, *, force, functions): captured["bucket"] = bucket captured["force"] = force - captured["handles"] = handles + captured["handles"] = vars(functions) return "digest" monkeypatch.setattr( From 314b3f11985a34cc8ee8446aed5d1f62cc5d68de Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:25:41 +0400 Subject: [PATCH 23/66] Pin every input the national reduce passes to the output builder Co-Authored-By: Claude Opus 5.5 --- .../tests/test_segmented_national_reduce.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py index 1a9b8db4d..9aec2e99e 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py @@ -158,6 +158,68 @@ def test__rejects_non_us(self): year=2026, ) + @pytest.mark.parametrize( + ("resolved_data_version", "runtime"), + [(None, None), ("data-1", "runtime-sentinel")], + ) + def test__passes_every_builder_input( + self, monkeypatch, resolved_data_version, runtime + ): + captured = {} + + class FakeBuilder: + def __init__(self, **kwargs): + captured.update(kwargs) + + def serialize(self): + return {"budget": {}} + + monkeypatch.setattr(reduce_mod, "SimulationOutputBuilder", FakeBuilder) + monkeypatch.setattr( + reduce_mod, + "PrecomputedSimulation", + lambda **kwargs: SimpleNamespace(output_dataset=None, **kwargs), + ) + monkeypatch.setattr( + reduce_mod, + "_us_dataset_from_frames", + lambda frames, *, year: SimpleNamespace(frames=frames, year=year), + ) + country_module = SimpleNamespace(model=object()) + simulation_params = {"country": "us", "scope": "macro"} + + build_national_output( + [_child([1], [10]), _child([2], [20])], + country="us", + simulation_params=simulation_params, + country_module=country_module, + year=2027, + resolved_data_version=resolved_data_version, + runtime=runtime, + ) + + assert list(captured) == [ + "country", + "simulation_params", + "country_module", + "dataset", + "baseline", + "reform", + "resolved_data_version", + "resolved_region_code", + "runtime", + ] + assert captured["country"] == "us" + assert captured["simulation_params"] is simulation_params + assert captured["country_module"] is country_module + assert captured["dataset"].year == 2027 + assert captured["baseline"].dataset is captured["dataset"] + assert captured["reform"].dataset.year == 2027 + assert captured["reform"].dataset is not captured["dataset"] + assert captured["resolved_data_version"] == resolved_data_version + assert captured["resolved_region_code"] == "us" + assert captured["runtime"] == runtime + def test__propagates_child_versions_over_builder_fallbacks(self, monkeypatch): # The rebuilt dataset has no artifact metadata, so the builder's # version fallbacks misreport provenance; the children loaded the From fb2e6a9a499bb7274ac7db3afc9c0cf59a4d768f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:26:41 +0400 Subject: [PATCH 24/66] Pass the national reduce's run facts to build_national_output as one group Co-Authored-By: Claude Opus 5.5 --- .../src/modal/segmented_national.py | 13 +++--- .../segmented_national_reduce.py | 26 +++++++---- .../tests/test_segmented_national.py | 9 ++-- .../tests/test_segmented_national_reduce.py | 43 +++++++++++-------- 4 files changed, 58 insertions(+), 33 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 37a8f4634..f44c19d10 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -37,6 +37,7 @@ national_region_groups, ) from policyengine_simulation_executor.segmented_national_reduce import ( + NationalReduceInputs, build_national_output, ) from policyengine_simulation_observability.errors import log_and_redact_exception @@ -542,11 +543,13 @@ def _reduce( ): output = build_national_output( child_results, - country=self.country, - simulation_params=simulation_params, - country_module=country_module, - year=_parse_year(simulation_params), - resolved_data_version=None, + NationalReduceInputs( + country=self.country, + simulation_params=simulation_params, + country_module=country_module, + year=_parse_year(simulation_params), + resolved_data_version=None, + ), runtime=self.runtime, ) output.update( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py index 8a75b99ef..06435548d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py @@ -25,6 +25,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import TYPE_CHECKING, Any import pandas as pd @@ -121,14 +122,21 @@ def _us_dataset_from_frames( ) +@dataclass(frozen=True) +class NationalReduceInputs: + """The national run's facts that the reduce passes to the output builder.""" + + country: str + simulation_params: dict[str, Any] + country_module: ModuleType + year: int + resolved_data_version: str | None = None + + def build_national_output( child_outputs: list[dict], + inputs: NationalReduceInputs, *, - country: str, - simulation_params: dict[str, Any], - country_module: ModuleType, - year: int, - resolved_data_version: str | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: """Reduce the children's microdata into the full national macro output dict. @@ -137,13 +145,15 @@ def build_national_output( with ``resolved_region_code=country`` so the national congressional district table populates. No model recompute occurs. """ + country = inputs.country + country_module = inputs.country_module if country.lower() != "us": msg = f"No segmented-national reduce for country {country!r} (US only)" raise ValueError(msg) sides = concatenate_microdata(child_outputs) datasets = { - side: _us_dataset_from_frames(frames, year=year) + side: _us_dataset_from_frames(frames, year=inputs.year) for side, frames in sides.items() } @@ -158,12 +168,12 @@ def stand_in(dataset: PolicyEngineUSDataset) -> PrecomputedSimulation: builder = SimulationOutputBuilder( country=country, - simulation_params=simulation_params, + simulation_params=inputs.simulation_params, country_module=country_module, dataset=datasets["baseline"], baseline=stand_in(datasets["baseline"]), reform=stand_in(datasets["reform"]), - resolved_data_version=resolved_data_version, + resolved_data_version=inputs.resolved_data_version, resolved_region_code=country, runtime=runtime, ) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 958adb623..8ee822723 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -342,9 +342,10 @@ def test__uncovered_registry_states_ride_in_the_last_group(self, monkeypatch): def test__reduce_receives_clean_params_and_all_children(self, monkeypatch): captured = {} - def fake_build_national_output(child_results, **kwargs): + def fake_build_national_output(child_results, inputs, *, runtime): captured["children"] = child_results - captured.update(kwargs) + captured["runtime"] = runtime + captured.update(vars(inputs)) return {"budget": {}} monkeypatch.setattr(sn, "build_national_output", fake_build_national_output) @@ -850,7 +851,9 @@ def stub_reduce(monkeypatch): The reduce itself is covered by test_segmented_national_reduce. """ monkeypatch.setattr( - sn, "build_national_output", lambda _children, **_kwargs: {"budget": {}} + sn, + "build_national_output", + lambda _children, _inputs, **_kwargs: {"budget": {}}, ) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py index 9aec2e99e..686bae6e4 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national_reduce.py @@ -12,6 +12,7 @@ import pytest from policyengine_simulation_executor import segmented_national_reduce as reduce_mod from policyengine_simulation_executor.segmented_national_reduce import ( + NationalReduceInputs, PrecomputedSimulation, _us_dataset_from_frames, build_national_output, @@ -129,10 +130,12 @@ def __init__(self, **kwargs): country_module = SimpleNamespace(model=object()) result = build_national_output( [_child([1], [10]), _child([2], [20])], - country="us", - simulation_params={"country": "us", "scope": "macro"}, - country_module=country_module, - year=2026, + NationalReduceInputs( + country="us", + simulation_params={"country": "us", "scope": "macro"}, + country_module=country_module, + year=2026, + ), ) assert result == {"budget": {"budgetary_impact": -42.0}} @@ -152,10 +155,12 @@ def test__rejects_non_us(self): with pytest.raises(ValueError, match="US only"): build_national_output( [_child([1], [10])], - country="uk", - simulation_params={}, - country_module=SimpleNamespace(model=object()), - year=2026, + NationalReduceInputs( + country="uk", + simulation_params={}, + country_module=SimpleNamespace(model=object()), + year=2026, + ), ) @pytest.mark.parametrize( @@ -190,11 +195,13 @@ def serialize(self): build_national_output( [_child([1], [10]), _child([2], [20])], - country="us", - simulation_params=simulation_params, - country_module=country_module, - year=2027, - resolved_data_version=resolved_data_version, + NationalReduceInputs( + country="us", + simulation_params=simulation_params, + country_module=country_module, + year=2027, + resolved_data_version=resolved_data_version, + ), runtime=runtime, ) @@ -252,10 +259,12 @@ def serialize(self): ] result = build_national_output( children, - country="us", - simulation_params={"country": "us", "scope": "macro"}, - country_module=SimpleNamespace(model=object()), - year=2026, + NationalReduceInputs( + country="us", + simulation_params={"country": "us", "scope": "macro"}, + country_module=SimpleNamespace(model=object()), + year=2026, + ), ) assert result["model_version"] == "1.764.6" assert result["data_version"] == "buildm-x" From 539f4eba87c807ab71d18411d22cdfc7ebdb3d93 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:27:58 +0400 Subject: [PATCH 25/66] Pass fan-out poll timing to the batch and segmented runners as one backoff Also name the segmented runner's consecutive poll-error limit. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/budget_window_scheduler.py | 18 +++++++++------ .../src/modal/fanout.py | 14 ++++++++++++ .../src/modal/segmented_national.py | 22 +++++++++++++------ .../tests/test_budget_window_batch.py | 9 +++----- .../tests/test_segmented_national.py | 8 +++++-- 5 files changed, 49 insertions(+), 22 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py index 0947b09d4..2c2f7be68 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -30,6 +30,7 @@ build_budget_window_result, extract_annual_impact, ) +from src.modal.fanout import PollBackoff import modal @@ -56,6 +57,11 @@ # Retained for backward compatibility with callers that imported the original # constant; new code should use the initial/max pair above. POLL_INTERVAL_SECONDS = POLL_INTERVAL_INITIAL_SECONDS +DEFAULT_POLL_BACKOFF = PollBackoff( + initial_seconds=POLL_INTERVAL_INITIAL_SECONDS, + maximum_seconds=POLL_INTERVAL_MAX_SECONDS, + factor=POLL_INTERVAL_BACKOFF_FACTOR, +) def serialize_batch_status( @@ -100,9 +106,7 @@ def __init__( *, runtime: ObservabilityRuntime, modal_module: ModuleType | None = None, - poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, - poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, - poll_interval_backoff_factor: float = POLL_INTERVAL_BACKOFF_FACTOR, + poll_backoff: PollBackoff = DEFAULT_POLL_BACKOFF, ) -> None: """Load or seed the batch state and resolve a lazy child-function handle. @@ -112,11 +116,11 @@ def __init__( self.context = context self.runtime = runtime self.modal = modal if modal_module is None else modal_module - self.poll_interval_initial_seconds = poll_interval_seconds - self.poll_interval_max_seconds = poll_interval_max_seconds - self.poll_interval_backoff_factor = poll_interval_backoff_factor + self.poll_interval_initial_seconds = poll_backoff.initial_seconds + self.poll_interval_max_seconds = poll_backoff.maximum_seconds + self.poll_interval_backoff_factor = poll_backoff.factor # Kept for tests that still read this attribute. - self.poll_interval_seconds = poll_interval_seconds + self.poll_interval_seconds = poll_backoff.initial_seconds self.runtime.set_context( batch_job_id=context.batch_job_id, resolved_app_name=context.resolved_app_name, diff --git a/projects/policyengine-simulation-executor/src/modal/fanout.py b/projects/policyengine-simulation-executor/src/modal/fanout.py index 42a9d4f12..d9a037c2e 100644 --- a/projects/policyengine-simulation-executor/src/modal/fanout.py +++ b/projects/policyengine-simulation-executor/src/modal/fanout.py @@ -10,6 +10,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -36,6 +37,19 @@ def build_child_payload( return payload +@dataclass(frozen=True) +class PollBackoff: + """How a parent job spaces its polls of child calls that are still running. + + Sleeps start at ``initial_seconds`` and grow by ``factor``, up to + ``maximum_seconds``, while no child finishes. + """ + + initial_seconds: float + maximum_seconds: float + factor: float + + def next_backoff(current: float, *, factor: float, maximum: float) -> float: """Return the next sleep in an exponential backoff walk, capped at ``maximum``.""" return min(current * factor, maximum) diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index f44c19d10..2a9c22ae9 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -46,7 +46,7 @@ Stage, ) from policyengine_simulation_observability.telemetry import split_internal_payload -from src.modal.fanout import build_child_payload, next_backoff +from src.modal.fanout import PollBackoff, build_child_payload, next_backoff import modal @@ -61,6 +61,14 @@ POLL_INTERVAL_INITIAL_SECONDS = 0.5 POLL_INTERVAL_MAX_SECONDS = 8.0 POLL_INTERVAL_BACKOFF_FACTOR = 2.0 +DEFAULT_POLL_BACKOFF = PollBackoff( + initial_seconds=POLL_INTERVAL_INITIAL_SECONDS, + maximum_seconds=POLL_INTERVAL_MAX_SECONDS, + factor=POLL_INTERVAL_BACKOFF_FACTOR, +) +# A child's poll may fail this many times in a row before the job fails: +# one transient poll-RPC blip is retried, a real child failure repeats. +CONSECUTIVE_POLL_ERROR_LIMIT = 2 # The dedicated children pool (see module docstring). SEGMENT_FUNCTION_NAME = "run_simulation_segment" @@ -184,8 +192,7 @@ def __init__( app_name: str, runtime: ObservabilityRuntime, modal_module: ModuleType | None = None, - poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, - poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, + poll_backoff: PollBackoff = DEFAULT_POLL_BACKOFF, child_deadline_seconds: float = CHILD_DEADLINE_SECONDS, parent_budget_seconds: float = ( PARENT_FUNCTION_TIMEOUT_SECONDS - PARENT_COMPLETION_MARGIN_SECONDS @@ -200,8 +207,9 @@ def __init__( self.app_name = app_name self.runtime = runtime self.modal = modal if modal_module is None else modal_module - self.poll_interval_initial_seconds = poll_interval_seconds - self.poll_interval_max_seconds = poll_interval_max_seconds + self.poll_interval_initial_seconds = poll_backoff.initial_seconds + self.poll_interval_max_seconds = poll_backoff.maximum_seconds + self.poll_interval_backoff_factor = poll_backoff.factor self.child_deadline_seconds = child_deadline_seconds self.parent_budget_seconds = parent_budget_seconds self.clock = clock @@ -346,7 +354,7 @@ def _collect(self, children: list[ChildExecution]) -> list[dict]: # real child failure re-raises on the next probe too. attempts = poll_errors.get(index, 0) + 1 poll_errors[index] = attempts - if attempts < 2: + if attempts < CONSECUTIVE_POLL_ERROR_LIMIT: logger.warning( "Poll error for group %s (attempt %d), retrying: %s", child.group_label, @@ -385,7 +393,7 @@ def _collect(self, children: list[ChildExecution]) -> list[dict]: time.sleep(current_sleep) current_sleep = next_backoff( current_sleep, - factor=POLL_INTERVAL_BACKOFF_FACTOR, + factor=self.poll_interval_backoff_factor, maximum=self.poll_interval_max_seconds, ) elif progress_made: diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py index 115493e4f..bb22abc98 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py @@ -13,6 +13,7 @@ PolicyEngineBundle, ) from src.modal.budget_window_batch import run_budget_window_batch_impl +from src.modal.fanout import PollBackoff class SequencedCall: @@ -350,9 +351,7 @@ def test_scheduler_sleep_exponentially_backs_off_then_resets_on_progress( raw_params=payload, ), runtime=observability_runtime, - poll_interval_seconds=0.5, - poll_interval_max_seconds=4.0, - poll_interval_backoff_factor=2.0, + poll_backoff=PollBackoff(initial_seconds=0.5, maximum_seconds=4.0, factor=2.0), ) runner.run() @@ -415,9 +414,7 @@ def test_scheduler_publishes_bounded_poll_aggregates( raw_params=payload, ), runtime=observability_runtime, - poll_interval_seconds=0.5, - poll_interval_max_seconds=4.0, - poll_interval_backoff_factor=2.0, + poll_backoff=PollBackoff(initial_seconds=0.5, maximum_seconds=4.0, factor=2.0), ) runner.run() diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 8ee822723..18d4bd54b 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -13,6 +13,7 @@ US_NATIONAL_REGION_GROUPS, ) from src.modal import segmented_national as sn +from src.modal.fanout import PollBackoff NATIONAL = {"country": "us", "scope": "macro", "time_period": "2026"} TEST_RUNTIME = NoOpObservabilityRuntime() @@ -222,8 +223,11 @@ def _runner(fake_modal, params=None, runtime=TEST_RUNTIME, **runner_options): app_name="policyengine-simulation-py-test", runtime=runtime, modal_module=fake_modal, - poll_interval_seconds=0.001, - poll_interval_max_seconds=0.002, + poll_backoff=PollBackoff( + initial_seconds=0.001, + maximum_seconds=0.002, + factor=sn.POLL_INTERVAL_BACKOFF_FACTOR, + ), **runner_options, ) From a43a5af10b501b6324c707ce8e8ab84f52ceb911 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:29:20 +0400 Subject: [PATCH 26/66] Pin the Stage 12 bundle schema version check Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_bundle.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index 21704607c..45d98f3f8 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -120,3 +120,13 @@ def test_stage12_extractor_does_not_inspect_installed_distribution_versions() -> ) assert "importlib.metadata" not in text assert "pkg_resources" not in text + + +@pytest.mark.parametrize("schema_version", [None, 1, 3, "2"]) +def test_only_bundle_schema_version_2_is_supported(schema_version) -> None: + bundle = {**_bundle(), "schema_version": schema_version} + + with pytest.raises( + Stage12BundleError, match=r"^PolicyEngine\.py bundle schema_version must be 2$" + ): + normalize_stage12_bundle(bundle) From 201d549ea63b1b1732f07b2ccf1fc89a267ee889 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:29:20 +0400 Subject: [PATCH 27/66] Name the supported Stage 12 bundle schema version Co-Authored-By: Claude Opus 5.5 --- .../stage12_bundle/manifest.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py index 20c06e6c5..90c24bca7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py @@ -25,6 +25,9 @@ if TYPE_CHECKING: from collections.abc import Sequence +# The packaged bundle layout this loader reads. +SUPPORTED_BUNDLE_SCHEMA_VERSION = 2 + def _normalize_stage12_bundle( raw_bundle: object, @@ -32,8 +35,11 @@ def _normalize_stage12_bundle( supported_countries: Sequence[CountryId] = SUPPORTED_COUNTRIES, ) -> Stage12BundleManifest: bundle = require_mapping(raw_bundle, "root") - if bundle.get("schema_version") != 2: - msg = "PolicyEngine.py bundle schema_version must be 2" + if bundle.get("schema_version") != SUPPORTED_BUNDLE_SCHEMA_VERSION: + msg = ( + "PolicyEngine.py bundle schema_version must be " + f"{SUPPORTED_BUNDLE_SCHEMA_VERSION}" + ) raise Stage12BundleError(msg) policyengine_version = require_text( bundle.get("policyengine_version"), From e98b123e288b05f82d600f34baf0b11bdf8e30fa Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:30:31 +0400 Subject: [PATCH 28/66] Pin every Stage 12 parity mismatch code, path and comparison order Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_parity.py | 248 ++++++++++++++++++ 1 file changed, 248 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_parity.py b/projects/policyengine-simulation-executor/tests/test_stage12_parity.py index 3a006766a..abd362b3d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_parity.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_parity.py @@ -90,3 +90,251 @@ def test_aggregate_parity_compares_every_field_and_bounds_diagnostics() -> None: candidate, tolerances={"budget.total": NumericalTolerance(absolute=0.01)}, ) + + +def _without_column(frames, entity, column): + return {**frames, entity: frames[entity].drop(columns=[column])} + + +def _with_column(frames, entity, column, values): + return {**frames, entity: frames[entity].assign(**{column: values})} + + +def _person_frames(**columns): + return { + "person": pd.DataFrame( + {"person_id": np.array([1, 2], dtype=np.int64), **columns} + ) + } + + +@pytest.mark.parametrize( + ("incumbent", "candidate", "code", "path"), + [ + pytest.param( + _without_column(_frames(), "person", "person_id"), + _frames(), + "missing_row_identity", + "person", + id="incumbent-missing-id", + ), + pytest.param( + _frames(), + { + **_frames(), + "person": pd.DataFrame( + { + "person_id": np.array([1, 1], dtype=np.int64), + "person_household_id": np.array([1, 1], dtype=np.int64), + } + ), + }, + "duplicate_row_identity", + "person", + id="candidate-duplicate-id", + ), + pytest.param( + _frames(), + {"household": _frames()["household"]}, + "entity_membership_mismatch", + "simulation", + id="missing-entity", + ), + pytest.param( + _frames(), + _with_column(_frames(), "household", "extra", [1, 2]), + "column_membership_mismatch", + "household", + id="extra-column", + ), + pytest.param( + _frames(), + _with_column( + _frames(), "person", "person_id", np.array([3, 1], dtype=np.int64) + ), + "row_identity_mismatch", + "person", + id="different-ids", + ), + pytest.param( + _person_frames(name=["a", "b"]), + _person_frames(name=["a", "c"]), + "value_mismatch", + "person.name", + id="text-value", + ), + pytest.param( + _person_frames(score=np.array([1.0, np.nan])), + _person_frames(score=np.array([1.0, 2.0])), + "numeric_value_mismatch", + "person.score", + id="nan-against-number", + ), + pytest.param( + { + "a": pd.DataFrame({"a_id": [1], "value": [1.0]}), + "b": pd.DataFrame({"b_id": [1], "value": [1.0]}), + }, + { + "a": pd.DataFrame({"a_id": [1], "value": [2.0]}), + "b": pd.DataFrame({"b_id": [1], "value": [3.0]}), + }, + "numeric_value_mismatch", + "a.value", + id="first-entity-reported", + ), + pytest.param( + _person_frames(name=["a", "b"], score=[1.0, 2.0]), + _person_frames(name=["a", "c"], score=[1.0, 3.0]), + "value_mismatch", + "person.name", + id="columns-checked-in-name-order", + ), + pytest.param( + _person_frames(score=[1.0, 2.0]), + _person_frames(score=[1, 2]), + "dtype_mismatch", + "person.score", + id="float-against-int", + ), + ], +) +def test_simulation_parity_reports_the_first_mismatch(incumbent, candidate, code, path): + with pytest.raises(ParityMismatchError) as error: + compare_simulation_frames(incumbent, candidate) + + assert (error.value.code, error.value.path) == (code, path) + assert str(error.value) == f"{code} at {path}" + + +def test_simulation_parity_treats_matching_nan_and_text_as_equal() -> None: + frames = _person_frames(score=np.array([np.nan, 1.0]), name=["a", None]) + + compare_simulation_frames( + frames, + _person_frames(score=np.array([np.nan, 1.0]), name=["a", None]), + ) + + +def test_simulation_parity_applies_relative_tolerances() -> None: + compare_simulation_frames( + _person_frames(score=[100.0, 200.0]), + _person_frames(score=[101.0, 202.0]), + tolerances={"person.score": NumericalTolerance(relative=0.01)}, + ) + with pytest.raises(ParityMismatchError, match="numeric_value_mismatch"): + compare_simulation_frames( + _person_frames(score=[100.0, 200.0]), + _person_frames(score=[101.0, 203.0]), + tolerances={"person.score": NumericalTolerance(relative=0.01)}, + ) + + +@pytest.mark.parametrize( + ("incumbent", "candidate", "code", "path"), + [ + pytest.param( + {"a": 1}, {"b": 1}, "field_membership_mismatch", "report", id="root-keys" + ), + pytest.param( + {"a": 1}, [1], "field_membership_mismatch", "report", id="root-type" + ), + pytest.param( + {"a": {"b": 1}}, + {"a": {"b": 1, "c": 2}}, + "field_membership_mismatch", + "a", + id="nested-keys", + ), + pytest.param( + {"a": [1]}, {"a": "1"}, "value_type_mismatch", "a", id="list-as-text" + ), + pytest.param( + {"a": [1]}, {"a": 1}, "value_type_mismatch", "a", id="list-as-number" + ), + pytest.param( + {"a": [1, 2]}, {"a": [1]}, "list_length_mismatch", "a", id="list-length" + ), + pytest.param( + {"a": [{"b": 1}, {"b": 2}]}, + {"a": [{"b": 1}, {"b": 3}]}, + "numeric_value_mismatch", + "a[1].b", + id="nested-index", + ), + pytest.param([1, 2], [1, 3], "numeric_value_mismatch", "[1]", id="root-list"), + pytest.param( + {"a": True}, {"a": 1}, "value_type_mismatch", "a", id="bool-as-number" + ), + pytest.param( + {"a": 1}, {"a": False}, "value_type_mismatch", "a", id="number-as-bool" + ), + pytest.param({"a": True}, {"a": False}, "value_mismatch", "a", id="bool-value"), + pytest.param({"a": "x"}, {"a": "y"}, "value_mismatch", "a", id="text-value"), + pytest.param( + {"a": None}, {"a": 0}, "value_type_mismatch", "a", id="none-as-number" + ), + pytest.param( + {"a": "x"}, {"a": b"x"}, "value_type_mismatch", "a", id="text-as-bytes" + ), + pytest.param( + {"b": 1, "a": 2}, + {"b": 2, "a": 3}, + "numeric_value_mismatch", + "a", + id="keys-checked-in-sorted-order", + ), + pytest.param( + {"a": float("nan")}, + {"a": float("nan")}, + "numeric_value_mismatch", + "a", + id="nan-never-matches", + ), + ], +) +def test_aggregate_parity_reports_the_first_mismatch(incumbent, candidate, code, path): + with pytest.raises(ParityMismatchError) as error: + compare_aggregate_reports(incumbent, candidate) + + assert (error.value.code, error.value.path) == (code, path) + + +@pytest.mark.parametrize( + ("incumbent", "candidate"), + [ + ({"a": 1}, {"a": 1.0}), + ({"a": [1, "x", None, True]}, {"a": (1, "x", None, True)}), + ({"a": {"b": [1.0]}}, {"a": {"b": [1]}}), + ("text", "text"), + (None, None), + ], +) +def test_aggregate_parity_accepts_equal_reports(incumbent, candidate): + compare_aggregate_reports(incumbent, candidate) + + +def test_aggregate_parity_applies_relative_tolerances_by_path() -> None: + compare_aggregate_reports( + {"deciles": [100.0]}, + {"deciles": [101.0]}, + tolerances={"deciles[0]": NumericalTolerance(relative=0.01)}, + ) + + +def test_parity_paths_are_truncated_to_255_characters() -> None: + long_key = "k" * 300 + + with pytest.raises(ParityMismatchError) as error: + compare_aggregate_reports({long_key: 1}, {long_key: 2}) + + assert error.value.path == long_key[:255] + assert str(error.value) == f"numeric_value_mismatch at {long_key[:255]}" + + +@pytest.mark.parametrize("tolerance", [{"absolute": -0.1}, {"relative": -0.1}]) +def test_negative_tolerances_are_rejected(tolerance) -> None: + with pytest.raises( + ValueError, match=r"^numerical tolerances must be non-negative$" + ): + NumericalTolerance(**tolerance) From c66fd2bcc8cdb20d42cb37476b8c79f0325154e9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:30:48 +0400 Subject: [PATCH 29/66] Split Stage 12 parity comparisons into entity, column and report-value checks Co-Authored-By: Claude Opus 5.5 --- .../stage12_parity.py | 214 +++++++++++------- 1 file changed, 137 insertions(+), 77 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py index 483360ab1..2e68cebb2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py @@ -78,38 +78,59 @@ def compare_simulation_frames( raise ParityMismatchError(error_code, "simulation") allowed = tolerances or {} for entity, expected_frame in expected.items(): - actual_frame = actual[entity] - if tuple(expected_frame.columns) != tuple(actual_frame.columns): - error_code = "column_membership_mismatch" - raise ParityMismatchError(error_code, entity) - if len(expected_frame) != len(actual_frame): - error_code = "row_membership_mismatch" - raise ParityMismatchError(error_code, entity) - identifier = f"{entity}_id" - if not expected_frame[identifier].equals(actual_frame[identifier]): - error_code = "row_identity_mismatch" - raise ParityMismatchError(error_code, entity) - for column in expected_frame.columns: - path = f"{entity}.{column}" - expected_values = expected_frame[column] - actual_values = actual_frame[column] - if str(expected_values.dtype) != str(actual_values.dtype): - error_code = "dtype_mismatch" - raise ParityMismatchError(error_code, path) - if pd.api.types.is_numeric_dtype(expected_values.dtype): - tolerance = allowed.get(path, NumericalTolerance()) - if not np.allclose( - expected_values.to_numpy(), - actual_values.to_numpy(), - rtol=tolerance.relative, - atol=tolerance.absolute, - equal_nan=True, - ): - error_code = "numeric_value_mismatch" - raise ParityMismatchError(error_code, path) - elif not expected_values.equals(actual_values): - error_code = "value_mismatch" - raise ParityMismatchError(error_code, path) + _compare_entity_frames(entity, expected_frame, actual[entity], allowed) + + +def _compare_entity_frames( + entity: str, + expected_frame: pd.DataFrame, + actual_frame: pd.DataFrame, + tolerances: Mapping[str, NumericalTolerance], +) -> None: + """Check one canonical entity's columns, row ids and then each column's values.""" + if tuple(expected_frame.columns) != tuple(actual_frame.columns): + error_code = "column_membership_mismatch" + raise ParityMismatchError(error_code, entity) + if len(expected_frame) != len(actual_frame): + error_code = "row_membership_mismatch" + raise ParityMismatchError(error_code, entity) + identifier = f"{entity}_id" + if not expected_frame[identifier].equals(actual_frame[identifier]): + error_code = "row_identity_mismatch" + raise ParityMismatchError(error_code, entity) + for column in expected_frame.columns: + path = f"{entity}.{column}" + _compare_column( + path, + expected_frame[column], + actual_frame[column], + tolerances.get(path, NumericalTolerance()), + ) + + +def _compare_column( + path: str, + expected_values: pd.Series, + actual_values: pd.Series, + tolerance: NumericalTolerance, +) -> None: + """Require one dtype, then values within ``tolerance`` (numeric) or equal.""" + if str(expected_values.dtype) != str(actual_values.dtype): + error_code = "dtype_mismatch" + raise ParityMismatchError(error_code, path) + if pd.api.types.is_numeric_dtype(expected_values.dtype): + if not np.allclose( + expected_values.to_numpy(), + actual_values.to_numpy(), + rtol=tolerance.relative, + atol=tolerance.absolute, + equal_nan=True, + ): + error_code = "numeric_value_mismatch" + raise ParityMismatchError(error_code, path) + elif not expected_values.equals(actual_values): + error_code = "value_mismatch" + raise ParityMismatchError(error_code, path) def compare_aggregate_reports( @@ -124,51 +145,90 @@ def compare_aggregate_reports( within the tolerance keyed by their path, such as ``budget.total`` or ``deciles[3]``; any other value must be equal and of the same type. """ - allowed = tolerances or {} + _compare_report_values(incumbent, candidate, "", tolerances or {}) - def compare(expected: object, actual: object, path: str) -> None: - if isinstance(expected, Mapping): - if not isinstance(actual, Mapping) or expected.keys() != actual.keys(): - error_code = "field_membership_mismatch" - raise ParityMismatchError(error_code, path or "report") - for key in sorted(expected): - compare(expected[key], actual[key], f"{path}.{key}".lstrip(".")) - return - if isinstance(expected, Sequence) and not isinstance( - expected, (str, bytes, bytearray) - ): - if not isinstance(actual, Sequence) or isinstance( - actual, (str, bytes, bytearray) - ): - error_code = "value_type_mismatch" - raise ParityMismatchError(error_code, path) - if len(expected) != len(actual): - error_code = "list_length_mismatch" - raise ParityMismatchError(error_code, path) - for index, value in enumerate(expected): - compare(value, actual[index], f"{path}[{index}]") - return - if ( - isinstance(expected, (int, float)) - and not isinstance(expected, bool) - and isinstance(actual, (int, float)) - and not isinstance(actual, bool) + +def _is_report_list(value: object) -> bool: + """Whether a report value is a list-like sequence rather than text or bytes.""" + return isinstance(value, Sequence) and not isinstance( + value, (str, bytes, bytearray) + ) + + +def _is_report_number(value: object) -> bool: + """Whether a report value is an int or float; booleans are not numbers here.""" + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _compare_report_values( + expected: object, + actual: object, + path: str, + tolerances: Mapping[str, NumericalTolerance], +) -> None: + """Compare two report values at ``path`` by the shape of the expected value.""" + if isinstance(expected, Mapping): + _compare_report_mappings(expected, actual, path, tolerances) + elif _is_report_list(expected): + _compare_report_lists(expected, actual, path, tolerances) + else: + _compare_report_scalars(expected, actual, path, tolerances) + + +def _compare_report_mappings( + expected: Mapping, + actual: object, + path: str, + tolerances: Mapping[str, NumericalTolerance], +) -> None: + """Require the same keys, then compare each field in sorted key order.""" + if not isinstance(actual, Mapping) or expected.keys() != actual.keys(): + error_code = "field_membership_mismatch" + raise ParityMismatchError(error_code, path or "report") + for key in sorted(expected): + _compare_report_values( + expected[key], actual[key], f"{path}.{key}".lstrip("."), tolerances + ) + + +def _compare_report_lists( + expected: Sequence, + actual: object, + path: str, + tolerances: Mapping[str, NumericalTolerance], +) -> None: + """Require a list of the same length, then compare each item in order.""" + if not _is_report_list(actual): + error_code = "value_type_mismatch" + raise ParityMismatchError(error_code, path) + if len(expected) != len(actual): + error_code = "list_length_mismatch" + raise ParityMismatchError(error_code, path) + for index, value in enumerate(expected): + _compare_report_values(value, actual[index], f"{path}[{index}]", tolerances) + + +def _compare_report_scalars( + expected: object, + actual: object, + path: str, + tolerances: Mapping[str, NumericalTolerance], +) -> None: + """Compare numbers within the path's tolerance; other values must be equal.""" + if _is_report_number(expected) and _is_report_number(actual): + tolerance = tolerances.get(path, NumericalTolerance()) + if not math.isclose( + float(expected), + float(actual), + rel_tol=tolerance.relative, + abs_tol=tolerance.absolute, ): - tolerance = allowed.get(path, NumericalTolerance()) - if not math.isclose( - float(expected), - float(actual), - rel_tol=tolerance.relative, - abs_tol=tolerance.absolute, - ): - error_code = "numeric_value_mismatch" - raise ParityMismatchError(error_code, path) - return - if type(expected) is not type(actual): - error_code = "value_type_mismatch" - raise ParityMismatchError(error_code, path) - if expected != actual: - error_code = "value_mismatch" + error_code = "numeric_value_mismatch" raise ParityMismatchError(error_code, path) - - compare(incumbent, candidate, "") + return + if type(expected) is not type(actual): + error_code = "value_type_mismatch" + raise ParityMismatchError(error_code, path) + if expected != actual: + error_code = "value_mismatch" + raise ParityMismatchError(error_code, path) From a52f75306f14ef231e81bcb3f2aa6a713e38e439 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:31:45 +0400 Subject: [PATCH 30/66] Pin parity qualification receipts, check order and default runner lifecycle Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_qualification.py | 199 ++++++++++++++++++ 1 file changed, 199 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index e2927ab51..da4a5e25a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -200,3 +200,202 @@ def existing_runner(_request): assert error.value.code == "country_package_provenance_mismatch" assert "different" not in str(error.value) + + +def _matching_existing_runner(report, baseline, reform, **overrides): + def existing_runner(_request): + return { + "model_version": report.baseline.bundle.country_package_version, + "data_version": report.baseline.bundle.dataset.artifact_revision, + "budget": {"budgetary_impact": 20.0}, + "_microdata": extract_output_microdata( + _Simulation(baseline), + _Simulation(reform), + ), + **overrides, + } + + return existing_runner + + +def _matching_single_runner(baseline, reform, calls=None): + def single_runner(simulation): + if calls is not None: + calls.append(simulation.role) + frames = baseline if simulation.role is SimulationRole.BASELINE else reform + return SimulationCalculation(frames=frames) + + return single_runner + + +def test_qualification_receipt_records_every_digest_and_tolerance_field() -> None: + from policyengine_simulation_executor.stage12_parity import NumericalTolerance + + report = _report() + baseline = _frames(100.0) + reform = _frames(120.0) + + receipt = qualify_report_parity( + _existing_request(), + report, + existing_runner=_matching_existing_runner(report, baseline, reform), + single_simulation_runner=_matching_single_runner(baseline, reform), + aggregate_builder=lambda **_: { + "result": { + "model_version": report.baseline.bundle.country_package_version, + "data_version": report.baseline.bundle.dataset.artifact_revision, + "budget": {"budgetary_impact": 20.0}, + } + }, + simulation_tolerances={ + "person.person_household_id": NumericalTolerance(), + "household.household_net_income": NumericalTolerance(absolute=0.1), + }, + aggregate_tolerances={"budget.budgetary_impact": NumericalTolerance()}, + ) + + dumped = receipt.model_dump(mode="json") + digests = {key: value for key, value in dumped.items() if key.endswith("sha256")} + assert set(digests) == { + "bundle_manifest_sha256", + "baseline_input_sha256", + "reform_input_sha256", + "existing_baseline_output_sha256", + "v2_baseline_output_sha256", + "existing_reform_output_sha256", + "v2_reform_output_sha256", + "existing_aggregate_output_sha256", + "v2_aggregate_output_sha256", + } + assert dumped["contract_version"] == 1 + assert dumped["evaluation_id"] == str(report.evaluation_id) + assert dumped["bundle_manifest_sha256"] == ( + report.baseline.bundle.bundle_manifest_sha256 + ) + assert dumped["baseline_input_sha256"] != dumped["reform_input_sha256"] + assert ( + dumped["existing_baseline_output_sha256"] + == (dumped["v2_baseline_output_sha256"]) + ) + assert ( + dumped["existing_baseline_output_sha256"] + != (dumped["existing_reform_output_sha256"]) + ) + assert dumped["simulation_tolerance_fields"] == [ + "household.household_net_income", + "person.person_household_id", + ] + assert dumped["aggregate_tolerance_fields"] == ["budget.budgetary_impact"] + + +@pytest.mark.parametrize( + ("overrides", "code", "path"), + [ + ( + {"data_version": "different"}, + "dataset_provenance_mismatch", + "report.data_version", + ), + ({"_microdata": None}, "missing_existing_microdata", "report._microdata"), + ( + {"_microdata": {"baseline": {}, "reform": {}}}, + "invalid_existing_microdata", + "report._microdata.baseline", + ), + ], +) +def test_existing_result_is_checked_before_v2_simulations_run( + overrides, code, path +) -> None: + report = _report() + baseline = _frames(100.0) + reform = _frames(120.0) + calls = [] + + with pytest.raises(ParityMismatchError) as error: + qualify_report_parity( + _existing_request(), + report, + existing_runner=_matching_existing_runner( + report, baseline, reform, **overrides + ), + single_simulation_runner=_matching_single_runner(baseline, reform, calls), + ) + + assert (error.value.code, error.value.path) == (code, path) + assert calls == [] + + +@pytest.mark.parametrize("v2_report", [{}, {"result": None}, {"result": []}]) +def test_v2_aggregate_without_a_result_object_is_rejected(v2_report) -> None: + report = _report() + baseline = _frames(100.0) + reform = _frames(120.0) + + with pytest.raises(ParityMismatchError) as error: + qualify_report_parity( + _existing_request(), + report, + existing_runner=_matching_existing_runner(report, baseline, reform), + single_simulation_runner=_matching_single_runner(baseline, reform), + aggregate_builder=lambda **_: v2_report, + ) + + assert (error.value.code, error.value.path) == ( + "invalid_v2_aggregate", + "report.result", + ) + + +@pytest.mark.parametrize("fails", [False, True]) +def test_default_existing_runner_uses_and_shuts_down_a_local_runtime( + monkeypatch, fails +) -> None: + from policyengine_simulation_executor import ( + simulation_runtime, + stage12_qualification, + ) + + report = _report() + baseline = _frames(100.0) + reform = _frames(120.0) + events = [] + + class LocalRuntime: + def shutdown(self): + events.append("shutdown") + + local_runtime = LocalRuntime() + + def run_simulation_impl(request, *, runtime): + events.append(("run", request, runtime)) + if fails: + msg = "existing executor failed" + raise RuntimeError(msg) + return _matching_existing_runner(report, baseline, reform)(request) + + monkeypatch.setattr( + stage12_qualification, "_local_qualification_runtime", lambda: local_runtime + ) + monkeypatch.setattr(simulation_runtime, "run_simulation_impl", run_simulation_impl) + + if fails: + with pytest.raises(RuntimeError, match="existing executor failed"): + qualify_report_parity(_existing_request(), report) + else: + with pytest.raises(ParityMismatchError): + qualify_report_parity( + _existing_request(), + report, + single_simulation_runner=_matching_single_runner(baseline, reform), + aggregate_builder=lambda **_: {"result": []}, + ) + + assert events == [ + ( + "run", + {**_existing_request(), "_emit_microdata": True}, + local_runtime, + ), + "shutdown", + ] From 5c4cfdbe179d614753addf6b40f97f977f11b970 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:33:35 +0400 Subject: [PATCH 31/66] Split parity qualification into existing-run, provenance, calculation and aggregate steps Group the injectable runners and the reviewed tolerances into two option objects, and update the qualification script and tests. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/utils/qualify_stage12_parity.py | 7 +- .../stage12_qualification.py | 192 ++++++++++++------ .../tests/test_stage12_qualification.py | 74 ++++--- 3 files changed, 183 insertions(+), 90 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py b/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py index 92db1debd..ac80360e0 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/qualify_stage12_parity.py @@ -9,6 +9,7 @@ from policyengine_simulation_executor.stage12_parity import NumericalTolerance from policyengine_simulation_executor.stage12_qualification import ( + ParityTolerances, qualify_report_parity, ) @@ -58,8 +59,10 @@ def main() -> None: receipt = qualify_report_parity( _json_object(arguments.existing_request), _json_object(arguments.report_input), - simulation_tolerances=_tolerances(arguments.simulation_tolerances), - aggregate_tolerances=_tolerances(arguments.aggregate_tolerances), + tolerances=ParityTolerances( + simulation=_tolerances(arguments.simulation_tolerances), + aggregate=_tolerances(arguments.aggregate_tolerances), + ), ) print(receipt.model_dump_json()) # noqa: T201 # command-line script: the receipt JSON on stdout is its output diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index ca02627f3..d8f6b14f1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -4,6 +4,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from hashlib import sha256 from typing import Any, Literal @@ -19,6 +20,7 @@ from policyengine_simulation_contract.stage12_execution import ( ArtifactMediaType, ArtifactReference, + BundleProvenance, PlannedSimulationExecutionInput, ReportExecutionInput, SimulationArtifactDescriptor, @@ -220,23 +222,59 @@ def _descriptor( ) -def qualify_report_parity( +@dataclass(frozen=True) +class QualificationRunners: + """Seams for the calculations a qualification runs; None selects the default. + + The defaults are the existing combined executor (on a local telemetry + runtime that sends nothing remotely), the Stage 12 single-simulation + calculation and the Stage 12 aggregate builder. + """ + + existing: ExistingRunner | None = None + single_simulation: SingleSimulationRunner | None = None + aggregate: AggregateBuilder | None = None + + +@dataclass(frozen=True) +class ParityTolerances: + """Reviewed numeric slack for the simulation and aggregate comparisons. + + Each mapping is keyed by field path; an absent path compares exactly. + """ + + simulation: Mapping[str, NumericalTolerance] | None = None + aggregate: Mapping[str, NumericalTolerance] | None = None + + +@dataclass(frozen=True) +class _QualifiedSide: + """One role's planned input, existing-executor frames and Stage 12 result.""" + + planned_input: PlannedSimulationExecutionInput + existing_frames: dict[str, pd.DataFrame] + calculation: SimulationCalculation + + def output_digests(self) -> tuple[str, str]: + """Return the SHA-256 of the existing and the Stage 12 output artifacts.""" + existing_payload, _ = serialize_simulation_frames(self.existing_frames) + v2_payload, _ = serialize_simulation_frames( + self.calculation.frames, + calculation_provenance=self.calculation.calculation_provenance, + uk_local_authority_metadata=self.calculation.uk_local_authority_metadata, + ) + return sha256(existing_payload).hexdigest(), sha256(v2_payload).hexdigest() + + +def _run_existing_executor( existing_request: Mapping[str, Any], - report_payload: object, - *, - existing_runner: ExistingRunner | None = None, - single_simulation_runner: SingleSimulationRunner | None = None, - aggregate_builder: AggregateBuilder | None = None, - simulation_tolerances: Mapping[str, NumericalTolerance] | None = None, - aggregate_tolerances: Mapping[str, NumericalTolerance] | None = None, -) -> Stage12ParityReceipt: - """Run both implementations and fail on every unexplained difference.""" - report = ReportExecutionInput.model_validate(report_payload) - _validate_matching_input(existing_request, report) - output_plan = resolve_report_output_plan(report) - baseline_input = plan_simulation_input(report.baseline, output_plan) - reform_input = plan_simulation_input(report.reform, output_plan) + existing_runner: ExistingRunner | None, +) -> dict[str, Any]: + """Run the existing executor with microdata emission on, and copy its result. + Without an injected runner this runs the real executor on a local runtime, + which is shut down even when the run fails. + """ local_runtime: ObservabilityRuntime | None = None if existing_runner is None: from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # default runner, bound only when none is injected @@ -250,23 +288,21 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: return run_simulation_impl(request, runtime=qualification_runtime) existing_runner = run_existing - if single_simulation_runner is None: - single_simulation_runner = calculate_simulation_frames - if aggregate_builder is None: - aggregate_builder = build_aggregate_report request_with_microdata = dict(existing_request) request_with_microdata["_emit_microdata"] = True try: - existing_result = dict(existing_runner(request_with_microdata)) + return dict(existing_runner(request_with_microdata)) finally: if local_runtime is not None: local_runtime.shutdown() - microdata = existing_result.pop("_microdata", None) - existing_baseline = _existing_frames(microdata, side="baseline") - existing_reform = _existing_frames(microdata, side="reform") - bundle = report.baseline.bundle + +def _validate_existing_provenance( + existing_result: Mapping[str, Any], + bundle: BundleProvenance, +) -> None: + """Require the existing result's model and data versions to be the bundle's.""" if existing_result.get("model_version") != bundle.country_package_version: error_code = "country_package_provenance_mismatch" raise ParityMismatchError( @@ -280,35 +316,40 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: "report.data_version", ) + +def _run_v2_simulations( + single_simulation_runner: SingleSimulationRunner, + baseline_input: PlannedSimulationExecutionInput, + reform_input: PlannedSimulationExecutionInput, +) -> tuple[SimulationCalculation, SimulationCalculation]: + """Run the Stage 12 baseline and reform calculations concurrently.""" with ThreadPoolExecutor(max_workers=2) as executor: baseline_future = executor.submit(single_simulation_runner, baseline_input) reform_future = executor.submit(single_simulation_runner, reform_input) baseline_calculation = _calculation(baseline_future.result()) reform_calculation = _calculation(reform_future.result()) + return baseline_calculation, reform_calculation - compare_simulation_frames( - existing_baseline, - baseline_calculation.frames, - tolerances=simulation_tolerances, - ) - compare_simulation_frames( - existing_reform, - reform_calculation.frames, - tolerances=simulation_tolerances, - ) - baseline_descriptor = _descriptor(baseline_input, baseline_calculation) - reform_descriptor = _descriptor(reform_input, reform_calculation) +def _v2_aggregate_result( + report: ReportExecutionInput, + aggregate_builder: AggregateBuilder, + baseline: _QualifiedSide, + reform: _QualifiedSide, +) -> Mapping[str, Any]: + """Build the Stage 12 aggregate report and return its ``result`` object.""" + baseline_descriptor = _descriptor(baseline.planned_input, baseline.calculation) + reform_descriptor = _descriptor(reform.planned_input, reform.calculation) v2_report = aggregate_builder( report=report, - baseline_frames=baseline_calculation.frames, - reform_frames=reform_calculation.frames, + baseline_frames=baseline.calculation.frames, + reform_frames=reform.calculation.frames, baseline_descriptor=baseline_descriptor, reform_descriptor=reform_descriptor, uk_local_authority_metadata=validate_uk_local_authority_metadata( report.baseline.geography.country, - baseline_calculation.uk_local_authority_metadata, - reform_calculation.uk_local_authority_metadata, + baseline.calculation.uk_local_authority_metadata, + reform.calculation.uk_local_authority_metadata, region_code=report.baseline.geography.region, ), ) @@ -316,37 +357,70 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: if not isinstance(v2_result, Mapping): error_code = "invalid_v2_aggregate" raise ParityMismatchError(error_code, "report.result") + return v2_result + + +def qualify_report_parity( + existing_request: Mapping[str, Any], + report_payload: object, + *, + runners: QualificationRunners | None = None, + tolerances: ParityTolerances | None = None, +) -> Stage12ParityReceipt: + """Run both implementations and fail on every unexplained difference.""" + runners = runners or QualificationRunners() + tolerances = tolerances or ParityTolerances() + report = ReportExecutionInput.model_validate(report_payload) + _validate_matching_input(existing_request, report) + output_plan = resolve_report_output_plan(report) + baseline_input = plan_simulation_input(report.baseline, output_plan) + reform_input = plan_simulation_input(report.reform, output_plan) + + existing_result = _run_existing_executor(existing_request, runners.existing) + microdata = existing_result.pop("_microdata", None) + existing_baseline = _existing_frames(microdata, side="baseline") + existing_reform = _existing_frames(microdata, side="reform") + bundle = report.baseline.bundle + _validate_existing_provenance(existing_result, bundle) + + baseline_calculation, reform_calculation = _run_v2_simulations( + runners.single_simulation or calculate_simulation_frames, + baseline_input, + reform_input, + ) + baseline = _QualifiedSide(baseline_input, existing_baseline, baseline_calculation) + reform = _QualifiedSide(reform_input, existing_reform, reform_calculation) + for side in (baseline, reform): + compare_simulation_frames( + side.existing_frames, + side.calculation.frames, + tolerances=tolerances.simulation, + ) + + v2_result = _v2_aggregate_result( + report, runners.aggregate or build_aggregate_report, baseline, reform + ) compare_aggregate_reports( existing_result, v2_result, - tolerances=aggregate_tolerances, + tolerances=tolerances.aggregate, ) - existing_baseline_payload, _ = serialize_simulation_frames(existing_baseline) - existing_reform_payload, _ = serialize_simulation_frames(existing_reform) - v2_baseline_payload, _ = serialize_simulation_frames( - baseline_calculation.frames, - calculation_provenance=baseline_calculation.calculation_provenance, - uk_local_authority_metadata=baseline_calculation.uk_local_authority_metadata, - ) - v2_reform_payload, _ = serialize_simulation_frames( - reform_calculation.frames, - calculation_provenance=reform_calculation.calculation_provenance, - uk_local_authority_metadata=reform_calculation.uk_local_authority_metadata, - ) + existing_baseline_sha256, v2_baseline_sha256 = baseline.output_digests() + existing_reform_sha256, v2_reform_sha256 = reform.output_digests() return Stage12ParityReceipt( evaluation_id=str(report.evaluation_id), bundle_manifest_sha256=bundle.bundle_manifest_sha256, baseline_input_sha256=simulation_input_sha256(baseline_input), reform_input_sha256=simulation_input_sha256(reform_input), - existing_baseline_output_sha256=sha256(existing_baseline_payload).hexdigest(), - v2_baseline_output_sha256=sha256(v2_baseline_payload).hexdigest(), - existing_reform_output_sha256=sha256(existing_reform_payload).hexdigest(), - v2_reform_output_sha256=sha256(v2_reform_payload).hexdigest(), + existing_baseline_output_sha256=existing_baseline_sha256, + v2_baseline_output_sha256=v2_baseline_sha256, + existing_reform_output_sha256=existing_reform_sha256, + v2_reform_output_sha256=v2_reform_sha256, existing_aggregate_output_sha256=sha256( canonical_json_bytes(existing_result) ).hexdigest(), v2_aggregate_output_sha256=sha256(canonical_json_bytes(v2_result)).hexdigest(), - simulation_tolerance_fields=tuple(sorted(simulation_tolerances or {})), - aggregate_tolerance_fields=tuple(sorted(aggregate_tolerances or {})), + simulation_tolerance_fields=tuple(sorted(tolerances.simulation or {})), + aggregate_tolerance_fields=tuple(sorted(tolerances.aggregate or {})), ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index da4a5e25a..5504faeb1 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -13,6 +13,8 @@ ) from policyengine_simulation_executor.stage12_parity import ParityMismatchError from policyengine_simulation_executor.stage12_qualification import ( + ParityTolerances, + QualificationRunners, qualify_report_parity, ) from policyengine_simulation_executor.stage12_runtime import SimulationCalculation @@ -102,9 +104,11 @@ def aggregate_builder(**arguments): receipt = qualify_report_parity( _existing_request(), report, - existing_runner=existing_runner, - single_simulation_runner=single_runner, - aggregate_builder=aggregate_builder, + runners=QualificationRunners( + existing=existing_runner, + single_simulation=single_runner, + aggregate=aggregate_builder, + ), ) assert {call.role for call in simulation_calls} == { @@ -149,9 +153,11 @@ def single_runner(simulation): qualify_report_parity( _existing_request(), report, - existing_runner=existing_runner, - single_simulation_runner=single_runner, - aggregate_builder=lambda **_: {"result": {}}, + runners=QualificationRunners( + existing=existing_runner, + single_simulation=single_runner, + aggregate=lambda **_: {"result": {}}, + ), ) assert error.value.code == "numeric_value_mismatch" @@ -195,7 +201,7 @@ def existing_runner(_request): qualify_report_parity( _existing_request(), report, - existing_runner=existing_runner, + runners=QualificationRunners(existing=existing_runner), ) assert error.value.code == "country_package_provenance_mismatch" @@ -238,20 +244,24 @@ def test_qualification_receipt_records_every_digest_and_tolerance_field() -> Non receipt = qualify_report_parity( _existing_request(), report, - existing_runner=_matching_existing_runner(report, baseline, reform), - single_simulation_runner=_matching_single_runner(baseline, reform), - aggregate_builder=lambda **_: { - "result": { - "model_version": report.baseline.bundle.country_package_version, - "data_version": report.baseline.bundle.dataset.artifact_revision, - "budget": {"budgetary_impact": 20.0}, - } - }, - simulation_tolerances={ - "person.person_household_id": NumericalTolerance(), - "household.household_net_income": NumericalTolerance(absolute=0.1), - }, - aggregate_tolerances={"budget.budgetary_impact": NumericalTolerance()}, + runners=QualificationRunners( + existing=_matching_existing_runner(report, baseline, reform), + single_simulation=_matching_single_runner(baseline, reform), + aggregate=lambda **_: { + "result": { + "model_version": report.baseline.bundle.country_package_version, + "data_version": report.baseline.bundle.dataset.artifact_revision, + "budget": {"budgetary_impact": 20.0}, + } + }, + ), + tolerances=ParityTolerances( + simulation={ + "person.person_household_id": NumericalTolerance(), + "household.household_net_income": NumericalTolerance(absolute=0.1), + }, + aggregate={"budget.budgetary_impact": NumericalTolerance()}, + ), ) dumped = receipt.model_dump(mode="json") @@ -316,10 +326,12 @@ def test_existing_result_is_checked_before_v2_simulations_run( qualify_report_parity( _existing_request(), report, - existing_runner=_matching_existing_runner( - report, baseline, reform, **overrides + runners=QualificationRunners( + existing=_matching_existing_runner( + report, baseline, reform, **overrides + ), + single_simulation=_matching_single_runner(baseline, reform, calls), ), - single_simulation_runner=_matching_single_runner(baseline, reform, calls), ) assert (error.value.code, error.value.path) == (code, path) @@ -336,9 +348,11 @@ def test_v2_aggregate_without_a_result_object_is_rejected(v2_report) -> None: qualify_report_parity( _existing_request(), report, - existing_runner=_matching_existing_runner(report, baseline, reform), - single_simulation_runner=_matching_single_runner(baseline, reform), - aggregate_builder=lambda **_: v2_report, + runners=QualificationRunners( + existing=_matching_existing_runner(report, baseline, reform), + single_simulation=_matching_single_runner(baseline, reform), + aggregate=lambda **_: v2_report, + ), ) assert (error.value.code, error.value.path) == ( @@ -387,8 +401,10 @@ def run_simulation_impl(request, *, runtime): qualify_report_parity( _existing_request(), report, - single_simulation_runner=_matching_single_runner(baseline, reform), - aggregate_builder=lambda **_: {"result": []}, + runners=QualificationRunners( + single_simulation=_matching_single_runner(baseline, reform), + aggregate=lambda **_: {"result": []}, + ), ) assert events == [ From 844b5a7b6a071be4ff7b0e9f159cd713188501a9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:37:45 +0400 Subject: [PATCH 32/66] Pin Stage 12 coordination validation, claims, failures, stages and aggregate output Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_coordination.py | 505 ++++++++++++++++++ .../tests/test_stage12_runtime.py | 81 +++ 2 files changed, 586 insertions(+) create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_coordination.py diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py new file mode 100644 index 000000000..ef41b323b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -0,0 +1,505 @@ +"""Characterization tests for Stage 12 report coordination state and order.""" + +from __future__ import annotations + +import re +from contextlib import contextmanager, nullcontext +from datetime import UTC, datetime, timedelta +from uuid import UUID + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, + ResultComparisonStatus, +) +from policyengine_simulation_executor.stage12_runtime import coordination +from test_stage12_runtime import ( + BASELINE_ID, + EVALUATION_ID, + NOW, + REFORM_ID, + ConcurrentInvoker, + FakeArtifacts, + FakeStore, + _aggregate_payload, + _child, + _context, + _output_plan, + _parent, + _report, + coordinate_report, +) + +CLAIMED_AT = datetime(2026, 9, 15, 1, 0, tzinfo=UTC) +PLACEHOLDER_ID = UUID("00000000-0000-4000-8000-0000000000ff") +# The redacted causes the report and its children record for these failures. +DISPATCH_FAILURE_SUMMARY = ( + "ConnectionError in stage12_child_dispatch (baseline): " + "dispatch rejected with sensitive detail" +) +RETAINED_DIGEST_FAILURE_SUMMARY = ( + "ValueError after stage12_child_state_create: " + "retained child artifact digest mismatch" +) + + +class FixedClock(datetime): + @classmethod + def now(cls, tz=None): + assert tz is UTC + return CLAIMED_AT + + +@pytest.fixture +def fixed_clock(monkeypatch): + monkeypatch.setattr(coordination, "datetime", FixedClock) + monkeypatch.setattr(coordination, "uuid4", lambda: PLACEHOLDER_ID) + + +class TracingRuntime: + def __init__(self) -> None: + self.events: list[tuple[str, str, dict[str, object]]] = [] + + @contextmanager + def operation(self, name, *, attributes=None, **_kwargs): + self.events.append(("operation", name, dict(attributes or {}))) + yield + + def span(self, name, *, attributes=None): + self.events.append(("span", name, dict(attributes or {}))) + return nullcontext() + + def capture_context(self): + self.events.append(("capture_context", "", {})) + return {"traceparent": "trace"} + + +class FailingSpawnInvoker: + def __init__(self, store=None, succeeded_child=None): + self.store = store + self.succeeded_child = succeeded_child + self.spawns = [] + + def spawn(self, **call): + self.spawns.append(call) + if self.succeeded_child is not None: + self.store.children[self.succeeded_child.simulation_execution_id] = ( + self.succeeded_child + ) + msg = "dispatch rejected with sensitive detail" + raise ConnectionError(msg) + + def restore(self, invocation_id): + msg = f"unexpected restore {invocation_id}" + raise AssertionError(msg) + + +def _coordinate(context=None, parent=None, **dependencies): + context = context or _context() + dependencies.setdefault("application_name", context.modal_application) + dependencies.setdefault("coordinator_invocation_id", "coordinator-1") + dependencies.setdefault("aggregator", lambda **_: _aggregate_payload()) + dependencies.setdefault("artifacts", FakeArtifacts()) + dependencies.setdefault("invoker", ConcurrentInvoker(dependencies["artifacts"])) + return coordinate_report( + _report().model_dump(mode="json"), + context.model_dump(mode="json"), + (parent or _parent()).model_dump(mode="json"), + **dependencies, + ) + + +def _succeeded_child(simulation, artifacts): + descriptor = artifacts.add_simulation(simulation) + return _child(simulation).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": descriptor.artifact.uri, + "output_sha256": descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": descriptor.row_identity.identifier_columns, + "row_count": descriptor.row_identity.row_count, + "row_identity_sha256": descriptor.row_identity.identity_sha256, + "completed_at": NOW, + } + ) + + +@pytest.mark.parametrize( + ("context_update", "application_name", "message"), + [ + ( + {}, + "another-application", + "report coordinator application differs from dispatch context", + ), + ( + {"worker_version": "9.9.9"}, + None, + "report worker version differs from dispatch context", + ), + ( + {"bundle_manifest_sha256": "c" * 64}, + None, + "report bundle digest differs from dispatch context", + ), + ], +) +def test_dispatch_context_must_match_the_report( + context_update, application_name, message +): + store = FakeStore() + context = _context().model_copy(update=context_update) + + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + _coordinate( + context, + store=store, + application_name=application_name or context.modal_application, + ) + + assert store.parent is None + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("evaluation_id", UUID("00000000-0000-0000-0000-0000000000aa")), + ("environment", "production"), + ("originating_request_id", "request-2"), + ("worker_version", "5.3.0"), + ("modal_application", "another-application"), + ("version_manifest_sha256", "c" * 64), + ("policyengine_version", "5.3.0"), + ("country_package_name", "policyengine-uk"), + ("country_package_version", "9.9.9"), + ("country", "uk"), + ("dataset_identity", "another-dataset"), + ("dataset_uri", "hf://policyengine/data/another.h5@revision"), + ("data_package_name", "another-data"), + ("data_package_version", "9.9.9"), + ("data_artifact_revision", "another-revision"), + ("created_at", NOW + timedelta(seconds=1)), + ("retention_expires_at", NOW + timedelta(days=29)), + ("status", ComparisonRunLifecycleStatus.RUNNING), + ("aggregation_status", ComparisonRunAggregationStatus.RUNNING), + ("coordinator_invocation_id", "coordinator-0"), + ("production_identity", "direct:another"), + ("incumbent_execution_id", "production-job-1"), + ("comparison_status", ResultComparisonStatus.PENDING), + ], +) +def test_submitted_direct_parent_must_describe_the_invocation(field, value): + store = FakeStore() + messages = { + "status": "submitted parent must be pending", + "aggregation_status": "submitted parent aggregation must not be started", + "coordinator_invocation_id": ( + "submitted parent must not contain an invocation identifier" + ), + "production_identity": "direct parent identity does not match its report", + "incumbent_execution_id": ( + "direct parent must not name a production execution" + ), + "comparison_status": "direct parent must not request result comparison", + } + + message = messages.get(field, f"submitted parent {field} does not match invocation") + + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + _coordinate(parent=_parent().model_copy(update={field: value}), store=store) + + assert store.parent is None + + +@pytest.mark.parametrize( + "update", + [ + {"production_identity": "production-job-2"}, + {"incumbent_execution_id": "production-job-2"}, + {"incumbent_execution_id": None}, + ], +) +def test_submitted_automatic_parent_must_name_the_production_job(update): + store = FakeStore() + context = _context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + parent = _parent().model_copy( + update={ + "production_identity": "production-job-1", + "incumbent_execution_id": "production-job-1", + "comparison_status": ResultComparisonStatus.PENDING, + **update, + } + ) + + with pytest.raises( + ValueError, + match=r"^automatic parent does not match the production execution$", + ): + _coordinate(context, parent, store=store) + + assert store.parent is None + + +def test_existing_parent_with_another_identifier_is_rejected_unwrapped(): + store = FakeStore() + store.parent = _parent().model_copy( + update={"evaluation_id": UUID("00000000-0000-0000-0000-0000000000aa")} + ) + + with pytest.raises( + ValueError, + match=r"^existing comparison identity has a different report identifier$", + ): + _coordinate(store=store) + + +@pytest.mark.usefixtures("fixed_clock") +@pytest.mark.parametrize("previous_start", [None, NOW + timedelta(minutes=5)]) +def test_failed_parent_is_reclaimed_keeping_its_first_start(previous_start): + store = FakeStore() + artifacts = FakeArtifacts() + store.parent = _parent().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "aggregation_status": ComparisonRunAggregationStatus.FAILED, + "coordinator_invocation_id": "coordinator-0", + "error_code": "report_coordination_failed", + "error_summary": "TimeoutError", + "started_at": previous_start, + "completed_at": NOW + timedelta(minutes=10), + } + ) + replaced = [] + replace_report = store.replace_report + + def record_replace(record): + replaced.append(record) + return replace_report(record) + + store.replace_report = record_replace + + _coordinate(store=store, artifacts=artifacts, invoker=ConcurrentInvoker(artifacts)) + + claimed = replaced[0] + assert claimed.status is ComparisonRunLifecycleStatus.RUNNING + assert claimed.aggregation_status is ComparisonRunAggregationStatus.RUNNING + assert claimed.coordinator_invocation_id == "coordinator-1" + assert claimed.error_code is None + assert claimed.error_summary is None + assert claimed.started_at == (previous_start or CLAIMED_AT) + assert claimed.updated_at == CLAIMED_AT + assert claimed.completed_at is None + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +@pytest.mark.parametrize( + "status", + [ + ComparisonRunLifecycleStatus.RUNNING, + ComparisonRunLifecycleStatus.SUCCEEDED, + ComparisonRunLifecycleStatus.SKIPPED, + ], +) +def test_claimed_or_finished_parent_is_deduplicated(status): + store = FakeStore() + existing = _parent().model_copy(update={"status": status, "error_code": "skipped"}) + if status is ComparisonRunLifecycleStatus.SUCCEEDED: + existing = existing.model_copy( + update={ + "aggregation_status": ComparisonRunAggregationStatus.SUCCEEDED, + "aggregate_output_uri": "gs://private/aggregate.json", + "aggregate_output_sha256": "d" * 64, + "aggregate_schema_version": 1, + "completed_at": NOW, + } + ) + store.parent = existing + invoker = ConcurrentInvoker(FakeArtifacts()) + + result = _coordinate(store=store, invoker=invoker) + + assert result == { + "deduplicated": True, + "evaluation_id": str(EVALUATION_ID), + "status": status.value, + } + assert invoker.events == [] + assert store.parent is existing + + +def test_retained_child_with_a_changed_artifact_fails_the_report(): + report = _report() + store = FakeStore() + artifacts = FakeArtifacts() + succeeded = _succeeded_child(report.baseline, artifacts) + store.children[BASELINE_ID] = succeeded + artifacts.payloads[succeeded.output_uri] = b"changed" + invoker = ConcurrentInvoker(artifacts) + + with pytest.raises(RuntimeError, match=r"^Stage 12 report coordination failed$"): + _coordinate(store=store, artifacts=artifacts, invoker=invoker) + + assert invoker.events == [] + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "report_coordination_failed" + assert store.parent.error_summary == RETAINED_DIGEST_FAILURE_SUMMARY + assert store.children[BASELINE_ID] == succeeded + + +@pytest.mark.usefixtures("fixed_clock") +def test_dispatch_failure_marks_the_child_and_report_failed(): + store = FakeStore() + invoker = FailingSpawnInvoker() + + with pytest.raises(RuntimeError) as error: + _coordinate(store=store, artifacts=FakeArtifacts(), invoker=invoker) + + assert str(error.value) == "Stage 12 report coordination failed" + assert error.value.__cause__ is None + assert error.value.__suppress_context__ is True + assert len(invoker.spawns) == 1 + failed_child = store.children[BASELINE_ID] + assert failed_child.status is ComparisonRunLifecycleStatus.FAILED + assert failed_child.modal_invocation_id == f"dispatch-pending-{PLACEHOLDER_ID}" + assert failed_child.error_code == "simulation_dispatch_failed" + assert failed_child.error_summary == DISPATCH_FAILURE_SUMMARY + assert failed_child.started_at == CLAIMED_AT + assert failed_child.completed_at == CLAIMED_AT + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.aggregation_status is ComparisonRunAggregationStatus.FAILED + assert store.parent.error_summary == DISPATCH_FAILURE_SUMMARY + assert store.parent.completed_at == CLAIMED_AT + + +def test_dispatch_failure_keeps_a_child_that_succeeded_meanwhile(): + report = _report() + store = FakeStore() + artifacts = FakeArtifacts() + succeeded = _succeeded_child(report.baseline, artifacts) + invoker = FailingSpawnInvoker(store, succeeded) + + with pytest.raises(RuntimeError, match=r"^Stage 12 report coordination failed$"): + _coordinate(store=store, artifacts=artifacts, invoker=invoker) + + assert store.children[BASELINE_ID] == succeeded + assert store.parent.error_summary == DISPATCH_FAILURE_SUMMARY + + +def test_default_output_plan_resolver_measures_its_own_stages(monkeypatch): + from test_stage12_runtime import coordinate_report_impl + + store = FakeStore() + artifacts = FakeArtifacts() + runtime = TracingRuntime() + resolved = [] + + def resolve(report, *, stage_scope): + with stage_scope(coordination.Stage.STAGE12_OUTPUT_PLANNING, {"probe": 1}): + resolved.append(report.evaluation_id) + return _output_plan() + + monkeypatch.setattr(coordination, "resolve_report_output_plan", resolve) + + coordinate_report_impl( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: _aggregate_payload(), + runtime=runtime, + ) + + assert resolved == [EVALUATION_ID] + assert ("operation", "stage12_output_planning", {"probe": 1}) in runtime.events + + +@pytest.mark.usefixtures("fixed_clock") +def test_successful_coordination_records_every_stage_and_state(): + store = FakeStore() + artifacts = FakeArtifacts() + runtime = TracingRuntime() + aggregated = [] + + def aggregate(**values): + aggregated.append(sorted(values)) + return _aggregate_payload({"budget": 1}) + + result = _coordinate( + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=aggregate, + runtime=runtime, + ) + + assert aggregated == [ + [ + "baseline_descriptor", + "baseline_frames", + "reform_descriptor", + "reform_frames", + "report", + "uk_local_authority_metadata", + ] + ] + assert runtime.events == [ + ("span", "stage12_coordinator_claim", {}), + ("operation", "stage12_coordinator_preparation", {}), + ("operation", "stage12_output_planning", {"country": "us"}), + ("operation", "stage12_child_input_planning", {"simulation_role": "baseline"}), + ("operation", "stage12_child_input_planning", {"simulation_role": "reform"}), + ("span", "stage12_child_state_create", {}), + ("span", "stage12_child_dispatch", {"simulation_role": "baseline"}), + ("capture_context", "", {}), + ("span", "stage12_child_dispatch", {"simulation_role": "reform"}), + ("capture_context", "", {}), + ("span", "stage12_child_wait", {"simulation_role": "baseline"}), + ("span", "stage12_child_wait", {"simulation_role": "reform"}), + ("span", "stage12_artifact_read", {}), + ("span", "stage12_aggregation", {}), + ("span", "stage12_aggregate_artifact_write", {}), + ] + aggregate_artifact = result["artifact"] + assert result == { + "contract_version": 1, + "evaluation_id": str(EVALUATION_ID), + "aggregate_schema_version": 1, + "artifact": aggregate_artifact, + "baseline_artifact_sha256": result["baseline_artifact_sha256"], + "reform_artifact_sha256": result["reform_artifact_sha256"], + "bundle": _report().baseline.bundle.model_dump(mode="json"), + } + aggregate_uri = f"gs://private/{_context().artifact_prefix}/reports/aggregate.json" + assert aggregate_artifact["uri"] == aggregate_uri + parent = store.parent.model_dump(mode="json") + expected_parent = _parent().model_dump(mode="json") + expected_parent.update( + { + "status": "succeeded", + "aggregation_status": "succeeded", + "coordinator_invocation_id": "coordinator-1", + "aggregate_output_uri": aggregate_uri, + "aggregate_output_sha256": aggregate_artifact["content_sha256"], + "aggregate_schema_version": 1, + "started_at": "2026-09-15T01:00:00Z", + "updated_at": "2026-09-15T01:00:00Z", + "completed_at": "2026-09-15T01:00:00Z", + } + ) + assert parent == expected_parent + for role, simulation_id in (("baseline", BASELINE_ID), ("reform", REFORM_ID)): + child = store.children[simulation_id] + assert child.status is ComparisonRunLifecycleStatus.RUNNING + assert child.modal_invocation_id == f"call-{role}" + assert child.started_at == CLAIMED_AT + assert child.updated_at == CLAIMED_AT + assert child.simulation_callable == "run_single_simulation_us" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 7da9ac18f..4d5d60308 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1157,6 +1157,87 @@ def serialize(self): assert json.loads(stored_aggregate)["result"]["spm_provenance"] is None +def test_aggregate_report_wraps_builder_output_with_report_provenance( + monkeypatch, +) -> None: + from policyengine_simulation_executor import simulation_output_builder + + report = _report() + report = report.model_copy( + update={"reform": report.reform.model_copy(update={"policy": {}})} + ) + entities = ("person", "marital_unit", "family", "spm_unit", "tax_unit") + + def frames(income): + return { + **{entity: pd.DataFrame({f"{entity}_id": [1]}) for entity in entities}, + "household": pd.DataFrame( + {"household_id": [1], "household_net_income": [income]} + ), + } + + observed = {} + + class CapturingBuilder: + def __init__(self, **values): + observed.update(values) + + def serialize(self): + return {"budget": {"budgetary_impact": 1.0}} + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder + ) + artifacts = FakeArtifacts() + + result = build_aggregate_report( + report=report, + baseline_frames=frames(100.0), + reform_frames=frames(120.0), + baseline_descriptor=artifacts.add_simulation(report.baseline), + reform_descriptor=artifacts.add_simulation(report.reform), + uk_local_authority_metadata=None, + ) + + assert result == { + "contract_version": 1, + "aggregate_schema_version": 1, + "evaluation_id": str(EVALUATION_ID), + "requested_aggregates": [item.value for item in report.requested_aggregates], + "bundle": report.baseline.bundle.model_dump(mode="json"), + "result": {"budget": {"budgetary_impact": 1.0}}, + } + assert list(observed) == [ + "country", + "simulation_params", + "country_module", + "dataset", + "baseline", + "reform", + "resolved_data_version", + "resolved_region_code", + "uk_local_authority_metadata", + ] + assert observed["country"] == "us" + assert observed["simulation_params"] == { + "country": "us", + "scope": "macro", + "data": report.baseline.bundle.dataset.uri, + "data_version": report.baseline.bundle.dataset.artifact_revision, + "time_period": str(report.baseline.year), + "region": report.baseline.geography.region, + } + assert observed["dataset"] is observed["baseline"].output_dataset + for role, income in (("baseline", 100.0), ("reform", 120.0)): + household = observed[role].output_dataset.data.household + assert household["household_net_income"].tolist() == [income] + assert observed["resolved_data_version"] == ( + report.baseline.bundle.dataset.artifact_revision + ) + assert observed["resolved_region_code"] == report.baseline.geography.region + assert observed["uk_local_authority_metadata"] is None + + def test_single_worker_rejects_combined_baseline_and_reform_input() -> None: payload = _planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") payload["baseline"] = {} From f4fa6fd9b92634260f3d3f165406d54e96ba4518 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:38:36 +0400 Subject: [PATCH 33/66] Pass each role's retained frames and descriptor to the aggregate builder together Co-Authored-By: Claude Opus 5.5 --- .../stage12_qualification.py | 9 ++-- .../stage12_runtime/__init__.py | 7 ++- .../stage12_runtime/aggregation.py | 23 ++++++---- .../stage12_runtime/coordination.py | 7 ++- .../tests/test_stage12_coordination.py | 9 +--- .../tests/test_stage12_qualification.py | 8 ++-- .../tests/test_stage12_runtime.py | 45 +++++++++++-------- 7 files changed, 61 insertions(+), 47 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index d8f6b14f1..dac0625e5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -42,6 +42,7 @@ compare_simulation_frames, ) from policyengine_simulation_executor.stage12_runtime import ( + RetainedSimulationOutput, SimulationCalculation, build_aggregate_report, calculate_simulation_frames, @@ -342,10 +343,10 @@ def _v2_aggregate_result( reform_descriptor = _descriptor(reform.planned_input, reform.calculation) v2_report = aggregate_builder( report=report, - baseline_frames=baseline.calculation.frames, - reform_frames=reform.calculation.frames, - baseline_descriptor=baseline_descriptor, - reform_descriptor=reform_descriptor, + baseline=RetainedSimulationOutput( + baseline.calculation.frames, baseline_descriptor + ), + reform=RetainedSimulationOutput(reform.calculation.frames, reform_descriptor), uk_local_authority_metadata=validate_uk_local_authority_metadata( report.baseline.geography.country, baseline.calculation.uk_local_authority_metadata, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index dd61bce17..c3dec4525 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -1,6 +1,10 @@ """Independent Stage 12 simulation execution and report coordination.""" -from .aggregation import build_aggregate_report, build_spm_result +from .aggregation import ( + RetainedSimulationOutput, + build_aggregate_report, + build_spm_result, +) from .coordination import coordinate_report from .simulation import ( SimulationCalculation, @@ -14,6 +18,7 @@ _build_spm_result = build_spm_result __all__ = [ + "RetainedSimulationOutput", "SimulationCalculation", "build_aggregate_report", "calculate_simulation_frames", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index e6e1b7281..5bcedc6ab 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from policyengine_simulation_contract.spm import build_spm_comparison_provenance @@ -143,13 +144,19 @@ def _dataset_from_frames( ) +@dataclass(frozen=True) +class RetainedSimulationOutput: + """One role's retained output frames and the artifact descriptor they came from.""" + + frames: Mapping[str, pd.DataFrame] + descriptor: SimulationArtifactDescriptor + + def build_aggregate_report( *, report: ReportExecutionInput, - baseline_frames: Mapping[str, pd.DataFrame], - reform_frames: Mapping[str, pd.DataFrame], - baseline_descriptor: SimulationArtifactDescriptor, - reform_descriptor: SimulationArtifactDescriptor, + baseline: RetainedSimulationOutput, + reform: RetainedSimulationOutput, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> dict[str, Any]: """Build the aggregate report payload from retained baseline and reform frames. @@ -172,12 +179,12 @@ def build_aggregate_report( country_module = _country_module(country) datasets = { "baseline": _dataset_from_frames( - baseline_frames, + baseline.frames, country=country, year=report.baseline.year, ), "reform": _dataset_from_frames( - reform_frames, + reform.frames, country=country, year=report.reform.year, ), @@ -222,8 +229,8 @@ def stand_in( output.update( build_spm_result( report=report, - baseline_descriptor=baseline_descriptor, - reform_descriptor=reform_descriptor, + baseline_descriptor=baseline.descriptor, + reform_descriptor=reform.descriptor, ) ) return { diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index d8b5f2ac8..627122ac0 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -36,6 +36,7 @@ ) from .aggregation import ( + RetainedSimulationOutput, build_aggregate_report, validate_aligned_outputs, validate_uk_local_authority_metadata, @@ -519,10 +520,8 @@ def coordinate_report( with stage_span(Stage.STAGE12_AGGREGATION): aggregate = aggregator( report=report, - baseline_frames=baseline_frames, - reform_frames=reform_frames, - baseline_descriptor=baseline, - reform_descriptor=reform, + baseline=RetainedSimulationOutput(baseline_frames, baseline), + reform=RetainedSimulationOutput(reform_frames, reform), uk_local_authority_metadata=uk_local_authority_metadata, ) with stage_span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index ef41b323b..e0373b6bf 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -442,14 +442,7 @@ def aggregate(**values): ) assert aggregated == [ - [ - "baseline_descriptor", - "baseline_frames", - "reform_descriptor", - "reform_frames", - "report", - "uk_local_authority_metadata", - ] + ["baseline", "reform", "report", "uk_local_authority_metadata"] ] assert runtime.events == [ ("span", "stage12_coordinator_claim", {}), diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index 5504faeb1..dd1f164f9 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -88,10 +88,10 @@ def single_runner(simulation): def aggregate_builder(**arguments): assert arguments["report"] == report - assert arguments["baseline_frames"] is baseline - assert arguments["reform_frames"] is reform - assert arguments["baseline_descriptor"].bundle == report.baseline.bundle - assert arguments["reform_descriptor"].bundle == report.reform.bundle + assert arguments["baseline"].frames is baseline + assert arguments["reform"].frames is reform + assert arguments["baseline"].descriptor.bundle == report.baseline.bundle + assert arguments["reform"].descriptor.bundle == report.reform.bundle return { "result": { "model_version": report.baseline.bundle.country_package_version, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 4d5d60308..16f29af73 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -74,6 +74,7 @@ serialize_simulation_frames, ) from policyengine_simulation_executor.stage12_runtime import ( + RetainedSimulationOutput, SimulationCalculation, _build_spm_result, build_aggregate_report, @@ -1005,10 +1006,12 @@ def serialize(self): result = build_aggregate_report( report=report, - baseline_frames=frames, - reform_frames=frames, - baseline_descriptor=artifacts.add_simulation(report.baseline), - reform_descriptor=artifacts.add_simulation(report.reform), + baseline=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.reform) + ), ) assert result["result"] == {"captured": True} @@ -1065,10 +1068,12 @@ def serialize(self): build_aggregate_report( report=report, - baseline_frames=frames, - reform_frames=frames, - baseline_descriptor=artifacts.add_simulation(report.baseline), - reform_descriptor=artifacts.add_simulation(report.reform), + baseline=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.reform) + ), ) for role, stop in (("baseline", "2026-12-31"), ("reform", "2027-12-31")): @@ -1141,10 +1146,12 @@ def serialize(self): aggregate = build_aggregate_report( report=report, - baseline_frames=uk_frames, - reform_frames=uk_frames, - baseline_descriptor=artifacts.add_simulation(report.baseline), - reform_descriptor=artifacts.add_simulation(report.reform), + baseline=RetainedSimulationOutput( + uk_frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + uk_frames, artifacts.add_simulation(report.reform) + ), ) assert aggregate["result"]["spm_provenance"] is None @@ -1192,10 +1199,12 @@ def serialize(self): result = build_aggregate_report( report=report, - baseline_frames=frames(100.0), - reform_frames=frames(120.0), - baseline_descriptor=artifacts.add_simulation(report.baseline), - reform_descriptor=artifacts.add_simulation(report.reform), + baseline=RetainedSimulationOutput( + frames(100.0), artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames(120.0), artifacts.add_simulation(report.reform) + ), uk_local_authority_metadata=None, ) @@ -1423,10 +1432,10 @@ def resolver(report): assert result["artifact"]["uri"] == ( f"gs://private/{context.artifact_prefix}/reports/aggregate.json" ) - assert observed["baseline_frames"]["household"][ + assert observed["baseline"].frames["household"][ "household_net_income" ].tolist() == [100.0] - assert observed["reform_frames"]["household"]["household_net_income"].tolist() == [ + assert observed["reform"].frames["household"]["household_net_income"].tolist() == [ 120.0 ] assert observed["uk_local_authority_metadata"] is None From 632b8e0b5317e25831ed7f2d5198e18f3a621758 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:41:19 +0400 Subject: [PATCH 34/66] Split Stage 12 report coordination into claim, plan, dispatch, wait and aggregate phases Group the coordinator's Modal identity and its injectable seams, and split submitted-parent validation into invocation, claim-state and production checks. The Modal coordinate_report function keeps its signature. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 12 +- .../stage12_runtime/__init__.py | 8 +- .../stage12_runtime/coordination.py | 641 ++++++++++++------ .../tests/test_stage12_coordination.py | 24 +- .../tests/test_stage12_runtime.py | 33 +- 5 files changed, 476 insertions(+), 242 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 5ff581fb8..080cffa8c 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -306,6 +306,10 @@ def coordinate_report( result is compared with that production job's. A report another invocation already claimed returns a ``deduplicated`` marker instead. """ + from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + CoordinatorDependencies, + CoordinatorInvocation, + ) from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time coordinate_report as run_report_coordinator, ) @@ -343,9 +347,11 @@ def coordinate_report( payload, context, parent, - application_name=APP_NAME, - coordinator_invocation_id=coordinator_invocation_id, - runtime=runtime, + invocation=CoordinatorInvocation( + application_name=APP_NAME, + invocation_id=coordinator_invocation_id, + ), + dependencies=CoordinatorDependencies(runtime=runtime), ) finally: runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index c3dec4525..d8c19fad4 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -5,7 +5,11 @@ build_aggregate_report, build_spm_result, ) -from .coordination import coordinate_report +from .coordination import ( + CoordinatorDependencies, + CoordinatorInvocation, + coordinate_report, +) from .simulation import ( SimulationCalculation, calculate_simulation_frames, @@ -18,6 +22,8 @@ _build_spm_result = build_spm_result __all__ = [ + "CoordinatorDependencies", + "CoordinatorInvocation", "RetainedSimulationOutput", "SimulationCalculation", "build_aggregate_report", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 627122ac0..2a4745369 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass from datetime import UTC, datetime from functools import partial from hashlib import sha256 @@ -18,6 +19,7 @@ ReportExecutionInput, ResultComparisonStatus, SimulationArtifactDescriptor, + SimulationRole, Stage12InvocationContext, Stage12OutputPlan, ) @@ -65,9 +67,11 @@ from .simulation import descriptor_from_record, simulation_input_sha256 if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping + from contextlib import AbstractContextManager from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 @@ -94,6 +98,23 @@ def _child_record( ) +def _validate_dispatch_context( + report: ReportExecutionInput, + context: Stage12InvocationContext, + application_name: str, +) -> None: + """Reject a dispatch context naming another application, worker or bundle.""" + if application_name != context.modal_application: + msg = "report coordinator application differs from dispatch context" + raise ValueError(msg) + if report.baseline.bundle.policyengine_version != context.worker_version: + msg = "report worker version differs from dispatch context" + raise ValueError(msg) + if report.baseline.bundle.bundle_manifest_sha256 != context.bundle_manifest_sha256: + msg = "report bundle digest differs from dispatch context" + raise ValueError(msg) + + def _validate_submitted_parent( *, report: ReportExecutionInput, @@ -101,6 +122,18 @@ def _validate_submitted_parent( parent: ComparisonReportRecord, ) -> None: """Reject dispatch metadata that does not describe this exact report.""" + _validate_parent_matches_invocation(report=report, context=context, parent=parent) + _validate_parent_is_unclaimed(parent) + _validate_parent_production_link(report=report, context=context, parent=parent) + + +def _validate_parent_matches_invocation( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + parent: ComparisonReportRecord, +) -> None: + """Require every parent field the dispatch fixes to match the report and context.""" bundle = report.baseline.bundle expected = { "evaluation_id": report.evaluation_id, @@ -125,6 +158,10 @@ def _validate_submitted_parent( if getattr(parent, field) != value: msg = f"submitted parent {field} does not match invocation" raise ValueError(msg) + + +def _validate_parent_is_unclaimed(parent: ComparisonReportRecord) -> None: + """Require a pending parent that no coordinator has started.""" if parent.status is not ComparisonRunLifecycleStatus.PENDING: msg = "submitted parent must be pending" raise ValueError(msg) @@ -134,6 +171,15 @@ def _validate_submitted_parent( if parent.coordinator_invocation_id is not None: msg = "submitted parent must not contain an invocation identifier" raise ValueError(msg) + + +def _validate_parent_production_link( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + parent: ComparisonReportRecord, +) -> None: + """Require a direct parent to stand alone and a shadow parent to name its job.""" production_job_id = context.production_function_call_id if production_job_id is None: if parent.production_identity != f"direct:{report.evaluation_id}": @@ -250,87 +296,103 @@ def _require_artifact_digest( raise ValueError(mismatch_message) -def coordinate_report( - payload: object, - context_payload: object, - parent_payload: object, - *, - application_name: str, - coordinator_invocation_id: str, - store: ComparisonStore | None = None, - artifacts: Stage12ArtifactStore | None = None, - invoker: ChildInvoker | None = None, - aggregator: Callable[..., dict[str, Any]] = build_aggregate_report, - output_plan_resolver: Callable[[ReportExecutionInput], Stage12OutputPlan] - | None = None, - runtime: ObservabilityRuntime | None = None, -) -> dict[str, Any]: - """Run one Stage 12 comparison report end to end; return its aggregate descriptor. +@dataclass(frozen=True) +class CoordinatorInvocation: + """The deployed coordinator call running a report: its Modal app and call ID.""" - Claims the parent report (one already running or finished returns a - ``deduplicated`` summary instead), starts or reattaches the baseline and - reform child simulations, and aggregates their digest-checked artifacts; - shadow runs are then compared with the production result. A dispatch - context that disagrees with the report raises ``ValueError`` before the - claim; any later failure marks the report FAILED and raises a - ``RuntimeError`` that carries no calculation detail. + application_name: str + invocation_id: str + + +@dataclass(frozen=True) +class CoordinatorDependencies: + """Seams the coordinator uses; an unset seam selects the deployed default. + + The defaults are the Stage 12 database store, the private artifact + bucket, Modal child invocation, the standard aggregate builder and output + plan resolution, and no telemetry runtime. """ - report = ReportExecutionInput.model_validate(payload) - context = Stage12InvocationContext.model_validate(context_payload) - submitted_parent = ComparisonReportRecord.model_validate(parent_payload) - if application_name != context.modal_application: - msg = "report coordinator application differs from dispatch context" - raise ValueError(msg) - if report.baseline.bundle.policyengine_version != context.worker_version: - msg = "report worker version differs from dispatch context" - raise ValueError(msg) - if report.baseline.bundle.bundle_manifest_sha256 != context.bundle_manifest_sha256: - msg = "report bundle digest differs from dispatch context" - raise ValueError(msg) - _validate_submitted_parent( - report=report, - context=context, - parent=submitted_parent, - ) - persistence = store or runtime_store(runtime) - artifact_storage = artifacts or artifact_store() - child_invoker = invoker or ModalChildInvoker() - stage_plan: StagePlan = ( - STAGE12_SHADOW_REPORT_STAGES - if context.production_function_call_id is not None - else STAGE12_CANONICAL_REPORT_STAGES - ) - progress = StageProgress() - stage_span = partial( - tracked_stage, - stage_plan=stage_plan, - progress=progress, - runtime=runtime, + + store: ComparisonStore | None = None + artifacts: Stage12ArtifactStore | None = None + invoker: ChildInvoker | None = None + aggregator: Callable[..., dict[str, Any]] = build_aggregate_report + output_plan_resolver: Callable[[ReportExecutionInput], Stage12OutputPlan] | None = ( + None ) - stage_scope = partial(stage_span, operation=True) - with stage_span(Stage.STAGE12_COORDINATOR_CLAIM): - parent, should_execute = _claim_parent( - submitted=submitted_parent, - coordinator_invocation_id=coordinator_invocation_id, - store=persistence, - ) - if not should_execute: - return { - "deduplicated": True, - "evaluation_id": str(parent.evaluation_id), - "status": parent.status.value, - } - descriptors: dict[str, SimulationArtifactDescriptor] = {} - calls: dict[str, ChildCall] = {} + runtime: ObservabilityRuntime | None = None + + +@dataclass +class _ReportCoordination: + """One claimed report's inputs and resolved seams, run through each phase.""" + + report: ReportExecutionInput + context: Stage12InvocationContext + invocation: CoordinatorInvocation + persistence: ComparisonStore + artifact_storage: Stage12ArtifactStore + child_invoker: ChildInvoker + dependencies: CoordinatorDependencies + stage_plan: StagePlan + # The last stage entered and whether it finished, for failure causes. + progress: StageProgress # A child failure is reported where it occurs and then propagates to the - # report handler below, which must not report the same error twice. + # report handler, which must not report the same error twice. reported_child_failure: Exception | None = None # A failed child records its own cause; the report keeps that cause # rather than the generic exception Modal returns to the coordinator. child_failure_summary: str | None = None - try: - with stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): - context = context.model_copy( + + @property + def runtime(self) -> ObservabilityRuntime | None: + """The telemetry runtime, or None when the coordinator records nothing.""" + return self.dependencies.runtime + + def span( + self, + stage: Stage, + attributes: Mapping[str, object] | None = None, + ) -> AbstractContextManager[object]: + """Track ``stage`` for failure causes and time it as a span.""" + return tracked_stage( + stage, + attributes, + stage_plan=self.stage_plan, + progress=self.progress, + runtime=self.runtime, + ) + + def stage_scope( + self, + stage: Stage, + attributes: Mapping[str, object] | None = None, + ) -> AbstractContextManager[object]: + """Track ``stage`` for failure causes and record it as an operation.""" + return tracked_stage( + stage, + attributes, + stage_plan=self.stage_plan, + progress=self.progress, + runtime=self.runtime, + operation=True, + ) + + def run(self, parent: ComparisonReportRecord) -> dict[str, Any]: + """Plan, start and await both children, then aggregate and complete.""" + output_plan, planned_simulations = self._prepare(parent) + children = self._create_children(planned_simulations) + descriptors, calls = self._start_children(planned_simulations, children) + # Every required call has been started before the coordinator waits. + self._wait_for_children(calls, descriptors, planned_simulations) + return self._aggregate_and_complete(parent, output_plan, descriptors) + + def _prepare( + self, parent: ComparisonReportRecord + ) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: + """Bind the run's artifact prefix, then plan the outputs and child inputs.""" + with self.stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): + self.context = self.context.model_copy( update={ "artifact_prefix": ( f"stage-12-runs/{parent.environment}/" @@ -341,162 +403,223 @@ def coordinate_report( "retention_expires_at": parent.retention_expires_at, } ) - function_name = context.simulation_callable - with stage_scope( + with self.stage_scope( Stage.STAGE12_OUTPUT_PLANNING, - {"country": report.baseline.geography.country}, + {"country": self.report.baseline.geography.country}, ): - if output_plan_resolver is None: + resolver = self.dependencies.output_plan_resolver + if resolver is None: output_plan = resolve_report_output_plan( - report, - stage_scope=stage_scope, + self.report, + stage_scope=self.stage_scope, ) else: - output_plan = output_plan_resolver(report) + output_plan = resolver(self.report) simulations = [] - for simulation in (report.baseline, report.reform): - with stage_scope( + for simulation in (self.report.baseline, self.report.reform): + with self.stage_scope( Stage.STAGE12_CHILD_INPUT_PLANNING, {"simulation_role": simulation.role.value}, ): simulations.append(plan_simulation_input(simulation, output_plan)) - planned_simulations = tuple(simulations) - with stage_span(Stage.STAGE12_CHILD_STATE_CREATE): - children = { - simulation.role: persistence.create_or_resolve_simulation( + return output_plan, tuple(simulations) + + def _create_children( + self, planned_simulations: tuple[PlannedSimulationExecutionInput, ...] + ) -> dict[SimulationRole, ComparisonSimulationRecord]: + """Create each child's row, or resolve the row an earlier attempt made.""" + with self.span(Stage.STAGE12_CHILD_STATE_CREATE): + return { + simulation.role: self.persistence.create_or_resolve_simulation( _child_record( simulation=simulation, - context=context, - function_name=function_name, + context=self.context, + function_name=self.context.simulation_callable, ) ).record for simulation in planned_simulations } + + def _start_children( + self, + planned_simulations: tuple[PlannedSimulationExecutionInput, ...], + children: dict[SimulationRole, ComparisonSimulationRecord], + ) -> tuple[dict[str, SimulationArtifactDescriptor], dict[str, ChildCall]]: + """Reuse finished children, reattach running ones and dispatch the rest.""" + descriptors: dict[str, SimulationArtifactDescriptor] = {} + calls: dict[str, ChildCall] = {} for simulation in planned_simulations: child = children[simulation.role] + role = simulation.role.value if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: - descriptor = descriptor_from_record(child, simulation) - retained = artifact_storage.read(descriptor.artifact.uri) - _require_artifact_digest( - retained, - descriptor.artifact.content_sha256, - mismatch_message="retained child artifact digest mismatch", - ) - descriptor = descriptor.model_copy( - update={ - "artifact": descriptor.artifact.model_copy( - update={"size_bytes": len(retained)} - ), - "calculation_provenance": ( - deserialize_calculation_provenance(retained) - ), - } - ) - descriptors[simulation.role.value] = descriptor - continue - if ( + descriptors[role] = self._retained_descriptor(child, simulation) + elif ( child.status is ComparisonRunLifecycleStatus.RUNNING and child.modal_invocation_id is not None and not child.modal_invocation_id.startswith("dispatch-pending-") ): - calls[simulation.role.value] = child_invoker.restore( - child.modal_invocation_id - ) - continue - dispatching_at = datetime.now(UTC) - dispatch_placeholder = f"dispatch-pending-{uuid4()}" - persistence.replace_simulation( - child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "modal_invocation_id": dispatch_placeholder, - "started_at": child.started_at or dispatching_at, - "updated_at": dispatching_at, - "error_code": None, - "error_summary": None, - "completed_at": None, - } - ) + calls[role] = self.child_invoker.restore(child.modal_invocation_id) + else: + calls[role] = self._dispatch_child(child, simulation) + return descriptors, calls + + def _retained_descriptor( + self, + child: ComparisonSimulationRecord, + simulation: PlannedSimulationExecutionInput, + ) -> SimulationArtifactDescriptor: + """Describe a succeeded child's digest-checked retained artifact.""" + descriptor = descriptor_from_record(child, simulation) + retained = self.artifact_storage.read(descriptor.artifact.uri) + _require_artifact_digest( + retained, + descriptor.artifact.content_sha256, + mismatch_message="retained child artifact digest mismatch", + ) + return descriptor.model_copy( + update={ + "artifact": descriptor.artifact.model_copy( + update={"size_bytes": len(retained)} + ), + "calculation_provenance": ( + deserialize_calculation_provenance(retained) + ), + } + ) + + def _dispatch_child( + self, + child: ComparisonSimulationRecord, + simulation: PlannedSimulationExecutionInput, + ) -> ChildCall: + """Mark the child dispatching, spawn it, and attach its Modal call ID. + + A spawn failure is reported, marks the child failed (unless it already + succeeded) and re-raises. + """ + dispatching_at = datetime.now(UTC) + dispatch_placeholder = f"dispatch-pending-{uuid4()}" + self.persistence.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": dispatch_placeholder, + "started_at": child.started_at or dispatching_at, + "updated_at": dispatching_at, + "error_code": None, + "error_summary": None, + "completed_at": None, + } ) - try: - with stage_span( - Stage.STAGE12_CHILD_DISPATCH, - {"simulation_role": simulation.role.value}, - ): - child_observability_context = ( - runtime.capture_context() if runtime is not None else None - ) - call = child_invoker.spawn( - application_name=application_name, - function_name=function_name, - environment=context.modal_environment, - simulation=simulation.model_dump(mode="json"), - context=context.model_dump(mode="json"), - observability_context=child_observability_context, - ) - except Exception as error: - report_original_failure( - error, - runtime=runtime, - failure_scope="child_dispatch", - attributes={"simulation_role": simulation.role.value}, + ) + try: + with self.span( + Stage.STAGE12_CHILD_DISPATCH, + {"simulation_role": simulation.role.value}, + ): + child_observability_context = ( + self.runtime.capture_context() if self.runtime is not None else None ) - reported_child_failure = error - persist_failure_record( - partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code="simulation_dispatch_failed", - error_summary=failure_summary(error, progress=progress), - ), - runtime=runtime, - failure_scope="child_dispatch", + call = self.child_invoker.spawn( + application_name=self.invocation.application_name, + function_name=self.context.simulation_callable, + environment=self.context.modal_environment, + simulation=simulation.model_dump(mode="json"), + context=self.context.model_dump(mode="json"), + observability_context=child_observability_context, ) - raise - calls[simulation.role.value] = call - persistence.attach_simulation_invocation( - simulation.simulation_execution_id, - expected_placeholder=dispatch_placeholder, - modal_invocation_id=call.object_id, - updated_at=datetime.now(UTC), + except Exception as error: + self._record_child_failure( + simulation, + error, + failure_scope="child_dispatch", + error_code="simulation_dispatch_failed", ) - # Every required call has been started before the coordinator waits. + raise + self.persistence.attach_simulation_invocation( + simulation.simulation_execution_id, + expected_placeholder=dispatch_placeholder, + modal_invocation_id=call.object_id, + updated_at=datetime.now(UTC), + ) + return call + + def _wait_for_children( + self, + calls: dict[str, ChildCall], + descriptors: dict[str, SimulationArtifactDescriptor], + planned_simulations: tuple[PlannedSimulationExecutionInput, ...], + ) -> None: + """Await each started child's descriptor. + + A failure is reported and marks that child failed before it is re-raised. + """ for role, call in calls.items(): try: - with stage_span(Stage.STAGE12_CHILD_WAIT, {"simulation_role": role}): + with self.span(Stage.STAGE12_CHILD_WAIT, {"simulation_role": role}): descriptors[role] = SimulationArtifactDescriptor.model_validate( call.get(timeout=SIMULATION_WAIT_TIMEOUT_SECONDS) ) except Exception as error: # noqa: PERF203 # per child: the failing child is marked failed before the error is re-raised - report_original_failure( - error, - runtime=runtime, - failure_scope="child_wait", - attributes={"simulation_role": role}, - ) - reported_child_failure = error simulation = next( - item for item in simulations if item.role.value == role + item for item in planned_simulations if item.role.value == role ) - child_failure_summary = persist_failure_record( - partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code="simulation_invocation_failed", - error_summary=failure_summary(error, progress=progress), - ), - runtime=runtime, + self.child_failure_summary = self._record_child_failure( + simulation, + error, failure_scope="child_wait", + error_code="simulation_invocation_failed", ) raise - baseline = descriptors["baseline"] - reform = descriptors["reform"] - validate_aligned_outputs(report, output_plan, baseline, reform) - with stage_span(Stage.STAGE12_ARTIFACT_READ): - baseline_payload = artifact_storage.read(baseline.artifact.uri) - reform_payload = artifact_storage.read(reform.artifact.uri) + + def _record_child_failure( + self, + simulation: PlannedSimulationExecutionInput, + error: Exception, + *, + failure_scope: str, + error_code: str, + ) -> str | None: + """Report a child failure, then try to mark the child failed with its cause. + + The original failure reaches telemetry before the database write, which + may itself fail without replacing it; a succeeded child is left as is. + Returns the redacted cause the child's record keeps, or ``None``. + """ + report_original_failure( + error, + runtime=self.runtime, + failure_scope=failure_scope, + attributes={"simulation_role": simulation.role.value}, + ) + self.reported_child_failure = error + return persist_failure_record( + partial( + _mark_child_failed, + self.persistence, + simulation.simulation_execution_id, + error_code=error_code, + error_summary=failure_summary(error, progress=self.progress), + ), + runtime=self.runtime, + failure_scope=failure_scope, + ) + + def _read_retained_outputs( + self, + output_plan: Stage12OutputPlan, + baseline: SimulationArtifactDescriptor, + reform: SimulationArtifactDescriptor, + ) -> tuple[ + RetainedSimulationOutput, + RetainedSimulationOutput, + UKLocalAuthorityMetadata | None, + ]: + """Read, digest-check and validate both children's retained artifacts.""" + validate_aligned_outputs(self.report, output_plan, baseline, reform) + with self.span(Stage.STAGE12_ARTIFACT_READ): + baseline_payload = self.artifact_storage.read(baseline.artifact.uri) + reform_payload = self.artifact_storage.read(reform.artifact.uri) _require_artifact_digest( baseline_payload, baseline.artifact.content_sha256, @@ -510,34 +633,53 @@ def coordinate_report( baseline_frames = deserialize_simulation_frames(baseline_payload) reform_frames = deserialize_simulation_frames(reform_payload) uk_local_authority_metadata = validate_uk_local_authority_metadata( - report.baseline.geography.country, + self.report.baseline.geography.country, deserialize_uk_local_authority_metadata(baseline_payload), deserialize_uk_local_authority_metadata(reform_payload), - region_code=report.baseline.geography.region, + region_code=self.report.baseline.geography.region, ) validate_output_frames(baseline_frames, output_plan) validate_output_frames(reform_frames, output_plan) - with stage_span(Stage.STAGE12_AGGREGATION): - aggregate = aggregator( - report=report, - baseline=RetainedSimulationOutput(baseline_frames, baseline), - reform=RetainedSimulationOutput(reform_frames, reform), + return ( + RetainedSimulationOutput(baseline_frames, baseline), + RetainedSimulationOutput(reform_frames, reform), + uk_local_authority_metadata, + ) + + def _aggregate_and_complete( + self, + parent: ComparisonReportRecord, + output_plan: Stage12OutputPlan, + descriptors: dict[str, SimulationArtifactDescriptor], + ) -> dict[str, Any]: + """Aggregate both outputs, store the aggregate and complete the parent. + + A shadow run's result is then compared with the production result. + """ + baseline, reform, uk_local_authority_metadata = self._read_retained_outputs( + output_plan, descriptors["baseline"], descriptors["reform"] + ) + with self.span(Stage.STAGE12_AGGREGATION): + aggregate = self.dependencies.aggregator( + report=self.report, + baseline=baseline, + reform=reform, uk_local_authority_metadata=uk_local_authority_metadata, ) - with stage_span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): - aggregate_artifact = artifact_storage.write_aggregate( - prefix=context.artifact_prefix, + with self.span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): + aggregate_artifact = self.artifact_storage.write_aggregate( + prefix=self.context.artifact_prefix, payload=aggregate, ) descriptor = AggregateReportArtifactDescriptor( - evaluation_id=report.evaluation_id, + evaluation_id=self.report.evaluation_id, artifact=aggregate_artifact, - baseline_artifact_sha256=baseline.artifact.content_sha256, - reform_artifact_sha256=reform.artifact.content_sha256, - bundle=report.baseline.bundle, + baseline_artifact_sha256=baseline.descriptor.artifact.content_sha256, + reform_artifact_sha256=reform.descriptor.artifact.content_sha256, + bundle=self.report.baseline.bundle, ) completed = datetime.now(UTC) - completed_parent = persistence.replace_report( + completed_parent = self.persistence.replace_report( parent.model_copy( update={ "status": ComparisonRunLifecycleStatus.SUCCEEDED, @@ -553,31 +695,90 @@ def coordinate_report( compare_completed_report( parent=completed_parent, aggregate=aggregate, - context=context, - store=persistence, - artifacts=artifact_storage, - invoker=child_invoker, - runtime=runtime, - stage_plan=stage_plan, + context=self.context, + store=self.persistence, + artifacts=self.artifact_storage, + invoker=self.child_invoker, + runtime=self.runtime, + stage_plan=self.stage_plan, ) return descriptor.model_dump(mode="json") + + +def coordinate_report( + payload: object, + context_payload: object, + parent_payload: object, + *, + invocation: CoordinatorInvocation, + dependencies: CoordinatorDependencies | None = None, +) -> dict[str, Any]: + """Run one Stage 12 comparison report end to end; return its aggregate descriptor. + + Claims the parent report (one already running or finished returns a + ``deduplicated`` summary instead), starts or reattaches the baseline and + reform child simulations, and aggregates their digest-checked artifacts; + shadow runs are then compared with the production result. A dispatch + context that disagrees with the report raises ``ValueError`` before the + claim; any later failure marks the report FAILED and raises a + ``RuntimeError`` that carries no calculation detail. + """ + dependencies = dependencies or CoordinatorDependencies() + report = ReportExecutionInput.model_validate(payload) + context = Stage12InvocationContext.model_validate(context_payload) + submitted_parent = ComparisonReportRecord.model_validate(parent_payload) + _validate_dispatch_context(report, context, invocation.application_name) + _validate_submitted_parent( + report=report, + context=context, + parent=submitted_parent, + ) + coordination = _ReportCoordination( + report=report, + context=context, + invocation=invocation, + persistence=dependencies.store or runtime_store(dependencies.runtime), + artifact_storage=dependencies.artifacts or artifact_store(), + child_invoker=dependencies.invoker or ModalChildInvoker(), + dependencies=dependencies, + stage_plan=( + STAGE12_SHADOW_REPORT_STAGES + if context.production_function_call_id is not None + else STAGE12_CANONICAL_REPORT_STAGES + ), + progress=StageProgress(), + ) + with coordination.span(Stage.STAGE12_COORDINATOR_CLAIM): + parent, should_execute = _claim_parent( + submitted=submitted_parent, + coordinator_invocation_id=invocation.invocation_id, + store=coordination.persistence, + ) + if not should_execute: + return { + "deduplicated": True, + "evaluation_id": str(parent.evaluation_id), + "status": parent.status.value, + } + try: + return coordination.run(parent) # Persist a redacted cause for any coordinator, child-call, aggregation, or # artifact failure before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 # every failure is persisted before a stable exception returns to Modal - if error is not reported_child_failure: + if error is not coordination.reported_child_failure: report_original_failure( error, - runtime=runtime, + runtime=dependencies.runtime, failure_scope="report_coordination", ) - report_failure_summary = child_failure_summary or failure_summary( - error, progress=progress + report_failure_summary = coordination.child_failure_summary or ( + failure_summary(error, progress=coordination.progress) ) - record_failure_cause(runtime, error, report_failure_summary) + record_failure_cause(dependencies.runtime, error, report_failure_summary) failed_at = datetime.now(UTC) persist_failure_record( partial( - persistence.replace_report, + coordination.persistence.replace_report, parent.model_copy( update={ "status": ComparisonRunLifecycleStatus.FAILED, @@ -589,7 +790,7 @@ def coordinate_report( } ), ), - runtime=runtime, + runtime=dependencies.runtime, failure_scope="report_coordination", ) msg = "Stage 12 report coordination failed" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index e0373b6bf..631c3a7fb 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -13,7 +13,11 @@ ComparisonRunLifecycleStatus, ResultComparisonStatus, ) -from policyengine_simulation_executor.stage12_runtime import coordination +from policyengine_simulation_executor.stage12_runtime import ( + CoordinatorDependencies, + CoordinatorInvocation, + coordination, +) from test_stage12_runtime import ( BASELINE_ID, EVALUATION_ID, @@ -409,13 +413,17 @@ def resolve(report, *, stage_scope): _report().model_dump(mode="json"), _context().model_dump(mode="json"), _parent().model_dump(mode="json"), - application_name=_context().modal_application, - coordinator_invocation_id="coordinator-1", - store=store, - artifacts=artifacts, - invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: _aggregate_payload(), - runtime=runtime, + invocation=CoordinatorInvocation( + application_name=_context().modal_application, + invocation_id="coordinator-1", + ), + dependencies=CoordinatorDependencies( + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: _aggregate_payload(), + runtime=runtime, + ), ) assert resolved == [EVALUATION_ID] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 16f29af73..fab5d113a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -74,6 +74,8 @@ serialize_simulation_frames, ) from policyengine_simulation_executor.stage12_runtime import ( + CoordinatorDependencies, + CoordinatorInvocation, RetainedSimulationOutput, SimulationCalculation, _build_spm_result, @@ -226,9 +228,16 @@ def _aggregate_payload(result=None) -> dict: ).model_dump(mode="json") -def coordinate_report(*args, **kwargs): - kwargs.setdefault("output_plan_resolver", lambda _: _output_plan()) - return coordinate_report_impl(*args, **kwargs) +def coordinate_report(*payloads, application_name, coordinator_invocation_id, **seams): + seams.setdefault("output_plan_resolver", lambda _: _output_plan()) + return coordinate_report_impl( + *payloads, + invocation=CoordinatorInvocation( + application_name=application_name, + invocation_id=coordinator_invocation_id, + ), + dependencies=CoordinatorDependencies(**seams), + ) class RecordingRuntime: @@ -1413,13 +1422,17 @@ def resolver(report): _report().model_dump(mode="json"), context.model_dump(mode="json"), parent.model_dump(mode="json"), - application_name=context.modal_application, - coordinator_invocation_id="coordinator-1", - store=store, - artifacts=artifacts, - invoker=invoker, - aggregator=aggregate, - output_plan_resolver=resolver, + invocation=CoordinatorInvocation( + application_name=context.modal_application, + invocation_id="coordinator-1", + ), + dependencies=CoordinatorDependencies( + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=aggregate, + output_plan_resolver=resolver, + ), ) assert invoker.events == ["spawn:baseline", "spawn:reform"] From 8a8b0b5affc07adbceafc786953f8c7d8da85b05 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:41:55 +0400 Subject: [PATCH 35/66] Pin how the Modal child invoker looks up, spawns and restores children Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_coordination.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index 631c3a7fb..4570b73b6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -504,3 +504,61 @@ def aggregate(**values): assert child.started_at == CLAIMED_AT assert child.updated_at == CLAIMED_AT assert child.simulation_callable == "run_single_simulation_us" + + +def test_modal_child_invoker_spawns_and_restores_deployed_functions(monkeypatch): + from types import SimpleNamespace + + from policyengine_simulation_executor.stage12_runtime import dependencies + + calls = [] + + def from_name(*args, **kwargs): + calls.append(("from_name", args, kwargs)) + + def spawn(*spawn_args, **spawn_kwargs): + calls.append(("spawn", spawn_args, spawn_kwargs)) + return "spawned-call" + + return SimpleNamespace(spawn=spawn) + + def from_id(invocation_id): + calls.append(("from_id", (invocation_id,), {})) + return "restored-call" + + fake_modal = SimpleNamespace( + Function=SimpleNamespace(from_name=from_name), + FunctionCall=SimpleNamespace(from_id=from_id), + ) + + def import_modal(module_name): + assert module_name == "modal" + return fake_modal + + monkeypatch.setattr(dependencies, "import_module", import_modal) + invoker = dependencies.ModalChildInvoker() + + spawned = invoker.spawn( + application_name="application", + function_name="run_single_simulation_us", + environment="staging", + simulation={"role": "baseline"}, + context={"request_id": "request-1"}, + observability_context={"traceparent": "trace"}, + ) + restored = invoker.restore("call-1") + + assert (spawned, restored) == ("spawned-call", "restored-call") + assert calls == [ + ( + "from_name", + ("application", "run_single_simulation_us"), + {"environment_name": "staging"}, + ), + ( + "spawn", + ({"role": "baseline"}, {"request_id": "request-1"}), + {"observability_context": {"traceparent": "trace"}}, + ), + ("from_id", ("call-1",), {}), + ] From 8cd00b2b6c9d3cefc39ebd4779ce0dd33736c5e6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:42:46 +0400 Subject: [PATCH 36/66] Pass the child invoker one spawn request instead of six keywords Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 15 +++--- .../stage12_runtime/dependencies.py | 47 +++++++++---------- .../tests/test_stage12_coordination.py | 18 +++---- .../tests/test_stage12_runtime.py | 22 +++++---- 4 files changed, 53 insertions(+), 49 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 2a4745369..6e9bda31b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -47,6 +47,7 @@ from .dependencies import ( ChildCall, ChildInvoker, + ChildSpawnRequest, ComparisonStore, ModalChildInvoker, artifact_store, @@ -521,12 +522,14 @@ def _dispatch_child( self.runtime.capture_context() if self.runtime is not None else None ) call = self.child_invoker.spawn( - application_name=self.invocation.application_name, - function_name=self.context.simulation_callable, - environment=self.context.modal_environment, - simulation=simulation.model_dump(mode="json"), - context=self.context.model_dump(mode="json"), - observability_context=child_observability_context, + ChildSpawnRequest( + application_name=self.invocation.application_name, + function_name=self.context.simulation_callable, + environment=self.context.modal_environment, + simulation=simulation.model_dump(mode="json"), + context=self.context.model_dump(mode="json"), + observability_context=child_observability_context, + ) ) except Exception as error: self._record_child_failure( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index 553ba4c76..3e45ff80e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from dataclasses import dataclass from functools import partial from importlib import import_module from typing import TYPE_CHECKING, Any, Protocol @@ -96,19 +97,22 @@ def get(self, *, timeout: float | None = None) -> object: ... +@dataclass(frozen=True) +class ChildSpawnRequest: + """One child simulation to start: the deployed function and its arguments.""" + + application_name: str + function_name: str + environment: str + simulation: dict[str, Any] + context: dict[str, Any] + observability_context: dict[str, Any] | None + + class ChildInvoker(Protocol): """Launcher for child simulation functions that can reattach to running calls.""" - def spawn( - self, - *, - application_name: str, - function_name: str, - environment: str, - simulation: dict[str, Any], - context: dict[str, Any], - observability_context: dict[str, Any] | None, - ) -> ChildCall: + def spawn(self, request: ChildSpawnRequest) -> ChildCall: """Start the named deployed function on one simulation without waiting.""" ... @@ -120,27 +124,18 @@ def restore(self, invocation_id: str) -> ChildCall: class ModalChildInvoker: """``ChildInvoker`` backed by functions deployed to Modal.""" - def spawn( - self, - *, - application_name: str, - function_name: str, - environment: str, - simulation: dict[str, Any], - context: dict[str, Any], - observability_context: dict[str, Any] | None, - ) -> ChildCall: + def spawn(self, request: ChildSpawnRequest) -> ChildCall: """Look up ``application_name.function_name`` in ``environment``; spawn it.""" modal: Any = import_module("modal") function = modal.Function.from_name( - application_name, - function_name, - environment_name=environment, + request.application_name, + request.function_name, + environment_name=request.environment, ) return function.spawn( - simulation, - context, - observability_context=observability_context, + request.simulation, + request.context, + observability_context=request.observability_context, ) def restore(self, invocation_id: str) -> ChildCall: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index 4570b73b6..69fe73c5f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -85,8 +85,8 @@ def __init__(self, store=None, succeeded_child=None): self.succeeded_child = succeeded_child self.spawns = [] - def spawn(self, **call): - self.spawns.append(call) + def spawn(self, request): + self.spawns.append(request) if self.succeeded_child is not None: self.store.children[self.succeeded_child.simulation_execution_id] = ( self.succeeded_child @@ -539,12 +539,14 @@ def import_modal(module_name): invoker = dependencies.ModalChildInvoker() spawned = invoker.spawn( - application_name="application", - function_name="run_single_simulation_us", - environment="staging", - simulation={"role": "baseline"}, - context={"request_id": "request-1"}, - observability_context={"traceparent": "trace"}, + dependencies.ChildSpawnRequest( + application_name="application", + function_name="run_single_simulation_us", + environment="staging", + simulation={"role": "baseline"}, + context={"request_id": "request-1"}, + observability_context={"traceparent": "trace"}, + ) ) restored = invoker.restore("call-1") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index fab5d113a..98e19f272 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1336,11 +1336,11 @@ def __init__( self.production_wait_timeouts = [] self.output_plans = [] - def spawn(self, *, simulation, environment, **_): - parsed = PlannedSimulationExecutionInput.model_validate(simulation) + def spawn(self, request): + parsed = PlannedSimulationExecutionInput.model_validate(request.simulation) role = parsed.role.value self.events.append(f"spawn:{role}") - self.environments.append(environment) + self.environments.append(request.environment) self.output_plans.append(parsed.output_plan) return self._submit(parsed, role) @@ -1866,8 +1866,10 @@ def __init__(self, dispatch_error): self.dispatch_error = dispatch_error self.events = [] - def spawn(self, *, simulation, **_): - role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role.value self.events.append(f"spawn:{role}") raise self.dispatch_error @@ -1951,8 +1953,8 @@ def __init__(self, store, worker_summary): self.store = store self.worker_summary = worker_summary - def spawn(self, *, simulation, **_): - planned = PlannedSimulationExecutionInput.model_validate(simulation) + def spawn(self, request): + planned = PlannedSimulationExecutionInput.model_validate(request.simulation) if planned.role is SimulationRole.BASELINE: return WorkerRecordedFailureCall( "call-baseline", @@ -2455,8 +2457,10 @@ class TimeoutInvoker: def __init__(self): self.events = [] - def spawn(self, *, simulation, **_): - role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role.value self.events.append(f"spawn:{role}") return TimeoutCall(f"call-{role}") From cff978ab64350db8dd0b7a35d3722462dd5b1889 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:43:21 +0400 Subject: [PATCH 37/66] Pin result comparison logs, spans and the failed state-update path Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_coordination.py | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index 69fe73c5f..327fbfcbb 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -564,3 +564,93 @@ def import_modal(module_name): ), ("from_id", ("call-1",), {}), ] + + +def _shadow_run(store, artifacts, invoker, runtime=None): + context = _context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + parent = _parent().model_copy( + update={ + "production_identity": "production-job-1", + "incumbent_execution_id": "production-job-1", + "comparison_status": ResultComparisonStatus.PENDING, + } + ) + return _coordinate( + context, + parent, + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: _aggregate_payload({"budget": {"total": 100.0}}), + runtime=runtime, + ) + + +def test_shadow_comparison_logs_its_outcome_and_times_its_stage(caplog): + import logging + + store = FakeStore() + artifacts = FakeArtifacts() + runtime = TracingRuntime() + invoker = ConcurrentInvoker( + artifacts, production_result={"budget": {"total": 100.0}} + ) + + with caplog.at_level(logging.INFO): + _shadow_run(store, artifacts, invoker, runtime) + + assert runtime.events[-1] == ("span", "stage12_result_comparison", {}) + assert all(name.startswith("stage12_") for _kind, name, _ in runtime.events if name) + [record] = [ + item + for item in caplog.records + if item.getMessage() == "stage12_result_comparison_completed" + ] + assert record.name == "policyengine_simulation_executor.stage12_runtime.comparison" + assert record.levelno == logging.INFO + assert record.evaluation_id == str(EVALUATION_ID) + assert record.production_job_id == "production-job-1" + assert record.comparison_status == "matched" + assert record.difference_count == 0 + assert store.parent.comparison_status is ResultComparisonStatus.MATCHED + + +def test_comparison_state_update_failure_is_logged_without_values(caplog): + import logging + + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker(artifacts, production_result="not-an-object") + + def unavailable_report(_evaluation_id): + msg = "database unavailable with sensitive detail" + raise ConnectionError(msg) + + store.get_report = unavailable_report + + with caplog.at_level(logging.INFO): + result = _shadow_run(store, artifacts, invoker) + + assert result["evaluation_id"] == str(EVALUATION_ID) + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert store.parent.comparison_status is ResultComparisonStatus.RUNNING + state_record, failure_record = [ + item + for item in caplog.records + if item.getMessage() + in { + "stage12_result_comparison_state_update_failed", + "stage12_result_comparison_failed", + } + ] + assert state_record.getMessage() == "stage12_result_comparison_state_update_failed" + assert state_record.levelno == logging.ERROR + assert state_record.exc_info is not None + assert state_record.evaluation_id == str(EVALUATION_ID) + assert state_record.production_job_id == "production-job-1" + assert failure_record.getMessage() == "stage12_result_comparison_failed" + assert failure_record.levelno == logging.ERROR + assert failure_record.exc_info is None + assert failure_record.error_type == "ValidationError" From c33af85f39666f84721a0b806206a3e40b9c8e9c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:43:59 +0400 Subject: [PATCH 38/66] Pass the result comparison its coordinator services as one group Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/comparison.py | 25 +++++++++++++------ .../stage12_runtime/coordination.py | 14 ++++++----- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py index 59954c130..cda7c608b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py @@ -4,6 +4,7 @@ import logging from contextlib import nullcontext +from dataclasses import dataclass from datetime import UTC, datetime from typing import TYPE_CHECKING, Any @@ -29,18 +30,28 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class ComparisonServices: + """The coordinator's store, artifacts, invoker and telemetry a comparison uses.""" + + store: ComparisonStore + artifacts: Stage12ArtifactStore + invoker: ChildInvoker + runtime: ObservabilityRuntime | None = None + stage_plan: StagePlan | None = None + + def compare_completed_report( *, parent: ComparisonReportRecord, aggregate: dict[str, Any], context: Stage12InvocationContext, - store: ComparisonStore, - artifacts: Stage12ArtifactStore, - invoker: ChildInvoker, - runtime: ObservabilityRuntime | None = None, - stage_plan: StagePlan | None = None, + services: ComparisonServices, ) -> None: """Persist an isolated production/Stage 12 comparison after aggregation.""" + store = services.store + runtime = services.runtime + stage_plan = services.stage_plan production_job_id = context.production_function_call_id if production_job_id is None: return @@ -66,7 +77,7 @@ def compare_completed_report( } ) ) - production_result = invoker.restore(production_job_id).get( + production_result = services.invoker.restore(production_job_id).get( timeout=PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS ) comparison = compare_results( @@ -75,7 +86,7 @@ def compare_completed_report( production_result=production_result, stage12_result=aggregate.get("result"), ) - comparison_artifact = artifacts.write_comparison( + comparison_artifact = services.artifacts.write_comparison( prefix=context.artifact_prefix, payload=comparison.model_dump(mode="json"), ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 6e9bda31b..296a76382 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -43,7 +43,7 @@ validate_aligned_outputs, validate_uk_local_authority_metadata, ) -from .comparison import compare_completed_report +from .comparison import ComparisonServices, compare_completed_report from .dependencies import ( ChildCall, ChildInvoker, @@ -699,11 +699,13 @@ def _aggregate_and_complete( parent=completed_parent, aggregate=aggregate, context=self.context, - store=self.persistence, - artifacts=self.artifact_storage, - invoker=self.child_invoker, - runtime=self.runtime, - stage_plan=self.stage_plan, + services=ComparisonServices( + store=self.persistence, + artifacts=self.artifact_storage, + invoker=self.child_invoker, + runtime=self.runtime, + stage_plan=self.stage_plan, + ), ) return descriptor.model_dump(mode="json") From ca01de631a2bd9e0aeb3c2dbc64b7b138ab82fb8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:44:49 +0400 Subject: [PATCH 39/66] Pin the Stage 12 worker's reuse path and calculator timing Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_simulation_worker.py | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py new file mode 100644 index 000000000..c3dd5da40 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py @@ -0,0 +1,128 @@ +"""Characterization tests for the Stage 12 single-simulation worker's phases.""" + +from __future__ import annotations + +from contextlib import nullcontext + +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + simulation as simulation_module, +) +from policyengine_simulation_executor.stage12_runtime.failure_causes import ( + StageProgress, +) +from test_stage12_runtime import ( + NOW, + FakeArtifacts, + FakeStore, + _child, + _context, + _frames, + _planned_simulation, +) + + +class SpanRecordingRuntime: + def __init__(self) -> None: + self.spans: list[str] = [] + + def span(self, name, *_args, **_kwargs): + self.spans.append(name) + return nullcontext() + + +def _run(store, artifacts, *, calculator=None, runtime=None): + simulation = _planned_simulation(SimulationRole.BASELINE) + seams = {"store": store, "artifacts": artifacts, "runtime": runtime} + if calculator is not None: + seams["calculator"] = calculator + return run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + **seams, + ) + + +def test_succeeded_child_returns_its_stored_descriptor_without_recalculating(): + store = FakeStore() + artifacts = FakeArtifacts() + simulation = _planned_simulation(SimulationRole.BASELINE) + descriptor = artifacts.add_simulation(simulation) + succeeded = _child(simulation).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": descriptor.artifact.uri, + "output_sha256": descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": descriptor.row_identity.identifier_columns, + "row_count": descriptor.row_identity.row_count, + "row_identity_sha256": descriptor.row_identity.identity_sha256, + "completed_at": NOW, + } + ) + store.children[simulation.simulation_execution_id] = succeeded + calculations = [] + + result = _run(store, artifacts, calculator=calculations.append) + + assert calculations == [] + assert artifacts.input_writes == [] + assert store.children[simulation.simulation_execution_id] == succeeded + assert result == simulation_module.descriptor_from_record( + succeeded, simulation + ).model_dump(mode="json") + + +def test_default_calculator_receives_the_runtime_and_times_itself(monkeypatch): + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + runtime = SpanRecordingRuntime() + calls = [] + + def calculate(planned, *, runtime=None, progress=None): + calls.append((planned.role, runtime, progress)) + return _frames() + + monkeypatch.setattr(simulation_module, "calculate_simulation_frames", calculate) + + _run(store, FakeArtifacts(), calculator=calculate, runtime=runtime) + + [(role, received_runtime, received_progress)] = calls + assert role is SimulationRole.BASELINE + assert received_runtime is runtime + # The worker's own stage tracker, so a failure names the calculation stage. + assert isinstance(received_progress, StageProgress) + assert received_progress.process_role == "baseline" + assert runtime.spans == [ + "stage12_input_write", + "stage12_simulation_artifact_write", + ] + + +def test_injected_calculator_is_timed_by_the_worker(): + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + runtime = SpanRecordingRuntime() + calls = [] + + def calculate(planned): + calls.append(planned.role) + return _frames() + + _run(store, FakeArtifacts(), calculator=calculate, runtime=runtime) + + assert calls == [SimulationRole.BASELINE] + assert runtime.spans == [ + "stage12_input_write", + "stage12_calculation", + "stage12_simulation_artifact_write", + ] From 7f6d5e97ba886ec4dbc55b91960845cdb782a1df Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:45:26 +0400 Subject: [PATCH 40/66] Pass the single-simulation worker its seams as one dependencies group The Modal run_single_simulation_us and _uk functions keep their signatures. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 6 +- .../stage12_runtime/__init__.py | 2 + .../stage12_runtime/simulation.py | 32 +++++--- .../tests/test_stage12_runtime.py | 81 ++++++++++++------- .../tests/test_stage12_simulation_worker.py | 3 +- 5 files changed, 81 insertions(+), 43 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 080cffa8c..4923f6882 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -207,6 +207,7 @@ def run_single_simulation_us( ) -> dict: """Calculate one planned US simulation and return its artifact descriptor.""" from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + SimulationWorkerDependencies, run_single_simulation, ) @@ -230,7 +231,7 @@ def run_single_simulation_us( payload, context, required_country="us", - runtime=runtime, + dependencies=SimulationWorkerDependencies(runtime=runtime), ) finally: runtime.shutdown() @@ -253,6 +254,7 @@ def run_single_simulation_uk( ) -> dict: """Calculate one planned UK simulation and return its artifact descriptor.""" from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + SimulationWorkerDependencies, run_single_simulation, ) @@ -276,7 +278,7 @@ def run_single_simulation_uk( payload, context, required_country="uk", - runtime=runtime, + dependencies=SimulationWorkerDependencies(runtime=runtime), ) finally: runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index d8c19fad4..a6eda98e9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -12,6 +12,7 @@ ) from .simulation import ( SimulationCalculation, + SimulationWorkerDependencies, calculate_simulation_frames, run_single_simulation, simulation_input_sha256, @@ -26,6 +27,7 @@ "CoordinatorInvocation", "RetainedSimulationOutput", "SimulationCalculation", + "SimulationWorkerDependencies", "build_aggregate_report", "calculate_simulation_frames", "coordinate_report", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 28a721055..d494aea6d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -283,18 +283,29 @@ def descriptor_from_record( ) +@dataclass(frozen=True) +class SimulationWorkerDependencies: + """Seams the single-simulation worker uses; unset seams use the defaults. + + The defaults are the Stage 12 database store, the private artifact bucket, + the standard country-package calculation and no telemetry runtime. + """ + + store: ComparisonStore | None = None + artifacts: Stage12ArtifactStore | None = None + calculator: Callable[ + [PlannedSimulationExecutionInput], + Mapping[str, pd.DataFrame] | SimulationCalculation, + ] = calculate_simulation_frames + runtime: ObservabilityRuntime | None = None + + def run_single_simulation( payload: object, context_payload: object, *, required_country: CountryId, - store: ComparisonStore | None = None, - artifacts: Stage12ArtifactStore | None = None, - calculator: Callable[ - [PlannedSimulationExecutionInput], - Mapping[str, pd.DataFrame] | SimulationCalculation, - ] = calculate_simulation_frames, - runtime: ObservabilityRuntime | None = None, + dependencies: SimulationWorkerDependencies | None = None, ) -> dict[str, Any]: """Calculate one planned simulation, store its artifact, and return its descriptor. @@ -302,11 +313,14 @@ def run_single_simulation( recalculating. A failure after the run is marked running marks it FAILED and raises a ``RuntimeError`` that carries no calculation detail. """ + dependencies = dependencies or SimulationWorkerDependencies() + calculator = dependencies.calculator + runtime = dependencies.runtime simulation = PlannedSimulationExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) - persistence = store or runtime_store(runtime) - artifact_storage = artifacts or artifact_store() + persistence = dependencies.store or runtime_store(runtime) + artifact_storage = dependencies.artifacts or artifact_store() child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: return descriptor_from_record(child, simulation).model_dump(mode="json") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 98e19f272..a0208a65b 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -78,6 +78,7 @@ CoordinatorInvocation, RetainedSimulationOutput, SimulationCalculation, + SimulationWorkerDependencies, _build_spm_result, build_aggregate_report, database_retry, @@ -683,9 +684,11 @@ def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=artifacts, - calculator=lambda _: _frames(), + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: _frames(), + ), ) assert result["role"] == "baseline" @@ -708,11 +711,13 @@ def test_single_worker_retains_detached_calculation_provenance() -> None: simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=artifacts, - calculator=lambda _: SimulationCalculation( - frames=_frames(), - calculation_provenance=provenance, + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: SimulationCalculation( + frames=_frames(), + calculation_provenance=provenance, + ), ), ) @@ -738,9 +743,11 @@ def fail(_simulation): simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=FakeArtifacts(), - calculator=fail, + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + ), ) assert "country package calculation failed" not in str(error.value) @@ -768,10 +775,12 @@ def fail(_simulation): simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=FakeArtifacts(), - calculator=fail, - runtime=runtime, + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=runtime, + ), ) child = store.children[simulation.simulation_execution_id] @@ -799,9 +808,11 @@ def write_simulation(self, **_kwargs): simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=FailingSimulationArtifacts(), - calculator=lambda _: _frames(), + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FailingSimulationArtifacts(), + calculator=lambda _: _frames(), + ), ) assert store.children[simulation.simulation_execution_id].error_summary == ( @@ -848,9 +859,11 @@ def record_store_runtime(runtime): simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - artifacts=FakeArtifacts(), - calculator=lambda _: _frames(), - runtime=runtime, + dependencies=SimulationWorkerDependencies( + artifacts=FakeArtifacts(), + calculator=lambda _: _frames(), + runtime=runtime, + ), ) # The runtime store reports each connection retry through this runtime. @@ -880,10 +893,12 @@ def fail(_simulation): simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=FakeArtifacts(), - calculator=fail, - runtime=TelemetryRecordingRuntime(ordered_events), + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=TelemetryRecordingRuntime(ordered_events), + ), ) original_report, failure_write, persistence_report = ordered_events @@ -920,9 +935,11 @@ def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> N simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: frames, + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: frames, + ), ) assert store.children[simulation.simulation_execution_id].status is ( @@ -1266,9 +1283,11 @@ def test_single_worker_rejects_combined_baseline_and_reform_input() -> None: payload, _context().model_dump(mode="json"), required_country="us", - store=FakeStore(), - artifacts=FakeArtifacts(), - calculator=lambda _: _frames(), + dependencies=SimulationWorkerDependencies( + store=FakeStore(), + artifacts=FakeArtifacts(), + calculator=lambda _: _frames(), + ), ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py index c3dd5da40..a326488c6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py @@ -9,6 +9,7 @@ SimulationRole, ) from policyengine_simulation_executor.stage12_runtime import ( + SimulationWorkerDependencies, run_single_simulation, ) from policyengine_simulation_executor.stage12_runtime import ( @@ -46,7 +47,7 @@ def _run(store, artifacts, *, calculator=None, runtime=None): simulation.model_dump(mode="json"), _context().model_dump(mode="json"), required_country="us", - **seams, + dependencies=SimulationWorkerDependencies(**seams), ) From a5e37334ab9ff9025004443acaa6a3dd1b944ea1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:47:05 +0400 Subject: [PATCH 41/66] Pin Stage 12 worker validation results, order and failures Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_worker_validation.py | 109 ++++++++++++++++++ 1 file changed, 109 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index 5dbbd88ea..f4748dfcc 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -226,3 +226,112 @@ def test_uk_dataset_check_admits_missing_codes_only_when_area_regions_are_routed routed["value"] = True _check_uk_local_authority_dataset(national) + + +def _validate(country="us", expected_digest=None, **seams): + seams.setdefault("environment", _environment()) + seams.setdefault("dataset_path_resolver", lambda _: f"/installed/{country}.h5") + seams.setdefault("dataset_check", lambda *_: None) + seams.setdefault("calculation_check", lambda _: None) + return validate_country_worker( + country=country, + expected_bundle_manifest_sha256=( + expected_digest or load_stage12_bundle().bundle_manifest_sha256 + ), + **seams, + ) + + +@pytest.mark.parametrize("country", ["us", "uk"]) +def test_validation_reports_the_complete_installed_bundle(country) -> None: + resolved = load_stage12_bundle() + country_bundle = next( + item for item in resolved.bundle.countries if item.country == country + ) + events = [] + + result = _validate( + country, + dataset_check=lambda path, _digest: events.append(("dataset", path)), + calculation_check=lambda name: events.append(("calculation", name)), + local_authority_resource_check=lambda path: events.append(("uk", path)), + ) + + assert result == { + "validated": True, + "country": country, + "policyengine_version": resolved.bundle.policyengine_version, + "country_package_name": country_bundle.country_package_name, + "country_package_version": country_bundle.country_package_version, + "dataset_identity": country_bundle.default_dataset, + "dataset_uri": country_bundle.default_dataset_uri, + "data_artifact_revision": country_bundle.data_artifact_revision, + "bundle_manifest_sha256": resolved.bundle_manifest_sha256, + } + expected_events = [("dataset", f"/installed/{country}.h5")] + if country == "uk": + expected_events.append(("uk", f"/installed/{country}.h5")) + expected_events.append(("calculation", country)) + assert events == expected_events + + +def test_validation_rejects_a_country_the_bundle_does_not_serve(monkeypatch) -> None: + from policyengine_simulation_executor import stage12_worker_validation + + resolved = load_stage12_bundle() + us_only = resolved.model_copy( + update={ + "bundle": resolved.bundle.model_copy( + update={ + "countries": tuple( + item + for item in resolved.bundle.countries + if item.country == "us" + ) + } + ) + } + ) + monkeypatch.setattr( + stage12_worker_validation, "load_stage12_bundle", lambda: us_only + ) + + with pytest.raises(RuntimeError) as error: + _validate("uk") + + assert str(error.value) == "installed PolicyEngine.py bundle does not support 'uk'" + assert error.value.__cause__ is None + assert error.value.__suppress_context__ is True + + +def test_validation_uses_the_process_environment_by_default(monkeypatch) -> None: + for name, value in _environment().items(): + monkeypatch.setenv(name, value) + + assert _validate(environment=None)["validated"] is True + + monkeypatch.delenv("STAGE12_ARTIFACT_BUCKET") + with pytest.raises( + RuntimeError, + match=r"^Stage 12 worker is missing required secret values: " + r"STAGE12_ARTIFACT_BUCKET$", + ): + _validate(environment=None) + + +def test_digest_mismatch_is_checked_before_any_dataset_access() -> None: + events = [] + + with pytest.raises( + RuntimeError, + match=( + r"^installed PolicyEngine\.py bundle digest differs " + r"from deployment input$" + ), + ): + _validate( + expected_digest="0" * 64, + dataset_path_resolver=lambda _: events.append("resolve"), + ) + + assert events == [] From 39a0ffd63b9618081753b2448688ce7dc200b044 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:48:01 +0400 Subject: [PATCH 42/66] Pass Stage 12 worker validation its environment and checks as one group Co-Authored-By: Claude Opus 5.5 --- .../stage12_worker_validation.py | 41 +++++++++----- .../tests/test_stage12_worker_validation.py | 55 +++++++++++-------- 2 files changed, 61 insertions(+), 35 deletions(-) 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 8bd9d0926..ab45f14ae 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 @@ -4,6 +4,7 @@ import os from collections.abc import Callable, Mapping +from dataclasses import dataclass from hashlib import sha256 from importlib import import_module from pathlib import Path @@ -99,21 +100,35 @@ def _check_uk_local_authority_dataset(dataset_path: str) -> None: detect_uk_local_authority_metadata_from_hdf(dataset_path) -def validate_country_worker( - *, - country: CountryId, - expected_bundle_manifest_sha256: str, - environment: Mapping[str, str] | None = None, +@dataclass(frozen=True) +class WorkerValidationSeams: + """What worker validation reads and runs; the defaults inspect this worker. + + ``environment`` defaults to the process environment. The checks default + to the installed dataset path, a streamed digest of that dataset, a + non-serving household calculation and the UK local-authority resources. + """ + + environment: Mapping[str, str] | None = None dataset_path_resolver: Callable[[CountryId], str | None] = ( _resolve_installed_dataset_path - ), - dataset_check: Callable[[str, str], None] = _check_dataset_access, - calculation_check: Callable[[CountryId], None] = _run_non_serving_calculation, + ) + dataset_check: Callable[[str, str], None] = _check_dataset_access + calculation_check: Callable[[CountryId], None] = _run_non_serving_calculation local_authority_resource_check: Callable[[str], None] = ( _check_uk_local_authority_dataset - ), + ) + + +def validate_country_worker( + *, + country: CountryId, + expected_bundle_manifest_sha256: str, + seams: WorkerValidationSeams | None = None, ) -> dict[str, Any]: """Validate imports, installed bundle, secrets, data, and calculation.""" + seams = seams or WorkerValidationSeams() + environment = seams.environment runtime_environment = environment if environment is not None else os.environ _require_secret_environment(runtime_environment) # Import the same typed persistence package and coordinator modules loaded @@ -132,7 +147,7 @@ def validate_country_worker( msg = f"installed PolicyEngine.py bundle does not support {country!r}" raise RuntimeError(msg) from None import_module(f"policyengine.tax_benefit_models.{country}") - installed_dataset_path = dataset_path_resolver(country) + installed_dataset_path = seams.dataset_path_resolver(country) if installed_dataset_path is None: msg = f"installed certified dataset is unavailable for {country!r}" raise RuntimeError(msg) @@ -141,10 +156,10 @@ def validate_country_worker( for dataset in country_bundle.datasets if dataset.identity == country_bundle.default_dataset ) - dataset_check(installed_dataset_path, selected_dataset.sha256) + seams.dataset_check(installed_dataset_path, selected_dataset.sha256) if country == "uk": - local_authority_resource_check(installed_dataset_path) - calculation_check(country) + seams.local_authority_resource_check(installed_dataset_path) + seams.calculation_check(country) return { "validated": True, "country": country, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index f4748dfcc..1319107ff 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -8,6 +8,7 @@ import pytest from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle from policyengine_simulation_executor.stage12_worker_validation import ( + WorkerValidationSeams, _run_non_serving_calculation, validate_country_worker, ) @@ -30,10 +31,12 @@ def test_validation_checks_dataset_and_non_serving_calculation() -> None: result = validate_country_worker( country="us", expected_bundle_manifest_sha256=resolved.bundle_manifest_sha256, - environment=_environment(), - dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", - dataset_check=lambda path, digest: datasets.append((path, digest)), - calculation_check=countries.append, + seams=WorkerValidationSeams( + environment=_environment(), + dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", + dataset_check=lambda path, digest: datasets.append((path, digest)), + calculation_check=countries.append, + ), ) assert result["validated"] is True @@ -82,11 +85,13 @@ def test_uk_validation_loads_packaged_local_authority_resources() -> None: result = validate_country_worker( country="uk", expected_bundle_manifest_sha256=resolved.bundle_manifest_sha256, - environment=_environment(), - dataset_path_resolver=lambda _: "/installed/enhanced_frs_2024_25.h5", - dataset_check=lambda *_: None, - calculation_check=lambda _: None, - local_authority_resource_check=checks.append, + seams=WorkerValidationSeams( + environment=_environment(), + dataset_path_resolver=lambda _: "/installed/enhanced_frs_2024_25.h5", + dataset_check=lambda *_: None, + calculation_check=lambda _: None, + local_authority_resource_check=checks.append, + ), ) assert result["validated"] is True @@ -99,10 +104,12 @@ def test_validation_rejects_digest_mismatch_before_dataset_access() -> None: validate_country_worker( country="us", expected_bundle_manifest_sha256="0" * 64, - environment=_environment(), - dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", - dataset_check=lambda path, _: accessed.append(path), - calculation_check=lambda _: None, + seams=WorkerValidationSeams( + environment=_environment(), + dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", + dataset_check=lambda path, _: accessed.append(path), + calculation_check=lambda _: None, + ), ) assert accessed == [] @@ -138,10 +145,12 @@ def test_validation_rejects_missing_secret_values(environment) -> None: validate_country_worker( country="us", expected_bundle_manifest_sha256=resolved.bundle_manifest_sha256, - environment=environment, - dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", - dataset_check=lambda *_: None, - calculation_check=lambda _: None, + seams=WorkerValidationSeams( + environment=environment, + dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", + dataset_check=lambda *_: None, + calculation_check=lambda _: None, + ), ) @@ -151,10 +160,12 @@ def test_validation_rejects_missing_installed_dataset() -> None: validate_country_worker( country="us", expected_bundle_manifest_sha256=resolved.bundle_manifest_sha256, - environment=_environment(), - dataset_path_resolver=lambda _: None, - dataset_check=lambda *_: None, - calculation_check=lambda _: None, + seams=WorkerValidationSeams( + environment=_environment(), + dataset_path_resolver=lambda _: None, + dataset_check=lambda *_: None, + calculation_check=lambda _: None, + ), ) @@ -238,7 +249,7 @@ def _validate(country="us", expected_digest=None, **seams): expected_bundle_manifest_sha256=( expected_digest or load_stage12_bundle().bundle_manifest_sha256 ), - **seams, + seams=WorkerValidationSeams(**seams), ) From bfda678205cf315c32647e5a2bedbeb645c6308d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:49:09 +0400 Subject: [PATCH 43/66] Pin routing state validation messages and check order Co-Authored-By: Claude Opus 5.5 --- .../tests/test_update_version_registry.py | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index 42b093598..9e5dcd48e 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import re import sys import pytest @@ -2052,3 +2053,100 @@ def test_pending_plan_can_be_inspected_locally_before_upload( _main(monkeypatch, "--print-pending-plan") assert json.loads(capsys.readouterr().out) == plan assert store.snapshot() == saved + + +def _routing_state(**overrides): + state = { + "schema_version": 1, + "generation": "4.19.1:app-4-19-1", + "latest": {"policyengine": "4.19.1", "us": "1.687.0", "uk": "2.88.14"}, + "routes": { + "policyengine": {"4.19.1": "app-4-19-1"}, + "us": {"1.687.0": "app-4-19-1"}, + "uk": {"2.88.14": "app-4-19-1"}, + }, + "bundles": { + "4.19.1": {"app_name": "app-4-19-1", "policyengine_version": "4.19.1"} + }, + } + for path, value in overrides.items(): + section, _, key = path.partition(".") + if key: + state[section] = {**state[section], key: value} + else: + state[section] = value + return state + + +@pytest.mark.parametrize( + ("overrides", "required", "message"), + [ + ({"schema_version": 2}, (), "Routing state schema_version must be 1"), + ({"latest.us": ""}, (), "Routing state latest.us is missing"), + ( + {"latest.uk": "9.9.9"}, + (), + "Routing state latest.uk='9.9.9' has no route", + ), + ( + {"latest.policyengine": None}, + (), + "Routing state latest.policyengine is missing", + ), + ( + {"routes.policyengine": {}, "latest.policyengine": None}, + ("4.19.1",), + "Routing state latest.policyengine is missing", + ), + ( + {"bundles": {}}, + ("4.19.1",), + "Routing state policyengine route '4.19.1' has no bundle manifest", + ), + ( + { + "bundles": { + "4.19.1": {"app_name": "app-4-19-1", "policyengine_version": "4"} + } + }, + (), + "Routing state bundle manifest key does not match policyengine_version " + "for '4.19.1'", + ), + ( + { + "bundles": { + "4.19.1": {"app_name": "other", "policyengine_version": "4.19.1"} + } + }, + (), + "Routing state bundle manifest app_name does not match route for '4.19.1'", + ), + ( + {"latest.us": "", "bundles": {}}, + ("4.19.1",), + "Routing state latest.us is missing", + ), + ], +) +def test_validate_routing_state_reports_the_first_inconsistency( + overrides, required, message +): + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + registry.validate_routing_state( + _routing_state(**overrides), + required_policyengine_manifests=required, + ) + + +@pytest.mark.parametrize( + "overrides", + [ + {}, + {"bundles": {}}, + {"bundles": {"4.19.1": "not-a-manifest"}}, + {"routes.policyengine": {}, "latest.policyengine": None, "bundles": {}}, + ], +) +def test_validate_routing_state_accepts_consistent_states(overrides): + registry.validate_routing_state(_routing_state(**overrides)) From 4fd407f7ee9f65b3ca094560b044e3742056ab3f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:49:52 +0400 Subject: [PATCH 44/66] Group a deployed app's versions in the routing publisher and split its validation The update_version_registry command-line options are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../modal/utils/update_version_registry.py | 79 +++++--- .../tests/test_canonical_spm.py | 10 +- .../tests/test_update_version_registry.py | 180 +++++++++++------- 3 files changed, 164 insertions(+), 105 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index 282830def..a2c30b2a1 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -13,6 +13,7 @@ import uuid from collections.abc import Iterable from copy import deepcopy +from dataclasses import dataclass from pathlib import Path from typing import NotRequired, TypedDict @@ -220,8 +221,19 @@ def validate_routing_state( raise ValueError(msg) required_manifests = set(required_policyengine_manifests) + _validate_latest_routes( + state, + require_policyengine=bool( + state["routes"]["policyengine"] or required_manifests + ), + ) + _validate_bundle_manifests(state, required_manifests) + + +def _validate_latest_routes(state: RoutingState, *, require_policyengine: bool) -> None: + """Require each US, UK and (when required) policyengine latest to be routed.""" required_latest_routes = ["us", "uk"] - if state["routes"]["policyengine"] or required_manifests: + if require_policyengine: required_latest_routes.append("policyengine") for route_name in required_latest_routes: @@ -233,6 +245,14 @@ def validate_routing_state( msg = f"Routing state latest.{route_name}={latest!r} has no route" raise ValueError(msg) + +def _validate_bundle_manifests( + state: RoutingState, required_manifests: set[str] +) -> None: + """Require each stored manifest to match its policyengine route. + + A route may lack a manifest unless its version is in ``required_manifests``. + """ for policyengine_version, app_name in state["routes"]["policyengine"].items(): manifest = state["bundles"].get(policyengine_version) if not isinstance(manifest, dict): @@ -271,13 +291,20 @@ def _assert_version_matches_manifest( raise ValueError(msg) +@dataclass(frozen=True) +class DeployedAppVersions: + """A deployed simulation app and the package versions it serves.""" + + app_name: str + policyengine_version: str + us_version: str + uk_version: str + + def build_next_routing_state( *, current_state: object, - app_name: str, - policyengine_version: str, - us_version: str, - uk_version: str, + deployed: DeployedAppVersions, force_latest: bool = False, ) -> RoutingState: """Return ``current_state`` with routes and the bundle manifest for a new app added. @@ -286,6 +313,10 @@ def build_next_routing_state( ``force_latest`` (an intentional rollback). Raises ValueError when the US or UK version differs from the installed bundle's model version. """ + app_name = deployed.app_name + policyengine_version = deployed.policyengine_version + us_version = deployed.us_version + uk_version = deployed.uk_version manifest = build_bundle_manifest_metadata( app_name=app_name, policyengine_version=policyengine_version, @@ -418,10 +449,7 @@ def _json_copy(value: object) -> object: def plan_routing_publication( *, current_state: object, - app_name: str, - policyengine_version: str, - us_version: str, - uk_version: str, + deployed: DeployedAppVersions, force_latest: bool = False, publication_id: str | None = None, ) -> dict: @@ -434,14 +462,13 @@ def plan_routing_publication( route to this release once the plan is applied. """ publication_id = publication_id or uuid.uuid4().hex + app_name = deployed.app_name + policyengine_version = deployed.policyengine_version before = _json_copy(current_state) if current_state is not None else None after = _json_copy( build_next_routing_state( current_state=deepcopy(current_state), - app_name=app_name, - policyengine_version=policyengine_version, - us_version=us_version, - uk_version=uk_version, + deployed=deployed, force_latest=force_latest, ) ) @@ -1023,10 +1050,7 @@ def seed_active_routing_state_from_legacy(*, environment: str) -> RoutingState: def publish_routing_state( *, environment: str, - app_name: str, - policyengine_version: str, - us_version: str, - uk_version: str, + deployed: DeployedAppVersions, force_latest: bool = False, ) -> RoutingState: """Add a deployed app to the active routing state and store it in one write.""" @@ -1039,10 +1063,7 @@ def publish_routing_state( _require_finished_publication(store) next_state = build_next_routing_state( current_state=store.get(ROUTING_STATE_ACTIVE_KEY), - app_name=app_name, - policyengine_version=policyengine_version, - us_version=us_version, - uk_version=uk_version, + deployed=deployed, force_latest=force_latest, ) store[ROUTING_STATE_ACTIVE_KEY] = next_state @@ -1384,14 +1405,17 @@ def active_state() -> object: missing_args = [name for name, value in required_args.items() if not value] if missing_args: parser.error(f"the following arguments are required: {', '.join(missing_args)}") + deployed = DeployedAppVersions( + app_name=args.app_name, + policyengine_version=args.policyengine_version, + us_version=args.us_version, + uk_version=args.uk_version, + ) if args.plan_out or args.dry_run: plan = plan_routing_publication( current_state=active_state(), - app_name=args.app_name, - policyengine_version=args.policyengine_version, - us_version=args.us_version, - uk_version=args.uk_version, + deployed=deployed, force_latest=args.force_latest, publication_id=args.publication_id, ) @@ -1412,10 +1436,7 @@ def active_state() -> object: state = publish_routing_state( environment=args.environment, - app_name=args.app_name, - policyengine_version=args.policyengine_version, - us_version=args.us_version, - uk_version=args.uk_version, + deployed=deployed, force_latest=args.force_latest, ) _report() diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index fa28cd15e..b48134a70 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -931,10 +931,12 @@ def test_versions_http_retains_capability_contract_version(monkeypatch): ) state = registry.build_next_routing_state( current_state=None, - app_name="test-app", - policyengine_version="9.9.9", - us_version="test-us", - uk_version="test-uk", + deployed=registry.DeployedAppVersions( + app_name="test-app", + policyengine_version="9.9.9", + us_version="test-us", + uk_version="test-uk", + ), ) monkeypatch.setattr(endpoints, "_active_routing_state", lambda: state) response = TestClient(create_gateway_app()).get("/versions") diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index 9e5dcd48e..dd29cfd30 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -131,10 +131,12 @@ def test__is_newer_version__does_not_advance_on_equal(): def test_validate_routing_state_accepts_complete_state(): state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) registry.validate_routing_state(state) @@ -170,10 +172,12 @@ def test_bundle_manifest_records_policyengine_6_spm_capability( def test_validate_routing_state_rejects_missing_latest_route(): state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) state["routes"]["us"].pop("1.687.0") @@ -185,10 +189,12 @@ def test_validate_routing_state_rejects_missing_latest_route(): def test_validate_routing_state_allows_migrated_policyengine_route_without_bundle(): state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) state["bundles"].pop("4.19.1") @@ -199,10 +205,12 @@ def test_validate_routing_state_allows_migrated_policyengine_route_without_bundl def test_validate_routing_state_rejects_required_policyengine_route_without_bundle(): state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) state["bundles"].pop("4.19.1") @@ -217,18 +225,22 @@ def test_validate_routing_state_rejects_required_policyengine_route_without_bund def test_build_next_routing_state_preserves_existing_routes(): current_state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-18-0", - policyengine_version="4.18.0", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-18-0", + policyengine_version="4.18.0", + us_version="1.687.0", + uk_version="2.88.14", + ), ) next_state = registry.build_next_routing_state( current_state=current_state, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) assert ( @@ -247,18 +259,22 @@ def test_build_next_routing_state_preserves_existing_routes(): def test_build_next_routing_state_keeps_existing_latest_when_incoming_older(): current_state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) next_state = registry.build_next_routing_state( current_state=current_state, - app_name="policyengine-simulation-py4-18-0", - policyengine_version="4.18.0", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-18-0", + policyengine_version="4.18.0", + us_version="1.687.0", + uk_version="2.88.14", + ), ) assert next_state["latest"]["policyengine"] == "4.19.1" @@ -272,18 +288,22 @@ def test_build_next_routing_state_keeps_existing_latest_when_incoming_older(): def test_build_next_routing_state_force_latest_allows_downgrade(): current_state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) next_state = registry.build_next_routing_state( current_state=current_state, - app_name="policyengine-simulation-py4-18-0", - policyengine_version="4.18.0", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-18-0", + policyengine_version="4.18.0", + us_version="1.687.0", + uk_version="2.88.14", + ), force_latest=True, ) @@ -295,10 +315,12 @@ def test_build_next_routing_state_rejects_country_version_mismatch(): with pytest.raises(ValueError, match="US version"): registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.999.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.999.0", + uk_version="2.88.14", + ), ) @@ -308,10 +330,12 @@ def test_publish_routing_state_writes_only_active_snapshot( ): registry.publish_routing_state( environment="main", - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) assert set(patched_modal) == {"main/simulation-api-routing-state"} @@ -430,10 +454,12 @@ def test_seed_active_routing_state_from_legacy_merges_existing_active( current_state = registry.build_next_routing_state( current_state=None, - app_name="policyengine-simulation-py4-19-1", - policyengine_version="4.19.1", - us_version="1.687.0", - uk_version="2.88.14", + deployed=registry.DeployedAppVersions( + app_name="policyengine-simulation-py4-19-1", + policyengine_version="4.19.1", + us_version="1.687.0", + uk_version="2.88.14", + ), ) patched_modal["main/simulation-api-routing-state"] = FakeDict( {"active": current_state} @@ -559,10 +585,12 @@ def _publish(monkeypatch, state, policyengine_version, us_version, uk_version): _manifests(monkeypatch, {policyengine_version: (us_version, uk_version)}) return registry.build_next_routing_state( current_state=state, - app_name=_app(policyengine_version), - policyengine_version=policyengine_version, - us_version=us_version, - uk_version=uk_version, + deployed=registry.DeployedAppVersions( + app_name=_app(policyengine_version), + policyengine_version=policyengine_version, + us_version=us_version, + uk_version=uk_version, + ), ) @@ -579,10 +607,12 @@ def _plan( _manifests(monkeypatch, {version: (us_version, uk_version)}) return registry.plan_routing_publication( current_state=state, - app_name=_app(version), - policyengine_version=version, - us_version=us_version, - uk_version=uk_version, + deployed=registry.DeployedAppVersions( + app_name=_app(version), + policyengine_version=version, + us_version=us_version, + uk_version=uk_version, + ), force_latest=force, publication_id=publication_id, ) @@ -937,10 +967,12 @@ def test_plans_carry_their_publication_id(monkeypatch): _manifests(monkeypatch, {"6.3.0": ("1.1.0", "2.2.0")}) plan = registry.plan_routing_publication( current_state=state, - app_name=_app("6.3.0"), - policyengine_version="6.3.0", - us_version="1.1.0", - uk_version="2.2.0", + deployed=registry.DeployedAppVersions( + app_name=_app("6.3.0"), + policyengine_version="6.3.0", + us_version="1.1.0", + uk_version="2.2.0", + ), publication_id="run-7-attempt-2-prod", ) assert plan["publication_id"] == "run-7-attempt-2-prod" @@ -1049,10 +1081,12 @@ def test_routing_guard_requires_every_condition(monkeypatch, variable, value): ( "publish_routing_state", { - "app_name": "policyengine-simulation-py6-3-0", - "policyengine_version": "6.3.0", - "us_version": "1.1.0", - "uk_version": "2.2.0", + "deployed": registry.DeployedAppVersions( + app_name="policyengine-simulation-py6-3-0", + policyengine_version="6.3.0", + us_version="1.1.0", + uk_version="2.2.0", + ), }, ), ("seed_active_routing_state_from_legacy", {}), @@ -1406,10 +1440,12 @@ def write_routing(): if writer == "publish": registry.publish_routing_state( environment="staging", - app_name=_app("6.2.0"), - policyengine_version="6.2.0", - us_version="1.2.0", - uk_version="2.3.0", + deployed=registry.DeployedAppVersions( + app_name=_app("6.2.0"), + policyengine_version="6.2.0", + us_version="1.2.0", + uk_version="2.3.0", + ), ) elif writer == "latest": registry.restore_routing_state( From df0591c8de95fef9e7159c0a0f0658ec5e8f7f67 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:52:03 +0400 Subject: [PATCH 45/66] Pin the artifact fetch layer's receipt, credential and download edge cases Co-Authored-By: Claude Opus 5.5 --- .../tests/test_image_setup_fetch.py | 158 ++++++++++++++++++ 1 file changed, 158 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py index 96b546f59..648f7935d 100644 --- a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py +++ b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py @@ -249,6 +249,27 @@ def test_falls_back_to_ambient_adc_without_a_blob( assert patched_google.infos == [] assert patched_google.client_calls == [{}] + @pytest.mark.usefixtures("fake_client", "data_folder") + def test_unwraps_an_outer_quoted_blob(self, monkeypatch, patched_google): + clean = json.dumps({"type": "service_account", "project_id": "pe-test"}) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", json.dumps(clean)) + + fetch_artifacts("test-bucket", _manifest()) + + assert patched_google.infos == [ + {"type": "service_account", "project_id": "pe-test"} + ] + + @pytest.mark.usefixtures("fake_client", "data_folder") + def test_rejects_a_blob_that_is_not_json(self, monkeypatch, patched_google): + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", "not-json") + + with pytest.raises(json.JSONDecodeError): + fetch_artifacts("test-bucket", _manifest()) + + assert patched_google.infos == [] + assert patched_google.client_calls == [] + def test_fetch_never_writes_credentials_to_disk(): """The credential shim must never materialize the key as a file. @@ -282,3 +303,140 @@ def test_image_setup_module_stays_self_contained(): if isinstance(node, ast.Import): for alias in node.names: assert not alias.name.startswith("policyengine_simulation_executor") + + +class TestCharacterization: + """Pin the remaining fetch branches before the function is split up.""" + + @pytest.mark.usefixtures("data_folder") + @pytest.mark.parametrize( + "receipt_text", + [ + "not json", + "[]", + '{"datasets": ["not-an-entry"]}', + '{"datasets": [{"country": "uk", "version": "9.9.9"}]}', + '{"datasets": [{"country": "us"}]}', + ], + ) + def test_unusable_receipts_skip_the_data_version_check_with_a_warning( + self, monkeypatch, tmp_path, fake_client, caplog, receipt_text + ): + import logging + + receipt_file = tmp_path / "bundle-receipt.json" + receipt_file.write_text(receipt_text) + monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_file)) + + with caplog.at_level(logging.INFO): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + assert len(fake_client.downloads) == 2 + warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + ] + assert warnings == [ + "No installed bundle receipt entry for us; " + "skipping the data-version freshness check" + ] + + @pytest.mark.usefixtures("data_folder") + @pytest.mark.parametrize( + ("versions", "passes"), + [(["9.9.9", "1.2.3"], True), (["1.2.3", "9.9.9"], False)], + ) + def test_the_last_receipt_entry_for_the_country_wins( + self, monkeypatch, tmp_path, fake_client, versions, passes + ): + receipt_file = tmp_path / "bundle-receipt.json" + receipt_file.write_text( + json.dumps( + { + "datasets": [ + {"country": "us", "version": version} for version in versions + ] + } + ) + ) + monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_file)) + + if passes: + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert len(fake_client.downloads) == 2 + else: + with pytest.raises(RuntimeError) as error: + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert str(error.value) == ( + "Stale artifact manifest: installed us data version is '9.9.9' " + "but the manifest was computed for '1.2.3'. Re-run the " + "precompute for the current version-set." + ) + + def test_stale_policyengine_version_message(self, monkeypatch, fake_client): + monkeypatch.setenv("POLICYENGINE_VERSION", "4.99.0") + monkeypatch.setenv("POLICYENGINE_US_VERSION", "9.9.9") + + with pytest.raises(RuntimeError) as error: + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + assert str(error.value) == ( + "Stale artifact manifest: the image has POLICYENGINE_VERSION='4.99.0' " + "but the manifest was computed for '4.22.0'. Re-run the precompute " + "for the current version-set." + ) + + @pytest.mark.usefixtures("data_folder") + def test_a_download_that_writes_nothing_fails_while_logging_it(self, fake_client): + class SilentBlob(FakeBlob): + def download_to_filename(self, filename): + self.store.downloads.append((self.path, Path(filename).name)) + + fake_client.bucket = lambda _name: SimpleNamespace( + blob=lambda path: SilentBlob(fake_client, path) + ) + + with pytest.raises(FileNotFoundError): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + assert fake_client.downloads == [ + ("datasets/us/d1/populace_year_2026.h5", "populace_year_2026.h5") + ] + + def test_files_missing_after_every_download_fail_the_build( + self, fake_client, data_folder + ): + class ReplacingBlob(FakeBlob): + def download_to_filename(self, filename): + for existing in data_folder.iterdir(): + existing.unlink() + super().download_to_filename(filename) + + fake_client.bucket = lambda _name: SimpleNamespace( + blob=lambda path: ReplacingBlob(fake_client, path) + ) + + with pytest.raises(RuntimeError) as error: + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + missing = str(data_folder / "populace_year_2026.h5") + assert str(error.value) == ( + f"Artifact fetch did not produce expected files: [{missing!r}]" + ) + + @pytest.mark.usefixtures("data_folder") + def test_logs_each_fetched_artifact(self, fake_client, caplog): + import logging + + with caplog.at_level(logging.INFO): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + fetched = [ + record.getMessage() + for record in caplog.records + if record.getMessage().startswith("Fetched ") + ] + assert len(fetched) == 2 + assert fetched[0].endswith("populace_year_2026.h5 (0.0 MB)") + assert fetched[1].endswith("bl1-aaaa.h5 (0.0 MB)") From ecb85cc5566c0f9b21331b993fe68ce7f85ded80 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:53:03 +0400 Subject: [PATCH 46/66] Split the image artifact fetch into freshness, credential and download steps The helpers stay in the self-contained image setup module with lazy imports. Changing fetch_artifacts's source changes its Modal build-layer definition, so the next simulation image build re-runs the artifact fetch layer once. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/_image_setup.py | 166 ++++++++++++------ 1 file changed, 111 insertions(+), 55 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/_image_setup.py b/projects/policyengine-simulation-executor/src/modal/_image_setup.py index 6019805c2..3abe3521b 100644 --- a/projects/policyengine-simulation-executor/src/modal/_image_setup.py +++ b/projects/policyengine-simulation-executor/src/modal/_image_setup.py @@ -11,7 +11,8 @@ # imports; the annotations that use them are therefore quoted. TYPE_CHECKING = False if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Mapping, Sequence + from pathlib import Path from typing import Any from google.cloud import storage @@ -39,7 +40,6 @@ def fetch_artifacts( ``client`` is a test seam; real builds construct a storage client from the layer's GCP secret (or ambient ADC for local drill runs). """ - import json # noqa: PLC0415 # image-build function: this module has no module-level imports by design import logging # noqa: PLC0415 # image-build function: this module has no module-level imports by design import os # noqa: PLC0415 # image-build function: this module has no module-level imports by design from pathlib import ( # noqa: PLC0415 # image-build function: this module has no module-level imports by design @@ -47,7 +47,6 @@ def fetch_artifacts( ) logging.basicConfig(level=logging.INFO) - logger = logging.getLogger(__name__) if not manifest: msg = ( @@ -58,8 +57,28 @@ def fetch_artifacts( ) raise RuntimeError(msg) - # Freshness gate: refuse to bake artifacts computed for a different - # version-set than the one this image installs. + _require_manifest_matches_image(manifest) + if client is None: + client = _storage_client_from_environment() + + data_folder = Path( + os.environ.get("POLICYENGINE_DATA_FOLDER", "/opt/policyengine/data") + ) + data_folder.mkdir(parents=True, exist_ok=True) + _download_manifest_artifacts( + client.bucket(bucket), manifest["artifacts"], data_folder + ) + + +def _require_manifest_matches_image(manifest: "Mapping[str, Any]") -> None: + """Freshness gate: refuse artifacts computed for another version-set. + + The image's policyengine and country model versions must equal the + manifest receipt's, and so must the installed data version when the + bundle receipt records one. + """ + import os # noqa: PLC0415 # image-build function: this module has no module-level imports by design + country = str(manifest["country"]) receipt = manifest["receipt"] expected_env = { @@ -77,20 +96,12 @@ def fetch_artifacts( ) raise RuntimeError(msg) - installed_data_version = None - receipt_path = os.environ.get("POLICYENGINE_BUNDLE_RECEIPT") - if receipt_path and Path(receipt_path).exists(): - try: - installed = json.loads(Path(receipt_path).read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - installed = None - if isinstance(installed, dict): - for dataset in installed.get("datasets", []): - if isinstance(dataset, dict) and dataset.get("country") == country: - installed_data_version = dataset.get("version") + installed_data_version = _installed_data_version(country) if installed_data_version is None: + import logging # noqa: PLC0415 # image-build function: this module has no module-level imports by design + # Lenient by design: local drill runs have no bundle receipt file. - logger.warning( + logging.getLogger(__name__).warning( "No installed bundle receipt entry for %s; " "skipping the data-version freshness check", country, @@ -104,49 +115,94 @@ def fetch_artifacts( ) raise RuntimeError(msg) - if client is None: - creds_kwargs = {} - blob = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON") - if blob: - # The secret payload is sometimes double-encoded (escaped - # interior quotes, with or without outer quotes) — same - # unwrap as the runtime's _normalize_credentials_blob. - try: - parsed = json.loads(blob) - except json.JSONDecodeError: - if blob.lstrip().startswith('"') or '\\"' in blob: - parsed = json.loads(f'"{blob}"') - else: - raise - info = json.loads(parsed) if isinstance(parsed, str) else parsed - - from google.oauth2 import ( # noqa: PLC0415 # image-build function: this module has no module-level imports by design - service_account, - ) - # Credentials built in memory, never written to disk: this - # layer's container filesystem is committed into the image, - # so a credentials file would ship inside it. - creds_kwargs = { - "credentials": service_account.Credentials.from_service_account_info( - info - ), - "project": info.get("project_id"), - } - # No blob: ambient ADC (the local fault-drill path). - from google.cloud import ( # noqa: PLC0415 # image-build function: this module has no module-level imports by design - storage, - ) +def _installed_data_version(country: str) -> object: + """Return the bundle receipt's data version for ``country``, or None. + + None also covers a missing or unreadable receipt; the last matching + receipt entry wins. + """ + import json # noqa: PLC0415 # image-build function: this module has no module-level imports by design + import os # noqa: PLC0415 # image-build function: this module has no module-level imports by design + from pathlib import ( # noqa: PLC0415 # image-build function: this module has no module-level imports by design + Path, + ) - client = storage.Client(**creds_kwargs) + receipt_path = os.environ.get("POLICYENGINE_BUNDLE_RECEIPT") + if not receipt_path or not Path(receipt_path).exists(): + return None + try: + installed = json.loads(Path(receipt_path).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + installed_data_version = None + if isinstance(installed, dict): + for dataset in installed.get("datasets", []): + if isinstance(dataset, dict) and dataset.get("country") == country: + installed_data_version = dataset.get("version") + return installed_data_version - data_folder = Path( - os.environ.get("POLICYENGINE_DATA_FOLDER", "/opt/policyengine/data") + +def _storage_client_from_environment() -> "storage.Client": + """Build a storage client from the layer's GCP secret, else ambient ADC.""" + import os # noqa: PLC0415 # image-build function: this module has no module-level imports by design + + creds_kwargs = {} + blob = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON") + if blob: + creds_kwargs = _service_account_credentials(blob) + # No blob: ambient ADC (the local fault-drill path). + from google.cloud import ( # noqa: PLC0415 # image-build function: this module has no module-level imports by design + storage, ) - data_folder.mkdir(parents=True, exist_ok=True) - bucket_handle = client.bucket(bucket) - artifacts = manifest["artifacts"] + return storage.Client(**creds_kwargs) + + +def _service_account_credentials(blob: str) -> "dict[str, Any]": + """Return storage client arguments for the service account in ``blob``. + + Credentials are built in memory, never written to disk: this layer's + container filesystem is committed into the image, so a credentials file + would ship inside it. + """ + import json # noqa: PLC0415 # image-build function: this module has no module-level imports by design + + # The secret payload is sometimes double-encoded (escaped interior + # quotes, with or without outer quotes) — same unwrap as the runtime's + # _normalize_credentials_blob. + try: + parsed = json.loads(blob) + except json.JSONDecodeError: + if blob.lstrip().startswith('"') or '\\"' in blob: + parsed = json.loads(f'"{blob}"') + else: + raise + info = json.loads(parsed) if isinstance(parsed, str) else parsed + + from google.oauth2 import ( # noqa: PLC0415 # image-build function: this module has no module-level imports by design + service_account, + ) + + return { + "credentials": service_account.Credentials.from_service_account_info(info), + "project": info.get("project_id"), + } + + +def _download_manifest_artifacts( + bucket_handle: "storage.Bucket", + artifacts: "Sequence[Mapping[str, Any]]", + data_folder: "Path", +) -> None: + """Download every artifact under its runtime filename, failing on any gap. + + Every store object is checked before anything downloads, and every + expected file is checked after. + """ + import logging # noqa: PLC0415 # image-build function: this module has no module-level imports by design + + logger = logging.getLogger(__name__) missing = [ artifact["path"] for artifact in artifacts From 3838e8b2db99a297fcc52f78e1f42d74b5ba551d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:19:45 +0400 Subject: [PATCH 47/66] Pin the release sentinel's polling, retry statuses and problem order Co-Authored-By: Claude Opus 5.5 --- .../tests/test_release_sentinel.py | 190 ++++++++++++++++++ 1 file changed, 190 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_release_sentinel.py b/projects/policyengine-simulation-executor/tests/test_release_sentinel.py index 5e4b4228d..719ce27ed 100644 --- a/projects/policyengine-simulation-executor/tests/test_release_sentinel.py +++ b/projects/policyengine-simulation-executor/tests/test_release_sentinel.py @@ -400,3 +400,193 @@ def test_main_check_exit_codes(tmp_path, monkeypatch, capsys): == 1 ) assert "does not reproduce" in capsys.readouterr().err + + +# --- Characterization: check_deployment's problem order and polling ---------- + + +class ScriptedAPI: + """Answer each call with the next scripted response and record the call.""" + + def __init__(self, *responses): + self.responses = list(responses) + self.calls = [] + + def __call__(self, method, url, *, token=None, body=None): + self.calls.append((method, url, token, body)) + return self.responses.pop(0) + + +SERVED_BUNDLE = { + "policyengine_version": "6.2.6", + "model_version": "2.125.1", + "dataset_sha256": "5" * 64, +} +SUBMITTED = (200, {}, {"job_id": "fc-1", "policyengine_bundle": SERVED_BUNDLE}) +COMPLETE = (200, {}, {"status": "complete", "result": RESULT}) + + +def _check_scripted(api, expected=None, **kwargs): + if expected is None: + expected, _ = _expected() + sleeps = [] + problems = sentinel.check_deployment( + "https://candidate.example/", + expected, + token="t0ken", # noqa: S106 # fake bearer token; the HTTP calls are stubbed + http=api, + sleep=sleeps.append, + **kwargs, + ) + return problems, sleeps + + +def test_check_submits_then_polls_the_job_url_with_the_token(): + api = ScriptedAPI(SUBMITTED, (202, {}, {"status": "submitted"}), COMPLETE) + + problems, sleeps = _check_scripted(api, poll_seconds=7) + + assert problems == [] + assert sleeps == [7] + expected, _ = _expected() + case = expected["cases"][0] + assert api.calls == [ + ( + "POST", + "https://candidate.example/simulate/economy/comparison", + "t0ken", + {**case["request"], "policyengine_version": "6.2.6"}, + ), + ("GET", "https://candidate.example/jobs/fc-1", "t0ken", None), + ("GET", "https://candidate.example/jobs/fc-1", "t0ken", None), + ] + + +@pytest.mark.parametrize( + ("headers", "slept"), + [ + ({"Retry-After": "3"}, 3.0), + ({"retry-after": "4.5"}, 4.5), + ({"RETRY-AFTER": "-2"}, 0.0), + ({"Retry-After": "soon"}, 7), + ({}, 7), + ], + ids=["header", "lowercase", "negative", "unparseable", "missing"], +) +@pytest.mark.parametrize("status", [202, 429, 503]) +def test_check_waits_for_a_pending_job_as_retry_after_says(status, headers, slept): + api = ScriptedAPI(SUBMITTED, (status, headers, {"status": "submitted"}), COMPLETE) + + problems, sleeps = _check_scripted(api, poll_seconds=7) + + assert problems == [] + assert sleeps == [slept] + + +def test_check_reports_identity_problems_before_value_problems(): + served = {"policyengine_version": "6.2.3", "model_version": "2.102.3"} + result = {**RESULT, "budget": {**RESULT["budget"], "budgetary_impact": 6.2e9}} + api = ScriptedAPI( + (200, {}, {"job_id": "fc-1", "policyengine_bundle": served}), + (200, {}, {"status": "complete", "result": result}), + ) + + problems, _ = _check_scripted(api) + + assert problems == [ + f"{LABEL}: served policyengine_version '6.2.3', expected '6.2.6'", + f"{LABEL}: served model_version '2.102.3', expected '2.125.1'", + f"{LABEL}: served dataset_sha256 None, expected '{'5' * 64}'", + f"{LABEL}: budget.budgetary_impact served 6200000000.0, computed 6100000000.0", + ] + + +def test_check_reports_a_submit_without_a_bundle_as_three_identity_problems(): + api = ScriptedAPI((200, {}, {"job_id": "fc-1"}), COMPLETE) + + problems, _ = _check_scripted(api) + + assert problems == [ + f"{LABEL}: served policyengine_version None, expected '6.2.6'", + f"{LABEL}: served model_version None, expected '2.125.1'", + f"{LABEL}: served dataset_sha256 None, expected '{'5' * 64}'", + ] + + +def test_check_reports_a_submit_that_returned_no_object(): + api = ScriptedAPI((200, {}, ["job"])) + + problems, sleeps = _check_scripted(api) + + assert problems == [f"{LABEL}: submit returned HTTP 200: ['job']"] + assert sleeps == [] + assert len(api.calls) == 1 + + +def test_check_reports_a_job_body_that_is_not_an_object(): + api = ScriptedAPI(SUBMITTED, (200, {}, "complete")) + + problems, _ = _check_scripted(api) + + assert problems == [f"{LABEL}: job returned HTTP 200: complete"] + + +def test_check_gives_up_once_the_deadline_has_passed(): + ticks = iter([0, 100, 200, 300]) + pending = (202, {}, {"status": "submitted"}) + api = ScriptedAPI(SUBMITTED, pending, pending, pending) + + problems, sleeps = _check_scripted( + api, timeout_seconds=250, poll_seconds=5, clock=lambda: next(ticks) + ) + + assert problems == [f"{LABEL}: no result within 250s"] + assert sleeps == [5, 5] + assert len(api.calls) == 4 + + +def test_check_still_compares_the_object_result_of_a_failed_job(): + api = ScriptedAPI(SUBMITTED, (200, {}, {"status": "failed", "result": {}})) + + problems, _ = _check_scripted(api) + + assert problems == [ + f"{LABEL}: job ended 'failed'", + *(f"{LABEL}: the result has no {path}" for path in COMPARED), + ] + + +def test_check_continues_with_the_next_case_after_a_failed_one(): + expected, _ = _expected() + second = {**expected["cases"][0], "case_id": "second"} + expected = {**expected, "cases": [expected["cases"][0], second]} + api = ScriptedAPI((500, {}, {"detail": "boom"}), SUBMITTED, COMPLETE) + + problems, _ = _check_scripted(api, expected=expected) + + assert problems == [f"{LABEL}: submit returned HTTP 500: {{'detail': 'boom'}}"] + assert [call[0] for call in api.calls] == ["POST", "POST", "GET"] + + +def test_check_refuses_another_schema_before_any_request(): + api = ScriptedAPI() + + with pytest.raises( + sentinel.SentinelError, match=r"^expected file schema is 'other'$" + ): + sentinel.check_deployment("https://x", {"schema": "other"}, http=api) + + assert api.calls == [] + + +def test_access_token_refuses_a_failed_or_empty_token_response(monkeypatch): + environ = dict.fromkeys(sentinel._AUTH_VARS, "x") + for response in ((401, {}, {"access_token": "minted"}), (200, {}, {})): + monkeypatch.setattr( + sentinel, "_http", lambda *_args, answer=response, **_kwargs: answer + ) + with pytest.raises( + sentinel.SentinelError, + match=rf"^token request failed with HTTP {response[0]}$", + ): + sentinel.access_token(environ) From d99395650feb18b2514ccbb64a2b972fc4a4736b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:20:43 +0400 Subject: [PATCH 48/66] Split the release sentinel's deployment check into submit, poll and compare steps Group the job timeout, poll interval, sleep and clock as one JobPolling, and name the HTTP statuses the sentinel compares with. The check command's options are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/utils/release_sentinel.py | 217 ++++++++++++------ .../tests/test_release_sentinel.py | 10 +- 2 files changed, 151 insertions(+), 76 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py index d7f00bfeb..51e54aef2 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py @@ -32,6 +32,8 @@ import urllib.error import urllib.request from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from http import HTTPStatus from importlib import metadata from pathlib import Path from typing import Any @@ -67,6 +69,11 @@ }, ) +#: Job poll statuses that mean the job has not finished yet: poll again. +_PENDING_JOB_STATUSES = frozenset( + {HTTPStatus.ACCEPTED, HTTPStatus.TOO_MANY_REQUESTS, HTTPStatus.SERVICE_UNAVAILABLE} +) + _AUTH_VARS = ( "SIMULATION_TEST_AUTH_ISSUER", "SIMULATION_TEST_AUTH_AUDIENCE", @@ -228,7 +235,7 @@ def access_token(environ: Mapping[str, str] = os.environ) -> str | None: }, ) token = payload.get("access_token") if isinstance(payload, Mapping) else None - if status != 200 or not token: + if status != HTTPStatus.OK or not token: msg = f"token request failed with HTTP {status}" raise SentinelError(msg) return str(token) @@ -244,86 +251,156 @@ def _retry_after(headers: Mapping[str, str], fallback: float) -> float: return fallback +HttpCall = Callable[..., tuple[int, Mapping[str, str], Any]] + + +@dataclass(frozen=True) +class JobPolling: + """How long to wait for each job, and the clock and sleep that measure it. + + ``poll_seconds`` is the wait between polls when the API sends no usable + ``Retry-After`` header. + """ + + timeout_seconds: float = 1800 + poll_seconds: float = 10 + sleep: Callable[[float], None] = time.sleep + clock: Callable[[], float] = time.monotonic + + +@dataclass +class _JobOutcome: + """What polling one job found: its problems, its result and whether it completed.""" + + problems: list[str] = field(default_factory=list) + result: object = None + completed: bool = False + + +@dataclass(frozen=True) +class _Deployment: + """The deployment under check: its URL, how to call it and how to poll it.""" + + base_url: str + token: str | None + http: HttpCall + polling: JobPolling + + def check_case( + self, case: Mapping[str, Any], expected: Mapping[str, Any] + ) -> list[str]: + """Submit one case and return its identity, job and value problems.""" + label = case["case_id"] + version = expected["policyengine_version"] + body = {**case["request"], "policyengine_version": version} + status, _, submitted = self.http( + "POST", + f"{self.base_url}/simulate/economy/comparison", + token=self.token, + body=body, + ) + if status != HTTPStatus.OK or not isinstance(submitted, Mapping): + return [f"{label}: submit returned HTTP {status}: {submitted}"] + problems = _identity_problems( + label, + submitted.get("policyengine_bundle") or {}, + version=version, + country_identity=expected[case["request"]["country"]], + ) + job = self._await_job(label, submitted["job_id"]) + problems += job.problems + if not isinstance(job.result, Mapping): + # A complete job must carry an object result; anything else + # (missing, null, a list or a string) compares nothing. + if job.completed: + problems.append( + f"{label}: complete job returned a " + f"{type(job.result).__name__} result, not an object" + ) + return problems + return problems + _value_problems(label, case["values"], job.result) + + def _await_job(self, label: str, job_id: str) -> _JobOutcome: + """Poll the job until it finishes, fails or runs out of time.""" + deadline = self.polling.clock() + self.polling.timeout_seconds + while True: + status, headers, job = self.http( + "GET", f"{self.base_url}/jobs/{job_id}", token=self.token + ) + if status == HTTPStatus.OK and isinstance(job, Mapping): + return _finished_job(label, job) + if status not in _PENDING_JOB_STATUSES: + return _JobOutcome([f"{label}: job returned HTTP {status}: {job}"]) + if self.polling.clock() >= deadline: + return _JobOutcome( + [f"{label}: no result within {self.polling.timeout_seconds:.0f}s"] + ) + self.polling.sleep(_retry_after(headers, self.polling.poll_seconds)) + + +def _finished_job(label: str, job: Mapping[str, Any]) -> _JobOutcome: + """Describe a finished job; one that did not complete is a problem.""" + completed = job.get("status") == "complete" + problems = [] if completed else [f"{label}: job ended {job.get('status')!r}"] + return _JobOutcome(problems, job.get("result"), completed) + + +def _identity_problems( + label: str, + served: Mapping[str, Any], + *, + version: str, + country_identity: Mapping[str, Any], +) -> list[str]: + """Compare the served wrapper version, country model and dataset digest.""" + return [ + f"{label}: served {name} {served.get(name)!r}, expected {wanted!r}" + for name, wanted in ( + ("policyengine_version", version), + ("model_version", country_identity["model_version"]), + ("dataset_sha256", country_identity["dataset_sha256"]), + ) + if served.get(name) != wanted + ] + + +def _value_problems( + label: str, values: Mapping[str, float], result: Mapping[str, Any] +) -> list[str]: + """Compare each computed value with the one the deployment served.""" + problems = [] + for path, wanted in values.items(): + try: + actual = value_at(result, path) + except SentinelError as exc: + problems.append(f"{label}: {exc}") + continue + if not agree(wanted, actual): + problems.append(f"{label}: {path} served {actual!r}, computed {wanted!r}") + return problems + + def check_deployment( base_url: str, expected: Mapping[str, Any], *, token: str | None = None, - timeout_seconds: float = 1800, - poll_seconds: float = 10, - http: Callable[..., tuple[int, Mapping[str, str], Any]] | None = None, - sleep: Callable[[float], None] = time.sleep, - clock: Callable[[], float] = time.monotonic, + http: HttpCall | None = None, + polling: JobPolling | None = None, ) -> list[str]: """Submit each expected case to ``base_url`` and return what disagrees.""" if expected.get("schema") != SCHEMA: msg = f"expected file schema is {expected.get('schema')!r}" raise SentinelError(msg) - http = http or _http - base_url = base_url.rstrip("/") - version = expected["policyengine_version"] + deployment = _Deployment( + base_url=base_url.rstrip("/"), + token=token, + http=http or _http, + polling=polling or JobPolling(), + ) problems: list[str] = [] for case in expected["cases"]: - label = case["case_id"] - body = {**case["request"], "policyengine_version": version} - status, _, submitted = http( - "POST", f"{base_url}/simulate/economy/comparison", token=token, body=body - ) - if status != 200 or not isinstance(submitted, Mapping): - problems.append(f"{label}: submit returned HTTP {status}: {submitted}") - continue - served = submitted.get("policyengine_bundle") or {} - want = expected[case["request"]["country"]] - for field, wanted in ( - ("policyengine_version", version), - ("model_version", want["model_version"]), - ("dataset_sha256", want["dataset_sha256"]), - ): - if served.get(field) != wanted: - problems.append( - f"{label}: served {field} {served.get(field)!r}, " - f"expected {wanted!r}" - ) - deadline = clock() + timeout_seconds - result = None - completed = False - while True: - status, headers, job = http( - "GET", f"{base_url}/jobs/{submitted['job_id']}", token=token - ) - if status == 200 and isinstance(job, Mapping): - if job.get("status") != "complete": - problems.append(f"{label}: job ended {job.get('status')!r}") - else: - completed = True - result = job.get("result") - break - if status not in (202, 429, 503): - problems.append(f"{label}: job returned HTTP {status}: {job}") - break - if clock() >= deadline: - problems.append(f"{label}: no result within {timeout_seconds:.0f}s") - break - sleep(_retry_after(headers, poll_seconds)) - if not isinstance(result, Mapping): - # A complete job must carry an object result; anything else - # (missing, null, a list or a string) compares nothing. - if completed: - problems.append( - f"{label}: complete job returned a {type(result).__name__} " - "result, not an object" - ) - continue - for path, wanted in case["values"].items(): - try: - actual = value_at(result, path) - except SentinelError as exc: - problems.append(f"{label}: {exc}") - continue - if not agree(wanted, actual): - problems.append( - f"{label}: {path} served {actual!r}, computed {wanted!r}" - ) + problems += deployment.check_case(case, expected) return problems @@ -358,7 +435,7 @@ def main(argv: list[str] | None = None) -> int: args.base_url, expected, token=access_token(), - timeout_seconds=args.timeout_seconds, + polling=JobPolling(timeout_seconds=args.timeout_seconds), ) if problems: _report("The deployed API does not reproduce the bundle's numbers:", error=True) diff --git a/projects/policyengine-simulation-executor/tests/test_release_sentinel.py b/projects/policyengine-simulation-executor/tests/test_release_sentinel.py index 719ce27ed..47f40c203 100644 --- a/projects/policyengine-simulation-executor/tests/test_release_sentinel.py +++ b/projects/policyengine-simulation-executor/tests/test_release_sentinel.py @@ -181,7 +181,7 @@ def __call__(self, method, url, *, token=None, body=None): return 200, {}, {"status": "complete", "result": self.result} -def _check(api, **kwargs): +def _check(api, **polling): expected, _ = _expected() sleeps = [] problems = sentinel.check_deployment( @@ -189,8 +189,7 @@ def _check(api, **kwargs): expected, token="t0ken", # noqa: S106 # fake bearer token; the HTTP calls are stubbed http=api, - sleep=sleeps.append, - **kwargs, + polling=sentinel.JobPolling(sleep=sleeps.append, **polling), ) return problems, sleeps @@ -426,7 +425,7 @@ def __call__(self, method, url, *, token=None, body=None): COMPLETE = (200, {}, {"status": "complete", "result": RESULT}) -def _check_scripted(api, expected=None, **kwargs): +def _check_scripted(api, expected=None, **polling): if expected is None: expected, _ = _expected() sleeps = [] @@ -435,8 +434,7 @@ def _check_scripted(api, expected=None, **kwargs): expected, token="t0ken", # noqa: S106 # fake bearer token; the HTTP calls are stubbed http=api, - sleep=sleeps.append, - **kwargs, + polling=sentinel.JobPolling(sleep=sleeps.append, **polling), ) return problems, sleeps From a6120fee21ca41cf9b875966a0c54b8a6f6b0b6b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:23:03 +0400 Subject: [PATCH 49/66] Pin the routing publisher CLI's mode order, outputs and argument errors Also pin how a publication id names its owning run, and the restore refusals for no plans and a missing routing store. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_update_version_registry.py | 395 ++++++++++++++++++ 1 file changed, 395 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index dd29cfd30..1e973bcc0 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -5,6 +5,7 @@ import json import re import sys +from types import SimpleNamespace import pytest from hypothesis import given, settings @@ -2186,3 +2187,397 @@ def test_validate_routing_state_reports_the_first_inconsistency( ) def test_validate_routing_state_accepts_consistent_states(overrides): registry.validate_routing_state(_routing_state(**overrides)) + + +# --- Characterization: the CLI's mode order, outputs and argument errors ----- + +PUBLISHING_ARGUMENTS = ( + "--app-name", + "policyengine-simulation-py6-3-0", + "--policyengine-version", + "6.3.0", + "--us-version", + "1.1.0", + "--uk-version", + "2.2.0", +) +DEPLOYED_6_3_0 = registry.DeployedAppVersions( + app_name="policyengine-simulation-py6-3-0", + policyengine_version="6.3.0", + us_version="1.1.0", + uk_version="2.2.0", +) +RECORDED_STATE = { + "latest": {"policyengine": "6.2.6", "us": "1.1.0", "uk": "2.1.0"}, + "routes": { + "policyengine": {"6.2.6": "app-6-2-6", "6.3.0": "app-6-3-0"}, + "us": {"1.1.0": "app-6-2-6"}, + "uk": {"2.1.0": "app-6-2-6", "2.2.0": "app-6-3-0", "2.0.0": "app-6-0-0"}, + }, + "bundles": {"6.2.6": {}}, +} + + +@pytest.fixture +def recorded_writers(monkeypatch, patched_modal, tmp_path): # noqa: ARG001 # patched_modal keeps every unrecorded Modal read empty + """Replace each routing operation with a recorder, and write one plan file.""" + calls = [] + plan = { + "publication_id": "p-1", + "policyengine_version": "6.3.0", + "after_sha256": "a" * 64, + } + plan_file = tmp_path / "plan.json" + plan_file.write_text(json.dumps(plan)) + + def recorder(name, result): + def record(**kwargs): + calls.append((name, kwargs)) + return result + + return record + + for attribute, name, result in ( + ("finalize_routing_publication", "finalize", "restored"), + ("record_routing_plan", "record", None), + ("complete_routing_plan", "complete", None), + ("apply_routing_plan", "apply", None), + ("restore_routing_plans", "restore-plans", "restored"), + ("restore_routing_state", "restore-latest", RECORDED_STATE), + ("seed_active_routing_state_from_legacy", "seed", RECORDED_STATE), + ( + "plan_routing_publication", + "plan", + { + "before_sha256": "b" * 64, + "after_sha256": "a" * 64, + "advances_latest": False, + }, + ), + ("publish_routing_state", "publish", RECORDED_STATE), + ): + monkeypatch.setattr(registry, attribute, recorder(name, result)) + return SimpleNamespace(calls=calls, plan=plan, plan_file=str(plan_file)) + + +@pytest.mark.parametrize( + "mode_case", + [ + ( + ("--print-latest", "--print-identity", "--apply-plan", "{plan}"), + None, + [""], + ), + ( + ("--print-identity", "--print-pending-plan", "--apply-plan", "{plan}"), + None, + [ + "policyengine_version=", + "app_name=", + "uk_model_version=", + "uk_dataset_sha256=", + ], + ), + (("--print-pending-plan", "--apply-plan", "{plan}"), None, ["null"]), + ( + ( + "--finalize-publication", + "--release-environment", + "beta", + "--gates-passed", + "false", + "--record-plan", + "{plan}", + ), + ( + "finalize", + { + "environment": "staging", + "release_environment": "beta", + "gates_passed": False, + "publication_run": None, + }, + ), + [ + "No pending routing publication; routing is unchanged.", + "Routing publication finalization: restored", + ], + ), + ( + ("--record-plan", "{plan}", "--complete-plan", "{plan}"), + ("record", {"environment": "staging", "plan": "{plan_body}"}), + ["Routing publication p-1: pending"], + ), + ( + ("--complete-plan", "{plan}", "--apply-plan", "{plan}"), + ("complete", {"environment": "staging", "plan": "{plan_body}"}), + ["Routing publication p-1: completed"], + ), + ( + ("--apply-plan", "{plan}", "--restore-plan", "{plan}"), + ("apply", {"environment": "staging", "plan": "{plan_body}"}), + [ + f"Published policyengine 6.3.0 routing ({'a' * 64}) in Modal " + "environment staging" + ], + ), + ( + ("--restore-plan", "{plan}", "{plan}", "--restore-latest", "6.2.6"), + ( + "restore-plans", + {"environment": "staging", "plans": ["{plan_body}", "{plan_body}"]}, + ), + ["Routing for policyengine 6.3.0: restored in Modal environment staging"], + ), + ( + ( + "--restore-latest", + "6.2.6", + "--expected-latest", + "6.3.0", + "--seed-active-from-legacy", + ), + ( + "restore-latest", + { + "environment": "staging", + "policyengine_version": "6.2.6", + "expected_latest": "6.3.0", + }, + ), + [ + "Restored routing state in Modal environment: staging", + " latest.policyengine: 6.2.6", + " latest.us: 1.1.0", + " latest.uk: 2.1.0", + ], + ), + ( + ("--seed-active-from-legacy", *PUBLISHING_ARGUMENTS), + ("seed", {"environment": "staging"}), + [ + "Seeding active routing state from legacy Modal dicts in " + "environment: staging", + "", + " simulation-api-routing-state[active]: updated", + " policyengine routes: 2", + " US routes: 1", + " UK routes: 3", + " bundle manifests: 1", + "", + "Legacy routing state seed completed successfully.", + ], + ), + ( + (*PUBLISHING_ARGUMENTS, "--dry-run", "--force-latest"), + ( + "plan", + { + "current_state": None, + "deployed": DEPLOYED_6_3_0, + "force_latest": True, + "publication_id": None, + }, + ), + [ + f"before_sha256={'b' * 64}", + f"after_sha256={'a' * 64}", + "advances_latest=false", + ], + ), + ( + PUBLISHING_ARGUMENTS, + ( + "publish", + { + "environment": "staging", + "deployed": DEPLOYED_6_3_0, + "force_latest": False, + }, + ), + [ + "Publishing routing state in Modal environment: staging", + " App name: policyengine-simulation-py6-3-0", + " policyengine.py version: 6.3.0", + " US version: 1.1.0", + " UK version: 2.2.0", + "", + "", + " simulation-api-routing-state[active]: updated", + " latest.policyengine: 6.2.6", + " latest.us: 1.1.0", + " latest.uk: 2.1.0", + "", + "Routing state published successfully.", + ], + ), + ], + ids=[ + "print-latest", + "print-identity", + "print-pending-plan", + "finalize", + "record", + "complete", + "apply", + "restore-plans", + "restore-latest", + "seed", + "plan", + "publish", + ], +) +def test_main_runs_the_first_requested_mode_and_reports_it( + monkeypatch, capsys, recorded_writers, mode_case +): + """``mode_case`` is the arguments, the operation they run and its output.""" + arguments, operation, output = mode_case + + def resolved(value): + if value == "{plan}": + return recorded_writers.plan_file + if value == "{plan_body}": + return recorded_writers.plan + if isinstance(value, list): + return [resolved(item) for item in value] + return value + + _main(monkeypatch, *(resolved(argument) for argument in arguments)) + + expected_calls = [] + if operation is not None: + name, kwargs = operation + expected_calls = [ + (name, {key: resolved(value) for key, value in kwargs.items()}) + ] + assert recorded_writers.calls == expected_calls + assert capsys.readouterr().out.splitlines() == output + + +def test_main_reports_a_foreign_publication_before_its_outcome( + monkeypatch, capsys, recorded_writers +): + monkeypatch.setattr( + registry, + "finalize_routing_publication", + lambda **_kwargs: "foreign-publication", + ) + + _main( + monkeypatch, + "--finalize-publication", + "--publication-run", + "74401", + "--release-environment", + "prod", + "--gates-passed", + "true", + ) + + assert recorded_writers.calls == [] + assert capsys.readouterr().out.splitlines() == [ + "No pending routing publication; routing is unchanged.", + "Pending routing publication belongs to another run or environment; " + "routing is unchanged.", + "Routing publication finalization: foreign-publication", + ] + + +@pytest.mark.parametrize( + ("arguments", "message"), + [ + ( + ("--finalize-publication", "--gates-passed", "false"), + "--finalize-publication requires --release-environment and --gates-passed", + ), + ( + ( + "--finalize-publication", + "--release-environment", + "beta", + "--gates-passed", + "true", + ), + "operator recovery requires --gates-passed false", + ), + ( + ("--expected-latest", "6.2.6", *PUBLISHING_ARGUMENTS), + "--expected-latest needs --restore-latest", + ), + ( + (), + "the following arguments are required: --app-name, " + "--policyengine-version, --us-version, --uk-version", + ), + ( + ("--app-name", "app", "--uk-version", "2.2.0", "--dry-run"), + "the following arguments are required: --policyengine-version, " + "--us-version", + ), + ], + ids=[ + "finalize-inputs", + "operator-recovery", + "expected-latest", + "publish-arguments", + "plan-arguments", + ], +) +def test_main_refuses_incomplete_arguments_before_any_operation( + monkeypatch, capsys, recorded_writers, arguments, message +): + with pytest.raises(SystemExit) as refusal: + _main(monkeypatch, *arguments) + + assert refusal.value.code == 2 + assert capsys.readouterr().err.splitlines()[-1] == ( + f"update_version_registry: error: {message}" + ) + assert recorded_writers.calls == [] + + +@pytest.mark.parametrize( + ("publication_id", "publication_run", "owned"), + [ + ("74401-1-beta", "74401", True), + ("74401-2-beta", "74401", True), + ("74401-1-prod", "74401", False), + ("74400-1-beta", "74401", False), + ("74401-x-beta", "74401", False), + ("74401-1-beta-extra", "74401", False), + ("74401-beta", "74401", False), + ("anything", None, True), + ], +) +def test_publication_ownership_reads_run_attempt_and_environment( + publication_id, publication_run, owned +): + assert ( + registry._publication_owned_by_release( + {"publication_id": publication_id}, + publication_run=publication_run, + release_environment="beta", + ) + is owned + ) + + +@pytest.mark.usefixtures("patched_modal") +def test_restoring_no_plans_is_refused(): + with pytest.raises(ValueError, match=r"^no routing plan to restore$"): + registry.restore_routing_plans(environment="staging", plans=[]) + + +def test_restoring_without_a_routing_store(monkeypatch, patched_modal): + first = _plan(monkeypatch, None, "6.1.0", "1.1.0", "2.2.0") + assert ( + registry.restore_routing_plans(environment="staging", plans=[first]) + == "unchanged" + ) + later = _plan(monkeypatch, first["after"], "6.2.0", "1.1.0", "2.3.0") + with pytest.raises( + registry.RoutingConflictError, + match=r"^there is no routing state to restore into$", + ): + registry.restore_routing_plans(environment="staging", plans=[first, later]) + assert patched_modal == {} From 16d8e9ed238b4f8b5375e6e459e6ad5b204fb0a9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:25:37 +0400 Subject: [PATCH 50/66] Split the routing publisher CLI into one command per mode The modes keep their precedence and output, and the help text is unchanged. Also split the plan-restore loop out of restore_routing_plans and name the parts of a workflow run's publication id. Co-Authored-By: Claude Opus 5.5 --- .../modal/utils/update_version_registry.py | 367 +++++++++++------- 1 file changed, 223 insertions(+), 144 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index a2c30b2a1..4bc6ad701 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -11,7 +11,7 @@ import json import os import uuid -from collections.abc import Iterable +from collections.abc import Callable, Iterable from copy import deepcopy from dataclasses import dataclass from pathlib import Path @@ -680,6 +680,18 @@ def restore_routing_plans(*, environment: str, plans: list[dict]) -> str: return "unchanged" msg = "there is no routing state to restore into" raise RoutingConflictError(msg) + outcome = _restore_live_plans(store, plans) + if pending is not None: + _write_publication_record(store, plan=pending, status="restored") + return outcome + + +def _restore_live_plans(store: "modal.Dict", plans: list[dict]) -> str: + """Restore each live plan's starting state, newest attempt first. + + Returns ``restored``, or ``unchanged`` when the live state is a starting + state no plan replaced; any other live state is a conflict. + """ remaining = [ plan for plan in plans if plan["before_sha256"] != plan["after_sha256"] ] @@ -691,16 +703,7 @@ def restore_routing_plans(*, environment: str, plans: list[dict]) -> str: break plan = live[-1] remaining.remove(plan) - if plan["before"] is None: - store.pop(ROUTING_STATE_ACTIVE_KEY) - else: - store[ROUTING_STATE_ACTIVE_KEY] = plan["before"] - if ( - routing_state_digest(store.get(ROUTING_STATE_ACTIVE_KEY)) - != plan["before_sha256"] - ): - msg = "the restored routing state reads back differently" - raise RoutingConflictError(msg) + _write_starting_state(store, plan) outcome = "restored" if outcome == "unchanged" and not any( plan["before_sha256"] == digest for plan in plans @@ -710,17 +713,33 @@ def restore_routing_plans(*, environment: str, plans: list[dict]) -> str: "it published; another change was made, so nothing was restored" ) raise RoutingConflictError(msg) - if pending is not None: - _write_publication_record(store, plan=pending, status="restored") return outcome +def _write_starting_state(store: "modal.Dict", plan: dict) -> None: + """Put back the state ``plan`` replaced, and check that it reads back.""" + if plan["before"] is None: + store.pop(ROUTING_STATE_ACTIVE_KEY) + else: + store[ROUTING_STATE_ACTIVE_KEY] = plan["before"] + if ( + routing_state_digest(store.get(ROUTING_STATE_ACTIVE_KEY)) + != plan["before_sha256"] + ): + msg = "the restored routing state reads back differently" + raise RoutingConflictError(msg) + + def restore_routing_plan(*, environment: str, plan: dict) -> str: """Undo one plan: ``restored``, ``unchanged`` (never applied) or a conflict.""" _require_serialized_routing_workflow() return restore_routing_plans(environment=environment, plans=[plan]) +# A workflow run's publication id is "--". +_RUN_PUBLICATION_ID_PARTS = 3 + + def _publication_owned_by_release( plan: dict, *, publication_run: str | None, release_environment: str ) -> bool: @@ -730,7 +749,7 @@ def _publication_owned_by_release( parts = str(plan.get("publication_id", "")).split("-") return ( - len(parts) == 3 + len(parts) == _RUN_PUBLICATION_ID_PARTS and parts[1].isdigit() and parts[2] == release_environment and parts[0] == publication_run @@ -1107,8 +1126,8 @@ def _report(line: str = "") -> None: print(line) # noqa: T201 # command-line script: stdout is its progress report -def main() -> None: - """Publish a deployed app's routes, or seed the state from legacy Dicts.""" +def _parser() -> argparse.ArgumentParser: + """Describe the routing publisher's modes and their options.""" parser = argparse.ArgumentParser( formatter_class=argparse.RawDescriptionHelpFormatter, description=( @@ -1255,29 +1274,40 @@ def main() -> None: metavar="POLICYENGINE_VERSION", help="With --restore-latest: refuse unless latest is still this version.", ) - args = parser.parse_args() + return parser - def active_state() -> object: - snapshot = _modal_dict_snapshot( - name=ROUTING_STATE_DICT_NAME, environment=args.environment - ) - return snapshot.get(ROUTING_STATE_ACTIVE_KEY) +class _UsageError(Exception): + """A command-line mistake, reported as an argparse usage error.""" + + +def _active_state(environment: str) -> object: + snapshot = _modal_dict_snapshot( + name=ROUTING_STATE_DICT_NAME, environment=environment + ) + return snapshot.get(ROUTING_STATE_ACTIVE_KEY) + + +def _print_requested_view(args: argparse.Namespace) -> bool: + """Print the requested read-only view, if any; return whether one was printed.""" if args.print_latest: _report(read_latest_policyengine_version(environment=args.environment) or "") - return - - if args.print_identity: - for key, value in latest_release_identity(active_state()).items(): + elif args.print_identity: + for key, value in latest_release_identity( + _active_state(args.environment) + ).items(): _report(f"{key}={value}") - return - - if args.print_pending_plan: + elif args.print_pending_plan: store = _routing_store(args.environment, create_if_missing=False) _report(json.dumps(_unfinished_publication(store), indent=2, sort_keys=True)) - return + else: + return False + return True - writes_routing = ( + +def _writes_routing(args: argparse.Namespace) -> bool: + """Whether the requested mode can write routing; only planning modes cannot.""" + return bool( args.apply_plan or args.record_plan or args.complete_plan @@ -1287,115 +1317,125 @@ def active_state() -> object: or args.seed_active_from_legacy or not (args.plan_out or args.dry_run) ) - if writes_routing: - try: - _require_serialized_routing_workflow() - except RuntimeError as exc: - parser.error(str(exc)) - - if args.finalize_publication: - if args.release_environment is None or args.gates_passed is None: - parser.error( - "--finalize-publication requires --release-environment and " - "--gates-passed" - ) - if args.publication_run is None and args.gates_passed != "false": - parser.error("operator recovery requires --gates-passed false") - store = _routing_store(args.environment, create_if_missing=False) - pending = _unfinished_publication(store) - if pending is None: - _report("No pending routing publication; routing is unchanged.") - else: - _report("Pending routing publication:") - _report( - json.dumps(store.get(ROUTING_PUBLICATION_KEY), indent=2, sort_keys=True) - ) - outcome = finalize_routing_publication( - environment=args.environment, - release_environment=args.release_environment, - gates_passed=args.gates_passed == "true", - publication_run=args.publication_run, - ) - if outcome == "foreign-publication": - _report( - "Pending routing publication belongs to another run or " - "environment; routing is unchanged." - ) - _report(f"Routing publication finalization: {outcome}") - if args.github_output: - with Path(args.github_output).open("a") as handle: - handle.write(f"outcome={outcome}\n") - return - if args.record_plan or args.complete_plan: - with Path(args.record_plan or args.complete_plan).open() as handle: - plan = json.load(handle) - operation = record_routing_plan if args.record_plan else complete_routing_plan - operation(environment=args.environment, plan=plan) - _report( - f"Routing publication {plan['publication_id']}: " - f"{'pending' if args.record_plan else 'completed'}" - ) - return - if args.apply_plan: - with Path(args.apply_plan).open() as handle: - plan = json.load(handle) - apply_routing_plan( - environment=args.environment, - plan=plan, - ) - _report( - f"Published policyengine {plan['policyengine_version']} routing " - f"({plan['after_sha256']}) in Modal environment {args.environment}" - ) +def _require_writes_allowed(args: argparse.Namespace) -> None: + """Refuse a writing mode outside the workflows that serialize routing writes.""" + if not _writes_routing(args): return + try: + _require_serialized_routing_workflow() + except RuntimeError as exc: + raise _UsageError(str(exc)) from exc + - if args.restore_plan: - plans = [] - for path in args.restore_plan: - with Path(path).open() as handle: - plans.append(json.load(handle)) - outcome = restore_routing_plans( - environment=args.environment, - plans=plans, +def _read_plan(path: str) -> dict: + with Path(path).open() as handle: + return json.load(handle) + + +def _finalize_command(args: argparse.Namespace) -> None: + if args.release_environment is None or args.gates_passed is None: + msg = "--finalize-publication requires --release-environment and --gates-passed" + raise _UsageError(msg) + if args.publication_run is None and args.gates_passed != "false": + msg = "operator recovery requires --gates-passed false" + raise _UsageError(msg) + store = _routing_store(args.environment, create_if_missing=False) + pending = _unfinished_publication(store) + if pending is None: + _report("No pending routing publication; routing is unchanged.") + else: + _report("Pending routing publication:") + _report( + json.dumps(store.get(ROUTING_PUBLICATION_KEY), indent=2, sort_keys=True) ) + outcome = finalize_routing_publication( + environment=args.environment, + release_environment=args.release_environment, + gates_passed=args.gates_passed == "true", + publication_run=args.publication_run, + ) + if outcome == "foreign-publication": _report( - f"Routing for policyengine {plans[-1]['policyengine_version']}: " - f"{outcome} in Modal environment {args.environment}" + "Pending routing publication belongs to another run or " + "environment; routing is unchanged." ) - return + _report(f"Routing publication finalization: {outcome}") + if args.github_output: + with Path(args.github_output).open("a") as handle: + handle.write(f"outcome={outcome}\n") + + +def _journal_command(args: argparse.Namespace) -> None: + plan = _read_plan(args.record_plan or args.complete_plan) + operation = record_routing_plan if args.record_plan else complete_routing_plan + operation(environment=args.environment, plan=plan) + _report( + f"Routing publication {plan['publication_id']}: " + f"{'pending' if args.record_plan else 'completed'}" + ) - if args.restore_latest: - state = restore_routing_state( - environment=args.environment, - policyengine_version=args.restore_latest, - expected_latest=args.expected_latest, - ) - _report(f"Restored routing state in Modal environment: {args.environment}") - _report(f" latest.policyengine: {state['latest']['policyengine']}") - _report(f" latest.us: {state['latest']['us']}") - _report(f" latest.uk: {state['latest']['uk']}") - return - if args.expected_latest: - parser.error("--expected-latest needs --restore-latest") - if args.seed_active_from_legacy: - _report( - "Seeding active routing state from legacy Modal dicts in " - f"environment: {args.environment}" - ) - state = seed_active_routing_state_from_legacy(environment=args.environment) - _report() - _report(f" {ROUTING_STATE_DICT_NAME}[{ROUTING_STATE_ACTIVE_KEY}]: updated") - _report(f" policyengine routes: {len(state['routes']['policyengine'])}") - _report(f" US routes: {len(state['routes']['us'])}") - _report(f" UK routes: {len(state['routes']['uk'])}") - _report(f" bundle manifests: {len(state['bundles'])}") - _report() - _report("Legacy routing state seed completed successfully.") - return +def _apply_plan_command(args: argparse.Namespace) -> None: + plan = _read_plan(args.apply_plan) + apply_routing_plan( + environment=args.environment, + plan=plan, + ) + _report( + f"Published policyengine {plan['policyengine_version']} routing " + f"({plan['after_sha256']}) in Modal environment {args.environment}" + ) + + +def _restore_plans_command(args: argparse.Namespace) -> None: + plans = [_read_plan(path) for path in args.restore_plan] + outcome = restore_routing_plans( + environment=args.environment, + plans=plans, + ) + _report( + f"Routing for policyengine {plans[-1]['policyengine_version']}: " + f"{outcome} in Modal environment {args.environment}" + ) + + +def _restore_latest_command(args: argparse.Namespace) -> None: + state = restore_routing_state( + environment=args.environment, + policyengine_version=args.restore_latest, + expected_latest=args.expected_latest, + ) + _report(f"Restored routing state in Modal environment: {args.environment}") + _report(f" latest.policyengine: {state['latest']['policyengine']}") + _report(f" latest.us: {state['latest']['us']}") + _report(f" latest.uk: {state['latest']['uk']}") + +def _expected_latest_without_restore(_args: argparse.Namespace) -> None: + msg = "--expected-latest needs --restore-latest" + raise _UsageError(msg) + + +def _seed_command(args: argparse.Namespace) -> None: + _report( + "Seeding active routing state from legacy Modal dicts in " + f"environment: {args.environment}" + ) + state = seed_active_routing_state_from_legacy(environment=args.environment) + _report() + _report(f" {ROUTING_STATE_DICT_NAME}[{ROUTING_STATE_ACTIVE_KEY}]: updated") + _report(f" policyengine routes: {len(state['routes']['policyengine'])}") + _report(f" US routes: {len(state['routes']['us'])}") + _report(f" UK routes: {len(state['routes']['uk'])}") + _report(f" bundle manifests: {len(state['bundles'])}") + _report() + _report("Legacy routing state seed completed successfully.") + + +def _deployed_app(args: argparse.Namespace) -> DeployedAppVersions: + """Read the deployed app's versions, which publishing and planning require.""" required_args = { "--app-name": args.app_name, "--policyengine-version": args.policyengine_version, @@ -1404,27 +1444,35 @@ def active_state() -> object: } missing_args = [name for name, value in required_args.items() if not value] if missing_args: - parser.error(f"the following arguments are required: {', '.join(missing_args)}") - deployed = DeployedAppVersions( + msg = f"the following arguments are required: {', '.join(missing_args)}" + raise _UsageError(msg) + return DeployedAppVersions( app_name=args.app_name, policyengine_version=args.policyengine_version, us_version=args.us_version, uk_version=args.uk_version, ) + +def _plan_command(args: argparse.Namespace, deployed: DeployedAppVersions) -> None: + plan = plan_routing_publication( + current_state=_active_state(args.environment), + deployed=deployed, + force_latest=args.force_latest, + publication_id=args.publication_id, + ) + if args.plan_out: + with Path(args.plan_out).open("w") as handle: + json.dump(plan, handle, indent=2, sort_keys=True) + _report(f"before_sha256={plan['before_sha256']}") + _report(f"after_sha256={plan['after_sha256']}") + _report(f"advances_latest={'true' if plan['advances_latest'] else 'false'}") + + +def _publish_command(args: argparse.Namespace) -> None: + deployed = _deployed_app(args) if args.plan_out or args.dry_run: - plan = plan_routing_publication( - current_state=active_state(), - deployed=deployed, - force_latest=args.force_latest, - publication_id=args.publication_id, - ) - if args.plan_out: - with Path(args.plan_out).open("w") as handle: - json.dump(plan, handle, indent=2, sort_keys=True) - _report(f"before_sha256={plan['before_sha256']}") - _report(f"after_sha256={plan['after_sha256']}") - _report(f"advances_latest={'true' if plan['advances_latest'] else 'false'}") + _plan_command(args, deployed) return _report(f"Publishing routing state in Modal environment: {args.environment}") @@ -1448,5 +1496,36 @@ def active_state() -> object: _report("Routing state published successfully.") +def _requested_command( + args: argparse.Namespace, +) -> Callable[[argparse.Namespace], None]: + """Return the first requested writing mode's command, or publishing.""" + commands = ( + (args.finalize_publication, _finalize_command), + (args.record_plan or args.complete_plan, _journal_command), + (args.apply_plan, _apply_plan_command), + (args.restore_plan, _restore_plans_command), + (args.restore_latest, _restore_latest_command), + (args.expected_latest, _expected_latest_without_restore), + (args.seed_active_from_legacy, _seed_command), + ) + return next( + (command for requested, command in commands if requested), _publish_command + ) + + +def main() -> None: + """Publish a deployed app's routes, or seed the state from legacy Dicts.""" + parser = _parser() + args = parser.parse_args() + if _print_requested_view(args): + return + try: + _require_writes_allowed(args) + _requested_command(args)(args) + except _UsageError as error: + parser.error(str(error)) + + if __name__ == "__main__": main() From 934110d9eb5c683294d2f97a632b6600f617497d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:27:39 +0400 Subject: [PATCH 51/66] Pin the package update helper's bundle-versions and verify-dispatch output Co-Authored-By: Claude Opus 5.5 --- ...est_policyengine_package_update_scripts.py | 337 ++++++++++++++++++ 1 file changed, 337 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py index 0594c6ee3..a97040588 100644 --- a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py @@ -739,3 +739,340 @@ def test_update_policyengine_package_accepts_a_matching_dispatch( assert result.returncode == 0, result.stderr assert "pr create" in gh_log.read_text(encoding="utf-8") + + +# --- Characterization: the helper's bundle-versions and verify-dispatch ------- + + +@pytest.fixture +def helper_module(): + """Load the update helper in this process, as the shell script runs it.""" + import importlib.util + + spec = importlib.util.spec_from_file_location("update_policyengine_package", HELPER) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _synthetic_bundle(**data_releases): + return { + "packages": { + "policyengine": {"version": "4.1.0"}, + "policyengine-core": {"version": "3.33.0"}, + "spm-calculator": {"version": "1.0.0"}, + "policyengine-us": {"version": "1.1.0"}, + "policyengine-uk": {"version": "2.1.0"}, + }, + "data_releases": { + "us": { + "version": "1.10.0", + "default_dataset": "cps_2024", + "datasets": {"cps_2024": {"sha256": " " + "B" * 64 + "\n"}}, + }, + "uk": { + "data_package": {"version": "1.20.0"}, + "default_dataset": "frs_2024", + "certified_data_artifact": {"dataset": "frs_2024", "sha256": "a" * 64}, + "datasets": {"frs_2024": {"sha256": "c" * 64}}, + }, + **data_releases, + }, + } + + +def _bundle_versions(monkeypatch, capsys, helper_module, bundle, *, receipt=None): + monkeypatch.setattr("policyengine.bundle.get_current_bundle", lambda: bundle) + # None makes the import fail, as with a wrapper that predates the helper. + monkeypatch.setitem(sys.modules, "policyengine.provenance.receipt", receipt) + helper_module.bundle_versions() + return capsys.readouterr().out.splitlines() + + +def test_bundle_versions_prints_every_version_and_digest_in_order( + monkeypatch, capsys, helper_module +): + lines = _bundle_versions(monkeypatch, capsys, helper_module, _synthetic_bundle()) + + assert lines == [ + "policyengine_version=4.1.0", + "policyengine_core_version=3.33.0", + "spm_calculator_version=1.0.0", + "us_version=1.1.0", + "us_data_version=1.10.0", + "uk_version=2.1.0", + "uk_data_version=1.20.0", + "uk_data_package_version=1.20.0", + "us_dataset_sha256=" + "b" * 64, + "uk_dataset_sha256=" + "a" * 64, + "uk_certificate_sha256=", + ] + + +@pytest.mark.parametrize( + ("uk_release", "printed"), + [ + ( + { + "default_dataset": "frs_2024", + "certified_data_artifact": {"dataset": "frs_2023", "sha256": "a" * 64}, + "datasets": {"frs_2024": {"sha256": "c" * 64}}, + }, + ["uk_data_version=2024.1", "uk_dataset_sha256=" + "c" * 64], + ), + ( + { + "default_dataset": "frs_2024", + "certified_data_artifact": {"dataset": "frs_2024", "sha256": ""}, + "datasets": {"frs_2024": {"sha256": "d" * 64}}, + }, + ["uk_data_version=2024.1", "uk_dataset_sha256=" + "d" * 64], + ), + ( + {"default_dataset": "frs_2024", "data_package": {"version": 7}}, + [ + "uk_data_version=2024.1", + "uk_dataset_sha256=", + "uk_data_package_version=", + ], + ), + ], + ids=["other-certified-dataset", "empty-certified-digest", "no-digest"], +) +def test_bundle_versions_falls_back_to_the_default_dataset_reference( + monkeypatch, capsys, helper_module, uk_release, printed +): + bundle = _synthetic_bundle(uk={"version": "2024.1", **uk_release}) + + lines = _bundle_versions(monkeypatch, capsys, helper_module, bundle) + + for line in printed: + assert line in lines + + +def test_bundle_versions_reports_the_receipt_digest_of_the_uk_record( + monkeypatch, capsys, helper_module +): + from types import SimpleNamespace + + digested = [] + + def certificate_sha256(record): + digested.append(record) + return "e" * 64 + + bundle = _synthetic_bundle() + lines = _bundle_versions( + monkeypatch, + capsys, + helper_module, + bundle, + receipt=SimpleNamespace(certificate_sha256=certificate_sha256), + ) + + assert lines[-1] == "uk_certificate_sha256=" + "e" * 64 + assert digested == [bundle["data_releases"]["uk"]] + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ( + {"packages": {"policyengine-uk": {}}}, + "Bundle has no version for policyengine-uk", + ), + ( + {"data_releases": {"us": {"data_package": {"version": ""}}}}, + "Bundle has no data release version for us", + ), + ( + { + "data_releases": { + "us": {"version": "1.10.0"}, + "uk": { + "data_package": {"version": "1.20.0"}, + "default_dataset": "frs_2024", + "datasets": {"frs_2024": {"sha256": "not-a-digest"}}, + }, + } + }, + "Bundle records a malformed uk dataset SHA-256", + ), + ( + { + "packages": {"policyengine-uk": {"version": ""}}, + "data_releases": {"us": {"version": ""}}, + }, + "Bundle has no data release version for us", + ), + ], + ids=["package", "data-release", "digest", "first-in-output-order"], +) +def test_bundle_versions_refuses_an_incomplete_bundle( + monkeypatch, capsys, helper_module, change, message +): + bundle = _synthetic_bundle() + for section, entries in change.items(): + bundle[section] = {**bundle[section], **entries} + + with pytest.raises(SystemExit) as refusal: + _bundle_versions(monkeypatch, capsys, helper_module, bundle) + + assert refusal.value.code == message + assert capsys.readouterr().out == "" + + +def _verify_in_process(monkeypatch, capsys, helper_module, payload, installed): + import io + + monkeypatch.setenv("RELEASE_PAYLOAD", payload) + monkeypatch.setattr(sys, "stdin", io.StringIO(installed)) + helper_module.verify_dispatch() + return capsys.readouterr().out.splitlines() + + +def test_verify_dispatch_prints_the_verdict_and_receipt_lines( + monkeypatch, capsys, helper_module +): + lines = _verify_in_process( + monkeypatch, capsys, helper_module, _payload(), INSTALLED + ) + + assert lines == [ + "dispatch_verified=true", + "uk_qualification_verdict=qualified", + "uk_qualification_receipt_sha256=" + "e" * 64, + ] + + +def test_verify_dispatch_strips_installed_lines_and_claims( + monkeypatch, capsys, helper_module +): + # One line per installed field; a line without a key is skipped. + installed = ( + " uk_version = 2.126.8 \n=ignored\nno-separator\npolicyengine_version=4.1.0" + ) + payload = _payload( + model_version=" 2.126.8 ", + core_version="", + data_release_version=None, + dataset_sha256=None, + certificate_sha256=None, + ) + + lines = _verify_in_process(monkeypatch, capsys, helper_module, payload, installed) + + assert lines[0] == "dispatch_verified=true" + + +def test_verify_dispatch_lowercases_only_digest_claims( + monkeypatch, capsys, helper_module +): + installed = INSTALLED.replace("uk_version=2.126.8", "uk_version=2.126.8a") + + with pytest.raises(SystemExit) as refusal: + _verify_in_process( + monkeypatch, + capsys, + helper_module, + _payload(model_version="2.126.8A", dataset_sha256="A" * 64), + installed, + ) + + assert refusal.value.code == ( + "The release dispatch names a different UK bundle from the installed " + "wheel: model_version: dispatch 2.126.8A, installed 2.126.8a" + ) + + +def test_verify_dispatch_lists_every_mismatch_in_check_order( + monkeypatch, capsys, helper_module +): + installed = INSTALLED.replace("uk_dataset_sha256=" + "a" * 64, "uk_dataset_sha256=") + + with pytest.raises(SystemExit) as refusal: + _verify_in_process( + monkeypatch, + capsys, + helper_module, + _payload(model_version="2.0.0", data_release_version="1.0.0"), + installed, + ) + + assert refusal.value.code == ( + "The release dispatch names a different UK bundle from the installed " + "wheel: model_version: dispatch 2.0.0, installed 2.126.8; " + "data_release_version: dispatch 1.0.0, installed 1.58.0; " + "dataset_sha256: dispatch " + "a" * 64 + ", installed bundle reports none" + ) + assert capsys.readouterr().out == "" + + +@pytest.mark.parametrize( + ("dispatch", "message"), + [ + ( + {"version": "4.2.0", "bundle": {"uk": {"model_version": "0"}}}, + "Dispatch is for policyengine 4.2.0, but 4.1.0 is installed.", + ), + ( + { + "bundle": {"uk": {"model_version": "0"}}, + "qualification": {"uk": {"verdict": "Rejected"}}, + }, + "The release dispatch names a different UK bundle from the installed " + "wheel: model_version: dispatch 0, installed 2.126.8", + ), + ( + {"qualification": {"uk": {"verdict": 5, "receipt_sha256": 7}}}, + "The release dispatch carries a malformed verdict", + ), + ( + {"qualification": {"uk": {"verdict": "qualified", "receipt_sha256": 7}}}, + "The release dispatch carries a malformed receipt digest", + ), + ( + {"qualification": {"uk": {"promotable": False}}}, + "The release dispatch says the UK pair was not promotable ()", + ), + ( + {"qualification": {"uk": {"verdict": "incomplete", "promotable": True}}}, + "The release dispatch says the UK pair was not promotable (incomplete)", + ), + ], + ids=[ + "wrapper-first", + "bundle-before-qualification", + "verdict-before-receipt", + "receipt", + "unpromotable-without-verdict", + "unpromotable-verdict", + ], +) +def test_verify_dispatch_refusals_and_their_order( + monkeypatch, capsys, helper_module, dispatch, message +): + import json + + with pytest.raises(SystemExit) as refusal: + _verify_in_process( + monkeypatch, capsys, helper_module, json.dumps(dispatch), INSTALLED + ) + + assert refusal.value.code == message + + +def test_verify_dispatch_without_a_payload_compares_nothing( + monkeypatch, capsys, helper_module +): + monkeypatch.delenv("RELEASE_PAYLOAD", raising=False) + import io + + monkeypatch.setattr(sys, "stdin", io.StringIO(INSTALLED)) + helper_module.verify_dispatch() + + assert capsys.readouterr().out.splitlines() == [ + "dispatch_verified=false", + "uk_qualification_verdict=", + "uk_qualification_receipt_sha256=", + ] From 96139060439ae721c24389009e59a04dd7f93512 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:29:09 +0400 Subject: [PATCH 52/66] Split the package update helper's bundle and dispatch checks into helpers The helper's subcommands, output lines and messages are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../scripts/update_policyengine_package.py | 203 +++++++++++------- 1 file changed, 122 insertions(+), 81 deletions(-) diff --git a/.github/scripts/update_policyengine_package.py b/.github/scripts/update_policyengine_package.py index 47af06036..b1a06e8f9 100644 --- a/.github/scripts/update_policyengine_package.py +++ b/.github/scripts/update_policyengine_package.py @@ -10,7 +10,7 @@ import tempfile from collections.abc import Mapping from pathlib import Path -from typing import Any +from typing import Any, NamedTuple import tomllib @@ -201,6 +201,56 @@ def latest_version(package: str) -> None: _write_result(version) +def _package_version(packages: Mapping[str, Any], name: str) -> str: + version = _mapping(packages.get(name)).get("version") + if not isinstance(version, str) or not version: + msg = f"Bundle has no version for {name}" + raise SystemExit(msg) + return version + + +def _data_release_version(release: Mapping[str, Any], country: str) -> str: + version = release.get("version") or _mapping(release.get("data_package")).get( + "version" + ) + if not isinstance(version, str) or not version: + msg = f"Bundle has no data release version for {country}" + raise SystemExit(msg) + return version + + +def _data_package_version(release: Mapping[str, Any]) -> str: + version = _mapping(release.get("data_package")).get("version") + return version if isinstance(version, str) else "" + + +def _dataset_sha256(release: Mapping[str, Any], country: str) -> str: + """Return the default dataset's digest, preferring the certified artifact's.""" + default = release.get("default_dataset") + certified = _mapping(release.get("certified_data_artifact")) + if certified.get("dataset") == default and certified.get("sha256"): + value = certified["sha256"] + else: + reference = _mapping(_mapping(release.get("datasets")).get(default)) + value = reference.get("sha256") or "" + digest = str(value).strip().lower() + if digest and not re.fullmatch(r"[0-9a-f]{64}", digest): + msg = f"Bundle records a malformed {country} dataset SHA-256" + raise SystemExit(msg) + return digest + + +def _certificate_sha256(release: object) -> str: + # Receipt-v1 certificate digest; older wrappers lack the helper. + try: + from policyengine.provenance.receipt import ( # noqa: PLC0415 # older policyengine wrappers lack the receipt helper, so the ImportError is handled here + certificate_sha256 as digest, + ) + except ImportError: + return "" + return digest(release) if isinstance(release, Mapping) else "" + + def bundle_versions() -> None: """Print the package and data versions selected by PolicyEngine.py.""" from policyengine.bundle import ( # noqa: PLC0415 # only bundle-versions runs under uv with policyengine installed; the other commands use the runner's python3 @@ -210,67 +260,20 @@ def bundle_versions() -> None: bundle = _mapping(get_current_bundle()) packages = _mapping(bundle.get("packages")) data_releases = _mapping(bundle.get("data_releases")) - - def package_version(name: str) -> str: - version = _mapping(packages.get(name)).get("version") - if not isinstance(version, str) or not version: - msg = f"Bundle has no version for {name}" - raise SystemExit(msg) - return version - - def data_release_version(country: str) -> str: - release = _mapping(data_releases.get(country)) - version = release.get("version") or _mapping(release.get("data_package")).get( - "version" - ) - if not isinstance(version, str) or not version: - msg = f"Bundle has no data release version for {country}" - raise SystemExit(msg) - return version - - def data_package_version(country: str) -> str: - release = _mapping(data_releases.get(country)) - version = _mapping(release.get("data_package")).get("version") - return version if isinstance(version, str) else "" - - def dataset_sha256(country: str) -> str: - release = _mapping(data_releases.get(country)) - default = release.get("default_dataset") - certified = _mapping(release.get("certified_data_artifact")) - if certified.get("dataset") == default and certified.get("sha256"): - value = certified["sha256"] - else: - reference = _mapping(_mapping(release.get("datasets")).get(default)) - value = reference.get("sha256") or "" - digest = str(value).strip().lower() - if digest and not re.fullmatch(r"[0-9a-f]{64}", digest): - msg = f"Bundle records a malformed {country} dataset SHA-256" - raise SystemExit(msg) - return digest - - def certificate_sha256(country: str) -> str: - # Receipt-v1 certificate digest; older wrappers lack the helper. - try: - from policyengine.provenance.receipt import ( # noqa: PLC0415 # older policyengine wrappers lack the receipt helper, so the ImportError is handled here - certificate_sha256 as digest, - ) - except ImportError: - return "" - release = data_releases.get(country) - return digest(release) if isinstance(release, Mapping) else "" - + us_release = _mapping(data_releases.get("us")) + uk_release = _mapping(data_releases.get("uk")) outputs = { - "policyengine_version": package_version("policyengine"), - "policyengine_core_version": package_version("policyengine-core"), - "spm_calculator_version": package_version("spm-calculator"), - "us_version": package_version("policyengine-us"), - "us_data_version": data_release_version("us"), - "uk_version": package_version("policyengine-uk"), - "uk_data_version": data_release_version("uk"), - "uk_data_package_version": data_package_version("uk"), - "us_dataset_sha256": dataset_sha256("us"), - "uk_dataset_sha256": dataset_sha256("uk"), - "uk_certificate_sha256": certificate_sha256("uk"), + "policyengine_version": _package_version(packages, "policyengine"), + "policyengine_core_version": _package_version(packages, "policyengine-core"), + "spm_calculator_version": _package_version(packages, "spm-calculator"), + "us_version": _package_version(packages, "policyengine-us"), + "us_data_version": _data_release_version(us_release, "us"), + "uk_version": _package_version(packages, "policyengine-uk"), + "uk_data_version": _data_release_version(uk_release, "uk"), + "uk_data_package_version": _data_package_version(uk_release), + "us_dataset_sha256": _dataset_sha256(us_release, "us"), + "uk_dataset_sha256": _dataset_sha256(uk_release, "uk"), + "uk_certificate_sha256": _certificate_sha256(data_releases.get("uk")), } for name, version in outputs.items(): _write_result(f"{name}={version}") @@ -293,25 +296,30 @@ def certificate_sha256(country: str) -> str: PROMOTABLE_VERDICTS = frozenset({"qualified", "needs_review"}) -def verify_dispatch() -> None: - """Check a policyengine-release payload against the installed bundle. +class _DispatchComparison(NamedTuple): + """What comparing a dispatch with the installed bundle found. - Reads the dispatch ``client_payload`` JSON from ``RELEASE_PAYLOAD`` and - ``bundle-versions`` output from stdin. A payload without ``bundle.uk`` - (older releases, manual runs) is accepted. A field the payload names - must equal the installed bundle's, and an installed bundle that cannot - report it is a mismatch, except for the certificate digest, which older - wrappers cannot compute. A qualification that says the pair was not - promotable is refused. ``dispatch_verified`` is true only when at least - one field was compared. Prints digests for the PR body. + ``mismatches`` lists the fields that disagree; ``compared`` counts the + fields both the dispatch and the installed bundle stated. """ - raw = os.environ.get("RELEASE_PAYLOAD", "").strip() - payload = json.loads(raw) if raw and raw != "null" else {} + + mismatches: list[str] + compared: int + + +def _installed_fields(bundle_output: str) -> dict[str, str]: + """Read ``bundle-versions`` output: one ``key=value`` line per field.""" installed = {} - for line in sys.stdin.read().splitlines(): + for line in bundle_output.splitlines(): key, _, value = line.partition("=") if key: installed[key.strip()] = value.strip() + return installed + + +def _require_dispatched_wrapper( + payload: Mapping[str, Any], installed: Mapping[str, str] +) -> None: if payload.get("version") and payload["version"] != installed.get( "policyengine_version" ): @@ -320,7 +328,12 @@ def verify_dispatch() -> None: f"{installed.get('policyengine_version')} is installed." ) raise SystemExit(msg) - expected = _mapping(_mapping(payload.get("bundle")).get("uk")) + + +def _compare_dispatch_bundle( + expected: Mapping[str, Any], installed: Mapping[str, str] +) -> _DispatchComparison: + """Compare each field the dispatch names with the installed bundle's value.""" mismatches = [] compared = 0 for payload_key, installed_key in DISPATCH_CHECKS: @@ -341,11 +354,11 @@ def verify_dispatch() -> None: claimed_text = claimed_text.lower() if claimed_text != actual: mismatches.append(f"{payload_key}: dispatch {claimed}, installed {actual}") - if mismatches: - raise SystemExit( - "The release dispatch names a different UK bundle from the installed " - "wheel: " + "; ".join(mismatches) - ) + return _DispatchComparison(mismatches, compared) + + +def _promotable_qualification(payload: Mapping[str, Any]) -> tuple[str, str]: + """Return the UK verdict and receipt digest, refusing malformed or rejected ones.""" qualification = _mapping(_mapping(payload.get("qualification")).get("uk")) verdict = qualification.get("verdict") or "" receipt_sha256 = qualification.get("receipt_sha256") or "" @@ -362,7 +375,35 @@ def verify_dispatch() -> None: ): msg = f"The release dispatch says the UK pair was not promotable ({verdict})" raise SystemExit(msg) - _write_result(f"dispatch_verified={'true' if compared else 'false'}") + return verdict, receipt_sha256 + + +def verify_dispatch() -> None: + """Check a policyengine-release payload against the installed bundle. + + Reads the dispatch ``client_payload`` JSON from ``RELEASE_PAYLOAD`` and + ``bundle-versions`` output from stdin. A payload without ``bundle.uk`` + (older releases, manual runs) is accepted. A field the payload names + must equal the installed bundle's, and an installed bundle that cannot + report it is a mismatch, except for the certificate digest, which older + wrappers cannot compute. A qualification that says the pair was not + promotable is refused. ``dispatch_verified`` is true only when at least + one field was compared. Prints digests for the PR body. + """ + raw = os.environ.get("RELEASE_PAYLOAD", "").strip() + payload = json.loads(raw) if raw and raw != "null" else {} + installed = _installed_fields(sys.stdin.read()) + _require_dispatched_wrapper(payload, installed) + comparison = _compare_dispatch_bundle( + _mapping(_mapping(payload.get("bundle")).get("uk")), installed + ) + if comparison.mismatches: + raise SystemExit( + "The release dispatch names a different UK bundle from the installed " + "wheel: " + "; ".join(comparison.mismatches) + ) + verdict, receipt_sha256 = _promotable_qualification(payload) + _write_result(f"dispatch_verified={'true' if comparison.compared else 'false'}") _write_result(f"uk_qualification_verdict={verdict}") _write_result(f"uk_qualification_receipt_sha256={receipt_sha256}") From f9e40dbac72a7df200ba2babc95aed440785226a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:29:34 +0400 Subject: [PATCH 53/66] Pin the Stage 12 connection retry's delay sequence and budget cut Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_database_retry.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py index 84fd33ec6..e00bdf0da 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py @@ -227,3 +227,42 @@ def connect_on_second_attempt() -> OpenedConnection: "retrying attempt 2 after 1.000 seconds" ) assert isinstance(record.exc_info[1], OperationalError) + + +def test_delays_double_up_to_the_maximum_until_the_budget_is_spent() -> None: + clock = ManualClock() + + def never_connect() -> None: + raise _connection_limit_error() + + with pytest.raises(OperationalError): + open_connection_with_retry( + never_connect, + operation_name="get_report", + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ) + + assert clock.sleeps == [1.0, 2.0, 4.0, 8.0, 15.0, 15.0, 15.0] + + +def test_the_last_delay_is_cut_to_the_remaining_budget() -> None: + clock = ManualClock() + seconds_per_refused_attempt = 12.0 + + def slowly_refuse() -> None: + clock.now += seconds_per_refused_attempt + raise _connection_limit_error() + + with pytest.raises(OperationalError): + open_connection_with_retry( + slowly_refuse, + operation_name="get_report", + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ) + + # 12 + 1 + 12 + 2 + 12 + 4 + 12 seconds leave 5 of the 60-second budget. + assert clock.sleeps == [1.0, 2.0, 4.0, 5.0] From 62c478c2d5ac720f28fc8464dc6c18575bcd118d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:30:09 +0400 Subject: [PATCH 54/66] Pass the Stage 12 connection retry its clock, sleep and jitter as one timing The persistence store still calls the opener with only the connect step and the operation name. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/database_retry.py | 28 ++++++--- .../tests/test_stage12_database_retry.py | 61 ++++++++++++------- 2 files changed, 60 insertions(+), 29 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py index 90432c2de..f8e8c8ca1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py @@ -6,6 +6,7 @@ import random import time from contextlib import suppress +from dataclasses import dataclass from typing import TYPE_CHECKING from sqlalchemy.exc import OperationalError @@ -27,14 +28,24 @@ RETRY_DELAY_GROWTH_FACTOR = 2.0 +@dataclass(frozen=True) +class RetryTiming: + """The clock, sleep and jitter a connection retry uses; tests replace them. + + ``choose_delay`` picks a delay between its lower and upper bound. + """ + + sleep: Callable[[float], None] = time.sleep + monotonic: Callable[[], float] = time.monotonic + choose_delay: Callable[[float, float], float] = random.uniform + + def open_connection_with_retry( connect: Callable[[], Connection], *, operation_name: str, runtime: ObservabilityRuntime | None = None, - sleep: Callable[[float], None] = time.sleep, - monotonic: Callable[[], float] = time.monotonic, - choose_delay: Callable[[float, float], float] = random.uniform, + timing: RetryTiming | None = None, ) -> Connection: """Open a Stage 12 database connection, retrying while it is refused. @@ -46,18 +57,21 @@ def open_connection_with_retry( Errors other than ``OperationalError`` propagate immediately, and retries stop once the retry budget is spent. """ - started = monotonic() + timing = timing or RetryTiming() + started = timing.monotonic() attempt = 1 delay_ceiling = INITIAL_RETRY_DELAY_SECONDS while True: try: return connect() except OperationalError as error: # noqa: PERF203 # each refused attempt is caught so the loop can retry it - remaining_budget = CONNECTION_RETRY_BUDGET_SECONDS - (monotonic() - started) + remaining_budget = CONNECTION_RETRY_BUDGET_SECONDS - ( + timing.monotonic() - started + ) if remaining_budget <= 0: raise delay = min( - choose_delay(delay_ceiling / 2, delay_ceiling), remaining_budget + timing.choose_delay(delay_ceiling / 2, delay_ceiling), remaining_budget ) _report_connection_retry( runtime, @@ -66,7 +80,7 @@ def open_connection_with_retry( delay_seconds=delay, error=error, ) - sleep(delay) + timing.sleep(delay) attempt += 1 delay_ceiling = min( delay_ceiling * RETRY_DELAY_GROWTH_FACTOR, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py index e00bdf0da..29cfbf255 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py @@ -7,6 +7,7 @@ import pytest from policyengine_simulation_executor.stage12_runtime import database_retry from policyengine_simulation_executor.stage12_runtime.database_retry import ( + RetryTiming, open_connection_with_retry, ) from sqlalchemy.exc import OperationalError @@ -73,9 +74,11 @@ def connect_after_two_refusals() -> OpenedConnection: connect_after_two_refusals, operation_name="create_or_resolve_simulation", runtime=runtime, - sleep=clock.sleep, - monotonic=clock.monotonic, - choose_delay=_upper_bound_delay, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ), ) assert result is opened_connection @@ -118,9 +121,11 @@ def connect_on_third_attempt() -> OpenedConnection: open_connection_with_retry( connect_on_third_attempt, operation_name="get_simulation", - sleep=clock.sleep, - monotonic=clock.monotonic, - choose_delay=lowest_delay, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=lowest_delay, + ), ) assert requested_ranges == [(0.5, 1.0), (1.0, 2.0)] @@ -140,8 +145,10 @@ def fail_with_invalid_connection_options() -> None: open_connection_with_retry( fail_with_invalid_connection_options, operation_name="get_report", - sleep=clock.sleep, - monotonic=clock.monotonic, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + ), ) assert attempts == [1] @@ -160,9 +167,11 @@ def never_connect() -> None: open_connection_with_retry( never_connect, operation_name="create_or_resolve_report", - sleep=clock.sleep, - monotonic=clock.monotonic, - choose_delay=_upper_bound_delay, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ), ) budget = database_retry.CONNECTION_RETRY_BUDGET_SECONDS @@ -192,8 +201,10 @@ def connect_on_second_attempt() -> OpenedConnection: connect_on_second_attempt, operation_name="create_or_resolve_report", runtime=FailingRuntime(), - sleep=clock.sleep, - monotonic=clock.monotonic, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + ), ) is opened_connection ) @@ -215,9 +226,11 @@ def connect_on_second_attempt() -> OpenedConnection: open_connection_with_retry( connect_on_second_attempt, operation_name="replace_simulation", - sleep=clock.sleep, - monotonic=clock.monotonic, - choose_delay=_upper_bound_delay, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ), ) [record] = caplog.records @@ -239,9 +252,11 @@ def never_connect() -> None: open_connection_with_retry( never_connect, operation_name="get_report", - sleep=clock.sleep, - monotonic=clock.monotonic, - choose_delay=_upper_bound_delay, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ), ) assert clock.sleeps == [1.0, 2.0, 4.0, 8.0, 15.0, 15.0, 15.0] @@ -259,9 +274,11 @@ def slowly_refuse() -> None: open_connection_with_retry( slowly_refuse, operation_name="get_report", - sleep=clock.sleep, - monotonic=clock.monotonic, - choose_delay=_upper_bound_delay, + timing=RetryTiming( + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ), ) # 12 + 1 + 12 + 2 + 12 + 4 + 12 seconds leave 5 of the 60-second budget. From a63fd9f6b45a3e8490bf9770772213c179ab9dd0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:31:11 +0400 Subject: [PATCH 55/66] Pin the workflow expression evaluator and routing command discovery Pin what the GitHub expression evaluator returns for each operator and status function, and the argv and refusal messages routing command discovery produces, before splitting both test helpers. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_deployment_assets.py | 40 +++++++++++ .../tests/test_routing_workflow_invariants.py | 71 +++++++++++++++++++ rules/registry.toml | 2 + 3 files changed, 113 insertions(+) diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index 370fcfddf..2e03ba852 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -983,6 +983,46 @@ def test_the_expression_evaluator_follows_github_semantics(): assert not _gh_eval("${{ !cancelled() }}", {**context, "cancelled": True}) +def test_the_expression_evaluator_returns_operands_and_job_statuses(): + context = { + "inputs": {"name": "beta", "empty": "", "zero": 0, "flag": False}, + "needs": {"a": {"result": "success"}, "b": {"result": "skipped"}}, + "steps": "not-a-mapping", + } + # || and && return one of their operands, as GitHub does, not a boolean. + assert _gh_eval("${{ inputs.empty || 'fallback' }}", context) == "fallback" + assert _gh_eval("${{ inputs.name || 'fallback' }}", context) == "beta" + assert _gh_eval("${{ inputs.name && inputs.zero }}", context) == 0 + assert _gh_eval("${{ inputs.flag && inputs.name }}", context) is False + assert _gh_eval("${{ inputs.name != 'BETA' }}", context) is False + assert _gh_eval("${{ !(inputs.empty || inputs.zero) }}", context) is True + assert _gh_eval("${{ steps.x.outputs }}", context) is None + assert _gh_eval("${{ inputs.missing.deeper }}", context) is None + assert _gh_eval("false", context) is False + # success() needs every needed job to succeed; failure() needs one failure. + assert _gh_eval("${{ success() }}", context) is False + assert _gh_eval("${{ failure() }}", context) is False + assert _gh_eval("${{ success() }}", {}) is True + assert _gh_eval("${{ success() }}", {"cancelled": True}) is False + assert _gh_eval("${{ cancelled() }}", {"cancelled": True}) is True + # Any status function in a job if replaces the implicit success(). + assert _gh_eval("${{ inputs.name == 'beta' }}", context, job_if=True) is False + assert ( + _gh_eval("${{ inputs.name == 'beta' || failure() }}", context, job_if=True) + is True + ) + + +@pytest.mark.parametrize( + "expression", + ["${{ unknown() }}", "${{ inputs.name inputs.name }}", "${{ inputs.name @ }}"], + ids=["unknown-function", "trailing-token", "unknown-character"], +) +def test_the_expression_evaluator_refuses_what_it_cannot_read(expression): + with pytest.raises(AssertionError): + _gh_eval(expression, {"inputs": {"name": "beta"}}) + + def _routing_context(**overrides): context = { "inputs": {"deploy_existing_stack": True, "promote_entrypoint": True}, diff --git a/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py b/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py index 2ab1ecbf2..6bbfa7885 100644 --- a/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py +++ b/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py @@ -710,3 +710,74 @@ def test_finalizer_reference_cannot_bypass_routing_cli_discovery(): def test_argument_value_cannot_hide_a_default_writer(): script = f"uv run python -m {REGISTRY_COMMAND} --app-name app--dry-run" assert list(_writing_steps({"steps": [{"run": script}]})) + + +# --- Characterization: routing command discovery's argv and refusals --------- + + +def test_routing_command_discovery_yields_each_foreground_argv_in_order(): + script = ( + f"# uv run python -m {REGISTRY_COMMAND} --apply-plan ignored.json\n" + 'base=(--environment "${MODAL_ENVIRONMENT}")\n' + "base+=(--dry-run)\n" + f"uv run python -m {REGISTRY_COMMAND} \\\n --print-latest && " + "python src/modal/utils/update_version_registry.py " + '"${base[@]}" --app-name app\n' + "echo done\n" + ) + + assert list(_routing_commands(script)) == [ + ("--print-latest",), + ("--environment", "${MODAL_ENVIRONMENT}", "--dry-run", "--app-name", "app"), + ] + + +def test_routing_command_discovery_ignores_scripts_without_routing(): + assert list(_routing_commands("echo update_version\nmake test\n")) == [] + + +@pytest.mark.parametrize( + ("script", "message"), + [ + ( + f'args+=(--dry-run)\nuv run python -m {REGISTRY_COMMAND} "${{args[@]}}"', + "Cannot append unresolved routing array args", + ), + ( + f"uv run python -m {REGISTRY_COMMAND} python -m {REGISTRY_COMMAND}", + "Unsupported nested routing invocation", + ), + ( + f"python3 -m {REGISTRY_COMMAND} --print-latest", + "Unsupported routing shell prefix: ['python3', '-m']", + ), + ( + "uv run python -m src/modal/utils/update_version_registry.py", + "Unsupported routing shell prefix: ['uv', 'run', 'python', '-m']", + ), + ( + f"echo update_version_registry; uv run python -m {REGISTRY_COMMAND}", + "Unsupported inline routing library invocation", + ), + ( + f"uv run python -m {REGISTRY_COMMAND} --bogus", + "Unsupported routing option or shell form: --bogus", + ), + ( + f"uv run python -m {REGISTRY_COMMAND} --print-latest`date`", + "Unsupported routing command substitution", + ), + ], + ids=[ + "append-before-declare", + "nested", + "module-prefix", + "file-prefix", + "inline-reference", + "unknown-option", + "backtick", + ], +) +def test_routing_command_discovery_names_each_refusal(script, message): + with pytest.raises(AssertionError, match=f"^{re.escape(message)}"): + list(_routing_commands(script)) diff --git a/rules/registry.toml b/rules/registry.toml index d42d1ca7c..c809cbb68 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -54,6 +54,8 @@ guard_test_globs = [ ] ignore_tests = [ { node = "projects/policyengine-simulation-entry/tests/test_deployment_assets.py::test_the_expression_evaluator_follows_github_semantics", reason = "Checks the GitHub Actions expression evaluator that the workflow-condition tests use; it enforces no repository rule itself." }, + { node = "projects/policyengine-simulation-entry/tests/test_deployment_assets.py::test_the_expression_evaluator_returns_operands_and_job_statuses", reason = "Pins the operand values and job statuses the GitHub Actions expression evaluator returns; it enforces no repository rule itself." }, + { node = "projects/policyengine-simulation-entry/tests/test_deployment_assets.py::test_the_expression_evaluator_refuses_what_it_cannot_read", reason = "Pins that the GitHub Actions expression evaluator refuses expressions it cannot read; it enforces no repository rule itself." }, ] [agents] From 2dd32a57584aae290769fff65cd9e1e3b2eb8f4f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:31:49 +0400 Subject: [PATCH 56/66] Split the workflow expression evaluator into a tokenizer and a parser Co-Authored-By: Claude Opus 5.5 --- .../tests/test_deployment_assets.py | 126 +++++++++++------- 1 file changed, 78 insertions(+), 48 deletions(-) diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index 2e03ba852..72853cd62 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -889,7 +889,8 @@ def _gh_equal(left, right) -> bool: return _gh_number(left) == _gh_number(right) -def _gh_eval(expression: str, context: dict, *, job_if: bool = False): +def _gh_tokens(expression: str) -> list[str]: + """Split an expression into tokens, dropping its ``${{`` and ``}}``.""" tokens = [] position = 0 text = expression.strip() @@ -898,77 +899,106 @@ def _gh_eval(expression: str, context: dict, *, job_if: bool = False): assert match, text[position:] tokens.append(match.group(1)) position = match.end() - tokens = [token for token in tokens if token not in ("${{", "}}")] - if job_if and not any( + return [token for token in tokens if token not in ("${{", "}}")] + + +def _gh_calls_a_status_function(tokens: list[str]) -> bool: + return any( token in _STATUS_FUNCTIONS and tokens[index + 1 : index + 2] == ["("] for index, token in enumerate(tokens) - ): - tokens = ["success", "(", ")", "&&", "(", *tokens, ")"] - results = [job.get("result") for job in (context.get("needs") or {}).values()] + ) + + +def _gh_truthy(value) -> bool: + return value not in (None, False, "", 0) + + +class _GitHubExpression: + """A recursive-descent evaluation of one tokenized expression in a context. + + Each ``_parse_*`` method reads from a token index and returns the value + it read and the index after it, from the loosest operator (``||``) to the + tightest (``!``, parentheses and operands). + """ - def status(name): - cancelled = bool(context.get("cancelled")) + def __init__(self, tokens: list[str], context: dict) -> None: + self.tokens = tokens + self.context = context + self.results = [ + job.get("result") for job in (context.get("needs") or {}).values() + ] + + def evaluate(self): + value, index = self._parse_or(0) + assert index == len(self.tokens), self.tokens[index:] + return value + + def _next_is(self, index: int, operators: tuple[str, ...]) -> bool: + return index < len(self.tokens) and self.tokens[index] in operators + + def _status(self, name: str) -> bool: + cancelled = bool(self.context.get("cancelled")) if name == "always": return True if name == "cancelled": return cancelled if name == "failure": - return "failure" in results - return not cancelled and all(result == "success" for result in results) + return "failure" in self.results + return not cancelled and all(result == "success" for result in self.results) - def truthy(value): - return value not in (None, False, "", 0) - - def lookup(path): + def _lookup(self, path: str): if path in ("true", "false"): return path == "true" - value = context + value = self.context for part in path.split("."): value = value.get(part) if isinstance(value, dict) else None return value - def parse_or(index): - left, index = parse_and(index) - while index < len(tokens) and tokens[index] == "||": - right, index = parse_and(index + 1) - left = left if truthy(left) else right + def _parse_or(self, index: int): + left, index = self._parse_and(index) + while self._next_is(index, ("||",)): + right, index = self._parse_and(index + 1) + left = left if _gh_truthy(left) else right return left, index - def parse_and(index): - left, index = parse_compare(index) - while index < len(tokens) and tokens[index] == "&&": - right, index = parse_compare(index + 1) - left = right if truthy(left) else left + def _parse_and(self, index: int): + left, index = self._parse_compare(index) + while self._next_is(index, ("&&",)): + right, index = self._parse_compare(index + 1) + left = right if _gh_truthy(left) else left return left, index - def parse_compare(index): - left, index = parse_unary(index) - while index < len(tokens) and tokens[index] in ("==", "!="): - operator = tokens[index] - right, index = parse_unary(index + 1) + def _parse_compare(self, index: int): + left, index = self._parse_unary(index) + while self._next_is(index, ("==", "!=")): + operator = self.tokens[index] + right, index = self._parse_unary(index + 1) equal = _gh_equal(left, right) left = equal if operator == "==" else not equal return left, index - def parse_unary(index): - if tokens[index] == "!": - value, index = parse_unary(index + 1) - return not truthy(value), index - if tokens[index] == "(": - value, index = parse_or(index + 1) - assert tokens[index] == ")" + def _parse_unary(self, index: int): + symbol = self.tokens[index] + if symbol == "!": + value, index = self._parse_unary(index + 1) + return not _gh_truthy(value), index + if symbol == "(": + value, index = self._parse_or(index + 1) + assert self.tokens[index] == ")" return value, index + 1 - token = tokens[index] - if token.startswith("'"): - return token[1:-1], index + 1 - if tokens[index + 1 : index + 3] == ["(", ")"]: - assert token in _STATUS_FUNCTIONS, token - return status(token), index + 3 - return lookup(token), index + 1 - - value, index = parse_or(0) - assert index == len(tokens), tokens[index:] - return value + if symbol.startswith("'"): + return symbol[1:-1], index + 1 + if self.tokens[index + 1 : index + 3] == ["(", ")"]: + assert symbol in _STATUS_FUNCTIONS, symbol + return self._status(symbol), index + 3 + return self._lookup(symbol), index + 1 + + +def _gh_eval(expression: str, context: dict, *, job_if: bool = False): + tokens = _gh_tokens(expression) + if job_if and not _gh_calls_a_status_function(tokens): + tokens = ["success", "(", ")", "&&", "(", *tokens, ")"] + return _GitHubExpression(tokens, context).evaluate() def test_the_expression_evaluator_follows_github_semantics(): From 55455cb3028d29b645526c59060a39f9f961f700 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:32:38 +0400 Subject: [PATCH 57/66] Split routing command discovery into script, array and per-command checks Co-Authored-By: Claude Opus 5.5 --- .../tests/test_routing_workflow_invariants.py | 183 ++++++++++-------- 1 file changed, 107 insertions(+), 76 deletions(-) diff --git a/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py b/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py index 6bbfa7885..43f47b669 100644 --- a/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py +++ b/projects/policyengine-simulation-entry/tests/test_routing_workflow_invariants.py @@ -98,20 +98,15 @@ def _shell_tokens(line: str) -> list[str]: ] -def _routing_commands(script: str): - """Inspect foreground CLI argv in execution order, refusing unknown forms. - - The supported subset is explicit CLI commands and previously declared - argument arrays. Conditional array appends are conservatively included. - Importing the library or using shell wrappers needs a deliberate extension - here, rather than silently escaping discovery or the ordering proof. - """ +def _without_comments(script: str) -> str: + """Join continued lines and drop comment lines.""" script = re.sub(r"\\\n\s*", " ", script) - script = "\n".join( + return "\n".join( line for line in script.splitlines() if not line.lstrip().startswith("#") ) - if not ROUTING_LIBRARY_REFERENCE.search(script): - return + + +def _assert_protected_environment_kept(script: str) -> None: for variable in PROTECTED_SHELL_ENVIRONMENT: # Shell builtins remove quotes from assignment operands, so both # export "MODAL_ENVIRONMENT"=main and export "MODAL_ENVIRONMENT=main" @@ -123,83 +118,119 @@ def _routing_commands(script: str): script, ), f"Shell overrides the proven routing environment or write guard: {variable}" + +def _logical_lines(script: str) -> list[str]: + """Return the script's lines with each array assignment joined onto one.""" # Preserve source order: a declaration after a command cannot prove that # command's argv. The repository uses only literal array assignments. - logical_lines = re.sub( + return re.sub( r"(\b\w+\+?=\()([^)]*)(\))", lambda match: "".join(match.groups()).replace("\n", " "), script, ).splitlines() + + +def _declare_array(assignment: re.Match[str], arrays: dict[str, list[str]]) -> None: + name, append, values = assignment.groups() + if append: + assert name in arrays, f"Cannot append unresolved routing array {name}" + arrays[name] += _shell_tokens(values) + else: + arrays[name] = _shell_tokens(values) + + +def _foreground_commands(line: str) -> list[list[str]]: + """Split a routing line into its commands, refusing substitution and grouping.""" + assert REGISTRY_INVOCATION.search(line), ( + f"Routing library calls must use the guarded foreground CLI: {line}" + ) + assert "$(" not in line, f"Unsupported routing command substitution: {line}" + assert "`" not in line, f"Unsupported routing command substitution: {line}" + tokens = _shell_tokens(line) + assert not any( + token in {"{", "}"} + or (re.fullmatch(r"[;&|()]+", token) and token not in {";", "&&", "||", "|"}) + for token in tokens + ), f"Routing commands must run in the foreground without shell grouping: {line}" + return [ + list(group) + for key, group in itertools.groupby( + tokens, key=lambda token: token in {";", "&&", "||", "|"} + ) + if not key + ] + + +def _assert_supported_prefix(command: list[str], index: int) -> None: + prefix = command[:index] + if command[index] == REGISTRY_COMMAND: + assert prefix in (["uv", "run", "python", "-m"], ["python", "-m"]), ( + f"Unsupported routing shell prefix: {prefix}" + ) + else: + assert prefix in (["uv", "run", "python"], ["python"]), ( + f"Unsupported routing shell prefix: {prefix}" + ) + + +def _expand_arrays(tokens: list[str], arrays: dict[str, list[str]]) -> tuple[str, ...]: + arguments = [] + for token in tokens: + array = re.fullmatch(r"\$\{(\w+)\[@\]\}", token) + if array: + name = array.group(1) + assert name in arrays, f"Cannot resolve routing argument array {name}" + arguments.extend(arrays[name]) + else: + arguments.append(token) + return tuple(arguments) + + +def _registry_arguments( + command: list[str], arrays: dict[str, list[str]], line: str +) -> tuple[str, ...] | None: + """Return a command's routing CLI argv, or None when it runs something else.""" + invocations = [ + index + for index, token in enumerate(command) + if REGISTRY_INVOCATION.fullmatch(token) + ] + if not invocations: + assert not ROUTING_LIBRARY_REFERENCE.search(" ".join(command)), ( + f"Unsupported inline routing library invocation: {line}" + ) + return None + assert len(invocations) == 1, f"Unsupported nested routing invocation: {line}" + index = invocations[0] + _assert_supported_prefix(command, index) + return _expand_arrays(command[index + 1 :], arrays) + + +def _routing_commands(script: str): + """Inspect foreground CLI argv in execution order, refusing unknown forms. + + The supported subset is explicit CLI commands and previously declared + argument arrays. Conditional array appends are conservatively included. + Importing the library or using shell wrappers needs a deliberate extension + here, rather than silently escaping discovery or the ordering proof. + """ + script = _without_comments(script) + if not ROUTING_LIBRARY_REFERENCE.search(script): + return + _assert_protected_environment_kept(script) arrays: dict[str, list[str]] = {} - for line in logical_lines: + for line in _logical_lines(script): assignment = re.fullmatch(r"\s*(\w+)(\+?)=\((.*?)\)\s*", line) if assignment: - name, append, values = assignment.groups() - if append: - assert name in arrays, f"Cannot append unresolved routing array {name}" - arrays[name] += _shell_tokens(values) - else: - arrays[name] = _shell_tokens(values) + _declare_array(assignment, arrays) continue if not ROUTING_LIBRARY_REFERENCE.search(line): continue - assert REGISTRY_INVOCATION.search(line), ( - f"Routing library calls must use the guarded foreground CLI: {line}" - ) - assert "$(" not in line, f"Unsupported routing command substitution: {line}" - assert "`" not in line, f"Unsupported routing command substitution: {line}" - tokens = _shell_tokens(line) - assert not any( - token in {"{", "}"} - or ( - re.fullmatch(r"[;&|()]+", token) and token not in {";", "&&", "||", "|"} - ) - for token in tokens - ), f"Routing commands must run in the foreground without shell grouping: {line}" - commands = [ - list(group) - for key, group in itertools.groupby( - tokens, key=lambda token: token in {";", "&&", "||", "|"} - ) - if not key - ] - for command in commands: - invocations = [ - index - for index, token in enumerate(command) - if REGISTRY_INVOCATION.fullmatch(token) - ] - if not invocations: - assert not ROUTING_LIBRARY_REFERENCE.search(" ".join(command)), ( - f"Unsupported inline routing library invocation: {line}" - ) - continue - assert len(invocations) == 1, ( - f"Unsupported nested routing invocation: {line}" - ) - index = invocations[0] - prefix = command[:index] - if command[index] == REGISTRY_COMMAND: - assert prefix in (["uv", "run", "python", "-m"], ["python", "-m"]), ( - f"Unsupported routing shell prefix: {prefix}" - ) - else: - assert prefix in (["uv", "run", "python"], ["python"]), ( - f"Unsupported routing shell prefix: {prefix}" - ) - arguments = [] - for token in command[index + 1 :]: - array = re.fullmatch(r"\$\{(\w+)\[@\]\}", token) - if array: - name = array.group(1) - assert name in arrays, ( - f"Cannot resolve routing argument array {name}" - ) - arguments.extend(arrays[name]) - else: - arguments.append(token) - _command_options(tuple(arguments)) - yield tuple(arguments) + for command in _foreground_commands(line): + arguments = _registry_arguments(command, arrays, line) + if arguments is not None: + _command_options(arguments) + yield arguments def _command_options(arguments: tuple[str, ...]) -> set[str]: From 8321ba79b16473c21aa3e6a463050e231ac2d6bf Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:35:54 +0400 Subject: [PATCH 58/66] Group the routing and package-update test parameters into cases Give the fake git its reported state as one FakeGitState, pass a planned release's versions to _plan together, and parametrize each multi-value test case as one argument. The parametrized test ids are unchanged. Co-Authored-By: Claude Opus 5.5 --- ...est_policyengine_package_update_scripts.py | 30 ++- ...est_policyengine_package_update_scripts.py | 87 ++++++--- .../tests/test_update_version_registry.py | 184 ++++++++---------- 3 files changed, 160 insertions(+), 141 deletions(-) diff --git a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py index 1043c9956..239876ca6 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_policyengine_package_update_scripts.py @@ -73,21 +73,33 @@ def write_executable(path: Path, content: str) -> None: path.chmod(0o755) +@dataclass(frozen=True) +class FakeGitState: + """What the fake ``git`` reports about the update branch and working tree. + + ``remote_branch_exists`` and ``diff_has_changes`` set the ``ls-remote`` + and ``diff`` exit codes, and ``show`` prints the update branch's pyproject + pinning ``branch_pin``. + """ + + remote_branch_exists: bool = False + diff_has_changes: bool = False + branch_pin: str = "4.1.0" + + def install_fake_git( fake_bin: Path, *, root: Path, log: Path, - remote_branch_exists: bool = False, - diff_has_changes: bool = False, - branch_pin: str = "4.1.0", + state: FakeGitState | None = None, ) -> None: """Install a fake ``git`` that appends each call to ``log``. - ``remote_branch_exists`` and ``diff_has_changes`` set the ``ls-remote`` - and ``diff`` exit codes, and ``show`` prints the update branch's pyproject - pinning ``branch_pin``; every other command succeeds silently. + It answers ``ls-remote``, ``diff`` and ``show`` as ``state`` says; every + other command succeeds silently. """ + state = state or FakeGitState() write_executable( fake_bin / "git", f"""#!/usr/bin/env bash @@ -100,21 +112,21 @@ def install_fake_git( fi if [[ "$1" == "ls-remote" ]]; then - if [[ "{int(remote_branch_exists)}" == "1" ]]; then + if [[ "{int(state.remote_branch_exists)}" == "1" ]]; then exit 0 fi exit 2 fi if [[ "$1" == "diff" ]]; then - if [[ "{int(diff_has_changes)}" == "1" ]]; then + if [[ "{int(state.diff_has_changes)}" == "1" ]]; then exit 1 fi exit 0 fi if [[ "$1" == "show" ]]; then - printf '[project]\ndependencies = ["policyengine[models]=={branch_pin}"]\n' + printf '[project]\ndependencies = ["policyengine[models]=={state.branch_pin}"]\n' exit 0 fi diff --git a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py index a97040588..b23348356 100644 --- a/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_policyengine_package_update_scripts.py @@ -10,6 +10,7 @@ from fixtures.test_policyengine_package_update_scripts import ( SCRIPT, BundledVersions, + FakeGitState, install_fake_gh, install_fake_git, install_fake_uv, @@ -115,7 +116,7 @@ def test_update_policyengine_package_dry_run_reports_existing_branch_recovery( fake_bin, root=fake_repo, log=git_log, - remote_branch_exists=True, + state=FakeGitState(remote_branch_exists=True), ) result = run_updater( @@ -159,7 +160,7 @@ def test_update_policyengine_package_opens_pr_for_existing_branch_without_open_p fake_bin, root=fake_repo, log=git_log, - remote_branch_exists=True, + state=FakeGitState(remote_branch_exists=True), ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv(fake_bin, log=uv_log) @@ -183,28 +184,28 @@ def test_update_policyengine_package_opens_pr_for_existing_branch_without_open_p @pytest.mark.parametrize( - ("uv_kwargs", "git_kwargs", "payload", "message"), + "recovery_case", [ ( - {"bundled": BundledVersions(uk_version="2.0.9")}, - {}, + BundledVersions(uk_version="2.0.9"), + "4.1.0", '{"version":"4.1.0","bundle":{"uk":{"model_version":"2.1.0"}}}', "does not match the release dispatch", ), ( - {}, - {}, + BundledVersions(), + "4.1.0", '{"version":"4.1.0","qualification":{"uk":{"verdict":"rejected",' '"promotable":false}}}', "does not match the release dispatch", ), ( - {"bundled": BundledVersions(policyengine_version="4.0.9")}, - {}, + BundledVersions(policyengine_version="4.0.9"), + "4.1.0", "", "reports bundle 4.0.9, expected 4.1.0", ), - ({}, {"branch_pin": "4.0.5"}, "", "pins policyengine 4.0.5, expected 4.1.0"), + (BundledVersions(), "4.0.5", "", "pins policyengine 4.0.5, expected 4.1.0"), ], ids=["bundle-mismatch", "rejected-verdict", "wrong-wrapper", "stale-branch"], ) @@ -212,22 +213,23 @@ def test_a_recovered_branch_gets_no_pr_without_verification( fake_bin: Path, fake_repo: Path, tmp_path: Path, - uv_kwargs: dict, - git_kwargs: dict, - payload: str, - message: str, + recovery_case: tuple[BundledVersions, str, str, str], ) -> None: - """#743 round-2 review: branch recovery may not bypass validation.""" + """#743 round-2 review: branch recovery may not bypass validation. + + ``recovery_case`` is the bundle the fake ``uv`` reports, the policyengine + pin on the recovered branch, the release payload and the expected error. + """ + bundled, branch_pin, payload, message = recovery_case gh_log = tmp_path / "gh.log" install_fake_git( fake_bin, root=fake_repo, log=tmp_path / "git.log", - remote_branch_exists=True, - **git_kwargs, + state=FakeGitState(remote_branch_exists=True, branch_pin=branch_pin), ) install_fake_gh(fake_bin, log=gh_log) - install_fake_uv(fake_bin, log=tmp_path / "uv.log", **uv_kwargs) + install_fake_uv(fake_bin, log=tmp_path / "uv.log", bundled=bundled) result = run_updater( env=updater_env(fake_bin, LATEST_OVERRIDE="4.1.0", RELEASE_PAYLOAD=payload), @@ -246,7 +248,12 @@ def test_update_policyengine_package_updates_only_wrapper_requirements( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv(fake_bin, log=uv_log) @@ -286,7 +293,12 @@ def test_update_policyengine_package_creates_and_links_missing_issue( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log, open_issue="") install_fake_uv(fake_bin, log=uv_log) @@ -309,7 +321,12 @@ def test_update_policyengine_package_stops_when_loaded_bundle_is_not_target( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv( fake_bin, @@ -335,7 +352,12 @@ def test_update_policyengine_package_requires_modal_models_requirement( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv(fake_bin, log=uv_log) pyproject = fake_repo / "simulation" / "pyproject.toml" @@ -362,7 +384,12 @@ def test_update_policyengine_package_rejects_direct_component_requirements( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv(fake_bin, log=uv_log) pyproject = fake_repo / "simulation" / "pyproject.toml" @@ -692,7 +719,12 @@ def test_update_policyengine_package_stops_when_dispatch_names_another_dataset( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv(fake_bin, log=uv_log) @@ -721,7 +753,12 @@ def test_update_policyengine_package_accepts_a_matching_dispatch( git_log = tmp_path / "git.log" gh_log = tmp_path / "gh.log" uv_log = tmp_path / "uv.log" - install_fake_git(fake_bin, root=fake_repo, log=git_log, diff_has_changes=True) + install_fake_git( + fake_bin, + root=fake_repo, + log=git_log, + state=FakeGitState(diff_has_changes=True), + ) install_fake_gh(fake_bin, log=gh_log) install_fake_uv(fake_bin, log=uv_log) diff --git a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py index 1e973bcc0..5a4dcd5d5 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_version_registry.py +++ b/projects/policyengine-simulation-executor/tests/test_update_version_registry.py @@ -595,16 +595,9 @@ def _publish(monkeypatch, state, policyengine_version, us_version, uk_version): ) -def _plan( - monkeypatch, - state, - version, - us_version, - uk_version, - *, - force=False, - publication_id=None, -): +def _plan(monkeypatch, state, versions, *, force=False, publication_id=None): + """Plan publishing ``versions``: a policyengine, US and UK version.""" + version, us_version, uk_version = versions _manifests(monkeypatch, {version: (us_version, uk_version)}) return registry.plan_routing_publication( current_state=state, @@ -671,9 +664,9 @@ def from_name(name, environment_name, create_if_missing): state["routes"]["policyengine"]["5.9.0"] = _app("5.9.0") state["routes"]["uk"]["2.1.0"] = _app("5.9.0") for version, us, uk, force in releases[:-1]: - state = _plan(monkeypatch, state, version, us, uk, force=force)["after"] + state = _plan(monkeypatch, state, (version, us, uk), force=force)["after"] version, us, uk, force = releases[-1] - plan = _plan(monkeypatch, state, version, us, uk, force=force) + plan = _plan(monkeypatch, state, (version, us, uk), force=force) if state is not None: stores["staging/simulation-api-routing-state"] = FakeDict( {"active": json.loads(json.dumps(state))} @@ -699,7 +692,7 @@ def test_the_review_counterexample_restores_every_pin(monkeypatch, patched_modal state["latest"]["policyengine"] = "6.0.0" state = _publish(monkeypatch, state, "6.1.0", "1.1.0", "2.2.0") store = _store(patched_modal, state) - plan = _plan(monkeypatch, state, "6.2.0", "1.1.0", "2.1.0") + plan = _plan(monkeypatch, state, ("6.2.0", "1.1.0", "2.1.0")) registry.apply_routing_plan(environment="staging", plan=plan) assert store.get("active")["routes"]["uk"]["2.1.0"] == _app("6.2.0") assert registry.restore_routing_plan(environment="staging", plan=plan) == "restored" @@ -710,7 +703,7 @@ def test_the_review_counterexample_restores_every_pin(monkeypatch, patched_modal def test_a_first_publication_is_undone_by_removing_the_state( monkeypatch, patched_modal ): - plan = _plan(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") + plan = _plan(monkeypatch, None, ("6.2.6", "1.1.0", "2.1.0")) registry.apply_routing_plan(environment="staging", plan=plan) store = patched_modal["staging/simulation-api-routing-state"] assert store.get("active") == plan["after"] @@ -721,7 +714,7 @@ def test_a_first_publication_is_undone_by_removing_the_state( def test_a_plan_that_was_never_applied_restores_nothing(monkeypatch, patched_modal): state = _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") store = _store(patched_modal, state) - plan = _plan(monkeypatch, state, "6.3.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, state, ("6.3.0", "1.1.0", "2.2.0")) assert ( registry.restore_routing_plan(environment="staging", plan=plan) == "unchanged" ) @@ -730,7 +723,7 @@ def test_a_plan_that_was_never_applied_restores_nothing(monkeypatch, patched_mod def test_a_plan_is_applied_only_onto_its_own_starting_state(monkeypatch, patched_modal): state = _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") - plan = _plan(monkeypatch, state, "6.3.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, state, ("6.3.0", "1.1.0", "2.2.0")) moved = _publish(monkeypatch, state, "6.2.7", "1.1.0", "2.1.0") store = _store(patched_modal, moved) with pytest.raises(registry.RoutingConflictError, match="nothing was published"): @@ -741,7 +734,7 @@ def test_a_plan_is_applied_only_onto_its_own_starting_state(monkeypatch, patched def test_a_restore_never_overwrites_a_later_change(monkeypatch, patched_modal): state = _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") store = _store(patched_modal, state) - plan = _plan(monkeypatch, state, "6.3.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, state, ("6.3.0", "1.1.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=plan) later = _publish(monkeypatch, store.get("active"), "6.4.0", "1.2.0", "2.2.0") store["active"] = later @@ -754,7 +747,7 @@ def test_a_restore_never_overwrites_a_later_change(monkeypatch, patched_modal): def test_a_tampered_plan_is_refused(monkeypatch, patched_modal, side): state = _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") _store(patched_modal, state) - plan = _plan(monkeypatch, state, "6.3.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, state, ("6.3.0", "1.1.0", "2.2.0")) plan[side] = {**(plan[side] or {}), "generation": "edited"} with pytest.raises(ValueError, match=f"{side} state does not match"): registry.apply_routing_plan(environment="staging", plan=plan) @@ -776,7 +769,7 @@ def test_plans_say_whether_latest_routes_to_the_release( monkeypatch, version, force, advances ): state = _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") - plan = _plan(monkeypatch, state, version, "1.1.0", "2.1.0", force=force) + plan = _plan(monkeypatch, state, (version, "1.1.0", "2.1.0"), force=force) assert plan["advances_latest"] is advances @@ -896,9 +889,9 @@ def test_a_failed_retry_attempt_cannot_hide_the_publication(monkeypatch, patched """#744 review: attempt 1 published A to B; attempt 2 republished B.""" a = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, a) - first = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") + first = _plan(monkeypatch, a, ("6.1.0", "1.1.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=first) - second = _plan(monkeypatch, first["after"], "6.1.0", "1.1.0", "2.2.0") + second = _plan(monkeypatch, first["after"], ("6.1.0", "1.1.0", "2.2.0")) # Each publication is its own state, even of the same release. assert second["before_sha256"] != second["after_sha256"] registry.apply_routing_plan(environment="staging", plan=second) @@ -915,12 +908,12 @@ def test_an_old_plan_cannot_withdraw_a_healthy_republication( """#744 review: A to B failed and was restored; a healthy run republished B.""" a = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, a) - failed = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") + failed = _plan(monkeypatch, a, ("6.1.0", "1.1.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=failed) assert ( registry.restore_routing_plan(environment="staging", plan=failed) == "restored" ) - healthy = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") + healthy = _plan(monkeypatch, a, ("6.1.0", "1.1.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=healthy) with pytest.raises(registry.RoutingConflictError, match="nothing was restored"): registry.restore_routing_plan(environment="staging", plan=failed) @@ -930,9 +923,9 @@ def test_an_old_plan_cannot_withdraw_a_healthy_republication( def test_a_chain_of_attempts_is_undone_in_order(monkeypatch, patched_modal): a = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, a) - first = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") + first = _plan(monkeypatch, a, ("6.1.0", "1.1.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=first) - second = _plan(monkeypatch, first["after"], "6.1.1", "1.2.0", "2.2.0") + second = _plan(monkeypatch, first["after"], ("6.1.1", "1.2.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=second) assert ( registry.restore_routing_plans(environment="staging", plans=[first, second]) @@ -950,9 +943,9 @@ def test_a_chain_of_attempts_is_undone_in_order(monkeypatch, patched_modal): def test_main_restores_several_plans(patched_modal, monkeypatch, capsys, tmp_path): a = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, a) - first = _plan(monkeypatch, a, "6.1.0", "1.1.0", "2.2.0") + first = _plan(monkeypatch, a, ("6.1.0", "1.1.0", "2.2.0")) registry.apply_routing_plan(environment="staging", plan=first) - second = _plan(monkeypatch, first["after"], "6.1.0", "1.1.0", "2.2.0") + second = _plan(monkeypatch, first["after"], ("6.1.0", "1.1.0", "2.2.0")) files = [] for index, plan in enumerate((first, second), start=1): path = tmp_path / f"plan-{index}.json" @@ -1105,7 +1098,7 @@ def test_routing_guard_requires_every_condition(monkeypatch, variable, value): ], ) @pytest.mark.parametrize( - ("variable", "value"), + "environment_override", [ (None, None), ("GITHUB_ACTIONS", "false"), @@ -1119,8 +1112,10 @@ def test_routing_guard_requires_every_condition(monkeypatch, variable, value): ) @pytest.mark.usefixtures("serialized_workflow") def test_library_writers_refuse_unserialized_calls_before_validation_or_modal_access( - monkeypatch, patched_modal, writer, kwargs, variable, value + monkeypatch, patched_modal, writer, kwargs, environment_override ): + """``environment_override`` names a GitHub variable and its value, or None.""" + variable, value = environment_override if variable is None: for name in ("GITHUB_ACTIONS", "GITHUB_REPOSITORY", "GITHUB_WORKFLOW_REF"): monkeypatch.delenv(name, raising=False) @@ -1182,7 +1177,7 @@ def forbidden_access(*_args, **_kwargs): ], ) @pytest.mark.parametrize( - ("variable", "value"), + "environment_override", [ (None, None), ("GITHUB_ACTIONS", "false"), @@ -1196,8 +1191,10 @@ def forbidden_access(*_args, **_kwargs): ) @pytest.mark.usefixtures("serialized_workflow") def test_main_refuses_unserialized_writes_before_accessing_modal_or_plan_files( - monkeypatch, patched_modal, capsys, args, variable, value + monkeypatch, patched_modal, capsys, args, environment_override ): + """``environment_override`` names a GitHub variable and its value, or None.""" + variable, value = environment_override if variable is None: for name in ("GITHUB_ACTIONS", "GITHUB_REPOSITORY", "GITHUB_WORKFLOW_REF"): monkeypatch.delenv(name, raising=False) @@ -1277,7 +1274,7 @@ def test_cli_help_explains_serialized_writes_and_has_no_lease_option( def test_applying_a_plan_verifies_the_written_state(monkeypatch, patched_modal): state = _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") - plan = _plan(monkeypatch, state, "6.3.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, state, ("6.3.0", "1.1.0", "2.2.0")) class DroppedWrite(FakeDict): def __setitem__(self, key, value): @@ -1302,7 +1299,7 @@ def test_restoring_a_plan_verifies_the_written_or_removed_state( if first_publication else _publish(monkeypatch, None, "6.2.6", "1.1.0", "2.1.0") ) - plan = _plan(monkeypatch, state, "6.3.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, state, ("6.3.0", "1.1.0", "2.2.0")) class DroppedRestore(FakeDict): def __setitem__(self, key, value): @@ -1325,11 +1322,11 @@ def test_runner_loss_cannot_be_adopted_by_the_next_release(monkeypatch, patched_ """Failed A cannot become successful B's starting state or rollback target.""" original = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, original) - failed_a = _plan(monkeypatch, original, "6.1.0", "1.1.0", "2.2.0") + failed_a = _plan(monkeypatch, original, ("6.1.0", "1.1.0", "2.2.0")) registry.record_routing_plan(environment="staging", plan=failed_a) registry.apply_routing_plan(environment="staging", plan=failed_a) # Runner disappears here, before outputs or a release completion record. - attempted_b = _plan(monkeypatch, store.get("active"), "6.2.0", "1.2.0", "2.3.0") + attempted_b = _plan(monkeypatch, store.get("active"), ("6.2.0", "1.2.0", "2.3.0")) for operation in (registry.record_routing_plan, registry.apply_routing_plan): with pytest.raises( registry.RoutingConflictError, match="Unfinished routing publication" @@ -1345,7 +1342,7 @@ def test_runner_loss_cannot_be_adopted_by_the_next_release(monkeypatch, patched_ ) assert store.get("active") == original assert store.get("publication") == {"status": "restored", "plan": failed_a} - healthy_b = _plan(monkeypatch, original, "6.2.0", "1.2.0", "2.3.0") + healthy_b = _plan(monkeypatch, original, ("6.2.0", "1.2.0", "2.3.0")) registry.record_routing_plan(environment="staging", plan=healthy_b) registry.apply_routing_plan(environment="staging", plan=healthy_b) registry.complete_routing_plan(environment="staging", plan=healthy_b) @@ -1353,7 +1350,7 @@ def test_runner_loss_cannot_be_adopted_by_the_next_release(monkeypatch, patched_ with pytest.raises(registry.RoutingConflictError, match="nothing was restored"): registry.restore_routing_plan(environment="staging", plan=failed_a) # A subsequent failed release restores healthy B and cannot resurrect A. - failed_c = _plan(monkeypatch, healthy_b["after"], "6.3.0", "1.3.0", "2.4.0") + failed_c = _plan(monkeypatch, healthy_b["after"], ("6.3.0", "1.3.0", "2.4.0")) registry.record_routing_plan(environment="staging", plan=failed_c) registry.apply_routing_plan(environment="staging", plan=failed_c) assert ( @@ -1374,7 +1371,7 @@ def test_a_failed_same_version_redeploy_is_resolved_by_its_restore( """ live = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, live) - redeploy = _plan(monkeypatch, live, "6.0.0", "1.1.0", "2.1.0") + redeploy = _plan(monkeypatch, live, ("6.0.0", "1.1.0", "2.1.0")) assert redeploy["before_sha256"] != redeploy["after_sha256"] changed = { key @@ -1385,7 +1382,7 @@ def test_a_failed_same_version_redeploy_is_resolved_by_its_restore( registry.record_routing_plan(environment="staging", plan=redeploy) registry.apply_routing_plan(environment="staging", plan=redeploy) # A release gate fails here. Until the restore runs, nothing else may write. - later = _plan(monkeypatch, live, "6.1.0", "1.1.0", "2.2.0") + later = _plan(monkeypatch, live, ("6.1.0", "1.1.0", "2.2.0")) with pytest.raises( registry.RoutingConflictError, match="Unfinished routing publication" ): @@ -1409,7 +1406,7 @@ def test_recorded_but_unpublished_plan_blocks_writes_until_restored( else _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") ) store = _store(patched_modal, before) - pending = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + pending = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) registry.record_routing_plan(environment="staging", plan=pending) # Covers loss after recording, before upload or before apply. with pytest.raises( @@ -1431,11 +1428,11 @@ def test_every_other_routing_write_refuses_an_unfinished_release( ): before = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, before) - pending = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + pending = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) registry.record_routing_plan(environment="staging", plan=pending) registry.apply_routing_plan(environment="staging", plan=pending) saved = json.loads(json.dumps(store.snapshot())) - other = _plan(monkeypatch, pending["after"], "6.2.0", "1.2.0", "2.3.0") + other = _plan(monkeypatch, pending["after"], ("6.2.0", "1.2.0", "2.3.0")) def write_routing(): if writer == "publish": @@ -1467,7 +1464,7 @@ def write_routing(): def test_only_matching_live_publication_can_be_completed(monkeypatch, patched_modal): before = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, before) - pending = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + pending = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) with pytest.raises(registry.RoutingConflictError, match="no matching unfinished"): registry.complete_routing_plan(environment="staging", plan=pending) registry.record_routing_plan(environment="staging", plan=pending) @@ -1476,7 +1473,7 @@ def test_only_matching_live_publication_can_be_completed(monkeypatch, patched_mo ): registry.complete_routing_plan(environment="staging", plan=pending) registry.apply_routing_plan(environment="staging", plan=pending) - other = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + other = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) with pytest.raises(registry.RoutingConflictError, match="no matching unfinished"): registry.complete_routing_plan(environment="staging", plan=other) assert store.get("publication")["status"] == "pending" @@ -1491,7 +1488,7 @@ def test_publication_ownership_compares_the_entire_plan( """A reused publication id never grants permission to a changed plan.""" before = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") store = _store(patched_modal, before) - pending = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + pending = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) registry.record_routing_plan(environment="staging", plan=pending) if operation != "apply": registry.apply_routing_plan(environment="staging", plan=pending) @@ -1511,7 +1508,7 @@ def test_publication_ownership_compares_the_entire_plan( @pytest.mark.parametrize( - ("scenario", "applied", "gates_passed", "expected"), + "release_case", [ ("failed_release", True, False, "restored"), ("successful_release", True, True, "completed"), @@ -1521,18 +1518,18 @@ def test_publication_ownership_compares_the_entire_plan( ("failure_inside_deployment_ready", True, True, "completed"), ("passing_gates_but_apply_never_ran", False, True, "restored"), ], + ids=lambda release_case: "-".join(str(part) for part in release_case), ) @pytest.mark.parametrize("first_publication", [False, True]) def test_finalization_finishes_each_release_and_allows_the_next_one( - monkeypatch, - patched_modal, - scenario, - applied, - gates_passed, - expected, - first_publication, + monkeypatch, patched_modal, release_case, first_publication ): - """Upload loss and attempt 2 use the full journal, with no artifact reads.""" + """Upload loss and attempt 2 use the full journal, with no artifact reads. + + ``release_case`` is the scenario, whether its plan was applied, whether + its gates passed and the finalization outcome. + """ + scenario, applied, gates_passed, expected = release_case before = ( None if first_publication @@ -1542,9 +1539,7 @@ def test_finalization_finishes_each_release_and_allows_the_next_one( pending = _plan( monkeypatch, before, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) @@ -1570,9 +1565,7 @@ def test_finalization_finishes_each_release_and_allows_the_next_one( later = _plan( monkeypatch, expected_live, - "6.2.0", - "1.2.0", - "2.3.0", + ("6.2.0", "1.2.0", "2.3.0"), publication_id="74402-1-beta", ) registry.record_routing_plan(environment="staging", plan=later) @@ -1588,14 +1581,12 @@ def test_lost_finalizer_is_recovered_by_operator_without_an_artifact( pending = _plan( monkeypatch, before, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) registry.apply_routing_plan(environment="staging", plan=pending) - later = _plan(monkeypatch, pending["after"], "6.2.0", "1.2.0", "2.3.0") + later = _plan(monkeypatch, pending["after"], ("6.2.0", "1.2.0", "2.3.0")) with pytest.raises( registry.RoutingConflictError, match="Unfinished routing publication" ): @@ -1609,7 +1600,7 @@ def test_lost_finalizer_is_recovered_by_operator_without_an_artifact( ) assert store.get("active") == before assert store.get("publication") == {"status": "restored", "plan": pending} - recovered_later = _plan(monkeypatch, before, "6.2.0", "1.2.0", "2.3.0") + recovered_later = _plan(monkeypatch, before, ("6.2.0", "1.2.0", "2.3.0")) registry.record_routing_plan(environment="staging", plan=recovered_later) registry.apply_routing_plan(environment="staging", plan=recovered_later) assert store.get("active") == recovered_later["after"] @@ -1634,9 +1625,7 @@ def test_finalizer_leaves_a_foreign_pending_publication_untouched( pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id=publication_id, ) registry.record_routing_plan(environment="staging", plan=pending) @@ -1666,9 +1655,7 @@ def test_finalizer_conflict_is_loud_and_keeps_the_journal_pending( pending = _plan( monkeypatch, before, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) @@ -1697,9 +1684,7 @@ def test_finalizer_without_its_pending_record_changes_nothing( plan = _plan( monkeypatch, before, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74400-1-beta", ) registry.record_routing_plan(environment="staging", plan=plan) @@ -1729,9 +1714,7 @@ def test_retry_after_finalization_returns_the_prior_outcome_without_writing( pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) @@ -1772,9 +1755,7 @@ def test_completed_finalizer_retry_cannot_report_a_release_that_is_no_longer_liv pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) @@ -1801,9 +1782,7 @@ def test_finalizer_owns_prod_publications_in_the_main_modal_environment( pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-prod", ) registry.record_routing_plan(environment="main", plan=pending) @@ -1845,9 +1824,7 @@ def test_operator_restores_any_valid_pending_plan_in_the_selected_store( pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id=publication_id, ) registry.record_routing_plan(environment="staging", plan=pending) @@ -1868,7 +1845,7 @@ def test_operator_restores_any_valid_pending_plan_in_the_selected_store( @pytest.mark.parametrize("status", ["pending", "completed", "restored"]) def test_publication_records_are_read_back(monkeypatch, patched_modal, status): before = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") - plan = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) class DroppedJournalWrite(FakeDict): def __setitem__(self, key, value): @@ -1895,7 +1872,7 @@ def __setitem__(self, key, value): def test_recording_checks_the_live_starting_state(monkeypatch, patched_modal): before = _publish(monkeypatch, None, "6.0.0", "1.1.0", "2.1.0") - plan = _plan(monkeypatch, before, "6.1.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, before, ("6.1.0", "1.1.0", "2.2.0")) later = _publish(monkeypatch, before, "6.2.0", "1.2.0", "2.3.0") store = _store(patched_modal, later) with pytest.raises(registry.RoutingConflictError, match="changed before recording"): @@ -1911,7 +1888,7 @@ def test_malformed_publication_record_refuses_writes( ): store = _store(patched_modal, None) store["publication"] = record - plan = _plan(monkeypatch, None, "6.1.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, None, ("6.1.0", "1.1.0", "2.2.0")) with pytest.raises((registry.RoutingConflictError, ValueError)): registry.apply_routing_plan(environment="staging", plan=plan) assert store.snapshot() == {"publication": record} @@ -1920,7 +1897,7 @@ def test_malformed_publication_record_refuses_writes( def test_cli_records_and_completes_release_publications( monkeypatch, patched_modal, tmp_path, capsys ): - plan = _plan(monkeypatch, None, "6.1.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, None, ("6.1.0", "1.1.0", "2.2.0")) path = tmp_path / "plan.json" path.write_text(json.dumps(plan)) _main(monkeypatch, "--record-plan", str(path)) @@ -1935,30 +1912,25 @@ def test_cli_records_and_completes_release_publications( @pytest.mark.parametrize( - ("publication_run", "gates_passed", "expected"), + "finalization_case", [ ("74401", "true", "completed"), ("74401", "false", "restored"), ("74400", "false", "foreign-publication"), (None, "false", "restored"), ], + ids=lambda finalization_case: "-".join(str(part) for part in finalization_case), ) def test_cli_finalizes_the_printed_journal_and_reports_its_outcome( - monkeypatch, - patched_modal, - tmp_path, - capsys, - publication_run, - gates_passed, - expected, + monkeypatch, patched_modal, tmp_path, capsys, finalization_case ): + """``finalization_case`` is the publication run, the gates and the outcome.""" + publication_run, gates_passed, expected = finalization_case store = _store(patched_modal, None) pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) @@ -2032,9 +2004,7 @@ def test_cli_output_loss_is_retryable_after_the_journal_is_completed( pending = _plan( monkeypatch, None, - "6.1.0", - "1.1.0", - "2.2.0", + ("6.1.0", "1.1.0", "2.2.0"), publication_id="74401-1-beta", ) registry.record_routing_plan(environment="staging", plan=pending) @@ -2081,7 +2051,7 @@ def test_cli_finalizer_requires_explicit_safe_inputs(monkeypatch, patched_modal, def test_pending_plan_can_be_inspected_locally_before_upload( monkeypatch, patched_modal, capsys ): - plan = _plan(monkeypatch, None, "6.1.0", "1.1.0", "2.2.0") + plan = _plan(monkeypatch, None, ("6.1.0", "1.1.0", "2.2.0")) registry.record_routing_plan(environment="staging", plan=plan) store = patched_modal["staging/simulation-api-routing-state"] saved = json.loads(json.dumps(store.snapshot())) @@ -2569,12 +2539,12 @@ def test_restoring_no_plans_is_refused(): def test_restoring_without_a_routing_store(monkeypatch, patched_modal): - first = _plan(monkeypatch, None, "6.1.0", "1.1.0", "2.2.0") + first = _plan(monkeypatch, None, ("6.1.0", "1.1.0", "2.2.0")) assert ( registry.restore_routing_plans(environment="staging", plans=[first]) == "unchanged" ) - later = _plan(monkeypatch, first["after"], "6.2.0", "1.1.0", "2.3.0") + later = _plan(monkeypatch, first["after"], ("6.2.0", "1.1.0", "2.3.0")) with pytest.raises( registry.RoutingConflictError, match=r"^there is no routing state to restore into$", From 9c703572dfeef9cbae8103fe4f6b5b1795821ff0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:36:23 +0400 Subject: [PATCH 59/66] Take the Public API's group 8b rules tool refactor Copy the Public API's scripts/rules, whose group 8b commits split the command line, the registry reader, validation, the workflow job parser and the mermaid graph builder, and pin their behaviour with new tests. The two copies were identical apart from that refactor and its tests. Co-Authored-By: Claude Opus 5.5 --- scripts/rules/cli.py | 60 +-- scripts/rules/load.py | 147 +++++--- scripts/rules/model.py | 66 ++-- scripts/rules/render.py | 92 +++-- scripts/rules/resolve.py | 56 +-- .../tests/fixtures/validate/broken_rules.toml | 159 ++++++++ .../validate/broken_rules_violations.txt | 41 +++ scripts/rules/tests/test_cli.py | 239 ++++++++++++ scripts/rules/tests/test_load.py | 194 ++++++++++ scripts/rules/tests/test_render.py | 211 +++++++++++ scripts/rules/tests/test_resolve.py | 196 ++++++++++ scripts/rules/tests/test_validate.py | 40 ++ scripts/rules/tests/test_workflows.py | 157 +++++++- scripts/rules/validate.py | 341 ++++++++++++------ scripts/rules/workflows.py | 338 ++++++++++------- 15 files changed, 1928 insertions(+), 409 deletions(-) create mode 100644 scripts/rules/tests/fixtures/validate/broken_rules.toml create mode 100644 scripts/rules/tests/fixtures/validate/broken_rules_violations.txt create mode 100644 scripts/rules/tests/test_cli.py create mode 100644 scripts/rules/tests/test_load.py create mode 100644 scripts/rules/tests/test_render.py create mode 100644 scripts/rules/tests/test_resolve.py create mode 100644 scripts/rules/tests/test_validate.py diff --git a/scripts/rules/cli.py b/scripts/rules/cli.py index ac232ae35..7917b2f4c 100644 --- a/scripts/rules/cli.py +++ b/scripts/rules/cli.py @@ -28,6 +28,38 @@ def _print(violations: list, notes: list[str] | None = None) -> int: return 1 if violations else 0 +def _ci_scripts_command(root: Path, *, write_baselines: bool) -> int: + """Check the workflow scripts, or record today's findings as the baselines.""" + if not write_baselines: + violations, notes = ci_scripts.run_check(root) + return _print(violations, notes) + registry, violations = load_registry(root) + if registry is None or registry.ci_scripts is None: + return _print(violations or ["no [ci_scripts] table"]) + ci_scripts.write_baselines( + registry, ci_scripts.find(registry.ci_scripts, root), root + ) + _report("wrote ci_scripts baselines; update the rules' baseline counts") + return 0 + + +def _registry_command(command: str, root: Path) -> int: + """Validate the registry and, for ``render``, write its views once it is valid.""" + clear_caches() + registry, violations = load_registry(root) + if registry is None: + return _print(violations) + violations += validate(registry, root, date.today()) # noqa: DTZ011 # review and expiry dates are calendar days, checked against the local date + if command != "render": + return _print(violations) + if violations: + _report("render refused: fix the registry first") + return _print(violations) + for path in write_outputs(registry, root): + _report(f"wrote {path.relative_to(root)}") + return 0 + + def main(argv: list[str] | None = None) -> int: """Run one rules command and return 1 if it found violations. @@ -56,32 +88,8 @@ def main(argv: list[str] | None = None) -> int: root = args.root.resolve() if args.command == "ci-scripts": - if args.write_baselines: - registry, violations = load_registry(root) - if registry is None or registry.ci_scripts is None: - return _print(violations or ["no [ci_scripts] table"]) - ci_scripts.write_baselines( - registry, ci_scripts.find(registry.ci_scripts, root), root - ) - _report("wrote ci_scripts baselines; update the rules' baseline counts") - return 0 - violations, notes = ci_scripts.run_check(root) - return _print(violations, notes) - + return _ci_scripts_command(root, write_baselines=args.write_baselines) if args.command == "guard": result = run_guard(root, base_ref=args.base_ref) return _print(result.violations, result.notes) - - clear_caches() - registry, violations = load_registry(root) - if registry is None: - return _print(violations) - violations += validate(registry, root, date.today()) # noqa: DTZ011 # review and expiry dates are calendar days, checked against the local date - if args.command == "render": - if violations: - _report("render refused: fix the registry first") - return _print(violations) - for path in write_outputs(registry, root): - _report(f"wrote {path.relative_to(root)}") - return 0 - return _print(violations) + return _registry_command(args.command, root) diff --git a/scripts/rules/load.py b/scripts/rules/load.py index 936a3a9f4..0a235a2b7 100644 --- a/scripts/rules/load.py +++ b/scripts/rules/load.py @@ -230,23 +230,23 @@ def _ci_scripts(c: _Collector, raw: object) -> CiScripts | None: ) -def _rule(c: _Collector, raw: object, index: int) -> Rule | None: - if not isinstance(raw, dict): - c.add(None, f"rule #{index + 1} must be a table", "write [[rule]]") - return None - rule_id = raw.get("id") if isinstance(raw.get("id"), str) else None - label = rule_id or f"rule #{index + 1}" - c.unknown_keys(rule_id, raw, _RULE_KEYS, label) - +def _enforcer( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> Enforcer: + """Read ``[rule.enforcer]``; a missing table or unknown kind enforces nothing.""" enforcer_raw = raw.get("enforcer") - enforcer = Enforcer(Kind.NONE) - if isinstance(enforcer_raw, dict): - kind = c.enum(rule_id, Kind, enforcer_raw.get("kind"), f"{label}.enforcer.kind") - fields = {k: v for k, v in enforcer_raw.items() if k != "kind"} - enforcer = Enforcer(kind or Kind.NONE, fields) - else: + if not isinstance(enforcer_raw, dict): c.add(rule_id, f"{label}.enforcer must be a table", "add [rule.enforcer]") + return Enforcer(Kind.NONE) + kind = c.enum(rule_id, Kind, enforcer_raw.get("kind"), f"{label}.enforcer.kind") + fields = {k: v for k, v in enforcer_raw.items() if k != "kind"} + return Enforcer(kind or Kind.NONE, fields) + +def _ci_refs( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> tuple[CiRef, ...]: + """Read the ``[[rule.ci]]`` entries, skipping any that are not tables.""" ci_refs = [] ci_raw = raw.get("ci", []) for entry in ci_raw if isinstance(ci_raw, list) else [ci_raw]: @@ -261,29 +261,45 @@ def _rule(c: _Collector, raw: object, index: int) -> Rule | None: c.string(rule_id, entry, "invocation", f"{label}.ci"), ) ) + return tuple(ci_refs) + + +def _proof( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> Proof | None: + """Read ``[rule.proof]``, which is optional but must be a table when present.""" + if "proof" not in raw: + return None + if isinstance(raw["proof"], dict): + return Proof(dict(raw["proof"])) + c.add(rule_id, f"{label}.proof must be a table", "write [rule.proof]") + return None + + +def _baseline( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> Baseline | None: + """Read ``[rule.baseline]``, which needs a file and an integer count.""" + if "baseline" not in raw: + return None + braw = raw["baseline"] + if isinstance(braw, dict) and isinstance(braw.get("count"), int): + c.unknown_keys(rule_id, braw, {"file", "count"}, f"{label}.baseline") + return Baseline( + c.string(rule_id, braw, "file", f"{label}.baseline"), braw["count"] + ) + c.add( + rule_id, + f"{label}.baseline needs 'file' and an integer 'count'", + "write [rule.baseline] with file and count", + ) + return None - proof = None - if "proof" in raw: - if isinstance(raw["proof"], dict): - proof = Proof(dict(raw["proof"])) - else: - c.add(rule_id, f"{label}.proof must be a table", "write [rule.proof]") - - baseline = None - if "baseline" in raw: - braw = raw["baseline"] - if isinstance(braw, dict) and isinstance(braw.get("count"), int): - c.unknown_keys(rule_id, braw, {"file", "count"}, f"{label}.baseline") - baseline = Baseline( - c.string(rule_id, braw, "file", f"{label}.baseline"), braw["count"] - ) - else: - c.add( - rule_id, - f"{label}.baseline needs 'file' and an integer 'count'", - "write [rule.baseline] with file and count", - ) +def _exceptions( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> tuple[RuleException, ...]: + """Read the ``[[rule.exception]]`` entries; one without a date never expires.""" exceptions = [] exc_raw = raw.get("exception", []) for entry in exc_raw if isinstance(exc_raw, list) else [exc_raw]: @@ -305,24 +321,53 @@ def _rule(c: _Collector, raw: object, index: int) -> Rule | None: expires or date.max, ) ) + return tuple(exceptions) - downgrade_reason = None - if "downgrade" in raw: - draw = raw["downgrade"] - if isinstance(draw, dict) and isinstance(draw.get("reason"), str): - downgrade_reason = draw["reason"] - else: - c.add( - rule_id, - f"{label}.downgrade needs a 'reason'", - 'write [rule.downgrade] reason = "..."', - ) +def _downgrade_reason( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> str | None: + """Read the reason in ``[rule.downgrade]``, which is optional.""" + if "downgrade" not in raw: + return None + draw = raw["downgrade"] + if isinstance(draw, dict) and isinstance(draw.get("reason"), str): + reason: str = draw["reason"] + return reason + c.add( + rule_id, + f"{label}.downgrade needs a 'reason'", + 'write [rule.downgrade] reason = "..."', + ) + return None + + +def _autofix( + c: _Collector, rule_id: str | None, raw: dict[str, Any], label: str +) -> bool: + """Read ``autofix``, which must be a boolean; anything else counts as false.""" autofix = raw.get("autofix") - if not isinstance(autofix, bool): - c.add(rule_id, f"{label}.autofix must be true or false", "set autofix") - autofix = False + if isinstance(autofix, bool): + return autofix + c.add(rule_id, f"{label}.autofix must be true or false", "set autofix") + return False + + +def _rule(c: _Collector, raw: object, index: int) -> Rule | None: + if not isinstance(raw, dict): + c.add(None, f"rule #{index + 1} must be a table", "write [[rule]]") + return None + rule_id = raw.get("id") if isinstance(raw.get("id"), str) else None + label = rule_id or f"rule #{index + 1}" + c.unknown_keys(rule_id, raw, _RULE_KEYS, label) + enforcer = _enforcer(c, rule_id, raw, label) + ci_refs = _ci_refs(c, rule_id, raw, label) + proof = _proof(c, rule_id, raw, label) + baseline = _baseline(c, rule_id, raw, label) + exceptions = _exceptions(c, rule_id, raw, label) + downgrade_reason = _downgrade_reason(c, rule_id, raw, label) + autofix = _autofix(c, rule_id, raw, label) review_by = None if "review_by" in raw: review_by = c.day(rule_id, raw["review_by"], f"{label}.review_by") @@ -342,11 +387,11 @@ def _rule(c: _Collector, raw: object, index: int) -> Rule | None: status=status, autofix=autofix, enforcer=enforcer, - ci=tuple(ci_refs), + ci=ci_refs, review_by=review_by, proof=proof, baseline=baseline, - exceptions=tuple(exceptions), + exceptions=exceptions, downgrade_reason=downgrade_reason, ) diff --git a/scripts/rules/model.py b/scripts/rules/model.py index 401becaaa..95f65570f 100644 --- a/scripts/rules/model.py +++ b/scripts/rules/model.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Callable, Mapping from datetime import date @@ -43,6 +43,40 @@ class Kind(StrEnum): NONE = "none" +def _describe_pytest(fields: Mapping[str, Any]) -> str: + """Name the test files behind a pytest enforcer, under its project if any.""" + prefix = f"{fields['project'].rstrip('/')}/" if "project" in fields else "" + files = sorted({node.split("::", 1)[0] for node in fields.get("nodes", [])}) + return "pytest " + ", ".join(f"`{prefix}{path}`" for path in files) + + +def _describe_ruff(fields: Mapping[str, Any]) -> str: + """Describe a ruff enforcer by its mode: format, a per-file ignore or a check.""" + codes = ", ".join(fields.get("codes", [])) + mode = fields.get("mode", "") + if mode == "format": + return "ruff format" + if mode == "per-file-ignore": + return f"ruff per-file ignore of {codes} for `{fields.get('pattern', '')}`" + return f"ruff check ({codes})" + + +# How each kind of enforcer describes itself; a kind not listed enforces nothing. +_ENFORCER_DESCRIPTIONS: dict[Kind, Callable[[Mapping[str, Any]], str]] = { + Kind.PYTEST: _describe_pytest, + Kind.QUALITY_GUARD: lambda fields: f"quality guard `{fields.get('name', '')}`", + Kind.RUFF: _describe_ruff, + Kind.MYPY: lambda _fields: "mypy", + Kind.WORKFLOW_JOB: lambda _fields: "CI job", + Kind.PYRIGHT: lambda fields: f"{Kind.PYRIGHT.value} `{fields.get('config', '')}`", + Kind.IMPORT_LINTER: ( + lambda fields: f"{Kind.IMPORT_LINTER.value} `{fields.get('config', '')}`" + ), + Kind.DOCS_GUARD: lambda fields: f"docs guard on `{fields.get('doc', '')}`", + Kind.LLM_REVIEW: lambda fields: f"LLM review with `{fields.get('prompt', '')}`", +} + + @dataclass(frozen=True) class Violation: """One guard failure. ``fix`` always tells the reader what to do next.""" @@ -124,34 +158,8 @@ def get( @property def description(self) -> str: """Short human-readable description of what enforces the rule.""" - kind, fields = self.kind, self.fields - if kind is Kind.PYTEST: - prefix = f"{fields['project'].rstrip('/')}/" if "project" in fields else "" - files = sorted({node.split("::", 1)[0] for node in fields.get("nodes", [])}) - return "pytest " + ", ".join(f"`{prefix}{path}`" for path in files) - if kind is Kind.QUALITY_GUARD: - return f"quality guard `{fields.get('name', '')}`" - if kind is Kind.RUFF: - codes = ", ".join(fields.get("codes", [])) - mode = fields.get("mode", "") - if mode == "format": - return "ruff format" - if mode == "per-file-ignore": - return ( - f"ruff per-file ignore of {codes} for `{fields.get('pattern', '')}`" - ) - return f"ruff check ({codes})" - if kind is Kind.MYPY: - return "mypy" - if kind is Kind.WORKFLOW_JOB: - return "CI job" - if kind in (Kind.PYRIGHT, Kind.IMPORT_LINTER): - return f"{kind.value} `{fields.get('config', '')}`" - if kind is Kind.DOCS_GUARD: - return f"docs guard on `{fields.get('doc', '')}`" - if kind is Kind.LLM_REVIEW: - return f"LLM review with `{fields.get('prompt', '')}`" - return "nothing" + describe = _ENFORCER_DESCRIPTIONS.get(self.kind) + return "nothing" if describe is None else describe(self.fields) @property def expects_proof(self) -> bool: diff --git a/scripts/rules/render.py b/scripts/rules/render.py index a6d5b639d..e22985d5f 100644 --- a/scripts/rules/render.py +++ b/scripts/rules/render.py @@ -63,46 +63,64 @@ def _mermaid_id(prefix: str, text: str) -> str: return prefix + "_" + re.sub(r"[^A-Za-z0-9]", "_", text) -def _mermaid(rules: list[Rule]) -> list[str]: - lines = ["```mermaid", "flowchart LR"] - declared: set[str] = set() - edges: list[str] = [] - - def node(node_id: str, label: str, shape: str = "box") -> str: - if node_id not in declared: - declared.add(node_id) - safe = label.replace('"', "'") - if shape == "round": - lines.append(f' {node_id}("{safe}")') - elif shape == "flag": - lines.append(f' {node_id}>"{safe}"]') - else: - lines.append(f' {node_id}["{safe}"]') +# Mermaid node declarations by shape: a box, a rounded box, or a flag. +_MERMAID_SHAPES = { + "box": ' {node_id}["{label}"]', + "round": ' {node_id}("{label}")', + "flag": ' {node_id}>"{label}"]', +} + + +class _MermaidGraph: + """A flowchart under construction: each node is declared once, in first use.""" + + def __init__(self) -> None: + self.nodes: list[str] = [] + self.edges: list[str] = [] + self._declared: set[str] = set() + + def node(self, node_id: str, label: str, shape: str = "box") -> str: + """Declare the node the first time it is used, and return its id.""" + if node_id not in self._declared: + self._declared.add(node_id) + template = _MERMAID_SHAPES.get(shape, _MERMAID_SHAPES["box"]) + self.nodes.append( + template.format(node_id=node_id, label=label.replace('"', "'")) + ) return node_id + def edge_once(self, edge: str) -> None: + """Add an edge unless the graph already has it.""" + if edge not in self.edges: + self.edges.append(edge) + + +def _add_rule_to_graph(graph: _MermaidGraph, rule: Rule) -> None: + """Link a rule to its enforcer, and the enforcer to its CI jobs or a gap.""" + rule_node = graph.node( + _mermaid_id("r", rule.id), f"{rule.id} ({rule.status})", "round" + ) + if rule.enforcer.kind is Kind.NONE: + gap = graph.node("gap_unenforced", "not enforced", "flag") + graph.edges.append(f" {rule_node} -.-> {gap}") + return + label = _enforced_by(rule).replace("`", "") + enforcer_node = graph.node(_mermaid_id("e", label), label) + graph.edges.append(f" {rule_node} --> {enforcer_node}") + for ref in rule.ci: + job_label = f"CI {Path(ref.workflow).name}:{ref.job}" + job_node = graph.node(_mermaid_id("j", job_label), job_label) + graph.edge_once(f" {enforcer_node} --> {job_node}") + if not rule.ci: + local = graph.node("local_only", "not run in CI", "flag") + graph.edge_once(f" {enforcer_node} -.-> {local}") + + +def _mermaid(rules: list[Rule]) -> list[str]: + graph = _MermaidGraph() for rule in rules: - rule_node = node( - _mermaid_id("r", rule.id), f"{rule.id} ({rule.status})", "round" - ) - if rule.enforcer.kind is Kind.NONE: - gap = node("gap_unenforced", "not enforced", "flag") - edges.append(f" {rule_node} -.-> {gap}") - continue - label = _enforced_by(rule).replace("`", "") - enforcer_node = node(_mermaid_id("e", label), label) - edges.append(f" {rule_node} --> {enforcer_node}") - for ref in rule.ci: - job_label = f"CI {Path(ref.workflow).name}:{ref.job}" - job_node = node(_mermaid_id("j", job_label), job_label) - edge = f" {enforcer_node} --> {job_node}" - if edge not in edges: - edges.append(edge) - if not rule.ci: - local = node("local_only", "not run in CI", "flag") - edge = f" {enforcer_node} -.-> {local}" - if edge not in edges: - edges.append(edge) - return lines + edges + ["```"] + _add_rule_to_graph(graph, rule) + return ["```mermaid", "flowchart LR", *graph.nodes, *graph.edges, "```"] def _mermaid_chunks(rules: list[Rule]) -> list[list[str]]: diff --git a/scripts/rules/resolve.py b/scripts/rules/resolve.py index 6f548139c..f27e5f44b 100644 --- a/scripts/rules/resolve.py +++ b/scripts/rules/resolve.py @@ -10,6 +10,7 @@ import ast import re import tomllib +from dataclasses import dataclass from functools import cache from typing import TYPE_CHECKING, Any @@ -381,15 +382,13 @@ def resolve_ci(rule: Rule, root: Path) -> list[Violation]: job, problems = _job_for(rule.id, ref, root, "ci", "rule.ci") violations += problems if job is not None and rule.enforcer.kind is Kind.PYTEST: - violations += _test_coverage( - rule.id, - root_relative_nodes(rule), - rule.enforcer.get("project"), - ref.invocation, - job, - root, - "ci", + required = _RequiredTests( + rule_id=rule.id, + check="ci", + nodes=root_relative_nodes(rule), + project=rule.enforcer.get("project"), ) + violations += _test_coverage(required, ref.invocation, job, root) return violations @@ -429,22 +428,27 @@ def resolve_proof_ci(rule: Rule, root: Path) -> list[Violation]: job, problems = _job_for(rule.id, ref, root, "proof", "proof") violations += problems if job is not None: - violations += _test_coverage( - rule.id, [node], project, ref.invocation, job, root, "proof" + required = _RequiredTests( + rule_id=rule.id, check="proof", nodes=[node], project=project ) + violations += _test_coverage(required, ref.invocation, job, root) return violations +@dataclass(frozen=True) +class _RequiredTests: + """Root-relative test nodes a CI job must collect, and how to report a miss.""" + + rule_id: str + check: str + nodes: list[str] + project: str | None + + def _test_coverage( - rule_id: str, - nodes: list[str], - project: str | None, - invocation: str, - job: JobBlock, - root: Path, - check: str, + required: _RequiredTests, invocation: str, job: JobBlock, root: Path ) -> list[Violation]: - """Check that the job's test command collects these root-relative test nodes. + """Check that the job's test command collects the required test nodes. Only ``make`` targets and ``unittest discover`` commands are inspected; any other command is trusted. @@ -461,30 +465,30 @@ def _test_coverage( collected = [match.group(1)] if match else None described = f"'{invocation}'" if collected: - for node in nodes: + for node in required.nodes: path = split_node(node)[0] if not any( path == p or path.startswith(p.rstrip("/") + "/") for p in collected ): violations.append( Violation( - rule_id, - check, + required.rule_id, + required.check, f"{path} is not collected by {described} " f"(collects {collected})", "add it to what the job runs, or move the test", ) ) if ( - project + required.project and job.matrix - and not any(project in values for values in job.matrix.values()) + and not any(required.project in values for values in job.matrix.values()) ): violations.append( Violation( - rule_id, - check, - f"project {project} is not in the matrix of job '{job.id}'", + required.rule_id, + required.check, + f"project {required.project} is not in the matrix of job '{job.id}'", "add the project to the job matrix", ) ) diff --git a/scripts/rules/tests/fixtures/validate/broken_rules.toml b/scripts/rules/tests/fixtures/validate/broken_rules.toml new file mode 100644 index 000000000..24d4b4b33 --- /dev/null +++ b/scripts/rules/tests/fixtures/validate/broken_rules.toml @@ -0,0 +1,159 @@ +[[category]] +id = "layers" +title = "Layers again" + +[agents] +instructions = "AGENTS.md" +adapters = ["CLAUDE.md", "MISSING.md"] + +[ci_scripts] +workflows = ".github/workflows/*.yml" +scripts_dir = ".github/scripts" +test_globs = ["tests/**/test_*.py"] +multiline_baseline = "rules/baselines/ci-multiline-run.txt" +untested_baseline = "rules/baselines/missing-untested.txt" +shellcheck_baseline = "rules/baselines/missing-shellcheck.txt" + +[[rule]] +id = "ex.layers.no-sql-in-routes" +category = "layers" +statement = "Duplicate id." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "mechanical" +status = "doc-only" +autofix = false +[rule.enforcer] +kind = "none" + +[[rule]] +id = "Bad_ID" +category = "nowhere" +statement = "Doc-only rule that still claims everything." +rationale = "docs/missing.md#section" +scope = ["pkg/*.py", "nothing/**/*.py"] +tier = "judgment" +status = "doc-only" +autofix = false +review_by = 2026-05-01 +[rule.enforcer] +kind = "ruff" +mode = "per-file-ignore" +extra = 1 +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "lint" +invocation = "ruff check" +[rule.baseline] +file = "rules/baselines/missing.txt" +count = 1 +[[rule.exception]] +path = "nothing/*.py" +reason = "legacy" +expires = 2026-01-01 +[[rule.exception]] +path = "pkg/*.py" +reason = "current" +expires = 2027-01-01 + +[[rule]] +id = "ex.quality.unenforced" +category = "quality" +statement = "Enforced but has no enforcer or CI." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "mechanical" +status = "ratcheted" +autofix = false +[rule.enforcer] +kind = "none" + +[[rule]] +id = "ex.quality.review-only" +category = "quality" +statement = "LLM review of a mechanical rule." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "mechanical" +status = "advisory" +autofix = false +[rule.enforcer] +kind = "llm-review" +prompt = 5 +[rule.proof] +bad_fixture = "fixtures/bad.py" +good_fixture = "pkg/routes.py" +test = "tests/unit/test_boundaries.py::test_missing" +workflow = ".github/workflows/pr.yml" +job = "test" + +[[rule]] +id = "ex.quality.pytest-empty" +category = "quality" +statement = "Pytest enforcer without nodes." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "metric" +status = "advisory" +autofix = false +[rule.enforcer] +kind = "pytest" +nodes = [] +[rule.proof] +below_fixture = "pkg/routes.py" +above_fixture = "fixtures/above.py" +threshold = "high" +test = "test_boundaries.py::test_routes_have_no_sql" +project = "tests/unit/" + +[[rule]] +id = "ex.quality.ruff-modes" +category = "quality" +statement = "Ruff lint without codes." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "mechanical" +status = "enforced" +autofix = false +[rule.enforcer] +kind = "ruff" +mode = "lint" +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "lint" +invocation = "ruff check" +[rule.baseline] +file = "rules/baselines/ruff.txt" +count = 3 + +[[rule]] +id = "ex.quality.ruff-bad-mode" +category = "quality" +statement = "Ruff with an unknown mode." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "mechanical" +status = "advisory" +autofix = false +[rule.enforcer] +kind = "ruff" +mode = "fix" + +[[rule]] +id = "ex.quality.judged" +category = "quality" +statement = "Judgment rule with a proof." +rationale = "docs/guide.md" +scope = ["pkg/*.py"] +tier = "judgment" +status = "advisory" +autofix = false +[rule.enforcer] +kind = "llm-review" +prompt = "prompts/review.md" +model = "claude" +[rule.proof] +eval_set = "evals/missing.jsonl" +min_precision = 0.9 +measured_precision = "high" +measured_on = 2026-01-01 diff --git a/scripts/rules/tests/fixtures/validate/broken_rules_violations.txt b/scripts/rules/tests/fixtures/validate/broken_rules_violations.txt new file mode 100644 index 000000000..aae2c3fcc --- /dev/null +++ b/scripts/rules/tests/fixtures/validate/broken_rules_violations.txt @@ -0,0 +1,41 @@ +[registry] schema: agents file MISSING.md does not exist. Fix: fix [agents] or create the file. +[registry] schema: ci_scripts.scripts_dir .github/scripts does not exist. Fix: fix [ci_scripts]. +[registry] schema: ci_scripts baseline rules/baselines/missing-untested.txt does not exist. Fix: create it (it may be empty). +[registry] schema: ci_scripts baseline rules/baselines/missing-shellcheck.txt does not exist. Fix: create it (it may be empty). +[registry] schema: duplicate category 'layers'. Fix: give each category a unique id. +[ex.layers.no-sql-in-routes] schema: duplicate rule id. Fix: give each rule a unique id. +[Bad_ID] schema: id must be lowercase dotted, e.g. 'api.v2.routes-no-sql'. Fix: rename the rule. +[Bad_ID] schema: unknown category 'nowhere'. Fix: use one of ['layers', 'quality', 'layers'] or add a [[category]]. +[Bad_ID] schema: rationale docs/missing.md does not exist. Fix: link an existing doc. +[Bad_ID] schema: scope 'nothing/**/*.py' matches no files. Fix: fix the glob. +[Bad_ID] schema: unknown key 'extra' in enforcer (ruff). Fix: remove it; allowed: ['codes', 'config', 'mode', 'pattern']. +[Bad_ID] schema: ruff per-file-ignore needs codes. Fix: list the ruff codes. +[Bad_ID] schema: ruff per-file-ignore needs a pattern. Fix: set pattern. +[Bad_ID] schema: doc-only rules have no enforcer. Fix: set enforcer.kind = "none", or pick a stronger status. +[Bad_ID] schema: doc-only rules have no CI wiring. Fix: remove rule.ci, or pick a stronger status. +[Bad_ID] schema: only ratcheted rules may have a baseline. Fix: remove [rule.baseline] or set status = "ratcheted". +[Bad_ID] schema: ruff is never ratcheted. Fix: remove the baseline and fix every violation in a standalone PR that also changes the ruff select list. +[Bad_ID] schema: judgment rules are enforced by review, not by a mechanical checker. Fix: use enforcer.kind = "llm-review" or "none". +[Bad_ID] baseline: baseline file rules/baselines/missing.txt does not exist. Fix: create it. +[Bad_ID] expiry: review_by 2026-05-01 has passed. Fix: review the rule, then remove or extend review_by. +[Bad_ID] expiry: exception for nothing/*.py expired on 2026-01-01. Fix: remove the exception (fix the code) or extend it with a reason. +[Bad_ID] schema: exception path 'nothing/*.py' matches no files. Fix: remove the stale exception. +[ex.quality.unenforced] schema: status ratcheted needs an enforcer. Fix: add [rule.enforcer] or use doc-only/advisory. +[ex.quality.unenforced] schema: ratcheted rules must run in CI. Fix: add a [[rule.ci]] reference. +[ex.quality.unenforced] schema: ratcheted rules need a baseline. Fix: add [rule.baseline] or use enforced. +[ex.quality.review-only] schema: enforcer (llm-review).prompt has the wrong type. Fix: make it a . +[ex.quality.review-only] schema: enforcer (llm-review) is missing 'model'. Fix: add model. +[ex.quality.review-only] schema: llm-review only enforces judgment rules. Fix: set tier = "judgment". +[ex.quality.review-only] schema: proof names its CI job only partly. Fix: set all of proof.workflow, proof.job and proof.invocation, or none. +[ex.quality.review-only] proof: proof bad_fixture fixtures/bad.py does not exist. Fix: add the fixture or fix the path. +[ex.quality.review-only] proof: tests/unit/test_boundaries.py has no test 'test_missing'. Fix: point proof.test at the test that runs the fixtures. +[ex.quality.pytest-empty] schema: pytest enforcer has no nodes. Fix: list at least one test node. +[ex.quality.pytest-empty] schema: proof (metric).threshold has the wrong type. Fix: make it a (, ). +[ex.quality.pytest-empty] proof: proof above_fixture fixtures/above.py does not exist. Fix: add the fixture or fix the path. +[ex.quality.ruff-modes] schema: ruff lint needs codes. Fix: list the ruff codes. +[ex.quality.ruff-modes] schema: only ratcheted rules may have a baseline. Fix: remove [rule.baseline] or set status = "ratcheted". +[ex.quality.ruff-modes] schema: ruff is never ratcheted. Fix: remove the baseline and fix every violation in a standalone PR that also changes the ruff select list. +[ex.quality.ruff-modes] baseline: baseline count is 3 but rules/baselines/ruff.txt lists 2. Fix: make the count match the file. +[ex.quality.ruff-bad-mode] schema: ruff mode is 'fix'. Fix: use one of ['format', 'lint', 'per-file-ignore']. +[ex.quality.judged] schema: proof (judgment).measured_precision has the wrong type. Fix: make it a . +[ex.quality.judged] proof: proof eval_set evals/missing.jsonl does not exist. Fix: add the fixture or fix the path. diff --git a/scripts/rules/tests/test_cli.py b/scripts/rules/tests/test_cli.py new file mode 100644 index 000000000..a9862f54a --- /dev/null +++ b/scripts/rules/tests/test_cli.py @@ -0,0 +1,239 @@ +"""Tests for the rules command line: what each command prints and returns.""" + +from __future__ import annotations + +import contextlib +import io +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from scripts.rules.ci_scripts import Findings +from scripts.rules.cli import main +from scripts.rules.guard import GuardResult +from scripts.rules.model import Violation +from scripts.rules.tests.support import REGISTRY, Repo + +MISSING_REGISTRY = ( + " - [registry] schema: rules/registry.toml is missing. " + "Fix: create rules/registry.toml." +) +DUPLICATE_CATEGORY = ( + " - [registry] schema: duplicate category 'layers'. " + "Fix: give each category a unique id." +) +CI_SCRIPTS_TABLE = """ +[ci_scripts] +workflows = ".github/workflows/*.yml" +scripts_dir = ".github/scripts" +test_globs = ["tests/**/test_*.py"] +multiline_baseline = "rules/baselines/ci-multiline-run.txt" +untested_baseline = "rules/baselines/ci-untested-scripts.txt" +shellcheck_baseline = "rules/baselines/ci-shellcheck.txt" +""" + + +def run_cli(*argv: str) -> tuple[int, list[str]]: + """Run the command line and return its exit status and printed lines.""" + output = io.StringIO() + with contextlib.redirect_stdout(output): + status = main(list(argv)) + return status, output.getvalue().splitlines() + + +def with_duplicate_category(repo: Repo) -> None: + """Append a second category with an existing id, which validation rejects.""" + repo.write( + "rules/registry.toml", + repo.read("rules/registry.toml") + + '\n[[category]]\nid = "layers"\ntitle = "Layers again"\n', + ) + + +class ValidateCommandTest(unittest.TestCase): + def test_valid_registry_passes(self) -> None: + with Repo() as repo: + self.assertEqual( + run_cli("validate", "--root", str(repo.root)), (0, ["rules: passed"]) + ) + + def test_missing_registry_fails(self) -> None: + with tempfile.TemporaryDirectory() as root: + self.assertEqual( + run_cli("validate", "--root", root), + (1, [MISSING_REGISTRY, "rules: failed"]), + ) + + def test_violations_are_listed(self) -> None: + with Repo(render=False) as repo: + with_duplicate_category(repo) + self.assertEqual( + run_cli("validate", "--root", str(repo.root)), + (1, [DUPLICATE_CATEGORY, "rules: failed"]), + ) + + +class RenderCommandTest(unittest.TestCase): + def test_render_writes_both_views(self) -> None: + with Repo(render=False) as repo: + self.assertEqual( + run_cli("render", "--root", str(repo.root)), + (0, ["wrote docs/generated/rules.md", "wrote AGENTS.md"]), + ) + self.assertTrue((repo.root / "docs/generated/rules.md").is_file()) + self.assertIn("BEGIN GENERATED: rules", repo.read("AGENTS.md")) + + def test_render_refuses_an_invalid_registry(self) -> None: + with Repo(render=False) as repo: + with_duplicate_category(repo) + self.assertEqual( + run_cli("render", "--root", str(repo.root)), + ( + 1, + [ + "render refused: fix the registry first", + DUPLICATE_CATEGORY, + "rules: failed", + ], + ), + ) + self.assertFalse((repo.root / "docs/generated/rules.md").exists()) + + def test_render_reports_a_missing_registry(self) -> None: + with tempfile.TemporaryDirectory() as root: + self.assertEqual( + run_cli("render", "--root", root), + (1, [MISSING_REGISTRY, "rules: failed"]), + ) + self.assertFalse((Path(root) / "docs").exists()) + + +class GuardCommandTest(unittest.TestCase): + def test_guard_prints_notes_then_violations(self) -> None: + result = GuardResult( + violations=[Violation("ex.rule", "ci", "job is missing", "add the job")], + notes=["base ref unavailable"], + ) + with ( + Repo(render=False) as repo, + mock.patch("scripts.rules.cli.run_guard", return_value=result) as guard, + ): + self.assertEqual( + run_cli("guard", "--root", str(repo.root), "--base-ref", "origin/main"), + ( + 1, + [ + "note: base ref unavailable", + " - [ex.rule] ci: job is missing. Fix: add the job.", + "rules: failed", + ], + ), + ) + guard.assert_called_once_with(repo.root.resolve(), base_ref="origin/main") + + def test_clean_guard_passes(self) -> None: + with ( + Repo(render=False) as repo, + mock.patch( + "scripts.rules.cli.run_guard", return_value=GuardResult() + ) as guard, + ): + self.assertEqual( + run_cli("guard", "--root", str(repo.root)), (0, ["rules: passed"]) + ) + guard.assert_called_once_with(repo.root.resolve(), base_ref=None) + + +class CiScriptsCommandTest(unittest.TestCase): + def test_check_prints_notes_and_violations(self) -> None: + violation = Violation(None, "ci-scripts", "untested script x.sh", "test it") + with ( + Repo(render=False) as repo, + mock.patch( + "scripts.rules.ci_scripts.run_check", + return_value=([violation], ["shellcheck is not installed"]), + ) as run_check, + ): + self.assertEqual( + run_cli("ci-scripts", "--root", str(repo.root)), + ( + 1, + [ + "note: shellcheck is not installed", + " - [registry] ci-scripts: untested script x.sh. " + "Fix: test it.", + "rules: failed", + ], + ), + ) + run_check.assert_called_once_with(repo.root.resolve()) + + def test_check_without_findings_passes(self) -> None: + with ( + Repo(render=False) as repo, + mock.patch("scripts.rules.ci_scripts.run_check", return_value=([], [])), + ): + self.assertEqual( + run_cli("ci-scripts", "--root", str(repo.root)), (0, ["rules: passed"]) + ) + + def test_write_baselines_needs_a_registry(self) -> None: + with tempfile.TemporaryDirectory() as root: + self.assertEqual( + run_cli("ci-scripts", "--root", root, "--write-baselines"), + (1, [MISSING_REGISTRY, "rules: failed"]), + ) + + def test_write_baselines_needs_a_ci_scripts_table(self) -> None: + with Repo(render=False) as repo: + self.assertEqual( + run_cli("ci-scripts", "--root", str(repo.root), "--write-baselines"), + (1, [" - no [ci_scripts] table", "rules: failed"]), + ) + + def test_write_baselines_reports_schema_problems_first(self) -> None: + with Repo(render=False) as repo: + repo.edit("rules/registry.toml", 'repo = "example"', "repo = 7") + self.assertEqual( + run_cli("ci-scripts", "--root", str(repo.root), "--write-baselines"), + ( + 1, + [ + " - [registry] schema: registry.repo must be a non-empty " + "string. Fix: set repo.", + "rules: failed", + ], + ), + ) + + def test_write_baselines_records_the_findings(self) -> None: + findings = Findings( + multiline={"wf.yml :: build :: Build"}, + untested={".github/scripts/b.sh", ".github/scripts/a.sh"}, + shellcheck=None, + ) + with ( + Repo({"rules/registry.toml": REGISTRY + CI_SCRIPTS_TABLE}) as repo, + mock.patch("scripts.rules.ci_scripts.find", return_value=findings) as find, + ): + self.assertEqual( + run_cli("ci-scripts", "--root", str(repo.root), "--write-baselines"), + (0, ["wrote ci_scripts baselines; update the rules' baseline counts"]), + ) + find.assert_called_once() + self.assertEqual(find.call_args.args[1], repo.root.resolve()) + header = "# Tolerated existing violations. This list may only shrink.\n" + self.assertEqual( + repo.read("rules/baselines/ci-multiline-run.txt"), + header + "wf.yml :: build :: Build\n", + ) + self.assertEqual( + repo.read("rules/baselines/ci-untested-scripts.txt"), + header + ".github/scripts/a.sh\n.github/scripts/b.sh\n", + ) + self.assertFalse((repo.root / "rules/baselines/ci-shellcheck.txt").exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/rules/tests/test_load.py b/scripts/rules/tests/test_load.py new file mode 100644 index 000000000..a25b017c2 --- /dev/null +++ b/scripts/rules/tests/test_load.py @@ -0,0 +1,194 @@ +"""Tests for loading rule tables: every schema problem is reported, in order.""" + +from __future__ import annotations + +import tomllib +import unittest +from datetime import date + +from scripts.rules.load import parse +from scripts.rules.model import ( + Baseline, + CiRef, + Enforcer, + Kind, + Proof, + Rule, + RuleException, + Status, + Tier, +) +from scripts.rules.tests.support import REGISTRY + +BOUNDARY_NODE = "tests/unit/test_boundaries.py::test_routes_have_no_sql" +KNOWN_RULE_KEYS = ( + "['autofix', 'baseline', 'category', 'ci', 'downgrade', 'enforcer', " + "'exception', 'id', 'proof', 'rationale', 'review_by', 'scope', 'statement', " + "'status', 'tier']" +) +KINDS = ( + "['pytest', 'quality-guard', 'ruff', 'mypy', 'pyright', 'import-linter', " + "'workflow-job', 'docs-guard', 'llm-review', 'none']" +) +MALFORMED = "[ex.layers.malformed] schema: ex.layers.malformed" + +MALFORMED_RULE = { + "id": "ex.layers.malformed", + "category": "layers", + "statement": "Malformed.", + "rationale": "docs/guide.md", + "scope": ["pkg/*.py"], + "tier": "mechanical", + "status": "enforced", + "autofix": "yes", + "review_by": "someday", + "enforcer": "pytest", + "ci": [ + 5, + { + "workflow": ".github/workflows/pr.yml", + "job": "", + "invocation": "make test", + "extra": 1, + }, + ], + "proof": "proof.py", + "baseline": {"file": "rules/baselines/x.txt"}, + "exception": [ + 7, + {"path": "pkg/*.py", "reason": "r", "expires": "soon", "extra": 1}, + ], + "downgrade": {"why": "x"}, + "bogus": 1, +} +UNKNOWN_TIER_RULE = { + "category": "layers", + "tier": "unknown", + "status": "enforced", + "autofix": False, + "enforcer": {"kind": "flake8", "codes": ["E1"]}, +} +SINGLE_TABLES_RULE = { + "id": "ex.layers.single-tables", + "category": "layers", + "statement": "Single tables.", + "rationale": "docs/guide.md", + "scope": ["pkg/*.py"], + "tier": "metric", + "status": "ratcheted", + "autofix": True, + "review_by": date(2026, 7, 1), + "enforcer": {"kind": "pytest", "nodes": [BOUNDARY_NODE]}, + "ci": { + "workflow": ".github/workflows/pr.yml", + "job": "test", + "invocation": "make test", + }, + "proof": {"test": BOUNDARY_NODE}, + "baseline": {"file": "rules/baselines/x.txt", "count": 2, "extra": True}, + "exception": {"path": "pkg/*.py", "reason": "legacy", "expires": "2026-12-31"}, + "downgrade": {"reason": "too noisy"}, +} + + +class RuleTableTest(unittest.TestCase): + def setUp(self) -> None: + raw = tomllib.loads(REGISTRY) + raw["rule"] = [ + "not a table", + MALFORMED_RULE, + UNKNOWN_TIER_RULE, + SINGLE_TABLES_RULE, + ] + self.registry, self.violations = parse(raw) + + def test_every_problem_is_reported_in_order(self) -> None: + self.assertEqual( + [str(violation) for violation in self.violations], + [ + "[registry] schema: rule #1 must be a table. Fix: write [[rule]].", + "[ex.layers.malformed] schema: unknown key 'bogus' in " + "ex.layers.malformed. Fix: remove it or use one of " + f"{KNOWN_RULE_KEYS}.", + f"{MALFORMED}.enforcer must be a table. Fix: add [rule.enforcer].", + f"{MALFORMED}.ci entries must be tables. Fix: write [[rule.ci]].", + "[ex.layers.malformed] schema: unknown key 'extra' in " + "ex.layers.malformed.ci. Fix: remove it or use one of " + "['invocation', 'job', 'workflow'].", + f"{MALFORMED}.ci.job must be a non-empty string. Fix: set job.", + f"{MALFORMED}.proof must be a table. Fix: write [rule.proof].", + f"{MALFORMED}.baseline needs 'file' and an integer 'count'. " + "Fix: write [rule.baseline] with file and count.", + f"{MALFORMED}.exception entries must be tables. " + "Fix: write [[rule.exception]].", + "[ex.layers.malformed] schema: unknown key 'extra' in " + "ex.layers.malformed.exception. Fix: remove it or use one of " + "['expires', 'path', 'reason'].", + f"{MALFORMED}.exception.expires must be a date. " + "Fix: write it as YYYY-MM-DD.", + f"{MALFORMED}.downgrade needs a 'reason'. " + 'Fix: write [rule.downgrade] reason = "...".', + f"{MALFORMED}.autofix must be true or false. Fix: set autofix.", + f"{MALFORMED}.review_by must be a date. Fix: write it as YYYY-MM-DD.", + f"[registry] schema: rule #3.enforcer.kind is 'flake8'. " + f"Fix: use one of {KINDS}.", + "[registry] schema: rule #3.tier is 'unknown'. " + "Fix: use one of ['mechanical', 'metric', 'judgment'].", + "[ex.layers.single-tables] schema: unknown key 'extra' in " + "ex.layers.single-tables.baseline. Fix: remove it or use one of " + "['count', 'file'].", + ], + ) + + def test_malformed_fields_fall_back_to_safe_values(self) -> None: + assert self.registry is not None + self.assertEqual( + self.registry.rule("ex.layers.malformed"), + Rule( + id="ex.layers.malformed", + category="layers", + statement="Malformed.", + rationale="docs/guide.md", + scope=("pkg/*.py",), + tier=Tier.MECHANICAL, + status=Status.ENFORCED, + autofix=False, + enforcer=Enforcer(Kind.NONE), + ci=(CiRef(".github/workflows/pr.yml", "", "make test"),), + exceptions=(RuleException("pkg/*.py", "r", date.max),), + ), + ) + + def test_a_rule_with_an_unknown_tier_is_dropped(self) -> None: + assert self.registry is not None + self.assertEqual( + [rule.id for rule in self.registry.rules], + ["ex.layers.malformed", "ex.layers.single-tables"], + ) + + def test_single_tables_stand_for_one_entry_lists(self) -> None: + assert self.registry is not None + self.assertEqual( + self.registry.rule("ex.layers.single-tables"), + Rule( + id="ex.layers.single-tables", + category="layers", + statement="Single tables.", + rationale="docs/guide.md", + scope=("pkg/*.py",), + tier=Tier.METRIC, + status=Status.RATCHETED, + autofix=True, + enforcer=Enforcer(Kind.PYTEST, {"nodes": [BOUNDARY_NODE]}), + ci=(CiRef(".github/workflows/pr.yml", "test", "make test"),), + review_by=date(2026, 7, 1), + proof=Proof({"test": BOUNDARY_NODE}), + baseline=Baseline("rules/baselines/x.txt", 2), + exceptions=(RuleException("pkg/*.py", "legacy", date(2026, 12, 31)),), + downgrade_reason="too noisy", + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/rules/tests/test_render.py b/scripts/rules/tests/test_render.py new file mode 100644 index 000000000..dbed7c2dc --- /dev/null +++ b/scripts/rules/tests/test_render.py @@ -0,0 +1,211 @@ +"""Tests for the rendered views: enforcer descriptions and the mermaid graphs.""" + +from __future__ import annotations + +import unittest + +from scripts.rules.model import ( + Category, + CiRef, + Discovery, + Enforcer, + Kind, + Registry, + Rule, + Status, + Tier, +) +from scripts.rules.render import render_rules_md + +PR_WORKFLOW = ".github/workflows/pr.yml" + + +def make_rule( + rule_id: str, + enforcer: Enforcer, + ci: tuple[CiRef, ...] = (), + status: Status = Status.ENFORCED, +) -> Rule: + return Rule( + id=rule_id, + category="checks", + statement="A statement.", + rationale="docs/guide.md", + scope=("pkg/*.py",), + tier=Tier.MECHANICAL, + status=status, + autofix=False, + enforcer=enforcer, + ci=ci, + ) + + +def mermaid_blocks(rules: tuple[Rule, ...]) -> list[list[str]]: + """Render the rules in one category and return each mermaid block's lines.""" + registry = Registry( + schema_version=1, + repo="example", + categories=(Category("checks", "Checks"),), + discovery=Discovery(), + rules=rules, + ) + blocks: list[list[str]] = [] + current: list[str] | None = None + for line in render_rules_md(registry).splitlines(): + if line == "```mermaid": + current = [line] + elif current is not None: + current.append(line) + if line == "```": + blocks.append(current) + current = None + return blocks + + +def rule_node(number: int) -> str: + """Return the mermaid declaration of the numbered rule in the chunking test.""" + return f' r_ex_many_rule_{number:02d}("ex.many.rule-{number:02d} (enforced)")' + + +class EnforcerDescriptionTest(unittest.TestCase): + def test_each_kind_describes_its_checker(self) -> None: + cases: list[tuple[Kind, dict[str, object], str]] = [ + ( + Kind.PYTEST, + {"project": "proj/", "nodes": ["b.py::t", "a.py::u", "a.py::v"]}, + "pytest `proj/a.py`, `proj/b.py`", + ), + (Kind.PYTEST, {"nodes": ["x.py"]}, "pytest `x.py`"), + (Kind.PYTEST, {}, "pytest "), + ( + Kind.QUALITY_GUARD, + {"name": "rules-registry"}, + "quality guard `rules-registry`", + ), + (Kind.QUALITY_GUARD, {}, "quality guard ``"), + (Kind.RUFF, {"codes": ["E1"], "mode": "format"}, "ruff format"), + ( + Kind.RUFF, + {"codes": ["E1", "F2"], "mode": "per-file-ignore", "pattern": "t/**"}, + "ruff per-file ignore of E1, F2 for `t/**`", + ), + (Kind.RUFF, {"mode": "per-file-ignore"}, "ruff per-file ignore of for ``"), + (Kind.RUFF, {"codes": ["E1", "F2"], "mode": "lint"}, "ruff check (E1, F2)"), + (Kind.RUFF, {}, "ruff check ()"), + (Kind.MYPY, {"paths": ["pkg"]}, "mypy"), + (Kind.WORKFLOW_JOB, {"step": "Build"}, "CI job"), + ( + Kind.PYRIGHT, + {"config": "pyrightconfig.json"}, + "pyright `pyrightconfig.json`", + ), + (Kind.PYRIGHT, {}, "pyright ``"), + ( + Kind.IMPORT_LINTER, + {"config": ".importlinter"}, + "import-linter `.importlinter`", + ), + (Kind.DOCS_GUARD, {"doc": "docs/x.md"}, "docs guard on `docs/x.md`"), + (Kind.DOCS_GUARD, {}, "docs guard on ``"), + (Kind.LLM_REVIEW, {"prompt": "p.md"}, "LLM review with `p.md`"), + (Kind.LLM_REVIEW, {}, "LLM review with ``"), + (Kind.NONE, {"name": "ignored"}, "nothing"), + ] + for kind, fields, expected in cases: + with self.subTest(kind=kind, fields=fields): + self.assertEqual(Enforcer(kind, fields).description, expected) + + +class MermaidTest(unittest.TestCase): + def test_graph_declares_each_node_once_and_skips_repeated_edges(self) -> None: + guard = Enforcer(Kind.QUALITY_GUARD, {"name": 'say "hi"'}) + mypy = Enforcer(Kind.MYPY, {"paths": ["pkg"]}) + rules = ( + make_rule("ex.a.unenforced", Enforcer(Kind.NONE), status=Status.DOC_ONLY), + make_rule( + "ex.b.also-unenforced", Enforcer(Kind.NONE), status=Status.DOC_ONLY + ), + make_rule( + "ex.c.guard", + guard, + ( + CiRef(PR_WORKFLOW, "guards", "run"), + CiRef(".github/workflows/push.yml", "guards", "run"), + ), + ), + make_rule("ex.d.same-guard", guard, (CiRef(PR_WORKFLOW, "guards", "run"),)), + make_rule("ex.e.local", mypy), + make_rule("ex.f.local-too", mypy), + make_rule( + "ex.g.job", + Enforcer(Kind.WORKFLOW_JOB), + ( + CiRef(PR_WORKFLOW, "build", "make"), + CiRef(PR_WORKFLOW, "deploy", "make"), + ), + ), + ) + self.assertEqual( + mermaid_blocks(rules), + [ + [ + "```mermaid", + "flowchart LR", + ' r_ex_a_unenforced("ex.a.unenforced (doc-only)")', + ' gap_unenforced>"not enforced"]', + ' r_ex_b_also_unenforced("ex.b.also-unenforced (doc-only)")', + ' r_ex_c_guard("ex.c.guard (enforced)")', + " e_quality_guard_say__hi_[\"quality guard say 'hi'\"]", + ' j_CI_pr_yml_guards["CI pr.yml:guards"]', + ' j_CI_push_yml_guards["CI push.yml:guards"]', + ' r_ex_d_same_guard("ex.d.same-guard (enforced)")', + ' r_ex_e_local("ex.e.local (enforced)")', + ' e_mypy["mypy"]', + ' local_only>"not run in CI"]', + ' r_ex_f_local_too("ex.f.local-too (enforced)")', + ' r_ex_g_job("ex.g.job (enforced)")', + ' e_CI_job_build__deploy["CI job build, deploy"]', + ' j_CI_pr_yml_build["CI pr.yml:build"]', + ' j_CI_pr_yml_deploy["CI pr.yml:deploy"]', + " r_ex_a_unenforced -.-> gap_unenforced", + " r_ex_b_also_unenforced -.-> gap_unenforced", + " r_ex_c_guard --> e_quality_guard_say__hi_", + " e_quality_guard_say__hi_ --> j_CI_pr_yml_guards", + " e_quality_guard_say__hi_ --> j_CI_push_yml_guards", + " r_ex_d_same_guard --> e_quality_guard_say__hi_", + " r_ex_e_local --> e_mypy", + " e_mypy -.-> local_only", + " r_ex_f_local_too --> e_mypy", + " r_ex_g_job --> e_CI_job_build__deploy", + " e_CI_job_build__deploy --> j_CI_pr_yml_build", + " e_CI_job_build__deploy --> j_CI_pr_yml_deploy", + "```", + ] + ], + ) + + def test_large_graphs_are_split_into_chunks_of_rules(self) -> None: + rules = tuple( + make_rule( + f"ex.many.rule-{number:02d}", + Enforcer(Kind.QUALITY_GUARD, {"name": f"guard-{number:02d}"}), + (CiRef(PR_WORKFLOW, f"job-{number:02d}", "run"),), + ) + for number in range(20) + ) + blocks = mermaid_blocks(rules) + self.assertEqual([len(block) for block in blocks], [78, 28]) + self.assertEqual( + [ + [line for line in block if line.startswith(" r_") and "(" in line] + for block in blocks + ], + [ + [rule_node(number) for number in range(15)], + [rule_node(number) for number in range(15, 20)], + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/rules/tests/test_resolve.py b/scripts/rules/tests/test_resolve.py new file mode 100644 index 000000000..750f83341 --- /dev/null +++ b/scripts/rules/tests/test_resolve.py @@ -0,0 +1,196 @@ +"""Tests that a rule's CI job really collects its test nodes and matrix project.""" + +from __future__ import annotations + +import unittest + +from scripts.rules.model import CiRef, Enforcer, Kind, Proof, Rule, Status, Tier +from scripts.rules.resolve import resolve_ci, resolve_proof_ci +from scripts.rules.tests.support import Repo + +PR_WORKFLOW = ".github/workflows/pr.yml" +WORKFLOW = """ +jobs: + test: + strategy: + matrix: + project: [projects/alpha, projects/beta] + steps: + - run: make test + unit: + steps: + - run: python -m unittest discover -s tools/tests -t . + unit-default: + steps: + - run: python -m unittest discover + other: + steps: + - run: tox -e py +""" +MAKEFILE = ( + "test:\n" + "\tpython -m pytest tests/unit tests/contract/test_a.py --disable-pytest-warnings\n" +) +MAKE_TEST = "'make test' (collects ['tests/unit', 'tests/contract/test_a.py'])" +UNITTEST = "'unittest discover -s tools/tests' (collects ['tools/tests'])" + + +def make_rule( + fields: dict[str, object], + ci: tuple[CiRef, ...] = (), + proof: dict[str, object] | None = None, + kind: Kind = Kind.PYTEST, +) -> Rule: + return Rule( + id="ex.r", + category="checks", + statement="A statement.", + rationale="docs/guide.md", + scope=("pkg/*.py",), + tier=Tier.MECHANICAL, + status=Status.ENFORCED, + autofix=False, + enforcer=Enforcer(kind, fields), + ci=ci, + proof=Proof(proof) if proof is not None else None, + ) + + +def not_collected(check: str, path: str, command: str) -> str: + return ( + f"[ex.r] {check}: {path} is not collected by {command}. " + "Fix: add it to what the job runs, or move the test." + ) + + +def not_in_matrix(check: str, project: str) -> str: + return ( + f"[ex.r] {check}: project {project} is not in the matrix of job 'test'. " + "Fix: add the project to the job matrix." + ) + + +class TestCoverageTest(unittest.TestCase): + def setUp(self) -> None: + self.repo = Repo({PR_WORKFLOW: WORKFLOW, "Makefile": MAKEFILE}, render=False) + self.addCleanup(self.repo.__exit__) + + def ci_messages(self, rule: Rule) -> list[str]: + return [str(violation) for violation in resolve_ci(rule, self.repo.root)] + + def proof_messages(self, rule: Rule) -> list[str]: + return [str(violation) for violation in resolve_proof_ci(rule, self.repo.root)] + + def test_make_target_must_collect_every_node(self) -> None: + rule = make_rule( + { + "nodes": [ + "tests/unit/test_a.py::t", + "tests/other/test_b.py::t", + "tests/contract/test_a.py", + "tests/unitx/test_c.py", + ] + }, + (CiRef(PR_WORKFLOW, "test", "make test"),), + ) + self.assertEqual( + self.ci_messages(rule), + [ + not_collected("ci", "tests/other/test_b.py", MAKE_TEST), + not_collected("ci", "tests/unitx/test_c.py", MAKE_TEST), + ], + ) + + def test_project_must_be_in_the_job_matrix(self) -> None: + missing = make_rule( + {"project": "projects/gamma", "nodes": ["tests/test_g.py::t"]}, + (CiRef(PR_WORKFLOW, "test", "make test"),), + ) + self.assertEqual( + self.ci_messages(missing), + [ + not_collected("ci", "projects/gamma/tests/test_g.py", MAKE_TEST), + not_in_matrix("ci", "projects/gamma"), + ], + ) + listed = make_rule( + {"project": "projects/alpha", "nodes": ["tests/test_g.py::t"]}, + (CiRef(PR_WORKFLOW, "test", "make test"),), + ) + self.assertEqual( + self.ci_messages(listed), + [not_collected("ci", "projects/alpha/tests/test_g.py", MAKE_TEST)], + ) + + def test_unittest_start_directory_must_hold_every_node(self) -> None: + rule = make_rule( + { + "nodes": [ + "tools/tests/test_a.py::T::t", + "tools/testsx/test_b.py", + "other/test_c.py", + ] + }, + (CiRef(PR_WORKFLOW, "unit", "unittest discover -s tools/tests"),), + ) + self.assertEqual( + self.ci_messages(rule), + [ + not_collected("ci", "tools/testsx/test_b.py", UNITTEST), + not_collected("ci", "other/test_c.py", UNITTEST), + ], + ) + + def test_other_commands_are_trusted(self) -> None: + default_discovery = make_rule( + {"nodes": ["other/test_c.py"]}, + (CiRef(PR_WORKFLOW, "unit-default", "unittest discover"),), + ) + other_command = make_rule( + {"project": "projects/gamma", "nodes": ["other/test_c.py"]}, + (CiRef(PR_WORKFLOW, "other", "tox"),), + ) + self.assertEqual(self.ci_messages(default_discovery), []) + self.assertEqual(self.ci_messages(other_command), []) + + def test_named_proof_job_must_collect_the_proof(self) -> None: + rule = make_rule( + {"paths": ["pkg"]}, + proof={ + "test": "elsewhere/test_p.py::t", + "project": "projects/gamma/", + "workflow": PR_WORKFLOW, + "job": "test", + "invocation": "make test", + }, + kind=Kind.MYPY, + ) + self.assertEqual( + self.proof_messages(rule), + [ + not_collected("proof", "projects/gamma/elsewhere/test_p.py", MAKE_TEST), + not_in_matrix("proof", "projects/gamma/"), + ], + ) + + def test_pytest_proof_runs_in_each_ci_job(self) -> None: + rule = make_rule( + {"nodes": ["tests/unit/test_a.py::t"]}, + ( + CiRef(PR_WORKFLOW, "unit", "unittest discover -s tools/tests"), + CiRef(PR_WORKFLOW, "other", "tox"), + ), + proof={"test": "tests/unit/test_p.py::t"}, + ) + self.assertEqual( + self.ci_messages(rule), + [not_collected("ci", "tests/unit/test_a.py", UNITTEST)], + ) + self.assertEqual( + self.proof_messages(rule), + [not_collected("proof", "tests/unit/test_p.py", UNITTEST)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/rules/tests/test_validate.py b/scripts/rules/tests/test_validate.py new file mode 100644 index 000000000..b3c575dbc --- /dev/null +++ b/scripts/rules/tests/test_validate.py @@ -0,0 +1,40 @@ +"""Tests for registry validation: every problem is reported, in a stable order.""" + +from __future__ import annotations + +import unittest +from pathlib import Path + +from scripts.rules.load import load_text +from scripts.rules.tests.support import REGISTRY, TODAY, Repo +from scripts.rules.validate import validate + +FIXTURES = Path(__file__).parent / "fixtures" / "validate" + + +class BrokenRegistryTest(unittest.TestCase): + """One registry that breaks every validation check at least once.""" + + def test_every_violation_is_reported_in_order(self) -> None: + files = { + "rules/registry.toml": REGISTRY + + "\n" + + (FIXTURES / "broken_rules.toml").read_text(), + "CLAUDE.md": "Read AGENTS.md.\n", + "rules/baselines/ci-multiline-run.txt": "", + "rules/baselines/ruff.txt": ( + "# header\n\nentry one\n # indented comment\nentry two\n" + ), + } + with Repo(files, render=False) as repo: + registry, load_violations = load_text(repo.read("rules/registry.toml")) + self.assertEqual(load_violations, []) + assert registry is not None + self.assertEqual( + [str(violation) for violation in validate(registry, repo.root, TODAY)], + (FIXTURES / "broken_rules_violations.txt").read_text().splitlines(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/rules/tests/test_workflows.py b/scripts/rules/tests/test_workflows.py index bf03f566e..e89793932 100644 --- a/scripts/rules/tests/test_workflows.py +++ b/scripts/rules/tests/test_workflows.py @@ -8,7 +8,7 @@ import unittest from pathlib import Path -from scripts.rules.workflows import parse_jobs +from scripts.rules.workflows import JobBlock, RunBlock, parse_jobs FIXTURES = Path(__file__).parent / "fixtures" @@ -83,5 +83,160 @@ def test_anchors_are_reported(self) -> None: self.assertTrue(problems) +SYNTHETIC_WORKFLOW = """ +name: Synthetic +on: push +jobs: + test: + name: "Quoted name" + strategy: + matrix: + project: [a, "b", 'c'] + python: + - "3.11" + - 3.12 + + - '3.13' + include: + - project: x + steps: + - name: First + run: > + folded one + folded two + - uses: actions/checkout@v4 + - run: echo a && echo b + - "just a string" + - name: Third + run: | + echo "a; b" + + echo ${{ x || y }} + - name: Fourth + uses: org/action@v1 + with: + - name: not a step + reusable: + uses: org/repo/.github/workflows/x.yml@main + other: + steps: + - name: Unindented list + run: make test + - run: | + multi + line + name: Named after run +permissions: {} +""" + + +class JobParsingTest(unittest.TestCase): + """Pins every field the reader extracts from a job, quirks included.""" + + def setUp(self) -> None: + self.jobs, self.problems = parse_jobs(SYNTHETIC_WORKFLOW) + + def test_jobs_end_at_the_next_top_level_key(self) -> None: + self.assertEqual(self.problems, []) + self.assertEqual(list(self.jobs), ["test", "reusable", "other"]) + + def test_job_with_matrix_and_steps(self) -> None: + self.assertEqual( + self.jobs["test"], + JobBlock( + id="test", + name="Quoted name", + run_text=( + 'folded one\nfolded two\necho a && echo b\necho "a; b"\n\n' + "echo ${{ x || y }}" + ), + uses="", + step_names=("First", "Third", "Fourth", "not a step"), + step_uses=("actions/checkout@v4", "org/action@v1"), + # A blank line ends a block-style matrix list. + matrix={ + "project": ("a", "b", "c"), + "python": ("3.11", "3.12"), + "include": (), + }, + run_blocks=( + RunBlock( + step="First", + code_lines=2, + chained=False, + text="folded one\nfolded two", + ), + RunBlock( + step="step 3", + code_lines=1, + chained=True, + text="echo a && echo b", + ), + RunBlock( + step="Third", + code_lines=2, + chained=False, + text='echo "a; b"\n\necho ${{ x || y }}', + ), + ), + ), + ) + + def test_reusable_workflow_job(self) -> None: + self.assertEqual( + self.jobs["reusable"], + JobBlock( + id="reusable", + name="reusable", + run_text="", + uses="org/repo/.github/workflows/x.yml@main", + step_names=(), + ), + ) + + def test_steps_listed_at_the_key_indent(self) -> None: + self.assertEqual( + self.jobs["other"], + JobBlock( + id="other", + name="other", + run_text="make test\nmulti\nline", + uses="", + step_names=("Unindented list", "Named after run"), + run_blocks=( + RunBlock( + step="Unindented list", + code_lines=1, + chained=False, + text="make test", + ), + RunBlock( + step="Named after run", + code_lines=2, + chained=False, + text="multi\nline", + ), + ), + ), + ) + + def test_flow_style_jobs_are_reported(self) -> None: + self.assertEqual( + parse_jobs("jobs: {a: 1}\n"), + ({}, ["flow-style 'jobs:' mapping is not supported"]), + ) + + def test_workflow_without_jobs(self) -> None: + self.assertEqual(parse_jobs("name: x\n"), ({}, [])) + + def test_stray_job_level_line_is_reported(self) -> None: + jobs, problems = parse_jobs( + "jobs:\n a:\n steps: []\n - stray\n b:\n name: B\n" + ) + self.assertEqual(list(jobs), ["a", "b"]) + self.assertEqual(jobs["b"].name, "B") + self.assertEqual(problems, ["unexpected line at job level: '- stray'"]) + + if __name__ == "__main__": unittest.main() diff --git a/scripts/rules/validate.py b/scripts/rules/validate.py index 9fffbb5a4..67f17ec4e 100644 --- a/scripts/rules/validate.py +++ b/scripts/rules/validate.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any from scripts.rules.model import Kind, Registry, Rule, Status, Tier, Violation -from scripts.rules.resolve import node_exists +from scripts.rules.resolve import PROOF_CI_FIELDS, node_exists if TYPE_CHECKING: from pathlib import Path @@ -144,53 +144,100 @@ def _validate_enforcer(rule: Rule) -> list[Violation]: return violations -def _validate_cross_fields(rule: Rule) -> list[Violation]: - violations = [] +def _schema_violation(rule: Rule, message: str, fix: str) -> Violation: + """Report a schema problem with one rule.""" + return Violation(rule.id, "schema", message, fix) - def v(message: str, fix: str) -> None: - violations.append(Violation(rule.id, "schema", message, fix)) - kind, status, tier = rule.enforcer.kind, rule.status, rule.tier +def _status_conflicts(rule: Rule) -> list[Violation]: + """Check that the rule's status fits its enforcer, CI wiring and baseline.""" + kind, status = rule.enforcer.kind, rule.status + violations = [] if status is Status.DOC_ONLY and kind is not Kind.NONE: - v( - "doc-only rules have no enforcer", - 'set enforcer.kind = "none", or pick a stronger status', + violations.append( + _schema_violation( + rule, + "doc-only rules have no enforcer", + 'set enforcer.kind = "none", or pick a stronger status', + ) ) if status is Status.DOC_ONLY and rule.ci: - v( - "doc-only rules have no CI wiring", - "remove rule.ci, or pick a stronger status", + violations.append( + _schema_violation( + rule, + "doc-only rules have no CI wiring", + "remove rule.ci, or pick a stronger status", + ) ) if kind is Kind.NONE and status not in (Status.DOC_ONLY, Status.ADVISORY): - v( - f"status {status} needs an enforcer", - "add [rule.enforcer] or use doc-only/advisory", + violations.append( + _schema_violation( + rule, + f"status {status} needs an enforcer", + "add [rule.enforcer] or use doc-only/advisory", + ) ) if status in (Status.ENFORCED, Status.RATCHETED) and not rule.ci: - v(f"{status} rules must run in CI", "add a [[rule.ci]] reference") + violations.append( + _schema_violation( + rule, f"{status} rules must run in CI", "add a [[rule.ci]] reference" + ) + ) if status is Status.RATCHETED and rule.baseline is None: - v("ratcheted rules need a baseline", "add [rule.baseline] or use enforced") + violations.append( + _schema_violation( + rule, + "ratcheted rules need a baseline", + "add [rule.baseline] or use enforced", + ) + ) if rule.baseline is not None and status is not Status.RATCHETED: - v( - "only ratcheted rules may have a baseline", - 'remove [rule.baseline] or set status = "ratcheted"', + violations.append( + _schema_violation( + rule, + "only ratcheted rules may have a baseline", + 'remove [rule.baseline] or set status = "ratcheted"', + ) ) + return violations + + +def _checker_conflicts(rule: Rule) -> list[Violation]: + """Check that the rule's enforcer kind suits its tier and baseline.""" + kind, tier = rule.enforcer.kind, rule.tier + violations = [] if kind is Kind.RUFF and rule.baseline is not None: - v( - "ruff is never ratcheted", - "remove the baseline and fix every violation in a standalone PR that " - "also changes the ruff select list", + violations.append( + _schema_violation( + rule, + "ruff is never ratcheted", + "remove the baseline and fix every violation in a standalone PR " + "that also changes the ruff select list", + ) ) if tier is Tier.JUDGMENT and kind not in (Kind.LLM_REVIEW, Kind.NONE): - v( - "judgment rules are enforced by review, not by a mechanical checker", - 'use enforcer.kind = "llm-review" or "none"', + violations.append( + _schema_violation( + rule, + "judgment rules are enforced by review, not by a mechanical checker", + 'use enforcer.kind = "llm-review" or "none"', + ) ) if kind is Kind.LLM_REVIEW and tier is not Tier.JUDGMENT: - v("llm-review only enforces judgment rules", 'set tier = "judgment"') + violations.append( + _schema_violation( + rule, + "llm-review only enforces judgment rules", + 'set tier = "judgment"', + ) + ) return violations +def _validate_cross_fields(rule: Rule) -> list[Violation]: + return _status_conflicts(rule) + _checker_conflicts(rule) + + def _validate_proof(rule: Rule, root: Path) -> list[Violation]: if rule.proof is None: return [] @@ -203,8 +250,8 @@ def _validate_proof(rule: Rule, root: Path) -> list[Violation]: violations = _check_fields( rule, fields, PROOF_FIELDS[rule.tier], optional, f"proof ({rule.tier})" ) - named = [key for key in ("workflow", "job", "invocation") if key in fields] - if named and len(named) != 3: + named = [key for key in PROOF_CI_FIELDS if key in fields] + if named and len(named) != len(PROOF_CI_FIELDS): violations.append( Violation( rule.id, @@ -243,120 +290,198 @@ def _validate_proof(rule: Rule, root: Path) -> list[Violation]: return violations -def validate(registry: Registry, root: Path, today: date) -> list[Violation]: - """Check the registry's fields and file references against the tree at ``root``. +def _validate_agents(registry: Registry, root: Path) -> list[Violation]: + """Check that every agent instruction file named in ``[agents]`` exists.""" + if registry.agents is None: + return [] + return [ + Violation( + None, + "schema", + f"agents file {path} does not exist", + "fix [agents] or create the file", + ) + for path in (registry.agents.instructions, *registry.agents.adapters) + if path and not (root / path).is_file() + ] - ``today`` decides whether review dates and exceptions have expired. - """ - violations: list[Violation] = [] - - def v(rule_id: str | None, message: str, fix: str, check: str = "schema") -> None: - violations.append(Violation(rule_id, check, message, fix)) - - if registry.agents is not None: - for path in (registry.agents.instructions, *registry.agents.adapters): - if path and not (root / path).is_file(): - v( - None, - f"agents file {path} does not exist", - "fix [agents] or create the file", - ) - if registry.ci_scripts is not None: - config = registry.ci_scripts - if config.scripts_dir and not (root / config.scripts_dir).is_dir(): - v( + +def _validate_ci_scripts(registry: Registry, root: Path) -> list[Violation]: + """Check that the ``[ci_scripts]`` directory and baseline files exist.""" + config = registry.ci_scripts + if config is None: + return [] + violations = [] + if config.scripts_dir and not (root / config.scripts_dir).is_dir(): + violations.append( + Violation( None, + "schema", f"ci_scripts.scripts_dir {config.scripts_dir} does not exist", "fix [ci_scripts]", ) + ) + violations.extend( + Violation( + None, + "schema", + f"ci_scripts baseline {path} does not exist", + "create it (it may be empty)", + ) for path in ( config.multiline_baseline, config.untested_baseline, config.shellcheck_baseline, - ): - if path and not (root / path).is_file(): - v( - None, - f"ci_scripts baseline {path} does not exist", - "create it (it may be empty)", - ) + ) + if path and not (root / path).is_file() + ) + return violations - category_ids = [category.id for category in registry.categories] - for duplicate in sorted({c for c in category_ids if category_ids.count(c) > 1}): - v(None, f"duplicate category '{duplicate}'", "give each category a unique id") +def _duplicate_ids(registry: Registry, category_ids: list[str]) -> list[Violation]: + """Report each category id, then each rule id, that is used more than once.""" + violations = [ + Violation( + None, + "schema", + f"duplicate category '{duplicate}'", + "give each category a unique id", + ) + for duplicate in sorted({c for c in category_ids if category_ids.count(c) > 1}) + ] rule_ids = [rule.id for rule in registry.rules] - for duplicate in sorted({r for r in rule_ids if rule_ids.count(r) > 1}): - v(duplicate, "duplicate rule id", "give each rule a unique id") + violations.extend( + Violation( + duplicate, "schema", "duplicate rule id", "give each rule a unique id" + ) + for duplicate in sorted({r for r in rule_ids if rule_ids.count(r) > 1}) + ) + return violations - for rule in registry.rules: - if rule.id and not RULE_ID.match(rule.id): - v( - rule.id, + +def _validate_identity( + rule: Rule, root: Path, category_ids: list[str] +) -> list[Violation]: + """Check the rule's id, category, rationale link and scope globs.""" + violations = [] + if rule.id and not RULE_ID.match(rule.id): + violations.append( + _schema_violation( + rule, "id must be lowercase dotted, e.g. 'api.v2.routes-no-sql'", "rename the rule", ) - if rule.category not in category_ids: - v( - rule.id, + ) + if rule.category not in category_ids: + violations.append( + _schema_violation( + rule, f"unknown category '{rule.category}'", f"use one of {category_ids} or add a [[category]]", ) - - rationale_path = rule.rationale.split("#", 1)[0] - if rationale_path and not (root / rationale_path).exists(): - v( - rule.id, + ) + rationale_path = rule.rationale.split("#", 1)[0] + if rationale_path and not (root / rationale_path).exists(): + violations.append( + _schema_violation( + rule, f"rationale {rationale_path} does not exist", "link an existing doc", ) + ) + violations.extend( + _schema_violation(rule, f"scope '{pattern}' matches no files", "fix the glob") + for pattern in rule.scope + if not any(root.glob(pattern)) + ) + return violations - for pattern in rule.scope: - if not any(root.glob(pattern)): - v(rule.id, f"scope '{pattern}' matches no files", "fix the glob") - violations.extend(_validate_enforcer(rule)) - violations.extend(_validate_cross_fields(rule)) - violations.extend(_validate_proof(rule, root)) +def _validate_baseline(rule: Rule, root: Path) -> list[Violation]: + """Check that a ratcheted rule's baseline file exists and matches its count.""" + if rule.baseline is None: + return [] + path = root / rule.baseline.file + if not path.is_file(): + return [ + Violation( + rule.id, + "baseline", + f"baseline file {rule.baseline.file} does not exist", + "create it", + ) + ] + if baseline_entries(path) != rule.baseline.count: + return [ + Violation( + rule.id, + "baseline", + f"baseline count is {rule.baseline.count} but " + f"{rule.baseline.file} lists {baseline_entries(path)}", + "make the count match the file", + ) + ] + return [] - if rule.baseline is not None: - path = root / rule.baseline.file - if not path.is_file(): - v( - rule.id, - f"baseline file {rule.baseline.file} does not exist", - "create it", - "baseline", - ) - elif baseline_entries(path) != rule.baseline.count: - v( - rule.id, - f"baseline count is {rule.baseline.count} but " - f"{rule.baseline.file} lists {baseline_entries(path)}", - "make the count match the file", - "baseline", - ) - if rule.review_by is not None and rule.review_by < today: - v( +def _validate_expiry(rule: Rule, root: Path, today: date) -> list[Violation]: + """Check the review date and each exception's expiry and path on ``today``.""" + violations = [] + if rule.review_by is not None and rule.review_by < today: + violations.append( + Violation( rule.id, + "expiry", f"review_by {rule.review_by} has passed", "review the rule, then remove or extend review_by", - "expiry", ) - for exception in rule.exceptions: - if exception.expires < today: - v( + ) + for exception in rule.exceptions: + if exception.expires < today: + violations.append( + Violation( rule.id, + "expiry", f"exception for {exception.path} expired on {exception.expires}", "remove the exception (fix the code) or extend it with a reason", - "expiry", ) - if not any(root.glob(exception.path)): - v( - rule.id, + ) + if not any(root.glob(exception.path)): + violations.append( + _schema_violation( + rule, f"exception path '{exception.path}' matches no files", "remove the stale exception", ) + ) + return violations + + +def _validate_rule( + rule: Rule, root: Path, category_ids: list[str], today: date +) -> list[Violation]: + """Run every per-rule check, in the order the guard reports them.""" + return [ + *_validate_identity(rule, root, category_ids), + *_validate_enforcer(rule), + *_validate_cross_fields(rule), + *_validate_proof(rule, root), + *_validate_baseline(rule, root), + *_validate_expiry(rule, root, today), + ] + +def validate(registry: Registry, root: Path, today: date) -> list[Violation]: + """Check the registry's fields and file references against the tree at ``root``. + + ``today`` decides whether review dates and exceptions have expired. + """ + category_ids = [category.id for category in registry.categories] + violations = [ + *_validate_agents(registry, root), + *_validate_ci_scripts(registry, root), + *_duplicate_ids(registry, category_ids), + ] + for rule in registry.rules: + violations.extend(_validate_rule(rule, root, category_ids, today)) return violations diff --git a/scripts/rules/workflows.py b/scripts/rules/workflows.py index 71ae03b13..53a2803f1 100644 --- a/scripts/rules/workflows.py +++ b/scripts/rules/workflows.py @@ -15,10 +15,13 @@ import re from dataclasses import dataclass, field +from typing import NamedTuple _KEY = re.compile(r"^(\s*)(-\s+)?([A-Za-z0-9_-]+):\s*(.*)$") _LIST_ITEM = re.compile(r"^(\s*)-\s+(.*)$") _UNSUPPORTED = re.compile(r"(<<:|:\s*[&*][A-Za-z0-9_-]+\s*$)") +# A quoted scalar is at least its opening and closing quote. +_QUOTE_PAIR_LENGTH = 2 @dataclass(frozen=True) @@ -77,7 +80,7 @@ def _strip_comment(line: str) -> str: def _unquote(value: str) -> str: value = value.strip() - if len(value) >= 2 and value[0] == value[-1] and value[0] in "'\"": + if len(value) >= _QUOTE_PAIR_LENGTH and value[0] == value[-1] and value[0] in "'\"": return value[1:-1] return value @@ -134,149 +137,201 @@ def _inline_list(value: str) -> tuple[str, ...]: return tuple(_unquote(item) for item in inner.split(",") if item.strip()) -def _parse_job(job_id: str, lines: list[str]) -> JobBlock: - name = "" - uses = "" - step_uses: list[str] = [] - runs: list[str] = [] - steps: list[str] = [] - matrix: dict[str, tuple[str, ...]] = {} - body_indent = None - # Step tracking, so each run: block can be attributed to its step. - item_indent: int | None = None - in_steps = False - step_index = 0 - step_key_column = -1 - step_labels: dict[int, str] = {} - blocks: list[tuple[int, str, str]] = [] - index = 0 - while index < len(lines): - line = lines[index] - match = _KEY.match(line) - if not match: - item = _LIST_ITEM.match(line) - if ( - in_steps - and item - and item_indent is not None - and len(item.group(1)) == item_indent - ): - step_index += 1 - index += 1 - continue +class _KeyLine(NamedTuple): + """One ``key: value`` line of a job, with the column its key starts at.""" + + leading: str + dash: str | None + key: str + value: str + column: int + top_level: bool + + +class _JobReader: + """Reads one job's lines, attributing each ``run:`` block to its step.""" + + def __init__(self, lines: list[str]) -> None: + self.lines = lines + self.index = 0 + self.name = "" + self.uses = "" + self.step_uses: list[str] = [] + self.runs: list[str] = [] + self.steps: list[str] = [] + self.matrix: dict[str, tuple[str, ...]] = {} + self.body_indent: int | None = None + # Step tracking, so each run: block can be attributed to its step. + self.item_indent: int | None = None + self.in_steps = False + self.step_index = 0 + self.step_key_column = -1 + self.step_labels: dict[int, str] = {} + self.blocks: list[tuple[int, str, str]] = [] + + def read(self) -> None: + """Read every line of the job.""" + while self.index < len(self.lines): + match = _KEY.match(self.lines[self.index]) + if match: + self._read_key(match) + else: + self._count_bare_step(self.lines[self.index]) + self.index += 1 + + def job_block(self, job_id: str) -> JobBlock: + """Return what was read as a job named ``job_id`` unless it has a name.""" + return JobBlock( + id=job_id, + name=self.name or job_id, + run_text="\n".join(self.runs), + uses=self.uses, + step_uses=tuple(self.step_uses), + step_names=tuple(self.steps), + matrix=self.matrix, + run_blocks=tuple( + RunBlock( + step=self.step_labels.get(number) or f"step {number}", + code_lines=code_lines(text), + chained=chains_commands(text), + text=text, + ) + for number, text, _style in self.blocks + ), + ) + + def _still_inside(self, column: int) -> bool: + """Whether the current line is blank or indented past ``column``.""" + return self.index < len(self.lines) and ( + not self.lines[self.index].strip() + or _indent(self.lines[self.index]) > column + ) + + def _count_bare_step(self, line: str) -> None: + """Count a list item that is not a key (such as a quoted string) as a step.""" + item = _LIST_ITEM.match(line) + if ( + self.in_steps + and item + and self.item_indent is not None + and len(item.group(1)) == self.item_indent + ): + self.step_index += 1 + + def _read_key(self, match: re.Match[str]) -> None: leading, dash, key, value = match.groups() - key_column = len(leading) + (len(dash) if dash else 0) - if body_indent is None and not dash: - body_indent = key_column - top_level = not dash and key_column == body_indent - if top_level: - in_steps = key == "steps" and not value - item_indent = None - elif in_steps and dash: - if item_indent is None: - item_indent = len(leading) - if len(leading) == item_indent: - step_index += 1 - step_key_column = key_column - if in_steps and key == "name" and key_column == step_key_column: - step_labels[step_index] = _unquote(value) - - if key == "name" and top_level: - name = _unquote(value) - elif key == "name": - steps.append(_unquote(value)) - elif key == "uses" and top_level: - uses = _unquote(value) - elif key == "uses": - step_uses.append(_unquote(value)) - elif key == "run": - if value.strip()[:1] in ("|", ">"): - block = [] - index += 1 - while index < len(lines) and ( - not lines[index].strip() or _indent(lines[index]) > key_column - ): - block.append(lines[index].strip()) - index += 1 - runs.append("\n".join(block)) - blocks.append((step_index, "\n".join(block), value.strip()[:1])) - continue - runs.append(_unquote(value)) - blocks.append((step_index, _unquote(value), "|")) - elif key == "matrix" and not value: - index += 1 - while index < len(lines) and ( - not lines[index].strip() or _indent(lines[index]) > key_column - ): - entry = _KEY.match(lines[index]) - if entry and not entry.group(2): - entry_key, entry_value = entry.group(3), entry.group(4) - entry_column = len(entry.group(1)) - if entry_value.startswith("["): - matrix[entry_key] = _inline_list(entry_value) - index += 1 - continue - items = [] - index += 1 - while index < len(lines) and ( - not lines[index].strip() - or _indent(lines[index]) >= entry_column - ): - item = _LIST_ITEM.match(lines[index]) - if not item or _KEY.match(lines[index]): - break - items.append(_unquote(item.group(2))) - index += 1 - matrix[entry_key] = tuple(items) - continue - index += 1 - continue - index += 1 - return JobBlock( - id=job_id, - name=name or job_id, - run_text="\n".join(runs), - uses=uses, - step_uses=tuple(step_uses), - step_names=tuple(steps), - matrix=matrix, - run_blocks=tuple( - RunBlock( - step=step_labels.get(number) or f"step {number}", - code_lines=code_lines(text), - chained=chains_commands(text), - text=text, - ) - for number, text, _style in blocks - ), - ) + column = len(leading) + (len(dash) if dash else 0) + if self.body_indent is None and not dash: + self.body_indent = column + entry = _KeyLine( + leading=leading, + dash=dash, + key=key, + value=value, + column=column, + top_level=not dash and column == self.body_indent, + ) + self._track_step(entry) + if entry.key == "run" and entry.value.strip()[:1] in ("|", ">"): + self._read_block_scalar(entry) + elif entry.key == "matrix" and not entry.value: + self._read_matrix(entry.column) + else: + self._record_scalar(entry) + self.index += 1 + + def _track_step(self, entry: _KeyLine) -> None: + """Note when a new step starts, and the name it gives itself.""" + if entry.top_level: + self.in_steps = entry.key == "steps" and not entry.value + self.item_indent = None + elif self.in_steps and entry.dash: + if self.item_indent is None: + self.item_indent = len(entry.leading) + if len(entry.leading) == self.item_indent: + self.step_index += 1 + self.step_key_column = entry.column + if ( + self.in_steps + and entry.key == "name" + and entry.column == self.step_key_column + ): + self.step_labels[self.step_index] = _unquote(entry.value) + + def _record_scalar(self, entry: _KeyLine) -> None: + """Record a one-line ``name:``, ``uses:`` or ``run:`` value.""" + if entry.key == "name" and entry.top_level: + self.name = _unquote(entry.value) + elif entry.key == "name": + self.steps.append(_unquote(entry.value)) + elif entry.key == "uses" and entry.top_level: + self.uses = _unquote(entry.value) + elif entry.key == "uses": + self.step_uses.append(_unquote(entry.value)) + elif entry.key == "run": + self.runs.append(_unquote(entry.value)) + self.blocks.append((self.step_index, _unquote(entry.value), "|")) + + def _read_block_scalar(self, entry: _KeyLine) -> None: + """Read a ``|`` or ``>`` run: block, one stripped line per script line.""" + block = [] + self.index += 1 + while self._still_inside(entry.column): + block.append(self.lines[self.index].strip()) + self.index += 1 + self.runs.append("\n".join(block)) + self.blocks.append((self.step_index, "\n".join(block), entry.value.strip()[:1])) + + def _read_matrix(self, column: int) -> None: + """Read the ``strategy.matrix`` entries nested under the key at ``column``.""" + self.index += 1 + while self._still_inside(column): + entry = _KEY.match(self.lines[self.index]) + if entry and not entry.group(2): + self._read_matrix_entry(entry) + else: + self.index += 1 + + def _read_matrix_entry(self, entry: re.Match[str]) -> None: + """Read one matrix key's values, written inline or as a block list.""" + key, value = entry.group(3), entry.group(4) + self.index += 1 + if value.startswith("["): + self.matrix[key] = _inline_list(value) + return + # Block list items may sit at the key's own column. + list_column = len(entry.group(1)) - 1 + items = [] + while self._still_inside(list_column): + item = _LIST_ITEM.match(self.lines[self.index]) + if not item or _KEY.match(self.lines[self.index]): + break + items.append(_unquote(item.group(2))) + self.index += 1 + self.matrix[key] = tuple(items) -def parse_jobs(text: str) -> tuple[dict[str, JobBlock], list[str]]: - """Return the jobs in a workflow, plus descriptions of unsupported constructs.""" - lines = [_strip_comment(line) for line in text.splitlines()] - problems = [ - f"line {number}: YAML anchors/aliases/merge keys are not supported" - for number, line in enumerate(lines, 1) - if _UNSUPPORTED.search(line) - ] +def _parse_job(job_id: str, lines: list[str]) -> JobBlock: + reader = _JobReader(lines) + reader.read() + return reader.job_block(job_id) - start = next( - (i for i, line in enumerate(lines) if re.match(r"^jobs:\s*$", line)), None - ) - if start is None: - if any(re.match(r"^jobs:\s*\S", line) for line in lines): - problems.append("flow-style 'jobs:' mapping is not supported") - return {}, problems +def _jobs_body(lines: list[str], start: int) -> list[str]: + """Return the lines after ``jobs:`` up to the next top-level key.""" body: list[str] = [] for line in lines[start + 1 :]: if line.strip() and _indent(line) == 0: break body.append(line) + return body + +def _split_jobs(body: list[str]) -> tuple[dict[str, JobBlock], list[str]]: + """Parse each job under ``jobs:``, and describe any stray job-level line.""" job_indent = next((_indent(line) for line in body if line.strip()), 0) jobs: dict[str, JobBlock] = {} + problems: list[str] = [] current: str | None = None current_lines: list[str] = [] for line in body: @@ -292,3 +347,24 @@ def parse_jobs(text: str) -> tuple[dict[str, JobBlock], list[str]]: if current: jobs[current] = _parse_job(current, current_lines) return jobs, problems + + +def parse_jobs(text: str) -> tuple[dict[str, JobBlock], list[str]]: + """Return the jobs in a workflow, plus descriptions of unsupported constructs.""" + lines = [_strip_comment(line) for line in text.splitlines()] + problems = [ + f"line {number}: YAML anchors/aliases/merge keys are not supported" + for number, line in enumerate(lines, 1) + if _UNSUPPORTED.search(line) + ] + + start = next( + (i for i, line in enumerate(lines) if re.match(r"^jobs:\s*$", line)), None + ) + if start is None: + if any(re.match(r"^jobs:\s*\S", line) for line in lines): + problems.append("flow-style 'jobs:' mapping is not supported") + return {}, problems + + jobs, job_problems = _split_jobs(_jobs_body(lines, start)) + return jobs, problems + job_problems From 02e6e218cf856ed2743d170eabbe7906e0976e96 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:37:33 +0400 Subject: [PATCH 60/66] Select every ruff rule (ALL) Every rule family now passes, so ruff.toml selects ALL instead of the per-group list. It ignores COM812 and ISC001, which conflict with ruff format (ruff's documentation says to disable them with the formatter), and D203 and D213, which are mutually exclusive with the kept D211 and D212. The per-file target version, namespace packages and per-file ignores are unchanged. The sim.quality.ruff-lint rule's enforcer codes become ["ALL"], and its statement names the four ignores and why. The rules-registry skill says every rule applies apart from the listed and registered per-file ignores, that a ruff upgrade PR fixes any new rule's violations, and that an exception is an inline noqa with a reason. The rendered rules page and the AGENTS.md rules block are regenerated, and the changelog fragment describes the complexity refactors and the switch. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- changelog_entry.yaml | 1 + docs/engineering/skills/rules-registry.md | 9 ++++--- docs/generated/rules.md | 8 +++--- ruff.toml | 33 +++++++---------------- rules/registry.toml | 4 +-- 6 files changed, 24 insertions(+), 33 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e841eb157..e13cf1c39 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -139,7 +139,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. - `sim.quality.every-project-tested`: Every project and library with unit tests runs them in the PR test matrix. (doc-only, not checked: follow it yourself) - `sim.quality.ruff-format-src`: Source code under projects/*/src and libs/*/src is formatted with ruff format. (blocking: ruff format) - `sim.quality.ruff-format-tests`: Test code is formatted with ruff format too; CI currently checks only src directories. (doc-only, not checked: follow it yourself) -- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (A, ANN, ARG, ASYNC, B, BLE, C4, D, DTZ, E, E4, E7, E9, EM, ERA, EXE, F, FAST, FBT, FIX, FLY, FURB, G, I, INP, LOG, N, PD, PERF, PGH, PIE, PLC, PLW, PT, PTH, PYI, RET, RSE, RUF, S, SIM, SLF, T20, TC, TD, TRY, UP, W)) +- `sim.quality.ruff-lint`: All Python code passes ruff check with every ruff rule selected (ALL), with no ratchet, except the ignores in ruff.toml: COM812 and ISC001, which conflict with ruff format, and D203 and D213, which conflict with the kept D211 and D212. (blocking: ruff check (ALL)) - `sim.quality.tests-pass`: Each project's unit tests pass in its own locked environment. (blocking: CI job `test`) - `sim.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert, need no docstrings, annotations or named constants for magic numbers, and may reach into private members, live outside packages and import inside a test (S101, D1, ANN, PLR2004, SLF001, INP001, PLC0415); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, INP001, PLC0415, PLR2004, S101, SLF001 for `**/tests/**`) - `sim.quality.unittest-ruff-exemptions`: The shared rules tool's tests use unittest, because the tool needs only the standard library, so ruff's pytest-style rules (PT) don't apply to them. (blocking: ruff per-file ignore of PT for `scripts/rules/tests/**`) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index ae9d56042..81f0621a7 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -4,6 +4,7 @@ - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md changed: + - Enabled every ruff rule (select ALL), ignoring only COM812 and ISC001, which conflict with ruff format, and the D203 and D213 halves of the mutually exclusive docstring pairs; complex functions, long parameter lists and magic numbers were split into smaller helpers and grouped parameters, each behind new characterization tests, across the contract and persistence validation, entry settings, app factory and Stage 12 adapter, gateway route resolution, Stage 12 coordination, workers, parity and qualification, the precompute and national reduce, the routing publisher and its command line, the release sentinel, the package update helper and the shared rules tool; Modal function, FastAPI route, Pydantic validator and command-line signatures and every OpenAPI document are unchanged, and three library APIs that other projects import keep a reasoned PLR0913 suppression - Enabled ruff's argument, idiom and pytest-style rules (ARG, ASYNC, B, C4, FBT, PD, PERF, PGH, PLW, PT, PTH, RUF, SIM); unused arguments are removed, underscored or explained, FastAPI auth dependencies use Annotated, re-raised errors keep their cause, and boolean options of in-repo helpers are keyword-only; OpenAPI documents are unchanged, and the shared rules tool's unittest tests are exempt from pytest-style rules - Enabled ruff's structure and import rules (INP, PLC, SLF, N, A, FAST, E); long lines are wrapped, deferred imports that Modal images, test seams, import cycles or heavy dependencies need carry a reason, script directories are declared to ruff as namespace packages, four internal exceptions gained an Error suffix, and published field names, OpenAPI documents and logged error types are unchanged; tests are also exempt from SLF001, INP001 and PLC0415 - Enabled ruff's security, logging and debugging rules (S, T20, LOG, G, DTZ, ERA, TD, FIX, EXE); asserts outside tests became explicit checks, command-line scripts keep their printed output, test subprocess calls go through helpers, and every deliberate exception carries a reasoned suppression; the gateway's auth misconfiguration log no longer repeats its message, and the deployed-service tests share the unit-test ruff exemptions diff --git a/docs/engineering/skills/rules-registry.md b/docs/engineering/skills/rules-registry.md index b4b843a05..4f80218e5 100644 --- a/docs/engineering/skills/rules-registry.md +++ b/docs/engineering/skills/rules-registry.md @@ -115,9 +115,12 @@ fail the guard yet. ## Baselines -- **Ruff is never ratcheted.** Adopting a ruff rule family is one standalone PR - that changes the select list and fixes every resulting violation. Don't use - `--add-noqa` or blanket per-file ignores. +- **Ruff is never ratcheted.** `ruff.toml` selects `ALL`, so every ruff rule + applies except the four listed in `ignore` (two formatter conflicts and one + of each mutually exclusive docstring pair) and the registered per-file + ignores. When a ruff upgrade adds rules, the upgrade PR fixes every new + violation. Don't use `--add-noqa` or blanket per-file ignores; a justified + exception is an inline `# noqa: CODE # reason`. - Other enforcers (mypy, Pyright, Import Linter, custom guards) may start `ratcheted` with a baseline file of tolerated violations. Baselines may only shrink: the guard compares counts against the base branch. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 099577ff5..c9786e610 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -246,7 +246,7 @@ flowchart LR | `sim.quality.every-project-tested` | Every project and library with unit tests runs them in the PR test matrix. | mechanical | doc-only | no | nothing | — | — | | `sim.quality.ruff-format-src` | Source code under projects/*/src and libs/*/src is formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | | `sim.quality.ruff-format-tests` | Test code is formatted with ruff format too; CI currently checks only src directories. | mechanical | doc-only | no | nothing | — | — | -| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (A, ANN, ARG, ASYNC, B, BLE, C4, D, DTZ, E, E4, E7, E9, EM, ERA, EXE, F, FAST, FBT, FIX, FLY, FURB, G, I, INP, LOG, N, PD, PERF, PGH, PIE, PLC, PLW, PT, PTH, PYI, RET, RSE, RUF, S, SIM, SLF, T20, TC, TD, TRY, UP, W) | `pr.yml:lint` | — | +| `sim.quality.ruff-lint` | All Python code passes ruff check with every ruff rule selected (ALL), with no ratchet, except the ignores in ruff.toml: COM812 and ISC001, which conflict with ruff format, and D203 and D213, which conflict with the kept D211 and D212. | mechanical | enforced | yes | ruff check (ALL) | `pr.yml:lint` | — | | `sim.quality.tests-pass` | Each project's unit tests pass in its own locked environment. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | | `sim.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert, need no docstrings, annotations or named constants for magic numbers, and may reach into private members, live outside packages and import inside a test (S101, D1, ANN, PLR2004, SLF001, INP001, PLC0415); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, INP001, PLC0415, PLR2004, S101, SLF001 for `**/tests/**` | `pr.yml:lint` | — | | `sim.quality.unittest-ruff-exemptions` | The shared rules tool's tests use unittest, because the tool needs only the standard library, so ruff's pytest-style rules (PT) don't apply to them. | mechanical | enforced | yes | ruff per-file ignore of PT for `scripts/rules/tests/**` | `pr.yml:lint` | — | @@ -262,7 +262,7 @@ flowchart LR e_ruff_format["ruff format"] r_sim_quality_ruff_format_tests("sim.quality.ruff-format-tests (doc-only)") r_sim_quality_ruff_lint("sim.quality.ruff-lint (enforced)") - e_ruff_check__A__ANN__ARG__ASYNC__B__BLE__C4__D__DTZ__E__E4__E7__E9__EM__ERA__EXE__F__FAST__FBT__FIX__FLY__FURB__G__I__INP__LOG__N__PD__PERF__PGH__PIE__PLC__PLW__PT__PTH__PYI__RET__RSE__RUF__S__SIM__SLF__T20__TC__TD__TRY__UP__W_["ruff check (A, ANN, ARG, ASYNC, B, BLE, C4, D, DTZ, E, E4, E7, E9, EM, ERA, EXE, F, FAST, FBT, FIX, FLY, FURB, G, I, INP, LOG, N, PD, PERF, PGH, PIE, PLC, PLW, PT, PTH, PYI, RET, RSE, RUF, S, SIM, SLF, T20, TC, TD, TRY, UP, W)"] + e_ruff_check__ALL_["ruff check (ALL)"] r_sim_quality_tests_pass("sim.quality.tests-pass (enforced)") e_CI_job_test["CI job test"] j_CI_pr_yml_test["CI pr.yml:test"] @@ -276,8 +276,8 @@ flowchart LR r_sim_quality_ruff_format_src --> e_ruff_format e_ruff_format --> j_CI_pr_yml_lint r_sim_quality_ruff_format_tests -.-> gap_unenforced - r_sim_quality_ruff_lint --> e_ruff_check__A__ANN__ARG__ASYNC__B__BLE__C4__D__DTZ__E__E4__E7__E9__EM__ERA__EXE__F__FAST__FBT__FIX__FLY__FURB__G__I__INP__LOG__N__PD__PERF__PGH__PIE__PLC__PLW__PT__PTH__PYI__RET__RSE__RUF__S__SIM__SLF__T20__TC__TD__TRY__UP__W_ - e_ruff_check__A__ANN__ARG__ASYNC__B__BLE__C4__D__DTZ__E__E4__E7__E9__EM__ERA__EXE__F__FAST__FBT__FIX__FLY__FURB__G__I__INP__LOG__N__PD__PERF__PGH__PIE__PLC__PLW__PT__PTH__PYI__RET__RSE__RUF__S__SIM__SLF__T20__TC__TD__TRY__UP__W_ --> j_CI_pr_yml_lint + r_sim_quality_ruff_lint --> e_ruff_check__ALL_ + e_ruff_check__ALL_ --> j_CI_pr_yml_lint r_sim_quality_tests_pass --> e_CI_job_test e_CI_job_test --> j_CI_pr_yml_test r_sim_quality_tests_ruff_exemptions --> e_ruff_per_file_ignore_of_ANN__D1__INP001__PLC0415__PLR2004__S101__SLF001_for____tests___ diff --git a/ruff.toml b/ruff.toml index d995e6942..ac444bb2c 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,5 +1,5 @@ -# Rule families are enabled one themed group per PR, each fixing everything it -# finds (no ratchet), until the select list is "ALL". See +# Every ruff rule applies, with no ratchet, apart from the ignores in [lint] and +# the per-file ignores below (sim.quality.ruff-lint). See # docs/engineering/skills/rules-registry.md and rules/registry.toml. # There is no root pyproject.toml, so files outside the projects and libs would # be treated as Python 3.9. scripts/rules supports Python 3.11+ (and is shared @@ -21,28 +21,15 @@ namespace-packages = [ ] [lint] -select = [ - "E4", "E7", "E9", "F", - # Group 1: safe auto-fixes. - "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W", - # Group 2: typing-only imports. - "TC", - # Group 3: annotations. - "ANN", - # Group 4: docstrings. - "D", - # Group 5: exceptions and messages. - "BLE", "EM", "TRY", - # Group 6: security, logging and debugging leftovers. - "DTZ", "ERA", "EXE", "FIX", "G", "LOG", "S", "T20", "TD", - # Group 7: structure and imports. - "A", "E", "FAST", "INP", "N", "PLC", "SLF", - # Group 8a: arguments, idioms and pytest style. - "ARG", "ASYNC", "B", "C4", "FBT", "PD", "PERF", "PGH", "PLW", "PT", "PTH", "RUF", "SIM", +select = ["ALL"] +ignore = [ + # These conflict with ruff format; ruff's documentation says to disable + # them when using the formatter. + "COM812", "ISC001", + # Mutually exclusive pydocstyle pairs: keep D211 (no blank line before a + # class docstring) and D212 (summary on the first line). + "D203", "D213", ] -# Mutually exclusive pydocstyle pairs: keep D211 (no blank line before a class -# docstring) and D212 (summary on the first line). -ignore = ["D203", "D213"] [lint.per-file-ignores] # Tests may use pytest's bare assert, and need no docstrings, annotations or diff --git a/rules/registry.toml b/rules/registry.toml index c809cbb68..fb5b238d6 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -703,7 +703,7 @@ kind = "none" [[rule]] id = "sim.quality.ruff-lint" category = "quality" -statement = "All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL." +statement = "All Python code passes ruff check with every ruff rule selected (ALL), with no ratchet, except the ignores in ruff.toml: COM812 and ISC001, which conflict with ruff format, and D203 and D213, which conflict with the kept D211 and D212." rationale = "docs/engineering/skills/rules-registry.md#baselines" scope = ["projects/*/src/**/*.py", "libs/*/src/**/*.py", "projects/*/tests/**/*.py", "libs/*/tests/**/*.py"] tier = "mechanical" @@ -712,7 +712,7 @@ autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["A", "ANN", "ARG", "ASYNC", "B", "BLE", "C4", "D", "DTZ", "E", "E4", "E7", "E9", "EM", "ERA", "EXE", "F", "FAST", "FBT", "FIX", "FLY", "FURB", "G", "I", "INP", "LOG", "N", "PD", "PERF", "PGH", "PIE", "PLC", "PLW", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "S", "SIM", "SLF", "T20", "TC", "TD", "TRY", "UP", "W"] +codes = ["ALL"] [[rule.ci]] workflow = ".github/workflows/pr.yml" job = "lint" From bd58bca8d64ae97e498b988d2a3f09844a6909a1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:33:09 +0400 Subject: [PATCH 61/66] Pin the poll, deadline and budget timing the segmented runner #795 extended #795 landed on main after this stack branched and gave the segmented national runner a child deadline, a parent budget and a clock beside its poll backoff. PLR0913, which this PR enables, flags the eight-argument initializer; these tests pin the deployed defaults and how given values are taken before the timing is passed as one group. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_segmented_national.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 18d4bd54b..be58b557a 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -1,6 +1,7 @@ """Unit tests for the segmented-national fan-out runner (C5/C6 + review fixes).""" import ast +import time from pathlib import Path from types import SimpleNamespace @@ -430,6 +431,46 @@ def _calls_with_stuck_child(clock, respawned_call, **sibling_options): return [*calls, respawned_call] +class TestSegmentedNationalRunnerTiming: + def test__runner_defaults_to_the_deployed_poll_deadline_and_budget(self): + runner = sn.SegmentedNationalRunner( + dict(NATIONAL), + app_name="policyengine-simulation-py-test", + runtime=TEST_RUNTIME, + modal_module=FakeModal([]), + ) + + assert runner.poll_interval_initial_seconds == sn.POLL_INTERVAL_INITIAL_SECONDS + assert runner.poll_interval_max_seconds == sn.POLL_INTERVAL_MAX_SECONDS + assert runner.poll_interval_backoff_factor == sn.POLL_INTERVAL_BACKOFF_FACTOR + assert runner.child_deadline_seconds == sn.CHILD_DEADLINE_SECONDS + assert runner.parent_budget_seconds == ( + sn.PARENT_FUNCTION_TIMEOUT_SECONDS - sn.PARENT_COMPLETION_MARGIN_SECONDS + ) + assert runner.clock is time.monotonic + assert runner.budget_ends_at == ( + runner.started_at + runner.parent_budget_seconds + ) + + def test__runner_takes_the_given_deadline_budget_and_clock(self): + clock = ManualClock() + clock.advance(50) + runner = _runner( + FakeModal([]), + child_deadline_seconds=600, + parent_budget_seconds=900, + clock=clock, + ) + + assert runner.poll_interval_initial_seconds == 0.001 + assert runner.poll_interval_max_seconds == 0.002 + assert runner.child_deadline_seconds == 600 + assert runner.parent_budget_seconds == 900 + assert runner.clock is clock + assert runner.started_at == 50 + assert runner.budget_ends_at == 950 + + class TestSegmentedNationalChildDeadline: @pytest.mark.usefixtures("bare_country") def test__overdue_child_is_cancelled_and_respawned_once(self): From b4bd8d0026f9a89347b51137b0c85aaba569113e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:33:51 +0400 Subject: [PATCH 62/66] Pass the segmented runner the poll and deadline timing #795 added as one group #795 landed on main after this stack branched. PLR0913, which this PR enables, flags the segmented national runner's initializer once #795's child deadline, parent budget and clock sit beside the poll backoff. The four now travel as one SegmentedRunTiming with the same defaults, which the previous commit pins; run_segmented_national keeps its signature. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/segmented_national.py | 43 +++++++++++++------ .../tests/test_segmented_national.py | 14 +++--- 2 files changed, 37 insertions(+), 20 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 2a9c22ae9..bed85163b 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -88,6 +88,26 @@ # The first dispatch plus one respawn after a missed deadline. MAX_CHILD_ATTEMPTS = 2 + +@dataclass(frozen=True) +class SegmentedRunTiming: + """How the runner spaces its polls and how long it waits for children. + + ``child_deadline_seconds`` runs from each child's dispatch and + ``parent_budget_seconds`` from the runner's start, both read from + ``clock``. + """ + + poll_backoff: PollBackoff = DEFAULT_POLL_BACKOFF + child_deadline_seconds: float = CHILD_DEADLINE_SECONDS + parent_budget_seconds: float = ( + PARENT_FUNCTION_TIMEOUT_SECONDS - PARENT_COMPLETION_MARGIN_SECONDS + ) + clock: Callable[[], float] = time.monotonic + + +DEFAULT_RUN_TIMING = SegmentedRunTiming() + # Parent-scoped keys never forwarded to children. (The shared payload # builder owns _telemetry: stripped, then re-attached when present.) # `region` is stripped because v1-style national parents carry @@ -192,12 +212,7 @@ def __init__( app_name: str, runtime: ObservabilityRuntime, modal_module: ModuleType | None = None, - poll_backoff: PollBackoff = DEFAULT_POLL_BACKOFF, - child_deadline_seconds: float = CHILD_DEADLINE_SECONDS, - parent_budget_seconds: float = ( - PARENT_FUNCTION_TIMEOUT_SECONDS - PARENT_COMPLETION_MARGIN_SECONDS - ), - clock: Callable[[], float] = time.monotonic, + timing: SegmentedRunTiming = DEFAULT_RUN_TIMING, ) -> None: """Resolve the country's region groups and a lazy segment-worker handle. @@ -207,16 +222,16 @@ def __init__( self.app_name = app_name self.runtime = runtime self.modal = modal if modal_module is None else modal_module - self.poll_interval_initial_seconds = poll_backoff.initial_seconds - self.poll_interval_max_seconds = poll_backoff.maximum_seconds - self.poll_interval_backoff_factor = poll_backoff.factor - self.child_deadline_seconds = child_deadline_seconds - self.parent_budget_seconds = parent_budget_seconds - self.clock = clock + self.poll_interval_initial_seconds = timing.poll_backoff.initial_seconds + self.poll_interval_max_seconds = timing.poll_backoff.maximum_seconds + self.poll_interval_backoff_factor = timing.poll_backoff.factor + self.child_deadline_seconds = timing.child_deadline_seconds + self.parent_budget_seconds = timing.parent_budget_seconds + self.clock = timing.clock # The parent's Modal timeout started before this runner existed; the # completion margin inside the default budget covers that gap. - self.started_at = clock() - self.budget_ends_at = self.started_at + parent_budget_seconds + self.started_at = self.clock() + self.budget_ends_at = self.started_at + self.parent_budget_seconds self.child_retry_count = 0 # Overdue attempts replaced by a respawn. They are never polled # again, but every later cancel retries them: the first cancel can diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index be58b557a..397627f5f 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -218,18 +218,20 @@ def _spawn(self, payload, *, observability_context=None): return self._calls[len(self.spawned_payloads) - 1] -def _runner(fake_modal, params=None, runtime=TEST_RUNTIME, **runner_options): +def _runner(fake_modal, params=None, runtime=TEST_RUNTIME, **timing_options): return sn.SegmentedNationalRunner( params or dict(NATIONAL), app_name="policyengine-simulation-py-test", runtime=runtime, modal_module=fake_modal, - poll_backoff=PollBackoff( - initial_seconds=0.001, - maximum_seconds=0.002, - factor=sn.POLL_INTERVAL_BACKOFF_FACTOR, + timing=sn.SegmentedRunTiming( + poll_backoff=PollBackoff( + initial_seconds=0.001, + maximum_seconds=0.002, + factor=sn.POLL_INTERVAL_BACKOFF_FACTOR, + ), + **timing_options, ), - **runner_options, ) From a463a79f9d83cc7b0321e1b2d6ec2a15ea9f7862 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:34:03 +0400 Subject: [PATCH 63/66] Set the fake child call's cancel failures #795 added after construction #795 landed on main after this stack branched. PLR0913 flags the six-argument initializer of the segmented national tests' fake child call once #795 added before_each_poll and cancel_failures. No test passes cancel_failures to the initializer; the two that use it set it on the call, so it now starts at zero and stays an attribute. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_segmented_national.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 397627f5f..6b6170a9c 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -144,7 +144,10 @@ def test__strips_v1_style_region_us_from_children(self): class FakeCall: - """A FunctionCall whose result resolves after N polls (or raises).""" + """A FunctionCall whose result resolves after N polls (or raises). + + Setting ``cancel_failures`` makes that many cancels fail first. + """ def __init__( self, @@ -153,14 +156,13 @@ def __init__( error=None, errors_once=0, before_each_poll=None, - cancel_failures=0, ): self.result = result self.polls_until_ready = polls_until_ready self.error = error self.errors_once = errors_once self.before_each_poll = before_each_poll - self.cancel_failures = cancel_failures + self.cancel_failures = 0 self.cancelled = False self.cancel_attempts = 0 self.poll_count = 0 From 601987584f268007c695a475c62f65e45f16f359 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:34:29 +0400 Subject: [PATCH 64/66] Pin how the tracked_stage #791 added scopes and records each stage #791 landed on main after this stack branched. PLR0913, which this PR enables, flags its six-argument tracked_stage, which other tests cover only through the coordinator and worker. These tests pin it directly before its stage plan, progress and runtime are passed as one group: the span or operation it opens, the attributes it passes, the role it records and the failing stage it keeps without a runtime. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_failure_causes.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py index 304cb8db9..c2b0b06f7 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py @@ -3,6 +3,7 @@ from __future__ import annotations import time +from contextlib import contextmanager import pytest from policyengine_simulation_contract.spm import SPMComparisonProvenance @@ -16,6 +17,11 @@ failure_summary, record_failure_cause, redact_failure_message, + tracked_stage, +) +from policyengine_simulation_observability.stages import ( + STAGE12_CANONICAL_REPORT_STAGES, + Stage, ) from pydantic import BaseModel, TypeAdapter, ValidationError from sqlalchemy.exc import IntegrityError, OperationalError @@ -472,3 +478,64 @@ def _progress_in(stage_name: str) -> StageProgress: with pytest.raises(RuntimeError), progress.stage(stage_name): raise RuntimeError(msg) return progress + + +class ScopeRecordingRuntime: + def __init__(self) -> None: + self.scopes: list[tuple[str, str, object]] = [] + + @contextmanager + def span(self, name, *, attributes=None): + self.scopes.append(("span", name, attributes)) + yield + + @contextmanager + def operation(self, name, *, attributes=None): + self.scopes.append(("operation", name, attributes)) + yield + + +@pytest.mark.parametrize( + ("operation", "scope_kind"), [(False, "span"), (True, "operation")] +) +def test_tracked_stage_scopes_and_records_the_stage_with_its_role( + operation, scope_kind +) -> None: + runtime = ScopeRecordingRuntime() + progress = StageProgress(role="baseline") + attributes = {"simulation_role": "reform"} + + with tracked_stage( + Stage.STAGE12_CHILD_WAIT, + attributes, + stage_plan=STAGE12_CANONICAL_REPORT_STAGES, + progress=progress, + runtime=runtime, + operation=operation, + ): + assert progress.stage_name == "stage12_child_wait" + assert progress.stage_role == "reform" + assert not progress.stage_finished + + assert runtime.scopes == [(scope_kind, "stage12_child_wait", attributes)] + assert progress.stage_finished + + +def test_tracked_stage_without_a_runtime_still_records_the_failing_stage() -> None: + progress = StageProgress(role="baseline") + msg = "wait failed" + + with ( + pytest.raises(RuntimeError, match=msg), + tracked_stage( + Stage.STAGE12_CHILD_WAIT, + stage_plan=STAGE12_CANONICAL_REPORT_STAGES, + progress=progress, + runtime=None, + ), + ): + raise RuntimeError(msg) + + assert progress.stage_name == "stage12_child_wait" + assert progress.stage_role == "baseline" + assert not progress.stage_finished From ad59907b00943c63d6a685a43f358a25887ef818 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:35:28 +0400 Subject: [PATCH 65/66] Pass the tracked_stage #791 added its stage plan, progress and runtime as one group #791 landed on main after this stack branched. PLR0913, which this PR enables, flags its six-argument tracked_stage. Each process's stage plan, progress and runtime now travel as one StageTracking: the worker builds one per run, and the coordinator derives its own from the report it coordinates. The previous commit's pins keep their assertions. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 21 +++++-------- .../stage12_runtime/failure_causes.py | 30 ++++++++++++++----- .../stage12_runtime/simulation.py | 9 ++---- .../tests/test_stage12_failure_causes.py | 9 ++---- 4 files changed, 37 insertions(+), 32 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 296a76382..30d97af37 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -55,6 +55,7 @@ ) from .failure_causes import ( StageProgress, + StageTracking, failure_summary, record_failure_cause, tracked_stage, @@ -350,19 +351,18 @@ def runtime(self) -> ObservabilityRuntime | None: """The telemetry runtime, or None when the coordinator records nothing.""" return self.dependencies.runtime + @property + def stage_tracking(self) -> StageTracking: + """How this report's stages are named, recorded and timed.""" + return StageTracking(self.stage_plan, self.progress, self.runtime) + def span( self, stage: Stage, attributes: Mapping[str, object] | None = None, ) -> AbstractContextManager[object]: """Track ``stage`` for failure causes and time it as a span.""" - return tracked_stage( - stage, - attributes, - stage_plan=self.stage_plan, - progress=self.progress, - runtime=self.runtime, - ) + return tracked_stage(stage, attributes, tracking=self.stage_tracking) def stage_scope( self, @@ -371,12 +371,7 @@ def stage_scope( ) -> AbstractContextManager[object]: """Track ``stage`` for failure causes and record it as an operation.""" return tracked_stage( - stage, - attributes, - stage_plan=self.stage_plan, - progress=self.progress, - runtime=self.runtime, - operation=True, + stage, attributes, tracking=self.stage_tracking, operation=True ) def run(self, parent: ComparisonReportRecord) -> dict[str, Any]: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py index 4147c439e..c636653f3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py @@ -30,6 +30,7 @@ import os import re from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass from typing import TYPE_CHECKING from urllib.parse import urlsplit @@ -152,23 +153,35 @@ def stage(self, stage_name: str, *, role: str | None = None) -> Iterator[None]: self.stage_finished = True +@dataclass(frozen=True) +class StageTracking: + """How one process names, records and times its Stage 12 stages. + + ``stage_plan`` names each stage, ``progress`` records the last one so that + a failure names it, and ``runtime``, when set, times it. + """ + + stage_plan: StagePlan + progress: StageProgress + runtime: ObservabilityRuntime | None + + @contextmanager def tracked_stage( stage: Stage, attributes: Mapping[str, object] | None = None, *, - stage_plan: StagePlan, - progress: StageProgress, - runtime: ObservabilityRuntime | None, + tracking: StageTracking, operation: bool = False, ) -> Iterator[None]: """Run a block as one Stage 12 stage for telemetry and failure causes. The block runs inside the stage's span, or its operation when - ``operation`` is true, and ``progress`` records the stage so that a - failure names it. A ``simulation_role`` attribute sets the stage's role. + ``operation`` is true, and ``tracking.progress`` records the stage so that + a failure names it. A ``simulation_role`` attribute sets the stage's role. """ - stage_name = stage_plan.name(stage) + stage_name = tracking.stage_plan.name(stage) + runtime = tracking.runtime scope: AbstractContextManager[object] if runtime is None: scope = nullcontext() @@ -176,7 +189,10 @@ def tracked_stage( scope = runtime.operation(stage_name, attributes=attributes) else: scope = runtime.span(stage_name, attributes=attributes) - with progress.stage(stage_name, role=_simulation_role(attributes)), scope: + with ( + tracking.progress.stage(stage_name, role=_simulation_role(attributes)), + scope, + ): yield diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index d494aea6d..fc76177f5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -39,6 +39,7 @@ from .dependencies import ComparisonStore, artifact_store, runtime_store from .failure_causes import ( StageProgress, + StageTracking, failure_summary, record_failure_cause, tracked_stage, @@ -137,9 +138,7 @@ def calculate_simulation_frames( progress = StageProgress(role=simulation.role.value) stage_span = partial( tracked_stage, - stage_plan=STAGE12_SIMULATION_STAGES, - progress=progress, - runtime=runtime, + tracking=StageTracking(STAGE12_SIMULATION_STAGES, progress, runtime), ) if simulation.population.kind != "dataset": msg = "Stage 12 society-wide worker requires a dataset population" @@ -338,9 +337,7 @@ def run_single_simulation( progress = StageProgress(role=simulation.role.value) stage_span = partial( tracked_stage, - stage_plan=STAGE12_SIMULATION_STAGES, - progress=progress, - runtime=runtime, + tracking=StageTracking(STAGE12_SIMULATION_STAGES, progress, runtime), ) try: with stage_span(Stage.STAGE12_INPUT_WRITE): diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py index c2b0b06f7..f9b0058b7 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py @@ -14,6 +14,7 @@ MAX_REDACTED_MESSAGE_LENGTH, Stage12FailureCause, StageProgress, + StageTracking, failure_summary, record_failure_cause, redact_failure_message, @@ -508,9 +509,7 @@ def test_tracked_stage_scopes_and_records_the_stage_with_its_role( with tracked_stage( Stage.STAGE12_CHILD_WAIT, attributes, - stage_plan=STAGE12_CANONICAL_REPORT_STAGES, - progress=progress, - runtime=runtime, + tracking=StageTracking(STAGE12_CANONICAL_REPORT_STAGES, progress, runtime), operation=operation, ): assert progress.stage_name == "stage12_child_wait" @@ -529,9 +528,7 @@ def test_tracked_stage_without_a_runtime_still_records_the_failing_stage() -> No pytest.raises(RuntimeError, match=msg), tracked_stage( Stage.STAGE12_CHILD_WAIT, - stage_plan=STAGE12_CANONICAL_REPORT_STAGES, - progress=progress, - runtime=None, + tracking=StageTracking(STAGE12_CANONICAL_REPORT_STAGES, progress, None), ), ): raise RuntimeError(msg) From 4734c934e6ee6351b0004c6b9d9e31e650a4b798 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:35:43 +0400 Subject: [PATCH 66/66] Compare the word count in the redaction helper #791 added without a magic 2 #791 landed on main after this stack branched. PLR2004, which this PR enables, flags the len(words) == 2 test in its _before_last_word. Splitting once from the right yields at most two parts, so len(words) > 1 is the same test. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/failure_causes.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py index c636653f3..27aac2e99 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py @@ -312,8 +312,9 @@ def _without_url_query(url_match: re.Match[str]) -> str: def _before_last_word(text: str) -> str: + # rsplit with maxsplit=1 returns at most two parts. words = text.rsplit(maxsplit=1) - return words[0] if len(words) == 2 else "" + return words[0] if len(words) > 1 else "" def _secret_environment_values() -> tuple[str, ...]: