Fix Stage 12 comparison and UK geography metadata

.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"

.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."

.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:

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

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."""
 

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"]

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:

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")