From bae99c581a1f32cc2534608a2aab2056ece84aca Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 28 Sep 2026 21:12:06 +0400 Subject: [PATCH 1/7] Fix Stage 12 comparison and UK geography metadata --- .../simulation_output_geographic.py | 95 +++++++++++++--- .../stage12_result_comparison.py | 12 +- .../tests/test_simulation_output_builder.py | 30 ++++- .../test_simulation_output_geographic.py | 107 ++++++++++++++++++ .../tests/test_stage12_result_comparison.py | 31 +++++ 5 files changed, 249 insertions(+), 26 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index 20503919e..0a882ca16 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -4,7 +4,12 @@ from collections.abc import Mapping from functools import lru_cache -from typing import Any +from typing import TYPE_CHECKING, Any + +from policyengine.data.uk_geography_assets import ( + CONSTITUENCY_ASSET_SPEC, + LOCAL_AUTHORITY_ASSET_SPEC, +) from policyengine_simulation_executor.simulation_macro_output import ( CongressionalDistrictImpactOutput, @@ -18,6 +23,56 @@ _try_compute_output, ) +if TYPE_CHECKING: + from policyengine.data.uk_geography_assets import UKGeographyAssetSpec + + +def _required_uk_geography_lookup_csv_path(spec: UKGeographyAssetSpec) -> str: + """Resolve one required lookup under the runtime GCP identity.""" + + from policyengine.outputs.uk_geography_impact import ( + resolve_uk_geography_lookup_csv_path, + ) + from policyengine_simulation_executor.simulation_runtime import ( + setup_gcp_credentials, + ) + + with setup_gcp_credentials(): + path = resolve_uk_geography_lookup_csv_path( + spec, + download_missing_assets=True, + ) + if path is None: + raise FileNotFoundError( + f"Required UK {spec.geography_type} lookup CSV " + f"{spec.lookup_csv_filename!r} could not be resolved" + ) + return path + + +def _complete_uk_geography_output( + value: object, + *, + code_field: str, + name_field: str, +) -> GeographicImpactOutput: + """Require every UK geography record to contain lookup-owned metadata.""" + + output = build_geographic_impact_output(value) + if output is None: + raise ValueError("UK geography output did not contain result records") + for record in output.root: + record_values = record.model_dump(mode="python") + code = record_values.get(code_field) + name = record_values.get(name_field) + if not isinstance(code, str) or not code.strip(): + raise ValueError(f"UK geography output is missing {code_field}") + if not isinstance(name, str) or not name.strip() or name == code: + raise ValueError(f"UK geography {code!r} is missing a lookup name") + if record_values.get("x") is None or record_values.get("y") is None: + raise ValueError(f"UK geography {code!r} is missing lookup coordinates") + return output + @lru_cache(maxsize=1) def _policyengine_us_district_metadata() -> tuple[dict[int, str], frozenset[str]]: @@ -160,17 +215,20 @@ def build_uk_constituency_impact( if country != "uk": return None - impact = _try_compute_output( - "constituency impacts", - _output_module_function( - "constituency_impact", "compute_uk_constituency_impacts" - ), + lookup_csv_path = _required_uk_geography_lookup_csv_path(CONSTITUENCY_ASSET_SPEC) + impact = _output_module_function( + "constituency_impact", "compute_uk_constituency_impacts" + )( baseline, reform, + constituency_csv_path=lookup_csv_path, + download_missing_assets=False, + ) + return _complete_uk_geography_output( + getattr(impact, "constituency_results", None), + code_field="constituency_code", + name_field="constituency_name", ) - if impact is None: - return None - return build_geographic_impact_output(getattr(impact, "constituency_results", None)) def build_uk_local_authority_impact( @@ -179,16 +237,17 @@ def build_uk_local_authority_impact( if country != "uk": return None - impact = _try_compute_output( - "local authority impacts", - _output_module_function( - "local_authority_impact", "compute_uk_local_authority_impacts" - ), + lookup_csv_path = _required_uk_geography_lookup_csv_path(LOCAL_AUTHORITY_ASSET_SPEC) + impact = _output_module_function( + "local_authority_impact", "compute_uk_local_authority_impacts" + )( baseline, reform, + local_authority_csv_path=lookup_csv_path, + download_missing_assets=False, ) - if impact is None: - return None - return build_geographic_impact_output( - getattr(impact, "local_authority_results", None) + return _complete_uk_geography_output( + getattr(impact, "local_authority_results", None), + code_field="local_authority_code", + name_field="local_authority_name", ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py index 9b256bb73..51d6a2a8b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py @@ -1,4 +1,4 @@ -"""Exact, private comparison of production and Stage 12 aggregate results.""" +"""Private comparison of production and Stage 12 aggregate results.""" from __future__ import annotations @@ -17,6 +17,7 @@ type ResultObject = dict[str, JsonValue] _result_adapter = TypeAdapter(ResultObject) +NUMERIC_ABSOLUTE_TOLERANCE = 0.01 class _MissingValue: @@ -56,6 +57,13 @@ def _leaf_values(value: JsonValue, path: str = "") -> dict[str, JsonValue]: def _equal_json_values(first: JsonValue, second: JsonValue) -> bool: + if ( + not isinstance(first, bool) + and not isinstance(second, bool) + and isinstance(first, (int, float)) + and isinstance(second, (int, float)) + ): + return abs(float(second) - float(first)) <= NUMERIC_ABSOLUTE_TOLERANCE return canonical_json_bytes(first) == canonical_json_bytes(second) @@ -84,7 +92,7 @@ def compare_results( stage12_result: object, compared_at: datetime | None = None, ) -> ResultComparisonArtifactPayload: - """Return an exact digest and every differing scalar result leaf.""" + """Return exact digests and scalar differences outside the numeric tolerance.""" production = _result_adapter.validate_python(production_result) stage12 = _result_adapter.validate_python(stage12_result) 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 95107cecb..3dc5638d2 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -212,8 +212,8 @@ def _constituency_impact_record() -> ConstituencyImpactRecord: return ConstituencyImpactRecord( constituency_code="E14000530", constituency_name="Birmingham, Ladywood", - x=None, - y=None, + x=34, + y=-42, average_household_income_change=10.0, relative_household_income_change=0.01, population=1000.0, @@ -224,8 +224,8 @@ def _local_authority_impact_record() -> LocalAuthorityImpactRecord: return LocalAuthorityImpactRecord( local_authority_code="E06000001", local_authority_name="Hartlepool", - x=None, - y=None, + x=63, + y=-19, average_household_income_change=12.0, relative_household_income_change=0.02, population=900.0, @@ -1650,9 +1650,13 @@ def fake_output_module_function(module_name, name): assert module_name == "constituency_impact" assert name == "compute_uk_constituency_impacts" - def compute(baseline_simulation, reform_simulation): + def compute(baseline_simulation, reform_simulation, **kwargs): assert baseline_simulation is baseline assert reform_simulation is reform + assert kwargs == { + "constituency_csv_path": "/tmp/constituencies_2024.csv", + "download_missing_assets": False, + } return SimpleNamespace(constituency_results=expected) return compute @@ -1661,6 +1665,11 @@ def compute(baseline_simulation, reform_simulation): "policyengine_simulation_executor.simulation_output_geographic._output_module_function", fake_output_module_function, ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "_required_uk_geography_lookup_csv_path", + lambda spec: "/tmp/constituencies_2024.csv", + ) result = _simulation_output_builder( "uk", baseline, reform @@ -1684,9 +1693,13 @@ def fake_output_module_function(module_name, name): assert module_name == "local_authority_impact" assert name == "compute_uk_local_authority_impacts" - def compute(baseline_simulation, reform_simulation): + def compute(baseline_simulation, reform_simulation, **kwargs): assert baseline_simulation is baseline assert reform_simulation is reform + assert kwargs == { + "local_authority_csv_path": "/tmp/local_authorities_2021.csv", + "download_missing_assets": False, + } return SimpleNamespace(local_authority_results=expected) return compute @@ -1695,6 +1708,11 @@ def compute(baseline_simulation, reform_simulation): "policyengine_simulation_executor.simulation_output_geographic._output_module_function", fake_output_module_function, ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "_required_uk_geography_lookup_csv_path", + lambda spec: "/tmp/local_authorities_2021.csv", + ) result = _simulation_output_builder( "uk", baseline, reform diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py new file mode 100644 index 000000000..8fe1ea009 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -0,0 +1,107 @@ +"""Unit tests for required UK geography metadata handling.""" + +from __future__ import annotations + +from contextlib import nullcontext + +import pytest +from policyengine.data.uk_geography_assets import CONSTITUENCY_ASSET_SPEC + +from policyengine_simulation_executor import simulation_output_geographic + + +def _complete_constituency_record() -> dict[str, object]: + return { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "x": 56, + "y": -40, + "average_household_income_change": -150.95, + "relative_household_income_change": -0.0026, + "population": 40_530.65, + } + + +def test_required_uk_lookup_uses_gcp_credentials(monkeypatch) -> None: + observed: list[str] = [] + + def credentials(): + observed.append("credentials-entered") + return nullcontext() + + def resolve(spec, **kwargs): + assert spec is CONSTITUENCY_ASSET_SPEC + assert kwargs == {"download_missing_assets": True} + observed.append("lookup-resolved") + return "/tmp/constituencies_2024.csv" + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.setup_gcp_credentials", + credentials, + ) + monkeypatch.setattr( + "policyengine.outputs.uk_geography_impact." + "resolve_uk_geography_lookup_csv_path", + resolve, + ) + + result = simulation_output_geographic._required_uk_geography_lookup_csv_path( + CONSTITUENCY_ASSET_SPEC + ) + + assert result == "/tmp/constituencies_2024.csv" + assert observed == ["credentials-entered", "lookup-resolved"] + + +def test_required_uk_lookup_rejects_missing_asset(monkeypatch) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.setup_gcp_credentials", + nullcontext, + ) + monkeypatch.setattr( + "policyengine.outputs.uk_geography_impact." + "resolve_uk_geography_lookup_csv_path", + lambda *args, **kwargs: None, + ) + + with pytest.raises(FileNotFoundError, match="constituency lookup CSV"): + simulation_output_geographic._required_uk_geography_lookup_csv_path( + CONSTITUENCY_ASSET_SPEC + ) + + +def test_complete_uk_geography_output_accepts_lookup_metadata() -> None: + result = simulation_output_geographic._complete_uk_geography_output( + [_complete_constituency_record()], + code_field="constituency_code", + name_field="constituency_name", + ) + + assert result.root[0].constituency_name == "Aldershot" + assert result.root[0].x == 56 + assert result.root[0].y == -40 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("constituency_name", "E14001063", "missing a lookup name"), + ("constituency_name", "", "missing a lookup name"), + ("x", None, "missing lookup coordinates"), + ("y", None, "missing lookup coordinates"), + ], +) +def test_complete_uk_geography_output_rejects_fallback_metadata( + field: str, + value: object, + message: str, +) -> None: + record = _complete_constituency_record() + record[field] = value + + with pytest.raises(ValueError, match=message): + simulation_output_geographic._complete_uk_geography_output( + [record], + code_field="constituency_code", + name_field="constituency_name", + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py b/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py index 3c8aed9bd..ee1e6707e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py @@ -3,6 +3,8 @@ from datetime import UTC, datetime from uuid import UUID +import pytest + from policyengine_simulation_executor.stage12_result_comparison import compare_results RUN_ID = UUID("00000000-0000-0000-0000-000000000001") @@ -59,3 +61,32 @@ def test_numeric_zero_has_an_absolute_but_no_relative_delta() -> None: difference = receipt.differences[0] assert difference.absolute_delta == 5.0 assert difference.relative_delta is None + + +def test_numeric_values_within_one_cent_match() -> None: + receipt = compare_results( + evaluation_id=RUN_ID, + production_job_id="production-job-1", + production_result={"value": 10}, + stage12_result={"value": 10.01}, + compared_at=COMPARED_AT, + ) + + assert receipt.status == "matched" + assert receipt.difference_count == 0 + assert receipt.differences == () + assert receipt.production_result_sha256 != receipt.stage12_result_sha256 + + +def test_numeric_values_above_one_cent_differ() -> None: + receipt = compare_results( + evaluation_id=RUN_ID, + production_job_id="production-job-1", + production_result={"value": 10}, + stage12_result={"value": 10.011}, + compared_at=COMPARED_AT, + ) + + assert receipt.status == "different" + assert receipt.difference_count == 1 + assert receipt.differences[0].absolute_delta == pytest.approx(0.011) From b9b5c928ec5081508a761a550c45a9947e4a62a2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 28 Sep 2026 21:13:29 +0400 Subject: [PATCH 2/7] Tighten geographic output validation typing --- .../simulation_output_geographic.py | 4 ++-- .../tests/test_simulation_output_geographic.py | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index 0a882ca16..1b38db7e9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -92,11 +92,11 @@ def build_geographic_impact_output(value: Any) -> GeographicImpactOutput | None: return value records = _output_model_dump(value) if isinstance(records, list): - return GeographicImpactOutput( + return GeographicImpactOutput.model_validate( [dict(item) for item in records if isinstance(item, Mapping)] ) if isinstance(value, list): - return GeographicImpactOutput( + return GeographicImpactOutput.model_validate( [dict(item) for item in value if isinstance(item, Mapping)] ) return None diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py index 8fe1ea009..03b036e45 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -77,9 +77,10 @@ def test_complete_uk_geography_output_accepts_lookup_metadata() -> None: name_field="constituency_name", ) - assert result.root[0].constituency_name == "Aldershot" - assert result.root[0].x == 56 - assert result.root[0].y == -40 + record = result.root[0].model_dump(mode="python") + assert record["constituency_name"] == "Aldershot" + assert record["x"] == 56 + assert record["y"] == -40 @pytest.mark.parametrize( From 0db843cc4ec1f6dd8ee47fcdfd0389206176a08e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 28 Sep 2026 21:43:54 +0400 Subject: [PATCH 3/7] Verify Stage 12 UK lookup access --- .../scripts/stage12-validate-infrastructure.sh | 15 ++++++++++++++- .../migration/stage-12-modal-worker-foundation.md | 7 ++++++- .../tests/test_stage12_infrastructure.py | 10 ++++++++-- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/.github/scripts/stage12-validate-infrastructure.sh b/.github/scripts/stage12-validate-infrastructure.sh index 64b309172..dc63c1337 100755 --- a/.github/scripts/stage12-validate-infrastructure.sh +++ b/.github/scripts/stage12-validate-infrastructure.sh @@ -37,6 +37,10 @@ for account in \ exit 1 fi done +uk_geography_lookup_objects=( + "gs://policyengine-uk-data-private/constituencies_2024.csv" + "gs://policyengine-uk-data-private/local_authorities_2021.csv" +) for secret_name in \ "${STAGE12_MODAL_TOKEN_ID_SECRET_NAME}" \ "${STAGE12_MODAL_TOKEN_SECRET_SECRET_NAME}" \ @@ -161,4 +165,13 @@ runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ "${canary_object}" --quiet canary_object="" -echo "Pre-provisioned Stage 12 database, secret, and storage access is verified." +for lookup_object in "${uk_geography_lookup_objects[@]}"; do + runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ + "${lookup_object}" "${canary_download_file}" --quiet + if [[ ! -s "${canary_download_file}" ]]; then + echo "Stage 12 worker downloaded an empty UK geography lookup: ${lookup_object}" >&2 + exit 1 + fi +done + +echo "Pre-provisioned Stage 12 database, secret, artifact storage, and UK geography lookup access is verified." diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 4e3a47231..fff728157 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -60,7 +60,12 @@ with the live API-owned schema. its selected applications unchanged. - V2 worker object access uses the already-provisioned separate `stage12-evaluation-gcp-credentials` Modal secret. The worker application - does not receive the existing general GCP credential secret. + does not receive the existing general GCP credential secret. The credential's + service account can read and write its environment-specific Stage 12 artifact + bucket. A conditional IAM binding separately permits it to read only + `constituencies_2024.csv` and `local_authorities_2021.csv` from + `policyengine-uk-data-private`; deployment validation downloads both lookup + files using the worker credential. - `STAGE12_DATABASE_URL` is delivered from the environment-specific Secret Manager resource named by `STAGE12_DATABASE_URL_SECRET_NAME`. It authenticates as the existing shared `policyengine_v2_runtime` account; Stage 12 does not diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py index f29177582..356603594 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py @@ -29,7 +29,10 @@ def test_infrastructure_validation_is_bounded_and_valid_shell() -> None: source = SCRIPT.read_text(encoding="utf-8") assert "STAGE12_ENVIRONMENT" in source assert "STAGE12_ARTIFACT_BUCKET" in source - assert "Pre-provisioned Stage 12 database, secret, and storage access" in source + assert ( + "Pre-provisioned Stage 12 database, secret, artifact storage, and UK " + "geography lookup access" in source + ) assert "STAGE12_DATABASE_URL_SECRET_NAME" in source assert "policyengine_v2_runtime" in source assert r"${expected_role}\.[a-z0-9]{20}" in source @@ -87,8 +90,11 @@ def test_storage_validation_uses_only_the_modal_worker_credentials() -> None: assert "runtime_gcloud auth activate-service-account" in source modal_account = '--account="${STAGE12_MODAL_SERVICE_ACCOUNT}"' - assert source.count(f"runtime_gcloud {modal_account} storage cp") == 2 + assert source.count(f"runtime_gcloud {modal_account} storage cp") == 3 assert source.count(f"runtime_gcloud {modal_account} storage rm") == 2 + assert "gs://policyengine-uk-data-private/constituencies_2024.csv" in source + assert "gs://policyengine-uk-data-private/local_authorities_2021.csv" in source + assert "STAGE12_UK_GEOGRAPHY_BUCKET" not in source def test_modal_secret_sync_rejects_missing_configuration() -> None: From 7729e52eafece363fba31357c5c2f09b34c25930 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:07:20 +0400 Subject: [PATCH 4/7] Reject empty Stage 12 geography results --- .../simulation_output_geographic.py | 2 +- .../tests/test_simulation_output_geographic.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py index 1b38db7e9..4e8abce55 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py @@ -59,7 +59,7 @@ def _complete_uk_geography_output( """Require every UK geography record to contain lookup-owned metadata.""" output = build_geographic_impact_output(value) - if output is None: + if output is None or not output.root: raise ValueError("UK geography output did not contain result records") for record in output.root: record_values = record.model_dump(mode="python") diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py index 03b036e45..ec9f30f75 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_geographic.py @@ -83,6 +83,15 @@ def test_complete_uk_geography_output_accepts_lookup_metadata() -> None: assert record["y"] == -40 +def test_complete_uk_geography_output_rejects_empty_result() -> None: + with pytest.raises(ValueError, match="did not contain result records"): + simulation_output_geographic._complete_uk_geography_output( + [], + code_field="constituency_code", + name_field="constituency_name", + ) + + @pytest.mark.parametrize( ("field", "value", "message"), [ From 7b6567f99e4d485a8f75cb4c8b9bf7f820fd03e9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:38:12 +0400 Subject: [PATCH 5/7] Compare Stage 12 results with decimal tolerance --- .../stage12_result_comparison.py | 14 ++++++++++++-- .../tests/test_stage12_result_comparison.py | 13 +++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py index 51d6a2a8b..335943bb0 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_result_comparison.py @@ -3,6 +3,7 @@ from __future__ import annotations from datetime import UTC, datetime +from decimal import Decimal from hashlib import sha256 from typing import TypeGuard from uuid import UUID @@ -17,7 +18,7 @@ type ResultObject = dict[str, JsonValue] _result_adapter = TypeAdapter(ResultObject) -NUMERIC_ABSOLUTE_TOLERANCE = 0.01 +NUMERIC_ABSOLUTE_TOLERANCE = Decimal("0.01") class _MissingValue: @@ -63,7 +64,16 @@ def _equal_json_values(first: JsonValue, second: JsonValue) -> bool: and isinstance(first, (int, float)) and isinstance(second, (int, float)) ): - return abs(float(second) - float(first)) <= NUMERIC_ABSOLUTE_TOLERANCE + # JSON decimal literals must be compared as decimal values. Converting + # them to binary floats first can make an exact one-cent difference, + # such as 1.11 - 1.10, appear slightly greater than the tolerance. + first_decimal = ( + Decimal(first) if isinstance(first, int) else Decimal(str(first)) + ) + second_decimal = ( + Decimal(second) if isinstance(second, int) else Decimal(str(second)) + ) + return abs(second_decimal - first_decimal) <= NUMERIC_ABSOLUTE_TOLERANCE return canonical_json_bytes(first) == canonical_json_bytes(second) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py b/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py index ee1e6707e..643bd3ba3 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_result_comparison.py @@ -78,6 +78,19 @@ def test_numeric_values_within_one_cent_match() -> None: assert receipt.production_result_sha256 != receipt.stage12_result_sha256 +def test_exact_decimal_tolerance_boundary_matches() -> None: + receipt = compare_results( + evaluation_id=RUN_ID, + production_job_id="production-job-1", + production_result={"value": 1.1}, + stage12_result={"value": 1.11}, + compared_at=COMPARED_AT, + ) + + assert receipt.status == "matched" + assert receipt.difference_count == 0 + + def test_numeric_values_above_one_cent_differ() -> None: receipt = compare_results( evaluation_id=RUN_ID, From 284624f5d56a440b0b3af112e43d116b3da32cc0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:47:28 +0400 Subject: [PATCH 6/7] Validate UK Hugging Face access before PR tests --- .github/scripts/modal-image-smoke.sh | 7 ++ .github/workflows/pr.yml | 27 +++++ .../hf_dataset.py | 44 ++++++-- .../tests/test_hf_dataset.py | 51 +++++++-- .../src/modal/smoke_app.py | 21 +++- .../uk_hf_access.py | 105 ++++++++++++++++++ .../tests/test_uk_hf_access.py | 61 ++++++++++ 7 files changed, 294 insertions(+), 22 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py create mode 100644 projects/policyengine-simulation-executor/tests/test_uk_hf_access.py diff --git a/.github/scripts/modal-image-smoke.sh b/.github/scripts/modal-image-smoke.sh index e259e489b..78cf5166f 100755 --- a/.github/scripts/modal-image-smoke.sh +++ b/.github/scripts/modal-image-smoke.sh @@ -11,6 +11,13 @@ set -euo pipefail MODAL_ENV="${1:?Modal environment required}" REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +echo "=== UK Hugging Face dataset access (env: $MODAL_ENV) ===" +( + cd "$REPO_ROOT/projects/policyengine-simulation-executor" + uv run modal run --env="$MODAL_ENV" \ + src/modal/smoke_app.py --uk-hf-access-only +) + echo "=== Gateway image smoke (env: $MODAL_ENV) ===" ( cd "$REPO_ROOT/projects/policyengine-simulation-gateway" diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index f4b4916f1..527e9870a 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -9,8 +9,34 @@ concurrency: cancel-in-progress: true jobs: + validate-uk-hf-access: + name: Validate UK Hugging Face dataset access + runs-on: ubuntu-latest + if: github.event.pull_request.head.repo.full_name == github.repository + + steps: + - uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.13' + + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + with: + enable-cache: true + + - name: Validate every UK bundle dataset with the staging credential + working-directory: projects/policyengine-simulation-executor + env: + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + run: uv run modal run --env=staging src/modal/smoke_app.py --uk-hf-access-only + test: name: Test + needs: validate-uk-hf-access runs-on: ubuntu-latest strategy: matrix: @@ -50,6 +76,7 @@ jobs: lint: name: Lint + needs: validate-uk-hf-access runs-on: ubuntu-latest steps: 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 407803240..c85862ba5 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 @@ -13,7 +13,7 @@ import os from dataclasses import dataclass from functools import lru_cache -from typing import Any +from typing import Any, Literal from urllib.error import HTTPError, URLError from urllib.parse import quote from urllib.request import Request, urlopen @@ -26,6 +26,7 @@ "HUGGINGFACE_HUB_TOKEN", "HUGGINGFACE_TOKEN", ) +HFRepositoryType = Literal["dataset", "model"] class HuggingFaceDatasetReferenceError(ValueError): @@ -77,13 +78,15 @@ def parse_hf_dataset_uri(dataset_uri: str) -> HFDatasetReference | None: @lru_cache -def _fetch_hf_dataset_revision( +def _fetch_hf_repository_revision( repo_id: str, revision: str, + repo_type: HFRepositoryType, token: str | None, ) -> dict[str, Any]: + resource_name = "datasets" if repo_type == "dataset" else "models" url = ( - f"{HF_ENDPOINT}/api/datasets/" + f"{HF_ENDPOINT}/api/{resource_name}/" f"{quote(repo_id, safe='/')}/revision/{quote(revision, safe='')}" ) headers = {"Accept": "application/json"} @@ -97,12 +100,13 @@ def _fetch_hf_dataset_revision( except HTTPError as exc: detail = exc.reason or f"HTTP {exc.code}" raise HuggingFaceDatasetReferenceError( - f"Hugging Face dataset revision {repo_id}@{revision} was not found: " + f"Hugging Face {repo_type} repository revision " + f"{repo_id}@{revision} was not found: " f"{detail}" ) from exc except (OSError, URLError, json.JSONDecodeError) as exc: raise HuggingFaceDatasetReferenceError( - f"Unable to validate Hugging Face dataset revision " + f"Unable to validate Hugging Face {repo_type} repository revision " f"{repo_id}@{revision}: {exc}" ) from exc @@ -124,23 +128,43 @@ def _siblings_contain_path(payload: dict[str, Any], path: str) -> bool | None: return False if seen_file_listing else None -def validate_hf_dataset_uri(dataset_uri: str) -> str: - """Validate an explicit ``hf://`` dataset URI if it pins a revision.""" +def validate_hf_artifact_uri( + dataset_uri: str, + repo_type: HFRepositoryType, +) -> str: + """Validate a pinned artifact against its declared Hub repository type.""" parsed = parse_hf_dataset_uri(dataset_uri) if parsed is None or parsed.revision is None: return dataset_uri - payload = _fetch_hf_dataset_revision(parsed.repo_id, parsed.revision, _hf_token()) + payload = _fetch_hf_repository_revision( + parsed.repo_id, + parsed.revision, + repo_type, + _hf_token(), + ) contains_path = _siblings_contain_path(payload, parsed.path) + if contains_path is None: + raise HuggingFaceDatasetReferenceError( + f"Hugging Face {repo_type} repository revision " + f"{parsed.repo_id}@{parsed.revision} did not provide an artifact listing" + ) if contains_path is False: raise HuggingFaceDatasetReferenceError( - f"Hugging Face dataset revision {parsed.repo_id}@{parsed.revision} " - f"does not contain artifact {parsed.path!r}" + f"Hugging Face {repo_type} repository revision " + f"{parsed.repo_id}@{parsed.revision} does not contain artifact " + f"{parsed.path!r}" ) return dataset_uri +def validate_hf_dataset_uri(dataset_uri: str) -> str: + """Validate a pinned URI stored in a Hugging Face dataset repository.""" + + return validate_hf_artifact_uri(dataset_uri, "dataset") + + def with_hf_revision(dataset_uri: str, revision: str) -> str: """Return ``dataset_uri`` pinned to ``revision`` and validate it on the Hub.""" diff --git a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py index 95d2f5e3b..e5f2d4464 100644 --- a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py +++ b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py @@ -10,6 +10,7 @@ from policyengine_simulation_contract.hf_dataset import ( HuggingFaceDatasetReferenceError, parse_hf_dataset_uri, + validate_hf_artifact_uri, validate_hf_dataset_uri, with_hf_revision, ) @@ -40,8 +41,8 @@ def test_parse_hf_dataset_uri_extracts_repo_path_and_revision(): assert parsed.revision == "custom-v1" -def test_fetch_hf_dataset_revision_uses_dataset_revision_api(monkeypatch): - hf_dataset._fetch_hf_dataset_revision.cache_clear() +def test_fetch_hf_repository_revision_uses_declared_repository_type(monkeypatch): + hf_dataset._fetch_hf_repository_revision.cache_clear() seen = {} def fake_urlopen(request, timeout): @@ -52,16 +53,16 @@ def fake_urlopen(request, timeout): monkeypatch.setattr(hf_dataset, "urlopen", fake_urlopen) - payload = hf_dataset._fetch_hf_dataset_revision( + payload = hf_dataset._fetch_hf_repository_revision( "policyengine/populace-us", "custom-v1", + "model", "hf-token", ) assert payload == {"sha": "abc123", "siblings": []} assert seen["url"] == ( - "https://huggingface.co/api/datasets/" - "policyengine/populace-us/revision/custom-v1" + "https://huggingface.co/api/models/policyengine/populace-us/revision/custom-v1" ) assert seen["headers"]["Authorization"] == "Bearer hf-token" assert seen["timeout"] == hf_dataset.HF_REQUEST_TIMEOUT_SECONDS @@ -70,8 +71,10 @@ def fake_urlopen(request, timeout): def test_validate_hf_dataset_uri_rejects_revision_missing_artifact(monkeypatch): monkeypatch.setattr( hf_dataset, - "_fetch_hf_dataset_revision", - lambda repo_id, revision, token: {"siblings": [{"rfilename": "other_file.h5"}]}, + "_fetch_hf_repository_revision", + lambda repo_id, revision, repo_type, token: { + "siblings": [{"rfilename": "other_file.h5"}] + }, ) with pytest.raises( @@ -83,6 +86,36 @@ def test_validate_hf_dataset_uri_rejects_revision_missing_artifact(monkeypatch): ) +def test_validate_hf_dataset_uri_rejects_missing_artifact_listing(monkeypatch): + monkeypatch.setattr( + hf_dataset, + "_fetch_hf_repository_revision", + lambda repo_id, revision, repo_type, token: {"sha": "abc123"}, + ) + + with pytest.raises( + HuggingFaceDatasetReferenceError, + match="did not provide an artifact listing", + ): + validate_hf_dataset_uri( + "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" + ) + + +def test_validate_hf_artifact_uri_uses_model_repository_api(monkeypatch): + seen_repository_types = [] + + def fake_fetch(repo_id, revision, repo_type, token): + seen_repository_types.append(repo_type) + return {"siblings": [{"rfilename": "enhanced_frs.h5"}]} + + monkeypatch.setattr(hf_dataset, "_fetch_hf_repository_revision", fake_fetch) + + uri = "hf://policyengine/private-uk/enhanced_frs.h5@release-v1" + assert validate_hf_artifact_uri(uri, "model") == uri + assert seen_repository_types == ["model"] + + def test_with_hf_revision_validates_and_preserves_requested_revision(monkeypatch): calls = [] @@ -99,6 +132,4 @@ def fake_validate(dataset_uri): ) == "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" ) - assert calls == [ - "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" - ] + assert calls == ["hf://policyengine/populace-us/populace_us_2024.h5@custom-v1"] diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index 164a4971f..b052b9c08 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -16,7 +16,7 @@ import modal -from src.modal.app import build_runtime_simulation_image +from src.modal.app import build_runtime_simulation_image, hf_secret app = modal.App("policyengine-simulation-executor-smoke") @@ -29,6 +29,17 @@ ) +@app.function(image=smoke_image, secrets=[hf_secret], timeout=120) +def smoke_uk_hf_dataset_access() -> str: + """Exercise the deployed HF_TOKEN against every bundle-declared UK artifact.""" + + from policyengine_simulation_executor.uk_hf_access import ( + validate_uk_hf_dataset_access, + ) + + return validate_uk_hf_dataset_access().model_dump_json() + + @app.function(image=smoke_image, timeout=600, memory=8192) def smoke_import_executor() -> dict: import importlib @@ -71,7 +82,13 @@ def smoke_import_executor() -> dict: @app.local_entrypoint() -def main(): +def main(uk_hf_access_only: bool = False): + access_report = smoke_uk_hf_dataset_access.remote() + print(access_report) + print("UK Hugging Face dataset access OK") + if uk_hf_access_only: + return + report = smoke_import_executor.remote() print(report) print("executor image smoke OK") diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py new file mode 100644 index 000000000..b9958ee75 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py @@ -0,0 +1,105 @@ +"""Live access validation for every UK dataset in the release bundle.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +import os +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from policyengine_simulation_contract.hf_dataset import ( + HFRepositoryType, + parse_hf_dataset_uri, + validate_hf_artifact_uri, +) + +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle + + +class UKHFDatasetAccessReport(BaseModel): + """Typed result of validating the complete UK bundle dataset set.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + country: Literal["uk"] = "uk" + dataset_count: int = Field(gt=0) + dataset_identities: tuple[str, ...] + repository_ids: tuple[str, ...] + repository_types: tuple[HFRepositoryType, ...] + revisions: tuple[str, ...] + + +DatasetURIValidator = Callable[[str, HFRepositoryType], str] + + +def validate_uk_hf_dataset_access( + *, + environment: Mapping[str, str] | None = None, + uri_validator: DatasetURIValidator = validate_hf_artifact_uri, +) -> UKHFDatasetAccessReport: + """Use ``HF_TOKEN`` to validate every UK dataset URI in the bundle. + + This deliberately reads dataset identities, repository paths, and revisions + from the installed PolicyEngine.py bundle. CI must not maintain a second, + independently editable list of private UK artifacts. + """ + + runtime_environment = environment if environment is not None else os.environ + if not runtime_environment.get("HF_TOKEN"): + raise RuntimeError("HF_TOKEN is required to validate UK dataset access") + + resolved = load_stage12_bundle() + uk_bundle = next( + country for country in resolved.bundle.countries if country.country == "uk" + ) + if not uk_bundle.datasets: + raise RuntimeError("The PolicyEngine.py bundle declares no UK datasets") + + repository_ids: set[str] = set() + repository_types: set[HFRepositoryType] = set() + revisions: set[str] = set() + identities: list[str] = [] + for dataset in uk_bundle.datasets: + reference = parse_hf_dataset_uri(dataset.uri) + if reference is None: + raise RuntimeError( + f"UK dataset {dataset.identity!r} is not a Hugging Face artifact" + ) + if reference.revision is None: + raise RuntimeError( + f"UK dataset {dataset.identity!r} does not pin a Hugging Face revision" + ) + if reference.revision != dataset.artifact_revision: + raise RuntimeError( + f"UK dataset {dataset.identity!r} URI revision differs from its " + "bundle artifact revision" + ) + + if dataset.repo_type == "dataset": + repository_type: HFRepositoryType = "dataset" + elif dataset.repo_type == "model": + repository_type = "model" + else: + raise RuntimeError( + f"UK dataset {dataset.identity!r} declares unsupported Hugging Face " + f"repository type {dataset.repo_type!r}" + ) + + validated_uri = uri_validator(dataset.uri, repository_type) + if validated_uri != dataset.uri: + raise RuntimeError( + f"UK dataset validator changed the URI for {dataset.identity!r}" + ) + repository_ids.add(reference.repo_id) + repository_types.add(repository_type) + revisions.add(reference.revision) + identities.append(dataset.identity) + + return UKHFDatasetAccessReport( + dataset_count=len(identities), + dataset_identities=tuple(identities), + repository_ids=tuple(sorted(repository_ids)), + repository_types=tuple(sorted(repository_types)), + revisions=tuple(sorted(revisions)), + ) diff --git a/projects/policyengine-simulation-executor/tests/test_uk_hf_access.py b/projects/policyengine-simulation-executor/tests/test_uk_hf_access.py new file mode 100644 index 000000000..d0f911ad9 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_uk_hf_access.py @@ -0,0 +1,61 @@ +"""Tests for complete UK Hugging Face bundle access validation.""" + +from __future__ import annotations + +import pytest + +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.uk_hf_access import ( + validate_uk_hf_dataset_access, +) + + +def _uk_dataset_uris() -> tuple[str, ...]: + bundle = load_stage12_bundle().bundle + uk_bundle = next(country for country in bundle.countries if country.country == "uk") + return tuple(dataset.uri for dataset in uk_bundle.datasets) + + +def test_validation_exercises_every_uk_dataset_declared_by_bundle() -> None: + validated_uris: list[str] = [] + + def record_validation(dataset_uri: str, repository_type: str) -> str: + validated_uris.append(dataset_uri) + return dataset_uri + + report = validate_uk_hf_dataset_access( + environment={"HF_TOKEN": "test-token"}, + uri_validator=record_validation, + ) + + expected_uris = _uk_dataset_uris() + assert tuple(validated_uris) == expected_uris + assert report.country == "uk" + assert report.dataset_count == len(expected_uris) + assert report.dataset_count > 0 + assert report.repository_ids + assert report.repository_types == ("model",) + assert report.revisions + + +def test_validation_requires_the_exact_modal_hf_token_variable() -> None: + with pytest.raises(RuntimeError, match="HF_TOKEN is required"): + validate_uk_hf_dataset_access( + environment={"HUGGINGFACE_TOKEN": "wrong-variable"}, + uri_validator=lambda dataset_uri, repository_type: dataset_uri, + ) + + +def test_validation_propagates_an_inaccessible_artifact_failure() -> None: + first_uri = _uk_dataset_uris()[0] + + def reject_first_artifact(dataset_uri: str, repository_type: str) -> str: + if dataset_uri == first_uri: + raise RuntimeError("artifact is inaccessible") + return dataset_uri + + with pytest.raises(RuntimeError, match="artifact is inaccessible"): + validate_uk_hf_dataset_access( + environment={"HF_TOKEN": "test-token"}, + uri_validator=reject_first_artifact, + ) From e118aa3ce710eac1f8c7c328d9e4633f3658ed7c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 17:34:11 +0400 Subject: [PATCH 7/7] Use Modal required keys for Hugging Face secret --- .github/scripts/modal-image-smoke.sh | 7 -- .github/workflows/pr.yml | 27 ----- .../hf_dataset.py | 44 ++------ .../tests/test_hf_dataset.py | 51 ++------- .../src/modal/app.py | 5 +- .../src/modal/smoke_app.py | 21 +--- .../src/modal/v2_app.py | 5 +- .../uk_hf_access.py | 105 ------------------ .../tests/test_modal_bundle_image.py | 4 + .../tests/test_stage12_modal_app.py | 4 + .../tests/test_uk_hf_access.py | 61 ---------- 11 files changed, 38 insertions(+), 296 deletions(-) delete mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py delete mode 100644 projects/policyengine-simulation-executor/tests/test_uk_hf_access.py diff --git a/.github/scripts/modal-image-smoke.sh b/.github/scripts/modal-image-smoke.sh index 78cf5166f..e259e489b 100755 --- a/.github/scripts/modal-image-smoke.sh +++ b/.github/scripts/modal-image-smoke.sh @@ -11,13 +11,6 @@ set -euo pipefail MODAL_ENV="${1:?Modal environment required}" REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -echo "=== UK Hugging Face dataset access (env: $MODAL_ENV) ===" -( - cd "$REPO_ROOT/projects/policyengine-simulation-executor" - uv run modal run --env="$MODAL_ENV" \ - src/modal/smoke_app.py --uk-hf-access-only -) - echo "=== Gateway image smoke (env: $MODAL_ENV) ===" ( cd "$REPO_ROOT/projects/policyengine-simulation-gateway" diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 527e9870a..f4b4916f1 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -9,34 +9,8 @@ concurrency: cancel-in-progress: true jobs: - validate-uk-hf-access: - name: Validate UK Hugging Face dataset access - runs-on: ubuntu-latest - if: github.event.pull_request.head.repo.full_name == github.repository - - steps: - - uses: actions/checkout@v6 - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: '3.13' - - - name: Install uv - uses: astral-sh/setup-uv@v8.1.0 - with: - enable-cache: true - - - name: Validate every UK bundle dataset with the staging credential - working-directory: projects/policyengine-simulation-executor - env: - MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} - MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} - run: uv run modal run --env=staging src/modal/smoke_app.py --uk-hf-access-only - test: name: Test - needs: validate-uk-hf-access runs-on: ubuntu-latest strategy: matrix: @@ -76,7 +50,6 @@ jobs: lint: name: Lint - needs: validate-uk-hf-access runs-on: ubuntu-latest steps: 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 c85862ba5..407803240 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 @@ -13,7 +13,7 @@ import os from dataclasses import dataclass from functools import lru_cache -from typing import Any, Literal +from typing import Any from urllib.error import HTTPError, URLError from urllib.parse import quote from urllib.request import Request, urlopen @@ -26,7 +26,6 @@ "HUGGINGFACE_HUB_TOKEN", "HUGGINGFACE_TOKEN", ) -HFRepositoryType = Literal["dataset", "model"] class HuggingFaceDatasetReferenceError(ValueError): @@ -78,15 +77,13 @@ def parse_hf_dataset_uri(dataset_uri: str) -> HFDatasetReference | None: @lru_cache -def _fetch_hf_repository_revision( +def _fetch_hf_dataset_revision( repo_id: str, revision: str, - repo_type: HFRepositoryType, token: str | None, ) -> dict[str, Any]: - resource_name = "datasets" if repo_type == "dataset" else "models" url = ( - f"{HF_ENDPOINT}/api/{resource_name}/" + f"{HF_ENDPOINT}/api/datasets/" f"{quote(repo_id, safe='/')}/revision/{quote(revision, safe='')}" ) headers = {"Accept": "application/json"} @@ -100,13 +97,12 @@ def _fetch_hf_repository_revision( except HTTPError as exc: detail = exc.reason or f"HTTP {exc.code}" raise HuggingFaceDatasetReferenceError( - f"Hugging Face {repo_type} repository revision " - f"{repo_id}@{revision} was not found: " + f"Hugging Face dataset revision {repo_id}@{revision} was not found: " f"{detail}" ) from exc except (OSError, URLError, json.JSONDecodeError) as exc: raise HuggingFaceDatasetReferenceError( - f"Unable to validate Hugging Face {repo_type} repository revision " + f"Unable to validate Hugging Face dataset revision " f"{repo_id}@{revision}: {exc}" ) from exc @@ -128,43 +124,23 @@ def _siblings_contain_path(payload: dict[str, Any], path: str) -> bool | None: return False if seen_file_listing else None -def validate_hf_artifact_uri( - dataset_uri: str, - repo_type: HFRepositoryType, -) -> str: - """Validate a pinned artifact against its declared Hub repository type.""" +def validate_hf_dataset_uri(dataset_uri: str) -> str: + """Validate an explicit ``hf://`` dataset URI if it pins a revision.""" parsed = parse_hf_dataset_uri(dataset_uri) if parsed is None or parsed.revision is None: return dataset_uri - payload = _fetch_hf_repository_revision( - parsed.repo_id, - parsed.revision, - repo_type, - _hf_token(), - ) + payload = _fetch_hf_dataset_revision(parsed.repo_id, parsed.revision, _hf_token()) contains_path = _siblings_contain_path(payload, parsed.path) - if contains_path is None: - raise HuggingFaceDatasetReferenceError( - f"Hugging Face {repo_type} repository revision " - f"{parsed.repo_id}@{parsed.revision} did not provide an artifact listing" - ) if contains_path is False: raise HuggingFaceDatasetReferenceError( - f"Hugging Face {repo_type} repository revision " - f"{parsed.repo_id}@{parsed.revision} does not contain artifact " - f"{parsed.path!r}" + f"Hugging Face dataset revision {parsed.repo_id}@{parsed.revision} " + f"does not contain artifact {parsed.path!r}" ) return dataset_uri -def validate_hf_dataset_uri(dataset_uri: str) -> str: - """Validate a pinned URI stored in a Hugging Face dataset repository.""" - - return validate_hf_artifact_uri(dataset_uri, "dataset") - - def with_hf_revision(dataset_uri: str, revision: str) -> str: """Return ``dataset_uri`` pinned to ``revision`` and validate it on the Hub.""" diff --git a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py index e5f2d4464..95d2f5e3b 100644 --- a/libs/policyengine-simulation-contract/tests/test_hf_dataset.py +++ b/libs/policyengine-simulation-contract/tests/test_hf_dataset.py @@ -10,7 +10,6 @@ from policyengine_simulation_contract.hf_dataset import ( HuggingFaceDatasetReferenceError, parse_hf_dataset_uri, - validate_hf_artifact_uri, validate_hf_dataset_uri, with_hf_revision, ) @@ -41,8 +40,8 @@ def test_parse_hf_dataset_uri_extracts_repo_path_and_revision(): assert parsed.revision == "custom-v1" -def test_fetch_hf_repository_revision_uses_declared_repository_type(monkeypatch): - hf_dataset._fetch_hf_repository_revision.cache_clear() +def test_fetch_hf_dataset_revision_uses_dataset_revision_api(monkeypatch): + hf_dataset._fetch_hf_dataset_revision.cache_clear() seen = {} def fake_urlopen(request, timeout): @@ -53,16 +52,16 @@ def fake_urlopen(request, timeout): monkeypatch.setattr(hf_dataset, "urlopen", fake_urlopen) - payload = hf_dataset._fetch_hf_repository_revision( + payload = hf_dataset._fetch_hf_dataset_revision( "policyengine/populace-us", "custom-v1", - "model", "hf-token", ) assert payload == {"sha": "abc123", "siblings": []} assert seen["url"] == ( - "https://huggingface.co/api/models/policyengine/populace-us/revision/custom-v1" + "https://huggingface.co/api/datasets/" + "policyengine/populace-us/revision/custom-v1" ) assert seen["headers"]["Authorization"] == "Bearer hf-token" assert seen["timeout"] == hf_dataset.HF_REQUEST_TIMEOUT_SECONDS @@ -71,10 +70,8 @@ def fake_urlopen(request, timeout): def test_validate_hf_dataset_uri_rejects_revision_missing_artifact(monkeypatch): monkeypatch.setattr( hf_dataset, - "_fetch_hf_repository_revision", - lambda repo_id, revision, repo_type, token: { - "siblings": [{"rfilename": "other_file.h5"}] - }, + "_fetch_hf_dataset_revision", + lambda repo_id, revision, token: {"siblings": [{"rfilename": "other_file.h5"}]}, ) with pytest.raises( @@ -86,36 +83,6 @@ def test_validate_hf_dataset_uri_rejects_revision_missing_artifact(monkeypatch): ) -def test_validate_hf_dataset_uri_rejects_missing_artifact_listing(monkeypatch): - monkeypatch.setattr( - hf_dataset, - "_fetch_hf_repository_revision", - lambda repo_id, revision, repo_type, token: {"sha": "abc123"}, - ) - - with pytest.raises( - HuggingFaceDatasetReferenceError, - match="did not provide an artifact listing", - ): - validate_hf_dataset_uri( - "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" - ) - - -def test_validate_hf_artifact_uri_uses_model_repository_api(monkeypatch): - seen_repository_types = [] - - def fake_fetch(repo_id, revision, repo_type, token): - seen_repository_types.append(repo_type) - return {"siblings": [{"rfilename": "enhanced_frs.h5"}]} - - monkeypatch.setattr(hf_dataset, "_fetch_hf_repository_revision", fake_fetch) - - uri = "hf://policyengine/private-uk/enhanced_frs.h5@release-v1" - assert validate_hf_artifact_uri(uri, "model") == uri - assert seen_repository_types == ["model"] - - def test_with_hf_revision_validates_and_preserves_requested_revision(monkeypatch): calls = [] @@ -132,4 +99,6 @@ def fake_validate(dataset_uri): ) == "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" ) - assert calls == ["hf://policyengine/populace-us/populace_us_2024.h5@custom-v1"] + assert calls == [ + "hf://policyengine/populace-us/populace_us_2024.h5@custom-v1" + ] diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index ce7f8af22..f86d13b55 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -125,7 +125,10 @@ def get_app_name(policyengine_version: str) -> str: # GCP credentials are shared across environments (always from main) gcp_secret = modal.Secret.from_name("gcp-credentials", environment_name="main") data_secret = modal.Secret.from_name("policyengine-data-credentials") -hf_secret = modal.Secret.from_name("huggingface-token") +hf_secret = modal.Secret.from_name( + "huggingface-token", + required_keys=["HF_TOKEN"], +) OBSERVABILITY_ENV = modal_image_environment() diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index b052b9c08..164a4971f 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -16,7 +16,7 @@ import modal -from src.modal.app import build_runtime_simulation_image, hf_secret +from src.modal.app import build_runtime_simulation_image app = modal.App("policyengine-simulation-executor-smoke") @@ -29,17 +29,6 @@ ) -@app.function(image=smoke_image, secrets=[hf_secret], timeout=120) -def smoke_uk_hf_dataset_access() -> str: - """Exercise the deployed HF_TOKEN against every bundle-declared UK artifact.""" - - from policyengine_simulation_executor.uk_hf_access import ( - validate_uk_hf_dataset_access, - ) - - return validate_uk_hf_dataset_access().model_dump_json() - - @app.function(image=smoke_image, timeout=600, memory=8192) def smoke_import_executor() -> dict: import importlib @@ -82,13 +71,7 @@ def smoke_import_executor() -> dict: @app.local_entrypoint() -def main(uk_hf_access_only: bool = False): - access_report = smoke_uk_hf_dataset_access.remote() - print(access_report) - print("UK Hugging Face dataset access OK") - if uk_hf_access_only: - return - +def main(): report = smoke_import_executor.remote() print(report) print("executor image smoke OK") diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 4a9f0affd..b25008f2b 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -68,7 +68,10 @@ def _external_assertions(environment: dict[str, str]) -> dict[str, str]: app = modal.App(APP_NAME) gcp_secret = modal.Secret.from_name("stage12-evaluation-gcp-credentials") data_secret = modal.Secret.from_name("policyengine-data-credentials") -hf_secret = modal.Secret.from_name("huggingface-token") +hf_secret = modal.Secret.from_name( + "huggingface-token", + required_keys=["HF_TOKEN"], +) comparison_runtime_secret = modal.Secret.from_name("stage12-evaluation-runtime") worker_secrets = [ gcp_secret, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py deleted file mode 100644 index b9958ee75..000000000 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_hf_access.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Live access validation for every UK dataset in the release bundle.""" - -from __future__ import annotations - -from collections.abc import Callable, Mapping -import os -from typing import Literal - -from pydantic import BaseModel, ConfigDict, Field - -from policyengine_simulation_contract.hf_dataset import ( - HFRepositoryType, - parse_hf_dataset_uri, - validate_hf_artifact_uri, -) - -from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle - - -class UKHFDatasetAccessReport(BaseModel): - """Typed result of validating the complete UK bundle dataset set.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - - country: Literal["uk"] = "uk" - dataset_count: int = Field(gt=0) - dataset_identities: tuple[str, ...] - repository_ids: tuple[str, ...] - repository_types: tuple[HFRepositoryType, ...] - revisions: tuple[str, ...] - - -DatasetURIValidator = Callable[[str, HFRepositoryType], str] - - -def validate_uk_hf_dataset_access( - *, - environment: Mapping[str, str] | None = None, - uri_validator: DatasetURIValidator = validate_hf_artifact_uri, -) -> UKHFDatasetAccessReport: - """Use ``HF_TOKEN`` to validate every UK dataset URI in the bundle. - - This deliberately reads dataset identities, repository paths, and revisions - from the installed PolicyEngine.py bundle. CI must not maintain a second, - independently editable list of private UK artifacts. - """ - - runtime_environment = environment if environment is not None else os.environ - if not runtime_environment.get("HF_TOKEN"): - raise RuntimeError("HF_TOKEN is required to validate UK dataset access") - - resolved = load_stage12_bundle() - uk_bundle = next( - country for country in resolved.bundle.countries if country.country == "uk" - ) - if not uk_bundle.datasets: - raise RuntimeError("The PolicyEngine.py bundle declares no UK datasets") - - repository_ids: set[str] = set() - repository_types: set[HFRepositoryType] = set() - revisions: set[str] = set() - identities: list[str] = [] - for dataset in uk_bundle.datasets: - reference = parse_hf_dataset_uri(dataset.uri) - if reference is None: - raise RuntimeError( - f"UK dataset {dataset.identity!r} is not a Hugging Face artifact" - ) - if reference.revision is None: - raise RuntimeError( - f"UK dataset {dataset.identity!r} does not pin a Hugging Face revision" - ) - if reference.revision != dataset.artifact_revision: - raise RuntimeError( - f"UK dataset {dataset.identity!r} URI revision differs from its " - "bundle artifact revision" - ) - - if dataset.repo_type == "dataset": - repository_type: HFRepositoryType = "dataset" - elif dataset.repo_type == "model": - repository_type = "model" - else: - raise RuntimeError( - f"UK dataset {dataset.identity!r} declares unsupported Hugging Face " - f"repository type {dataset.repo_type!r}" - ) - - validated_uri = uri_validator(dataset.uri, repository_type) - if validated_uri != dataset.uri: - raise RuntimeError( - f"UK dataset validator changed the URI for {dataset.identity!r}" - ) - repository_ids.add(reference.repo_id) - repository_types.add(repository_type) - revisions.add(reference.revision) - identities.append(dataset.identity) - - return UKHFDatasetAccessReport( - dataset_count=len(identities), - dataset_identities=tuple(identities), - repository_ids=tuple(sorted(repository_ids)), - repository_types=tuple(sorted(repository_types)), - revisions=tuple(sorted(revisions)), - ) 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 8b84c2698..d31b223ef 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py @@ -60,6 +60,10 @@ def test_modal_image_uses_policyengine_bundle_install(monkeypatch): assert app.OBSERVABILITY_ENV["OTEL_EXPORTER_OTLP_ENDPOINT"] == ( "https://collector.test" ) + assert app.hf_secret == { + "args": ("huggingface-token",), + "kwargs": {"required_keys": ["HF_TOKEN"]}, + } assert command_calls[0][2]["secrets"] == [app.data_secret, app.hf_secret] uv_sync_calls = [ call for call in app.simulation_image.calls if call[0] == "uv_sync" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index f1841008c..e38390f96 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -97,6 +97,10 @@ def test_v2_app_name_and_images_are_separate_and_bundle_derived(monkeypatch) -> assert "--country us" not in uk_command assert module.RESOLVED_BUNDLE.bundle.policyengine_requirement in us_command assert module.gcp_secret["args"] == ("stage12-evaluation-gcp-credentials",) + assert module.hf_secret == { + "args": ("huggingface-token",), + "kwargs": {"required_keys": ["HF_TOKEN"]}, + } assert all( secret.get("args") != ("policyengine-logfire",) for secret in module.worker_secrets diff --git a/projects/policyengine-simulation-executor/tests/test_uk_hf_access.py b/projects/policyengine-simulation-executor/tests/test_uk_hf_access.py deleted file mode 100644 index d0f911ad9..000000000 --- a/projects/policyengine-simulation-executor/tests/test_uk_hf_access.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Tests for complete UK Hugging Face bundle access validation.""" - -from __future__ import annotations - -import pytest - -from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle -from policyengine_simulation_executor.uk_hf_access import ( - validate_uk_hf_dataset_access, -) - - -def _uk_dataset_uris() -> tuple[str, ...]: - bundle = load_stage12_bundle().bundle - uk_bundle = next(country for country in bundle.countries if country.country == "uk") - return tuple(dataset.uri for dataset in uk_bundle.datasets) - - -def test_validation_exercises_every_uk_dataset_declared_by_bundle() -> None: - validated_uris: list[str] = [] - - def record_validation(dataset_uri: str, repository_type: str) -> str: - validated_uris.append(dataset_uri) - return dataset_uri - - report = validate_uk_hf_dataset_access( - environment={"HF_TOKEN": "test-token"}, - uri_validator=record_validation, - ) - - expected_uris = _uk_dataset_uris() - assert tuple(validated_uris) == expected_uris - assert report.country == "uk" - assert report.dataset_count == len(expected_uris) - assert report.dataset_count > 0 - assert report.repository_ids - assert report.repository_types == ("model",) - assert report.revisions - - -def test_validation_requires_the_exact_modal_hf_token_variable() -> None: - with pytest.raises(RuntimeError, match="HF_TOKEN is required"): - validate_uk_hf_dataset_access( - environment={"HUGGINGFACE_TOKEN": "wrong-variable"}, - uri_validator=lambda dataset_uri, repository_type: dataset_uri, - ) - - -def test_validation_propagates_an_inaccessible_artifact_failure() -> None: - first_uri = _uk_dataset_uris()[0] - - def reject_first_artifact(dataset_uri: str, repository_type: str) -> str: - if dataset_uri == first_uri: - raise RuntimeError("artifact is inaccessible") - return dataset_uri - - with pytest.raises(RuntimeError, match="artifact is inaccessible"): - validate_uk_hf_dataset_access( - environment={"HF_TOKEN": "test-token"}, - uri_validator=reject_first_artifact, - )