From a4c6ad2b552996528f73db6479c833b67b538bf6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:23:45 +0400 Subject: [PATCH 01/18] Model bundle dataset source packages --- src/policyengine/data/bundle/manifest.json | 8 + src/policyengine/provenance/__init__.py | 12 ++ .../provenance/dataset_materialization.py | 150 ++++++++++++++++++ src/policyengine/provenance/manifest.py | 3 + tests/test_dataset_materialization.py | 86 ++++++++++ 5 files changed, 259 insertions(+) create mode 100644 src/policyengine/provenance/dataset_materialization.py create mode 100644 tests/test_dataset_materialization.py diff --git a/src/policyengine/data/bundle/manifest.json b/src/policyengine/data/bundle/manifest.json index 7aa72022..98c4d1f0 100644 --- a/src/policyengine/data/bundle/manifest.json +++ b/src/policyengine/data/bundle/manifest.json @@ -201,28 +201,36 @@ "dataset_overlays": { "uk": { "enhanced_frs_2023_24": { + "data_package_name": "policyengine-uk-data", "path": "enhanced_frs_2023_24.h5", "repo_id": "policyengine/policyengine-uk-data-private", + "repo_type": "model", "revision": "655dd07e4bb9c777b00dac044949611f1feb824f", "sha256": "584ae33d80ca0431254610a3f8254d132da73477d31966d6446282861ecae50d" }, "frs_2023_24": { + "data_package_name": "policyengine-uk-data", "path": "frs_2023_24.h5", "repo_id": "policyengine/policyengine-uk-data-private", + "repo_type": "model", "revision": "655dd07e4bb9c777b00dac044949611f1feb824f", "sha256": "df26d4d7af9d164aa2d064181b39290292d2f62bb26fee6126fc095fc06da292" }, "populace_uk_2023": { + "data_package_name": "populace-data", "path": "populace_uk_2023.h5", "repo_id": "policyengine/populace-uk-private", + "repo_type": "model", "revision": "populace-uk-2023-dd68c73-4aa4b14-20260619T023711Z", "sha256": "f17306ccb2aad7ff0130be3589b560afb2e2a12a943570911cd0c77f07934833" } }, "us": { "populace_us_2024_acs_local": { + "data_package_name": "populace-data", "path": "populace_us_2024_acs_local.h5", "repo_id": "policyengine/populace-us", + "repo_type": "dataset", "revision": "populace-us-2024-buildo-acs-local-77e2061-20260724T110908Z", "sha256": "71763290ded993789af0ad818fd833a6aabb9ca7ff7d14f78315340d8617c6f4" } diff --git a/src/policyengine/provenance/__init__.py b/src/policyengine/provenance/__init__.py index 5ef0789d..8afc8f55 100644 --- a/src/policyengine/provenance/__init__.py +++ b/src/policyengine/provenance/__init__.py @@ -24,6 +24,18 @@ from .certification import ( certify_data_release as certify_data_release, ) +from .dataset_materialization import ( + BundleDatasetPlan as BundleDatasetPlan, +) +from .dataset_materialization import ( + DatasetMaterializationError as DatasetMaterializationError, +) +from .dataset_materialization import ( + MaterializedDataset as MaterializedDataset, +) +from .dataset_materialization import ( + resolve_bundle_dataset_plan as resolve_bundle_dataset_plan, +) from .manifest import ( CertifiedDataArtifact as CertifiedDataArtifact, ) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py new file mode 100644 index 00000000..8204c222 --- /dev/null +++ b/src/policyengine/provenance/dataset_materialization.py @@ -0,0 +1,150 @@ +"""Resolve and materialize datasets certified by a PolicyEngine.py bundle.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Literal, Optional + +from pydantic import BaseModel + +from .manifest import ( + CountryReleaseManifest, + _artifact_revision, + build_hf_uri, + get_release_manifest, +) + +DEFAULT_DATA_DIR = Path("./data") + + +class DatasetMaterializationError(ValueError): + """Raised when a dataset cannot be resolved or materialized safely.""" + + +class BundleDatasetPlan(BaseModel): + """Exact bundle metadata needed to materialize one managed dataset.""" + + country_id: str + dataset: str + data_package_name: str + repo_id: str + repo_type: Literal["model", "dataset"] + path: str + revision: str + expected_sha256: str + source_uri: str + destination: Path + build_id: Optional[str] = None + + +class MaterializedDataset(BaseModel): + """Verified local representation of one bundle-managed dataset.""" + + country_id: str + dataset: str + data_package_name: str + repo_id: str + repo_type: Literal["model", "dataset"] + revision: str + source_uri: str + expected_sha256: str + actual_sha256: str + path: Path + cache_hit: bool + build_id: Optional[str] = None + + +class _DatasetPackageStrategy: + """Package-specific validation for a bundle dataset source.""" + + def validate(self, plan: BundleDatasetPlan) -> None: + raise NotImplementedError + + +class _CountryDataPackageStrategy(_DatasetPackageStrategy): + def validate(self, plan: BundleDatasetPlan) -> None: + package_name = plan.data_package_name + if not ( + package_name.startswith("policyengine-") and package_name.endswith("-data") + ): + raise DatasetMaterializationError( + f"Unsupported country data package: {package_name!r}." + ) + + +class _PopulaceDataPackageStrategy(_DatasetPackageStrategy): + def validate(self, plan: BundleDatasetPlan) -> None: + if plan.data_package_name != "populace-data": + raise DatasetMaterializationError( + f"Unsupported Populace data package: {plan.data_package_name!r}." + ) + + +def _dataset_package_strategy(data_package_name: str) -> _DatasetPackageStrategy: + if data_package_name == "populace-data": + return _PopulaceDataPackageStrategy() + if data_package_name.startswith("policyengine-") and data_package_name.endswith( + "-data" + ): + return _CountryDataPackageStrategy() + raise DatasetMaterializationError( + "Unsupported bundle data package " + f"{data_package_name!r}; expected 'populace-data' or " + "'policyengine--data'." + ) + + +def resolve_bundle_dataset_plan( + country_id: str, + dataset: Optional[str] = None, + *, + data_dir: Path = DEFAULT_DATA_DIR, + manifest: Optional[CountryReleaseManifest] = None, +) -> BundleDatasetPlan: + """Resolve one logical dataset to its exact bundle-certified source.""" + + country_manifest = manifest or get_release_manifest(country_id) + dataset_name = dataset or country_manifest.default_dataset + reference = country_manifest.datasets.get(dataset_name) + if reference is None: + raise DatasetMaterializationError( + f"Unknown managed dataset {dataset_name!r} for country " + f"{country_id!r}. Known datasets: " + f"{sorted(country_manifest.datasets)}" + ) + + data_package_name = ( + reference.data_package_name or country_manifest.data_package.name + ) + repo_id = reference.repo_id or country_manifest.data_package.repo_id + repo_type = reference.repo_type or country_manifest.data_package.repo_type + revision = reference.revision or _artifact_revision(country_manifest.data_package) + if repo_type not in {"model", "dataset"}: + raise DatasetMaterializationError( + f"Dataset {dataset_name!r} has unsupported Hugging Face repository " + f"type {repo_type!r}." + ) + if not reference.sha256: + raise DatasetMaterializationError( + f"Managed dataset {dataset_name!r} is missing a certified sha256." + ) + + plan = BundleDatasetPlan( + country_id=country_id, + dataset=dataset_name, + data_package_name=data_package_name, + repo_id=repo_id, + repo_type=repo_type, + path=reference.path, + revision=revision, + expected_sha256=reference.sha256, + source_uri=build_hf_uri(repo_id, reference.path, revision), + destination=data_dir / Path(reference.path).name, + build_id=( + country_manifest.certified_data_artifact.build_id + if country_manifest.certified_data_artifact is not None + else None + ), + ) + _dataset_package_strategy(data_package_name).validate(plan) + return plan diff --git a/src/policyengine/provenance/manifest.py b/src/policyengine/provenance/manifest.py index 088676fd..5ada2d33 100644 --- a/src/policyengine/provenance/manifest.py +++ b/src/policyengine/provenance/manifest.py @@ -61,6 +61,9 @@ class ArtifactPathReference(BaseModel): revision: Optional[str] = None sha256: Optional[str] = None metadata_sha256: Optional[str] = None + # Set when the artifact is supplied by a different package than the + # country release manifest's primary data package. + data_package_name: Optional[str] = None # Set when the artifact lives outside the data package's repo (inherited # datasets keep their original repo + revision pins). repo_id: Optional[str] = None diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py new file mode 100644 index 00000000..957b018a --- /dev/null +++ b/tests/test_dataset_materialization.py @@ -0,0 +1,86 @@ +from pathlib import Path + +import pytest + +from policyengine.provenance.dataset_materialization import ( + BundleDatasetPlan, + DatasetMaterializationError, + MaterializedDataset, + _dataset_package_strategy, + resolve_bundle_dataset_plan, +) +from policyengine.provenance.manifest import CountryReleaseManifest + + +def _manifest() -> CountryReleaseManifest: + return CountryReleaseManifest.model_validate( + { + "country_id": "uk", + "policyengine_version": "5.0.4", + "model_package": {"name": "policyengine-uk", "version": "2.90.2"}, + "data_package": { + "name": "policyengine-uk-data", + "version": "1.56.16", + "repo_id": "policyengine/policyengine-uk-data-private", + "repo_type": "model", + }, + "default_dataset": "enhanced_frs_2024_25", + "datasets": { + "enhanced_frs_2024_25": { + "path": "enhanced_frs_2024_25.h5", + "revision": "uk-release", + "sha256": "a" * 64, + }, + "populace_uk_2023": { + "data_package_name": "populace-data", + "path": "populace_uk_2023.h5", + "repo_id": "policyengine/populace-uk-private", + "repo_type": "model", + "revision": "populace-release", + "sha256": "b" * 64, + }, + }, + } + ) + + +def test_resolve_bundle_dataset_plan_inherits_primary_package(tmp_path): + plan = resolve_bundle_dataset_plan("uk", data_dir=tmp_path, manifest=_manifest()) + + assert plan.data_package_name == "policyengine-uk-data" + assert plan.repo_id == "policyengine/policyengine-uk-data-private" + assert plan.repo_type == "model" + assert plan.revision == "uk-release" + assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" + + +def test_resolve_bundle_dataset_plan_uses_cross_package_overlay(tmp_path): + plan = resolve_bundle_dataset_plan( + "uk", + "populace_uk_2023", + data_dir=tmp_path, + manifest=_manifest(), + ) + + assert plan.data_package_name == "populace-data" + assert plan.repo_id == "policyengine/populace-uk-private" + assert plan.repo_type == "model" + assert plan.revision == "populace-release" + + +def test_bundle_dataset_models_round_trip_json(): + plan = resolve_bundle_dataset_plan("uk", manifest=_manifest()) + assert BundleDatasetPlan.model_validate_json(plan.model_dump_json()) == plan + + result = MaterializedDataset( + **plan.model_dump(exclude={"path", "destination"}), + actual_sha256=plan.expected_sha256, + path=Path("data/enhanced_frs_2024_25.h5"), + cache_hit=True, + ) + assert MaterializedDataset.model_validate_json(result.model_dump_json()) == result + + +def test_unknown_data_package_is_rejected(): + with pytest.raises(DatasetMaterializationError, match="Unsupported bundle"): + _dataset_package_strategy("unknown-data") From 5fd35ab4bdf156d34d1cbb029767a4836f149a8c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:28:39 +0400 Subject: [PATCH 02/18] Add canonical bundle dataset materialization --- src/policyengine/provenance/__init__.py | 6 + .../provenance/dataset_materialization.py | 317 +++++++++++++++++- tests/test_dataset_materialization.py | 201 +++++++++++ 3 files changed, 521 insertions(+), 3 deletions(-) diff --git a/src/policyengine/provenance/__init__.py b/src/policyengine/provenance/__init__.py index 8afc8f55..47010765 100644 --- a/src/policyengine/provenance/__init__.py +++ b/src/policyengine/provenance/__init__.py @@ -33,6 +33,12 @@ from .dataset_materialization import ( MaterializedDataset as MaterializedDataset, ) +from .dataset_materialization import ( + materialize_bundle_dataset as materialize_bundle_dataset, +) +from .dataset_materialization import ( + materialize_unmanaged_dataset_source as materialize_unmanaged_dataset_source, +) from .dataset_materialization import ( resolve_bundle_dataset_plan as resolve_bundle_dataset_plan, ) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 8204c222..5aff298a 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -2,9 +2,16 @@ from __future__ import annotations +import hashlib +import os +import shutil +import tempfile +from datetime import datetime, timezone from pathlib import Path -from typing import Literal, Optional +from typing import Literal, Optional, Union, cast +from urllib.parse import quote +import requests from pydantic import BaseModel from .manifest import ( @@ -15,6 +22,8 @@ ) DEFAULT_DATA_DIR = Path("./data") +BACKUP_DIR_NAME = ".policyengine-bundle-backups" +DOWNLOAD_TIMEOUT_SECONDS = 60 class DatasetMaterializationError(ValueError): @@ -55,11 +64,19 @@ class MaterializedDataset(BaseModel): class _DatasetPackageStrategy: - """Package-specific validation for a bundle dataset source.""" + """Package-specific behavior for a bundle dataset source.""" def validate(self, plan: BundleDatasetPlan) -> None: raise NotImplementedError + def materialize( + self, + plan: BundleDatasetPlan, + *, + session=requests, + ) -> MaterializedDataset: + raise NotImplementedError + class _CountryDataPackageStrategy(_DatasetPackageStrategy): def validate(self, plan: BundleDatasetPlan) -> None: @@ -71,6 +88,14 @@ def validate(self, plan: BundleDatasetPlan) -> None: f"Unsupported country data package: {package_name!r}." ) + def materialize( + self, + plan: BundleDatasetPlan, + *, + session=requests, + ) -> MaterializedDataset: + return _materialize_country_data_package(plan, session=session) + class _PopulaceDataPackageStrategy(_DatasetPackageStrategy): def validate(self, plan: BundleDatasetPlan) -> None: @@ -79,6 +104,14 @@ def validate(self, plan: BundleDatasetPlan) -> None: f"Unsupported Populace data package: {plan.data_package_name!r}." ) + def materialize( + self, + plan: BundleDatasetPlan, + *, + session=requests, + ) -> MaterializedDataset: + return _materialize_populace_data_package(plan, session=session) + def _dataset_package_strategy(data_package_name: str) -> _DatasetPackageStrategy: if data_package_name == "populace-data": @@ -124,6 +157,7 @@ def resolve_bundle_dataset_plan( f"Dataset {dataset_name!r} has unsupported Hugging Face repository " f"type {repo_type!r}." ) + validated_repo_type = cast(Literal["model", "dataset"], repo_type) if not reference.sha256: raise DatasetMaterializationError( f"Managed dataset {dataset_name!r} is missing a certified sha256." @@ -134,7 +168,7 @@ def resolve_bundle_dataset_plan( dataset=dataset_name, data_package_name=data_package_name, repo_id=repo_id, - repo_type=repo_type, + repo_type=validated_repo_type, path=reference.path, revision=revision, expected_sha256=reference.sha256, @@ -148,3 +182,280 @@ def resolve_bundle_dataset_plan( ) _dataset_package_strategy(data_package_name).validate(plan) return plan + + +def materialize_bundle_dataset( + country_id: str, + dataset: Optional[str] = None, + *, + data_dir: Path = DEFAULT_DATA_DIR, + manifest: Optional[CountryReleaseManifest] = None, + session=requests, +) -> MaterializedDataset: + """Download and verify one dataset certified by the release bundle.""" + + plan = resolve_bundle_dataset_plan( + country_id, + dataset, + data_dir=data_dir, + manifest=manifest, + ) + strategy = _dataset_package_strategy(plan.data_package_name) + return strategy.materialize(plan, session=session) + + +def _materialize_country_data_package( + plan: BundleDatasetPlan, + *, + session=requests, +) -> MaterializedDataset: + return _materialize_managed_hf_dataset(plan, session=session) + + +def _materialize_populace_data_package( + plan: BundleDatasetPlan, + *, + session=requests, +) -> MaterializedDataset: + return _materialize_managed_hf_dataset(plan, session=session) + + +def _materialize_managed_hf_dataset( + plan: BundleDatasetPlan, + *, + session=requests, +) -> MaterializedDataset: + destination = plan.destination + if destination.is_file(): + actual_sha256 = _sha256_file(destination) + if actual_sha256 == plan.expected_sha256: + return _materialized_result( + plan, + actual_sha256=actual_sha256, + cache_hit=True, + ) + + url = _hf_download_url( + repo_id=plan.repo_id, + repo_type=plan.repo_type, + path=plan.path, + revision=plan.revision, + ) + downloaded = _download_to_temp( + url, + destination=destination, + source_description=(f"{plan.country_id.upper()} dataset {plan.dataset!r}"), + session=session, + ) + try: + actual_sha256 = _sha256_file(downloaded) + if actual_sha256 != plan.expected_sha256: + raise DatasetMaterializationError( + f"Downloaded {plan.country_id.upper()} dataset {plan.dataset!r} " + f"has sha256 {actual_sha256}, expected {plan.expected_sha256}." + ) + _backup_existing(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + os.replace(downloaded, destination) + finally: + downloaded.unlink(missing_ok=True) + + return _materialized_result( + plan, + actual_sha256=actual_sha256, + cache_hit=False, + ) + + +def _materialized_result( + plan: BundleDatasetPlan, + *, + actual_sha256: str, + cache_hit: bool, +) -> MaterializedDataset: + return MaterializedDataset( + country_id=plan.country_id, + dataset=plan.dataset, + data_package_name=plan.data_package_name, + repo_id=plan.repo_id, + repo_type=plan.repo_type, + revision=plan.revision, + source_uri=plan.source_uri, + expected_sha256=plan.expected_sha256, + actual_sha256=actual_sha256, + path=plan.destination, + cache_hit=cache_hit, + build_id=plan.build_id, + ) + + +class _UnmanagedHFReference(BaseModel): + repo_id: str + path: str + revision: str + + +class _DatasetNotFoundError(DatasetMaterializationError): + pass + + +def materialize_unmanaged_dataset_source( + dataset_source: Union[str, Path], + *, + version: Optional[str] = None, + data_dir: Path = DEFAULT_DATA_DIR, + repo_type: Optional[Literal["model", "dataset"]] = None, + session=requests, +) -> str: + """Return a local path for an explicitly unmanaged local or HF source.""" + + source = str(dataset_source) + if source.startswith("gs://"): + raise DatasetMaterializationError( + "GCS dataset sources are no longer supported. Publish the dataset " + "on Hugging Face and reference that artifact instead." + ) + if not source.startswith("hf://"): + if "://" in source: + raise DatasetMaterializationError( + f"Unsupported unmanaged dataset URI: {source!r}." + ) + return source + + reference = _parse_unmanaged_hf_reference(source, version=version) + destination = data_dir / Path(reference.path).name + repo_types: list[Literal["model", "dataset"]] = ( + [repo_type] if repo_type is not None else ["model", "dataset"] + ) + for index, candidate_repo_type in enumerate(repo_types): + try: + downloaded = _download_to_temp( + _hf_download_url( + repo_id=reference.repo_id, + repo_type=candidate_repo_type, + path=reference.path, + revision=reference.revision, + ), + destination=destination, + source_description=f"unmanaged dataset {source!r}", + session=session, + ) + except _DatasetNotFoundError: + if index + 1 < len(repo_types): + continue + raise + destination.parent.mkdir(parents=True, exist_ok=True) + os.replace(downloaded, destination) + return str(destination) + + raise DatasetMaterializationError(f"Could not materialize dataset {source!r}.") + + +def _parse_unmanaged_hf_reference( + uri: str, + *, + version: Optional[str], +) -> _UnmanagedHFReference: + path_with_repo, uri_revision = ( + uri[5:].rsplit("@", maxsplit=1) if "@" in uri[5:] else (uri[5:], None) + ) + if uri_revision is not None and version is not None and uri_revision != version: + raise DatasetMaterializationError( + "Conflicting dataset versions: " + f"URI requests {uri_revision!r} but version is {version!r}." + ) + parts = path_with_repo.split("/", maxsplit=2) + if len(parts) != 3 or not all(parts): + raise DatasetMaterializationError( + "Invalid Hugging Face dataset URI. Expected format " + f"'hf://owner/repo/path/to/file[@revision]', got {uri!r}." + ) + return _UnmanagedHFReference( + repo_id=f"{parts[0]}/{parts[1]}", + path=parts[2], + revision=uri_revision or version or "main", + ) + + +def _hf_download_url( + *, + repo_id: str, + repo_type: Literal["model", "dataset"], + path: str, + revision: str, +) -> str: + prefix = "datasets/" if repo_type == "dataset" else "" + return ( + f"https://huggingface.co/{prefix}{repo_id}/resolve/" + f"{quote(revision, safe='')}/{quote(path)}" + ) + + +def _download_to_temp( + url: str, + *, + destination: Path, + source_description: str, + session=requests, +) -> Path: + destination.parent.mkdir(parents=True, exist_ok=True) + suffix = destination.suffix or ".download" + file_descriptor, temp_name = tempfile.mkstemp( + prefix=".policyengine-download-", + suffix=suffix, + dir=destination.parent, + ) + os.close(file_descriptor) + temp_path = Path(temp_name) + try: + with session.get( + url, + headers=_hugging_face_auth_headers(), + stream=True, + timeout=DOWNLOAD_TIMEOUT_SECONDS, + ) as response: + if response.status_code in {401, 403}: + raise DatasetMaterializationError( + f"Could not download {source_description}: Hugging Face " + "rejected the configured credentials. Set HUGGING_FACE_TOKEN " + "to a token with access to the certified repository." + ) + if response.status_code == 404: + raise _DatasetNotFoundError( + f"Could not find {source_description} at {url}." + ) + response.raise_for_status() + with temp_path.open("wb") as output: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + output.write(chunk) + except Exception: + temp_path.unlink(missing_ok=True) + raise + return temp_path + + +def _hugging_face_auth_headers() -> dict[str, str]: + token = ( + os.environ.get("HUGGING_FACE_TOKEN") + or os.environ.get("HF_TOKEN") + or os.environ.get("HUGGINGFACE_HUB_TOKEN") + ) + return {"Authorization": f"Bearer {token}"} if token else {} + + +def _backup_existing(path: Path) -> None: + if not path.exists(): + return + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") + backup_dir = path.parent / BACKUP_DIR_NAME / timestamp + backup_dir.mkdir(parents=True, exist_ok=True) + shutil.move(str(path), str(backup_dir / path.name)) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 957b018a..5e81b9f9 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -1,3 +1,4 @@ +import hashlib from pathlib import Path import pytest @@ -7,6 +8,8 @@ DatasetMaterializationError, MaterializedDataset, _dataset_package_strategy, + materialize_bundle_dataset, + materialize_unmanaged_dataset_source, resolve_bundle_dataset_plan, ) from policyengine.provenance.manifest import CountryReleaseManifest @@ -84,3 +87,201 @@ def test_bundle_dataset_models_round_trip_json(): def test_unknown_data_package_is_rejected(): with pytest.raises(DatasetMaterializationError, match="Unsupported bundle"): _dataset_package_strategy("unknown-data") + + +def _sha256(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +class _Response: + def __init__(self, payload: bytes = b"dataset", status_code: int = 200): + self.payload = payload + self.status_code = status_code + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + return False + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + def iter_content(self, chunk_size): + yield self.payload + + +class _Session: + def __init__(self, *responses: _Response): + self.responses = list(responses or [_Response()]) + self.calls = [] + + def get(self, url, **kwargs): + self.calls.append((url, kwargs)) + return self.responses.pop(0) + + +def _manifest_with_hash( + sha256: str, + *, + data_package_name: str = "policyengine-uk-data", + repo_type: str = "model", +) -> CountryReleaseManifest: + manifest = _manifest() + manifest.data_package.name = data_package_name + manifest.data_package.repo_type = repo_type + manifest.datasets[manifest.default_dataset].sha256 = sha256 + return manifest + + +def test_materialize_country_data_package_uses_model_repo_url(tmp_path): + session = _Session(_Response(b"country-data")) + manifest = _manifest_with_hash(_sha256(b"country-data")) + + result = materialize_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest, session=session + ) + + assert result.path.read_bytes() == b"country-data" + assert result.cache_hit is False + assert session.calls[0][0].startswith( + "https://huggingface.co/policyengine/policyengine-uk-data-private/resolve/" + ) + + +def test_materialize_populace_package_uses_dataset_repo_url(tmp_path): + session = _Session(_Response(b"populace-data")) + manifest = _manifest_with_hash( + _sha256(b"populace-data"), + data_package_name="populace-data", + repo_type="dataset", + ) + + result = materialize_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest, session=session + ) + + assert result.path.read_bytes() == b"populace-data" + assert session.calls[0][0].startswith( + "https://huggingface.co/datasets/policyengine/" + ) + + +def test_materialize_reuses_only_hash_verified_cache(tmp_path): + payload = b"certified" + manifest = _manifest_with_hash(_sha256(payload)) + destination = tmp_path / "enhanced_frs_2024_25.h5" + destination.write_bytes(payload) + session = _Session() + + result = materialize_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest, session=session + ) + + assert result.cache_hit is True + assert session.calls == [] + + +def test_materialize_replaces_and_backs_up_mismatched_cache(tmp_path): + payload = b"certified" + manifest = _manifest_with_hash(_sha256(payload)) + destination = tmp_path / "enhanced_frs_2024_25.h5" + destination.write_bytes(b"old") + + materialize_bundle_dataset( + "uk", + data_dir=tmp_path, + manifest=manifest, + session=_Session(_Response(payload)), + ) + + assert destination.read_bytes() == payload + backups = list( + (tmp_path / ".policyengine-bundle-backups").glob("*/enhanced_frs_2024_25.h5") + ) + assert len(backups) == 1 + assert backups[0].read_bytes() == b"old" + + +def test_hash_failure_does_not_replace_existing_cache(tmp_path): + manifest = _manifest_with_hash(_sha256(b"expected")) + destination = tmp_path / "enhanced_frs_2024_25.h5" + destination.write_bytes(b"old") + + with pytest.raises(DatasetMaterializationError, match="sha256"): + materialize_bundle_dataset( + "uk", + data_dir=tmp_path, + manifest=manifest, + session=_Session(_Response(b"wrong")), + ) + + assert destination.read_bytes() == b"old" + assert not (tmp_path / ".policyengine-bundle-backups").exists() + + +def test_materialize_passes_hugging_face_token(monkeypatch, tmp_path): + monkeypatch.setenv("HUGGING_FACE_TOKEN", "secret-token") + payload = b"certified" + session = _Session(_Response(payload)) + + materialize_bundle_dataset( + "uk", + data_dir=tmp_path, + manifest=_manifest_with_hash(_sha256(payload)), + session=session, + ) + + assert session.calls[0][1]["headers"] == {"Authorization": "Bearer secret-token"} + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_managed_auth_failure_does_not_retry_repo_type(tmp_path, status_code): + session = _Session(_Response(status_code=status_code)) + + with pytest.raises(DatasetMaterializationError, match="credentials"): + materialize_bundle_dataset( + "uk", + data_dir=tmp_path, + manifest=_manifest_with_hash(_sha256(b"certified")), + session=session, + ) + + assert len(session.calls) == 1 + + +def test_unmanaged_hf_retries_dataset_repo_only_after_not_found(tmp_path): + session = _Session(_Response(status_code=404), _Response(b"dataset")) + + result = materialize_unmanaged_dataset_source( + "hf://policyengine/example/data.h5@release", + data_dir=tmp_path, + session=session, + ) + + assert Path(result).read_bytes() == b"dataset" + assert "/policyengine/example/" in session.calls[0][0] + assert "/datasets/policyengine/example/" in session.calls[1][0] + + +def test_unmanaged_auth_failure_does_not_retry(tmp_path): + session = _Session(_Response(status_code=403)) + + with pytest.raises(DatasetMaterializationError, match="credentials"): + materialize_unmanaged_dataset_source( + "hf://policyengine/example/data.h5@release", + data_dir=tmp_path, + session=session, + ) + + assert len(session.calls) == 1 + + +def test_unmanaged_local_path_is_preserved(): + assert materialize_unmanaged_dataset_source("/tmp/custom.h5") == ("/tmp/custom.h5") + + +def test_unmanaged_gcs_source_is_rejected(): + with pytest.raises(DatasetMaterializationError, match="no longer supported"): + materialize_unmanaged_dataset_source("gs://bucket/data.h5@release") From adfc9a8c0005c6fea73492a197f7d472032abab3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:32:40 +0400 Subject: [PATCH 03/18] Route bundle installs through dataset materialization --- src/policyengine/bundle.py | 476 +++++++++---------------------------- tests/test_bundle.py | 138 ++++------- 2 files changed, 169 insertions(+), 445 deletions(-) diff --git a/src/policyengine/bundle.py b/src/policyengine/bundle.py index eeb45376..02f9dec6 100644 --- a/src/policyengine/bundle.py +++ b/src/policyengine/bundle.py @@ -8,30 +8,36 @@ from __future__ import annotations -import hashlib import json import os import shutil import subprocess -import tempfile import venv as venv_module -from dataclasses import dataclass from datetime import datetime, timezone from importlib import metadata from importlib.resources import files from pathlib import Path from typing import Any, Iterable, Mapping, Optional, Sequence -from urllib.parse import quote import requests +from policyengine.provenance.dataset_materialization import ( + BACKUP_DIR_NAME, + BundleDatasetPlan, + DatasetMaterializationError, + MaterializedDataset, + _sha256_file, + materialize_bundle_dataset, + resolve_bundle_dataset_plan, +) +from policyengine.provenance.manifest import CountryReleaseManifest + BUNDLE_MANIFEST_RESOURCE = ("data", "bundle", "manifest.json") BUNDLE_HISTORY_RESOURCE = ("data", "bundles") DEFAULT_COUNTRIES = ("us", "uk") DEFAULT_DATA_DIR = Path("./data") DEFAULT_VENV = Path(".venv") RECEIPT_FILENAME = ".policyengine-bundle-receipt.json" -BACKUP_DIR_NAME = ".policyengine-bundle-backups" DOWNLOAD_TIMEOUT_SECONDS = 60 @@ -39,168 +45,6 @@ class BundleError(ValueError): """Raised when bundle metadata or local installation state is invalid.""" -@dataclass(frozen=True) -class DatasetPlan: - country: str - dataset: str - uri: str - filename: str - data_version: Optional[str] - release_manifest_uri: Optional[str] - data_producer: str - repo_type: str - destination: Path - expected_sha256: Optional[str] - build_id: Optional[str] - - -class DataProducerRuntimeStrategy: - """Runtime install/verification behavior for a certified data producer.""" - - data_producer = "legacy" - - def dataset_plan( - self, - *, - country: str, - release: Mapping[str, Any], - data_dir: Path, - ) -> Optional[DatasetPlan]: - uri = release.get("default_dataset_uri") - dataset = release.get("default_dataset") - if not uri or not dataset: - return None - data_package = release.get("data_package", {}) - repo_type = ( - data_package.get("repo_type", "model") - if isinstance(data_package, Mapping) - else "model" - ) - dataset_name = str(dataset) - filename = _filename_from_uri(str(uri)) - return DatasetPlan( - country=country, - dataset=dataset_name, - uri=str(uri), - filename=filename, - data_version=( - str(release["version"]) if release.get("version") is not None else None - ), - release_manifest_uri=( - str(release["release_manifest_uri"]) - if release.get("release_manifest_uri") - else None - ), - data_producer=str(release.get("data_producer") or self.data_producer), - repo_type=str(repo_type), - destination=data_dir / filename, - expected_sha256=self.expected_sha256(release, dataset_name), - build_id=self.build_id(release), - ) - - def expected_sha256( - self, - release: Mapping[str, Any], - dataset: str, - ) -> Optional[str]: - dataset_artifact = self.default_dataset_artifact(release, dataset) - if dataset_artifact is not None and dataset_artifact.get("sha256"): - return str(dataset_artifact["sha256"]) - certified_artifact = release.get("certified_data_artifact") - if isinstance(certified_artifact, Mapping) and certified_artifact.get("sha256"): - return str(certified_artifact["sha256"]) - return None - - def default_dataset_artifact( - self, - release: Mapping[str, Any], - dataset: str, - ) -> Optional[Mapping[str, Any]]: - datasets = release.get("datasets") - if isinstance(datasets, Mapping): - artifact = datasets.get(dataset) - if isinstance(artifact, Mapping): - return artifact - return None - - def build_id(self, release: Mapping[str, Any]) -> Optional[str]: - build_id = release.get("build_id") - if build_id: - return str(build_id) - certified_artifact = release.get("certified_data_artifact") - if isinstance(certified_artifact, Mapping) and certified_artifact.get( - "build_id" - ): - return str(certified_artifact["build_id"]) - version = release.get("version") - return str(version) if version is not None else None - - def verify_download(self, plan: DatasetPlan, path: Path) -> str: - actual_sha256 = _sha256_file(path) - if plan.expected_sha256 and actual_sha256 != plan.expected_sha256: - raise BundleError( - f"Downloaded {plan.country.upper()} dataset {plan.dataset} " - f"has sha256 {actual_sha256}, expected {plan.expected_sha256}." - ) - return actual_sha256 - - def dataset_check( - self, - plan: DatasetPlan, - receipt_dataset: Optional[Mapping[str, Any]], - ) -> dict[str, Any]: - check: dict[str, Any] = { - "country": plan.country, - "dataset": plan.dataset, - "expected_version": plan.data_version, - "expected_path": str(plan.destination), - } - if plan.expected_sha256: - check["expected_sha256"] = plan.expected_sha256 - if receipt_dataset is None: - check["status"] = "missing_receipt" - return check - if receipt_dataset.get("version") != plan.data_version: - check["status"] = "mismatch" - check["installed_version"] = receipt_dataset.get("version") - return check - path = Path(str(receipt_dataset.get("path", plan.destination))) - if not path.exists(): - check["status"] = "missing_file" - return check - check["installed_version"] = receipt_dataset.get("version") - check["path"] = str(path) - if plan.expected_sha256: - actual_sha256 = _sha256_file(path) - check["installed_sha256"] = actual_sha256 - if actual_sha256 != plan.expected_sha256: - check["status"] = "sha256_mismatch" - return check - check["status"] = "ok" - return check - - -class LegacyDataProducerRuntimeStrategy(DataProducerRuntimeStrategy): - data_producer = "legacy" - - -class PopulaceDataProducerRuntimeStrategy(DataProducerRuntimeStrategy): - data_producer = "populace" - - def expected_sha256( - self, - release: Mapping[str, Any], - dataset: str, - ) -> str: - expected = super().expected_sha256(release, dataset) - if not expected: - raise BundleError( - f"Populace data release for dataset {dataset!r} is missing " - "a certified sha256." - ) - return expected - - def _bundle_resource_path(): path = files("policyengine") for part in BUNDLE_MANIFEST_RESOURCE: @@ -355,13 +199,6 @@ def _requirement(component: Mapping[str, Any]) -> str: return requirement -def runtime_strategy(data_producer: Optional[str]) -> DataProducerRuntimeStrategy: - producer = data_producer or "legacy" - if producer == "populace": - return PopulaceDataProducerRuntimeStrategy() - return LegacyDataProducerRuntimeStrategy() - - def resolve_target_python( *, python: Optional[str] = None, @@ -432,84 +269,14 @@ def install_package_scaffold( subprocess.run(command, check=True) -def dataset_plans( - manifest: Optional[Mapping[str, Any]] = None, - *, - countries: Optional[Sequence[str]] = None, - data_dir: Path = DEFAULT_DATA_DIR, -) -> list[DatasetPlan]: - bundle = _normalise_manifest(manifest or get_current_bundle()) - releases = bundle.get("data_releases") or _data_releases_from_countries(bundle) - plans: list[DatasetPlan] = [] - for country in normalise_countries(countries, bundle): - release = releases.get(country, {}) if isinstance(releases, Mapping) else {} - strategy = runtime_strategy(str(release.get("data_producer") or "legacy")) - plan = strategy.dataset_plan( - country=country, release=release, data_dir=data_dir - ) - if plan is None: - continue - plans.append(plan) - return plans - - -def _filename_from_uri(uri: str) -> str: - without_revision = uri.rsplit("@", 1)[0] - if without_revision.startswith("hf://"): - return ( - without_revision.removeprefix("hf://").split("/", 2)[2].rsplit("/", 1)[-1] - ) - if without_revision.startswith("gs://"): - return ( - without_revision.removeprefix("gs://").split("/", 1)[1].rsplit("/", 1)[-1] - ) - return Path(without_revision).name - - -def install_datasets( - manifest: Mapping[str, Any], - *, - countries: Optional[Sequence[str]] = None, - data_dir: Path = DEFAULT_DATA_DIR, - yes: bool = False, - dry_run: bool = False, - session=requests, -) -> list[dict[str, Any]]: - plans = dataset_plans(manifest, countries=countries, data_dir=data_dir) - if not plans: - return [] - _confirm_dataset_install(plans, data_dir=data_dir, yes=yes, dry_run=dry_run) - installed = [] - for plan in plans: - if dry_run: - print(f"download {plan.uri} -> {plan.destination}") - installed.append(_receipt_dataset(plan)) - continue - downloaded = _download_to_temp(plan, data_dir=data_dir, session=session) - installed_sha256 = None - try: - installed_sha256 = runtime_strategy(plan.data_producer).verify_download( - plan, - downloaded, - ) - _backup_existing(plan.destination) - plan.destination.parent.mkdir(parents=True, exist_ok=True) - shutil.move(str(downloaded), str(plan.destination)) - finally: - if downloaded.exists(): - downloaded.unlink() - installed.append(_receipt_dataset(plan, installed_sha256=installed_sha256)) - return installed - - def _confirm_dataset_install( - plans: Sequence[DatasetPlan], + plans: Sequence[BundleDatasetPlan], *, data_dir: Path, yes: bool, dry_run: bool, ) -> None: - countries = ", ".join(plan.country for plan in plans) + countries = ", ".join(plan.country_id for plan in plans) print( "This will download certified PolicyEngine datasets for " f"{countries} into {data_dir}." @@ -525,114 +292,60 @@ def _confirm_dataset_install( raise BundleError("Dataset installation cancelled.") -def _download_to_temp(plan: DatasetPlan, *, data_dir: Path, session=requests) -> Path: - data_dir.mkdir(parents=True, exist_ok=True) - url = _download_url(plan.uri, repo_type=plan.repo_type) - headers = _auth_headers(plan.uri) - suffix = Path(plan.filename).suffix or ".download" - fd, temp_name = tempfile.mkstemp( - prefix=".policyengine-download-", suffix=suffix, dir=data_dir - ) - os.close(fd) - temp_path = Path(temp_name) - try: - with session.get( - url, - headers=headers, - stream=True, - timeout=DOWNLOAD_TIMEOUT_SECONDS, - ) as response: - if response.status_code in {401, 403}: - raise BundleError( - f"Could not download {plan.country.upper()} dataset. " - "If this is a private Hugging Face dataset, set HUGGING_FACE_TOKEN." - ) - response.raise_for_status() - with temp_path.open("wb") as stream: - for chunk in response.iter_content(chunk_size=1024 * 1024): - if chunk: - stream.write(chunk) - except Exception: - if temp_path.exists(): - temp_path.unlink() - raise - return temp_path - - -def _download_url(uri: str, *, repo_type: str = "model") -> str: - without_revision, revision = _split_revision(uri) - if without_revision.startswith("hf://"): - parts = without_revision.removeprefix("hf://").split("/", 2) - if len(parts) != 3: - raise BundleError(f"Invalid Hugging Face dataset URI: {uri}") - repo_id = f"{parts[0]}/{parts[1]}" - path = parts[2] - if not revision: - raise BundleError(f"Hugging Face dataset URI must pin a revision: {uri}") - prefix = "datasets/" if repo_type == "dataset" else "" - return ( - f"https://huggingface.co/{prefix}{repo_id}/resolve/{quote(revision)}/{path}" - ) - if without_revision.startswith("gs://"): - bucket_and_path = without_revision.removeprefix("gs://") - bucket, _, path = bucket_and_path.partition("/") - return f"https://storage.googleapis.com/{bucket}/{quote(path)}" - if without_revision.startswith(("http://", "https://")): - return uri - return uri - - -def _split_revision(uri: str) -> tuple[str, Optional[str]]: - if "@" not in uri: - return uri, None - without_revision, revision = uri.rsplit("@", 1) - return without_revision, revision - - -def _auth_headers(uri: str) -> dict[str, str]: - if not uri.startswith(("hf://", "https://huggingface.co/")): - return {} - token = ( - os.environ.get("HUGGING_FACE_TOKEN") - or os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_HUB_TOKEN") - ) - return {"Authorization": f"Bearer {token}"} if token else {} - - -def _backup_existing(path: Path) -> None: - if not path.exists(): - return - timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") - backup_dir = path.parent / BACKUP_DIR_NAME / timestamp - backup_dir.mkdir(parents=True, exist_ok=True) - shutil.move(str(path), str(backup_dir / path.name)) - - def _receipt_dataset( - plan: DatasetPlan, + plan: BundleDatasetPlan, + release: Mapping[str, Any], *, - installed_sha256: Optional[str] = None, + materialized: Optional[MaterializedDataset] = None, ) -> dict[str, Any]: receipt = { - "country": plan.country, + "country": plan.country_id, "dataset": plan.dataset, - "version": plan.data_version, - "uri": plan.uri, + "version": release.get("version") or plan.build_id, + "uri": plan.source_uri, "path": str(plan.destination), - "release_manifest_uri": plan.release_manifest_uri, - "data_producer": plan.data_producer, + "release_manifest_uri": release.get("release_manifest_uri"), + "data_package_name": plan.data_package_name, "repo_type": plan.repo_type, } if plan.build_id: receipt["build_id"] = plan.build_id - if plan.expected_sha256: - receipt["expected_sha256"] = plan.expected_sha256 - if installed_sha256: - receipt["installed_sha256"] = installed_sha256 + receipt["expected_sha256"] = plan.expected_sha256 + if materialized is not None: + receipt["installed_sha256"] = materialized.actual_sha256 return receipt +def _selected_dataset_plans( + manifest: Mapping[str, Any], + countries: Sequence[str], + *, + data_dir: Path, +) -> list[tuple[BundleDatasetPlan, CountryReleaseManifest, Mapping[str, Any]]]: + releases = manifest.get("data_releases") + if not isinstance(releases, Mapping): + raise BundleError("Bundle manifest does not contain data releases.") + + selected = [] + for country in countries: + release = releases.get(country) + if not isinstance(release, Mapping): + raise BundleError( + f"Bundle manifest does not contain a {country.upper()} data release." + ) + try: + country_manifest = CountryReleaseManifest.model_validate(release) + plan = resolve_bundle_dataset_plan( + country, + data_dir=data_dir, + manifest=country_manifest, + ) + except (ValueError, DatasetMaterializationError) as exc: + raise BundleError(str(exc)) from exc + selected.append((plan, country_manifest, release)) + return selected + + def write_receipt( manifest: Mapping[str, Any], *, @@ -692,13 +405,34 @@ def install_bundle( install_package_scaffold(target_python, requirements, dry_run=dry_run) installed_datasets: list[dict[str, Any]] = [] if not no_datasets: - installed_datasets = install_datasets( - manifest, - countries=selected_countries, - data_dir=data_dir, - yes=yes, - dry_run=dry_run, + dataset_entries = _selected_dataset_plans( + manifest, selected_countries, data_dir=data_dir ) + plans = [entry[0] for entry in dataset_entries] + if plans: + _confirm_dataset_install( + plans, + data_dir=data_dir, + yes=yes, + dry_run=dry_run, + ) + for plan, country_manifest, release in dataset_entries: + if dry_run: + print(f"download {plan.source_uri} -> {plan.destination}") + installed_datasets.append(_receipt_dataset(plan, release)) + continue + try: + materialized = materialize_bundle_dataset( + plan.country_id, + plan.dataset, + data_dir=data_dir, + manifest=country_manifest, + ) + except DatasetMaterializationError as exc: + raise BundleError(str(exc)) from exc + installed_datasets.append( + _receipt_dataset(plan, release, materialized=materialized) + ) if not dry_run: write_receipt( manifest, @@ -904,17 +638,43 @@ def _dataset_checks( if isinstance(dataset, Mapping) and dataset.get("country"): receipt_datasets[str(dataset["country"])] = dataset checks = [] - for plan in dataset_plans(manifest, countries=countries, data_dir=data_dir): - receipt_dataset = receipt_datasets.get(plan.country) - checks.append( - runtime_strategy(plan.data_producer).dataset_check(plan, receipt_dataset) - ) + for plan, _, release in _selected_dataset_plans( + manifest, countries, data_dir=data_dir + ): + receipt_dataset = receipt_datasets.get(plan.country_id) + checks.append(_dataset_check(plan, release, receipt_dataset)) return checks -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() +def _dataset_check( + plan: BundleDatasetPlan, + release: Mapping[str, Any], + receipt_dataset: Optional[Mapping[str, Any]], +) -> dict[str, Any]: + expected_version = release.get("version") or plan.build_id + check: dict[str, Any] = { + "country": plan.country_id, + "dataset": plan.dataset, + "expected_version": expected_version, + "expected_path": str(plan.destination), + "expected_sha256": plan.expected_sha256, + } + if receipt_dataset is None: + check["status"] = "missing_receipt" + return check + if receipt_dataset.get("version") != expected_version: + check["status"] = "mismatch" + check["installed_version"] = receipt_dataset.get("version") + return check + path = Path(str(receipt_dataset.get("path", plan.destination))) + if not path.exists(): + check["status"] = "missing_file" + return check + actual_sha256 = _sha256_file(path) + check["installed_version"] = receipt_dataset.get("version") + check["installed_sha256"] = actual_sha256 + check["path"] = str(path) + check["status"] = ( + "ok" if actual_sha256 == plan.expected_sha256 else "sha256_mismatch" + ) + return check diff --git a/tests/test_bundle.py b/tests/test_bundle.py index 8e26cd86..f6806c68 100644 --- a/tests/test_bundle.py +++ b/tests/test_bundle.py @@ -7,6 +7,7 @@ from policyengine import bundle from policyengine.cli import main as cli_main +from policyengine.provenance.dataset_materialization import MaterializedDataset def _sha256(payload: bytes) -> str: @@ -67,43 +68,17 @@ def test_bundle_install_requirements_are_country_scoped(): assert not any("policyengine-us-data" in req for req in us_requirements) -def test_dataset_plans_use_certified_release_metadata(tmp_path): - plans = bundle.dataset_plans( - bundle.get_current_bundle(), - countries=["uk"], - data_dir=tmp_path, - ) - - assert len(plans) == 1 - assert plans[0].country == "uk" - assert plans[0].data_version.startswith("policyengine-uk-data-") - assert plans[0].data_producer == "populace" - assert plans[0].repo_type == "model" - assert plans[0].destination == tmp_path / "enhanced_frs_2024_25.h5" - assert ( - plans[0].expected_sha256 - == bundle.get_current_bundle()["data_releases"]["uk"]["datasets"][ - "enhanced_frs_2024_25" - ]["sha256"] +def test_selected_dataset_plan_uses_certified_release_metadata(tmp_path): + entries = bundle._selected_dataset_plans( + bundle.get_current_bundle(), ["uk"], data_dir=tmp_path ) - -def test_runtime_strategy_selects_populace(): - assert isinstance( - bundle.runtime_strategy("populace"), - bundle.PopulaceDataProducerRuntimeStrategy, - ) - - -def test_populace_runtime_strategy_requires_certified_hash(): - manifest = json.loads(json.dumps(bundle.get_current_bundle())) - release = manifest["data_releases"]["uk"] - dataset = release["default_dataset"] - del release["datasets"][dataset]["sha256"] - del release["certified_data_artifact"]["sha256"] - - with pytest.raises(bundle.BundleError, match="certified sha256"): - bundle.dataset_plans(manifest, countries=["uk"]) + plan, _, release = entries[0] + assert plan.country_id == "uk" + assert plan.data_package_name == "policyengine-uk-data" + assert plan.repo_type == "model" + assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" + assert plan.expected_sha256 == release["datasets"][plan.dataset]["sha256"] def test_install_bundle_package_only_uses_explicit_python(monkeypatch, tmp_path): @@ -167,66 +142,55 @@ def test_resolve_target_python_uses_local_venv_from_runner_env(monkeypatch, tmp_ assert not (tmp_path / ".venv").exists() -class FakeResponse: - status_code = 200 - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def raise_for_status(self): - return None - - def iter_content(self, chunk_size): - yield b"new-data" - +def test_install_bundle_materializes_defaults_and_records_receipt( + monkeypatch, tmp_path +): + calls = [] -class FakeSession: - def get(self, *args, **kwargs): - return FakeResponse() + monkeypatch.setattr( + bundle, "install_package_scaffold", lambda *args, **kwargs: None + ) + def fake_materialize(country_id, dataset, *, data_dir, manifest): + plan = bundle.resolve_bundle_dataset_plan( + country_id, + dataset, + data_dir=data_dir, + manifest=manifest, + ) + calls.append(plan) + plan.destination.parent.mkdir(parents=True, exist_ok=True) + plan.destination.write_bytes(b"materialized") + return MaterializedDataset( + country_id=plan.country_id, + dataset=plan.dataset, + data_package_name=plan.data_package_name, + repo_id=plan.repo_id, + repo_type=plan.repo_type, + revision=plan.revision, + source_uri=plan.source_uri, + expected_sha256=plan.expected_sha256, + actual_sha256=plan.expected_sha256, + path=plan.destination, + cache_hit=False, + build_id=plan.build_id, + ) -def test_install_datasets_downloads_then_backs_up_existing_file(tmp_path): - manifest = _manifest_with_dataset_sha("us", _sha256(b"new-data")) - existing = tmp_path / "populace_us_2024.h5" - existing.write_bytes(b"old-data") + monkeypatch.setattr(bundle, "materialize_bundle_dataset", fake_materialize) - installed = bundle.install_datasets( - manifest, - countries=["us"], + result = bundle.install_bundle( + python=sys.executable, + countries=["uk"], data_dir=tmp_path, yes=True, - session=FakeSession(), ) - assert installed[0]["country"] == "us" - assert installed[0]["expected_sha256"] == _sha256(b"new-data") - assert installed[0]["installed_sha256"] == _sha256(b"new-data") - assert installed[0]["build_id"] == manifest["data_releases"]["us"]["build_id"] - assert existing.read_bytes() == b"new-data" - backups = list((tmp_path / bundle.BACKUP_DIR_NAME).glob("*/populace_us_2024.h5")) - assert len(backups) == 1 - assert backups[0].read_bytes() == b"old-data" - - -def test_install_datasets_rejects_downloaded_hash_mismatch(tmp_path): - manifest = _manifest_with_dataset_sha("us", _sha256(b"expected-data")) - existing = tmp_path / "populace_us_2024.h5" - existing.write_bytes(b"old-data") - - with pytest.raises(bundle.BundleError, match="sha256"): - bundle.install_datasets( - manifest, - countries=["us"], - data_dir=tmp_path, - yes=True, - session=FakeSession(), - ) - - assert existing.read_bytes() == b"old-data" - assert not (tmp_path / bundle.BACKUP_DIR_NAME).exists() + assert [plan.country_id for plan in calls] == ["uk"] + assert result["datasets"][0]["data_package_name"] == "policyengine-uk-data" + assert result["datasets"][0]["installed_sha256"] == calls[0].expected_sha256 + receipt = bundle.read_receipt(tmp_path) + assert receipt is not None + assert receipt["datasets"] == result["datasets"] def test_status_matches_receipt_and_packages(monkeypatch, tmp_path): From bd63ab34cea1643eaf0d3c445044643f4b75346b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:42:41 +0400 Subject: [PATCH 04/18] Route model datasets through bundle materialization --- .../provenance/dataset_materialization.py | 113 +++++++++++++++++- src/policyengine/provenance/manifest.py | 9 +- .../tax_benefit_models/uk/datasets.py | 38 +++++- .../tax_benefit_models/uk/model.py | 61 +++++++--- .../tax_benefit_models/us/datasets.py | 88 +++++++++----- .../tax_benefit_models/us/model.py | 61 +++++++--- tests/test_dataset_materialization.py | 60 ++++++++++ tests/test_dataset_sources.py | 46 +++++-- tests/test_release_manifests.py | 59 ++++++++- tests/test_us_long_term_datasets.py | 104 +++++++++------- 10 files changed, 507 insertions(+), 132 deletions(-) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 5aff298a..02a351d7 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -19,6 +19,7 @@ _artifact_revision, build_hf_uri, get_release_manifest, + resolve_local_managed_dataset_source, ) DEFAULT_DATA_DIR = Path("./data") @@ -44,6 +45,9 @@ class BundleDatasetPlan(BaseModel): source_uri: str destination: Path build_id: Optional[str] = None + metadata_expected_sha256: Optional[str] = None + metadata_source_uri: Optional[str] = None + metadata_destination: Optional[Path] = None class MaterializedDataset(BaseModel): @@ -61,6 +65,9 @@ class MaterializedDataset(BaseModel): path: Path cache_hit: bool build_id: Optional[str] = None + metadata_expected_sha256: Optional[str] = None + metadata_actual_sha256: Optional[str] = None + metadata_path: Optional[Path] = None class _DatasetPackageStrategy: @@ -179,6 +186,17 @@ def resolve_bundle_dataset_plan( if country_manifest.certified_data_artifact is not None else None ), + metadata_expected_sha256=reference.metadata_sha256, + metadata_source_uri=( + build_hf_uri(repo_id, f"{reference.path}.metadata.json", revision) + if reference.metadata_sha256 + else None + ), + metadata_destination=( + data_dir / f"{Path(reference.path).name}.metadata.json" + if reference.metadata_sha256 + else None + ), ) _dataset_package_strategy(data_package_name).validate(plan) return plan @@ -190,6 +208,7 @@ def materialize_bundle_dataset( *, data_dir: Path = DEFAULT_DATA_DIR, manifest: Optional[CountryReleaseManifest] = None, + allow_local_mirror: bool = True, session=requests, ) -> MaterializedDataset: """Download and verify one dataset certified by the release bundle.""" @@ -200,6 +219,35 @@ def materialize_bundle_dataset( data_dir=data_dir, manifest=manifest, ) + local_source = resolve_local_managed_dataset_source( + country_id, + plan.source_uri, + allow_local_mirror=allow_local_mirror, + ) + if local_source != plan.source_uri: + local_path = Path(local_source).expanduser() + if local_path.is_file(): + actual_sha256 = _sha256_file(local_path) + if actual_sha256 == plan.expected_sha256: + metadata_path = Path(f"{local_path}.metadata.json") + if plan.metadata_expected_sha256 is None: + return _materialized_result( + plan, + actual_sha256=actual_sha256, + cache_hit=True, + path=local_path, + ) + if metadata_path.is_file(): + metadata_actual_sha256 = _sha256_file(metadata_path) + if metadata_actual_sha256 == plan.metadata_expected_sha256: + return _materialized_result( + plan, + actual_sha256=actual_sha256, + cache_hit=True, + path=local_path, + metadata_actual_sha256=metadata_actual_sha256, + metadata_path=metadata_path, + ) strategy = _dataset_package_strategy(plan.data_package_name) return strategy.materialize(plan, session=session) @@ -229,10 +277,16 @@ def _materialize_managed_hf_dataset( if destination.is_file(): actual_sha256 = _sha256_file(destination) if actual_sha256 == plan.expected_sha256: + metadata_actual_sha256, metadata_path = _materialize_metadata( + plan, + session=session, + ) return _materialized_result( plan, actual_sha256=actual_sha256, cache_hit=True, + metadata_actual_sha256=metadata_actual_sha256, + metadata_path=metadata_path, ) url = _hf_download_url( @@ -260,10 +314,16 @@ def _materialize_managed_hf_dataset( finally: downloaded.unlink(missing_ok=True) + metadata_actual_sha256, metadata_path = _materialize_metadata( + plan, + session=session, + ) return _materialized_result( plan, actual_sha256=actual_sha256, cache_hit=False, + metadata_actual_sha256=metadata_actual_sha256, + metadata_path=metadata_path, ) @@ -272,6 +332,9 @@ def _materialized_result( *, actual_sha256: str, cache_hit: bool, + path: Optional[Path] = None, + metadata_actual_sha256: Optional[str] = None, + metadata_path: Optional[Path] = None, ) -> MaterializedDataset: return MaterializedDataset( country_id=plan.country_id, @@ -283,12 +346,60 @@ def _materialized_result( source_uri=plan.source_uri, expected_sha256=plan.expected_sha256, actual_sha256=actual_sha256, - path=plan.destination, + path=path or plan.destination, cache_hit=cache_hit, build_id=plan.build_id, + metadata_expected_sha256=plan.metadata_expected_sha256, + metadata_actual_sha256=metadata_actual_sha256, + metadata_path=metadata_path, ) +def _materialize_metadata( + plan: BundleDatasetPlan, + *, + session=requests, +) -> tuple[Optional[str], Optional[Path]]: + if plan.metadata_expected_sha256 is None: + return None, None + if plan.metadata_destination is None: + raise DatasetMaterializationError( + f"Managed dataset {plan.dataset!r} has a metadata hash but no " + "metadata destination." + ) + + destination = plan.metadata_destination + if destination.is_file(): + actual_sha256 = _sha256_file(destination) + if actual_sha256 == plan.metadata_expected_sha256: + return actual_sha256, destination + + url = _hf_download_url( + repo_id=plan.repo_id, + repo_type=plan.repo_type, + path=f"{plan.path}.metadata.json", + revision=plan.revision, + ) + downloaded = _download_to_temp( + url, + destination=destination, + source_description=f"metadata for {plan.dataset!r}", + session=session, + ) + try: + actual_sha256 = _sha256_file(downloaded) + if actual_sha256 != plan.metadata_expected_sha256: + raise DatasetMaterializationError( + f"Downloaded metadata for dataset {plan.dataset!r} has sha256 " + f"{actual_sha256}, expected {plan.metadata_expected_sha256}." + ) + _backup_existing(destination) + os.replace(downloaded, destination) + finally: + downloaded.unlink(missing_ok=True) + return actual_sha256, destination + + class _UnmanagedHFReference(BaseModel): repo_id: str path: str diff --git a/src/policyengine/provenance/manifest.py b/src/policyengine/provenance/manifest.py index 5ada2d33..b2795c0a 100644 --- a/src/policyengine/provenance/manifest.py +++ b/src/policyengine/provenance/manifest.py @@ -588,7 +588,14 @@ def resolve_managed_dataset_reference( "bypass bundle enforcement." ) - return resolve_dataset_reference(country_id, dataset) + if allow_unmanaged: + return resolve_dataset_reference(country_id, dataset) + raise ValueError( + f"Unknown managed dataset {dataset!r} for country {country_id!r}. " + f"Known bundled datasets: {sorted(manifest.datasets)}. Set " + "`allow_unmanaged=True` only if you intend to resolve a dataset " + "outside the policyengine.py release bundle." + ) def resolve_local_managed_dataset_source( diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index 82522111..086c96a2 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -6,10 +6,15 @@ from pydantic import ConfigDict from policyengine.core import Dataset, YearData -from policyengine.provenance.dataset_sources import materialize_dataset_source +from policyengine.provenance.dataset_materialization import ( + materialize_bundle_dataset, + materialize_unmanaged_dataset_source, +) from policyengine.provenance.manifest import ( dataset_logical_name, + get_release_manifest, resolve_dataset_reference, + resolve_managed_dataset_reference, ) @@ -118,12 +123,31 @@ def create_datasets( ], years: list[int] = [2026, 2027, 2028, 2029, 2030], data_folder: str = "./data", + allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUKDataset]: result = {} for dataset in datasets: - resolved_dataset = resolve_dataset_reference("uk", dataset) + manifest = get_release_manifest("uk") + managed_dataset = dataset if dataset in manifest.datasets else None + if managed_dataset is not None: + materialized = materialize_bundle_dataset( + "uk", + managed_dataset, + data_dir=Path(data_folder), + ) + resolved_dataset = materialized.source_uri + runtime_dataset = str(materialized.path) + else: + resolved_dataset = resolve_managed_dataset_reference( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + ) + runtime_dataset = materialize_unmanaged_dataset_source( + resolved_dataset, + data_dir=Path(data_folder), + ) dataset_stem = dataset_logical_name(resolved_dataset) - runtime_dataset = materialize_dataset_source(resolved_dataset) from policyengine_uk import Microsimulation sim = Microsimulation(dataset=runtime_dataset) @@ -226,6 +250,7 @@ def ensure_datasets( ], years: list[int] = [2026, 2027, 2028, 2029, 2030], data_folder: str = "./data", + allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUKDataset]: """Ensure datasets exist, loading if available or creating if not. @@ -253,4 +278,9 @@ def ensure_datasets( if all_exist: return load_datasets(datasets=datasets, years=years, data_folder=data_folder) else: - return create_datasets(datasets=datasets, years=years, data_folder=data_folder) + return create_datasets( + datasets=datasets, + years=years, + data_folder=data_folder, + allow_unmanaged=allow_unmanaged, + ) diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index dfa54985..12c8b93f 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -1,14 +1,18 @@ import datetime -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional import pandas as pd from microdf import MicroDataFrame from policyengine.core import TaxBenefitModel -from policyengine.provenance.dataset_sources import materialize_dataset_source +from policyengine.provenance.dataset_materialization import ( + MaterializedDataset, + materialize_bundle_dataset, + materialize_unmanaged_dataset_source, +) from policyengine.provenance.manifest import ( dataset_logical_name, - resolve_local_managed_dataset_source, + get_release_manifest, resolve_managed_dataset_reference, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion @@ -267,12 +271,24 @@ def run(self, simulation: "Simulation") -> "Simulation": def _managed_release_bundle( dataset_uri: str, dataset_source: Optional[str] = None, -) -> dict[str, Optional[str]]: - bundle = dict(uk_latest.release_bundle) + materialized: Optional[MaterializedDataset] = None, +) -> dict[str, Any]: + bundle: dict[str, Any] = dict(uk_latest.release_bundle) bundle["runtime_dataset"] = dataset_logical_name(dataset_uri) bundle["runtime_dataset_uri"] = dataset_uri if dataset_source: bundle["runtime_dataset_source"] = dataset_source + if materialized is not None: + bundle.update( + { + "runtime_dataset_data_package": materialized.data_package_name, + "runtime_dataset_repo_type": materialized.repo_type, + "runtime_dataset_revision": materialized.revision, + "runtime_dataset_expected_sha256": materialized.expected_sha256, + "runtime_dataset_sha256": materialized.actual_sha256, + "runtime_dataset_cache_hit": materialized.cache_hit, + } + ) bundle["managed_by"] = "policyengine.py" return bundle @@ -298,19 +314,27 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - dataset_uri = resolve_managed_dataset_reference( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - ) - dataset_source = resolve_local_managed_dataset_source( - "uk", - dataset_uri, - allow_local_mirror=not ( - allow_unmanaged and dataset is not None and "://" in dataset - ), - ) - runtime_dataset_source = materialize_dataset_source(dataset_source) + manifest = get_release_manifest("uk") + managed_dataset = None + if dataset is None: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + elif dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + + materialized = None + if managed_dataset is not None: + materialized = materialize_bundle_dataset("uk", managed_dataset) + dataset_uri = materialized.source_uri + runtime_dataset_source = str(materialized.path) + else: + dataset_uri = resolve_managed_dataset_reference( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + ) + runtime_dataset_source = materialize_unmanaged_dataset_source(dataset_uri) runtime_dataset = runtime_dataset_source if isinstance(runtime_dataset_source, str) and "://" not in runtime_dataset_source: from policyengine_uk.data.dataset_schema import ( @@ -326,6 +350,7 @@ def managed_microsimulation( microsim.policyengine_bundle = _managed_release_bundle( dataset_uri, runtime_dataset_source, + materialized, ) return microsim diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 4b2dc2ba..e3691868 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -12,12 +12,15 @@ from pydantic import ConfigDict, Field from policyengine.core import Dataset, YearData -from policyengine.provenance.dataset_sources import materialize_dataset_source +from policyengine.provenance.dataset_materialization import ( + MaterializedDataset, + materialize_bundle_dataset, + materialize_unmanaged_dataset_source, +) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, - resolve_local_managed_dataset_source, resolve_managed_dataset_reference, ) @@ -279,6 +282,7 @@ def create_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2024, 2025, 2026, 2027, 2028], data_folder: str = "./data", + allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUSDataset]: """Create PolicyEngineUSDataset instances from logical dataset names or URLs. @@ -295,9 +299,27 @@ def create_datasets( datasets = datasets or [get_release_manifest("us").default_dataset] result = {} for dataset in datasets: - resolved_dataset = resolve_dataset_reference("us", dataset) + manifest = get_release_manifest("us") + managed_dataset = dataset if dataset in manifest.datasets else None + if managed_dataset is not None: + materialized = materialize_bundle_dataset( + "us", + managed_dataset, + data_dir=Path(data_folder), + ) + resolved_dataset = materialized.source_uri + runtime_dataset = str(materialized.path) + else: + resolved_dataset = resolve_managed_dataset_reference( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + ) + runtime_dataset = materialize_unmanaged_dataset_source( + resolved_dataset, + data_dir=Path(data_folder), + ) dataset_stem = dataset_logical_name(resolved_dataset) - runtime_dataset = materialize_dataset_source(resolved_dataset) sim = Microsimulation(dataset=runtime_dataset) for year in years: @@ -836,6 +858,7 @@ def _build_long_term_dataset( metadata: dict, metadata_path: Optional[Path], dataset_uri: Optional[str] = None, + materialized: Optional[MaterializedDataset] = None, ) -> PolicyEngineUSDataset: dataset = PolicyEngineUSDataset( id=_long_term_dataset_key(dataset_name, year), @@ -855,6 +878,17 @@ def _build_long_term_dataset( "runtime_dataset_uri": dataset_uri, } ) + if materialized is not None: + dataset.metadata["policyengine_bundle"].update( + { + "runtime_dataset_data_package": (materialized.data_package_name), + "runtime_dataset_repo_type": materialized.repo_type, + "runtime_dataset_revision": materialized.revision, + "runtime_dataset_expected_sha256": (materialized.expected_sha256), + "runtime_dataset_sha256": materialized.actual_sha256, + "runtime_dataset_cache_hit": materialized.cache_hit, + } + ) return dataset @@ -1013,6 +1047,7 @@ def load_long_term_datasets( def load_managed_long_term_datasets( years: list[int], dataset_name: str = "long_term_cps", + data_folder: str = "./data", require_metadata: bool = True, required_profile: Optional[str] = None, required_target_source: Optional[str] = None, @@ -1037,12 +1072,10 @@ def load_managed_long_term_datasets( """Load bundled long-term US datasets from the managed release manifest. Each requested year must have a logical dataset entry named - ``{dataset_name}_{year}`` in the bundled US manifest. For local development, - policyengine.py first checks for the corresponding sibling data-repo mirror - before falling back to the managed URI. Long-term H5 files are large, so this - helper intentionally refuses to stream remote files directly; callers should - either provide the published local mirror or use ``load_long_term_datasets`` - with an explicit local data folder. + ``{dataset_name}_{year}`` in the bundled US manifest. A verified local mirror + is reused when available; otherwise the exact bundle-certified Hugging Face + artifact and its certified metadata sidecar are materialized into + ``data_folder``. """ manifest = get_release_manifest("us") @@ -1070,24 +1103,14 @@ def load_managed_long_term_datasets( f"Managed long-term dataset {key!r} is missing a sha256 in " "the bundled US release manifest." ) - dataset_uri = resolve_managed_dataset_reference("us", key) - dataset_source = resolve_local_managed_dataset_source("us", dataset_uri) - if "://" in dataset_source: - raise FileNotFoundError( - f"Managed long-term dataset {key!r} resolves to {dataset_uri}, " - "but no local mirror exists. Download the bundled artifact into " - "the sibling policyengine-us-data storage mirror or call " - "load_long_term_datasets(..., data_folder=...) with an explicit " - "local directory." - ) - - path = Path(dataset_source).expanduser() - actual_sha256 = _sha256_file(path) - if actual_sha256 != path_reference.sha256: - raise ValueError( - f"Managed long-term dataset {key!r} at {path} has sha256 " - f"{actual_sha256}, expected {path_reference.sha256}." - ) + materialized = materialize_bundle_dataset( + "us", + key, + data_dir=Path(data_folder), + manifest=manifest, + ) + dataset_uri = materialized.source_uri + path = materialized.path metadata, metadata_path = _load_dataset_metadata(path, require_metadata) if path_reference.metadata_sha256: if metadata_path is None: @@ -1144,6 +1167,7 @@ def load_managed_long_term_datasets( metadata=metadata, metadata_path=metadata_path, dataset_uri=dataset_uri, + materialized=materialized, ) return result @@ -1153,6 +1177,7 @@ def ensure_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2024, 2025, 2026, 2027, 2028], data_folder: str = "./data", + allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUSDataset]: """Ensure datasets exist, loading if available or creating if not. @@ -1182,4 +1207,9 @@ def ensure_datasets( if all_exist: return load_datasets(datasets=datasets, years=years, data_folder=data_folder) else: - return create_datasets(datasets=datasets, years=years, data_folder=data_folder) + return create_datasets( + datasets=datasets, + years=years, + data_folder=data_folder, + allow_unmanaged=allow_unmanaged, + ) diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 090ac514..402396c3 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -1,14 +1,18 @@ import datetime -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional import pandas as pd from microdf import MicroDataFrame from policyengine.core import TaxBenefitModel -from policyengine.provenance.dataset_sources import materialize_dataset_source +from policyengine.provenance.dataset_materialization import ( + MaterializedDataset, + materialize_bundle_dataset, + materialize_unmanaged_dataset_source, +) from policyengine.provenance.manifest import ( dataset_logical_name, - resolve_local_managed_dataset_source, + get_release_manifest, resolve_managed_dataset_reference, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion @@ -410,12 +414,24 @@ def _build_simulation_from_dataset(self, microsim, dataset, system): def _managed_release_bundle( dataset_uri: str, dataset_source: Optional[str] = None, -) -> dict[str, Optional[str]]: - bundle = dict(us_latest.release_bundle) + materialized: Optional[MaterializedDataset] = None, +) -> dict[str, Any]: + bundle: dict[str, Any] = dict(us_latest.release_bundle) bundle["runtime_dataset"] = dataset_logical_name(dataset_uri) bundle["runtime_dataset_uri"] = dataset_uri if dataset_source: bundle["runtime_dataset_source"] = dataset_source + if materialized is not None: + bundle.update( + { + "runtime_dataset_data_package": materialized.data_package_name, + "runtime_dataset_repo_type": materialized.repo_type, + "runtime_dataset_revision": materialized.revision, + "runtime_dataset_expected_sha256": materialized.expected_sha256, + "runtime_dataset_sha256": materialized.actual_sha256, + "runtime_dataset_cache_hit": materialized.cache_hit, + } + ) bundle["managed_by"] = "policyengine.py" return bundle @@ -441,23 +457,32 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - dataset_uri = resolve_managed_dataset_reference( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) - dataset_source = resolve_local_managed_dataset_source( - "us", - dataset_uri, - allow_local_mirror=not ( - allow_unmanaged and dataset is not None and "://" in dataset - ), - ) - runtime_dataset_source = materialize_dataset_source(dataset_source) + manifest = get_release_manifest("us") + managed_dataset = None + if dataset is None: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + elif dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + + materialized = None + if managed_dataset is not None: + materialized = materialize_bundle_dataset("us", managed_dataset) + dataset_uri = materialized.source_uri + runtime_dataset_source = str(materialized.path) + else: + dataset_uri = resolve_managed_dataset_reference( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + ) + runtime_dataset_source = materialize_unmanaged_dataset_source(dataset_uri) microsim = Microsimulation(dataset=runtime_dataset_source, **kwargs) microsim.policyengine_bundle = _managed_release_bundle( dataset_uri, runtime_dataset_source, + materialized, ) return microsim diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 5e81b9f9..633bf537 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -168,6 +168,24 @@ def test_materialize_populace_package_uses_dataset_repo_url(tmp_path): ) +def test_materialize_downloads_and_verifies_metadata_sidecar(tmp_path): + dataset_payload = b"long-term-data" + metadata_payload = b'{"year": 2100}' + manifest = _manifest_with_hash(_sha256(dataset_payload)) + reference = manifest.datasets[manifest.default_dataset] + reference.metadata_sha256 = _sha256(metadata_payload) + session = _Session(_Response(dataset_payload), _Response(metadata_payload)) + + result = materialize_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest, session=session + ) + + assert result.metadata_path == (tmp_path / "enhanced_frs_2024_25.h5.metadata.json") + assert result.metadata_path.read_bytes() == metadata_payload + assert result.metadata_actual_sha256 == reference.metadata_sha256 + assert session.calls[1][0].endswith("/enhanced_frs_2024_25.h5.metadata.json") + + def test_materialize_reuses_only_hash_verified_cache(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) @@ -183,6 +201,48 @@ def test_materialize_reuses_only_hash_verified_cache(tmp_path): assert session.calls == [] +def test_materialize_reuses_hash_verified_local_mirror(monkeypatch, tmp_path): + payload = b"certified-local-mirror" + manifest = _manifest_with_hash(_sha256(payload)) + mirror = tmp_path / "mirror" / "enhanced_frs_2024_25.h5" + mirror.parent.mkdir() + mirror.write_bytes(payload) + monkeypatch.setattr( + "policyengine.provenance.dataset_materialization.resolve_local_managed_dataset_source", + lambda *args, **kwargs: str(mirror), + ) + session = _Session() + + result = materialize_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest, session=session + ) + + assert result.path == mirror + assert result.cache_hit is True + assert session.calls == [] + + +def test_materialize_ignores_mismatched_local_mirror(monkeypatch, tmp_path): + payload = b"certified-download" + manifest = _manifest_with_hash(_sha256(payload)) + mirror = tmp_path / "mirror" / "enhanced_frs_2024_25.h5" + mirror.parent.mkdir() + mirror.write_bytes(b"wrong") + monkeypatch.setattr( + "policyengine.provenance.dataset_materialization.resolve_local_managed_dataset_source", + lambda *args, **kwargs: str(mirror), + ) + session = _Session(_Response(payload)) + + result = materialize_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest, session=session + ) + + assert result.path == tmp_path / "enhanced_frs_2024_25.h5" + assert result.path.read_bytes() == payload + assert mirror.read_bytes() == b"wrong" + + def test_materialize_replaces_and_backs_up_mismatched_cache(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) diff --git a/tests/test_dataset_sources.py b/tests/test_dataset_sources.py index ff50865c..4b7ff960 100644 --- a/tests/test_dataset_sources.py +++ b/tests/test_dataset_sources.py @@ -7,6 +7,10 @@ import pytest from policyengine.provenance import dataset_sources +from policyengine.provenance.dataset_materialization import ( + MaterializedDataset, + resolve_bundle_dataset_plan, +) from policyengine.provenance.dataset_sources import ( materialize_dataset_source, parse_gs_uri, @@ -26,6 +30,24 @@ def _load_module_from_path(module_name: str, path: Path): return module +def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDataset: + plan = resolve_bundle_dataset_plan(country_id, dataset) + return MaterializedDataset( + country_id=plan.country_id, + dataset=plan.dataset, + data_package_name=plan.data_package_name, + repo_id=plan.repo_id, + repo_type=plan.repo_type, + revision=plan.revision, + source_uri=plan.source_uri, + expected_sha256=plan.expected_sha256, + actual_sha256=plan.expected_sha256, + path=Path(path), + cache_hit=False, + build_id=plan.build_id, + ) + + def test_parse_gs_uri_extracts_bucket_path_and_version(): reference = parse_gs_uri("gs://policyengine-us-data/states/CA.h5@1.77.0") @@ -101,9 +123,11 @@ def test_us_create_datasets_passes_materialized_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/us/datasets.py", ) - materialize = Mock(return_value="/tmp/enhanced_cps_2024.h5") + materialize = Mock( + return_value=_materialized("us", "populace_us_2024", "/tmp/populace_us_2024.h5") + ) microsimulation = Mock() - monkeypatch.setattr(us_datasets, "materialize_dataset_source", materialize) + monkeypatch.setattr(us_datasets, "materialize_bundle_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_us", @@ -111,14 +135,14 @@ def test_us_create_datasets_passes_materialized_source_to_country_package( ) us_datasets.create_datasets( - datasets=["gs://policyengine-us-data/enhanced_cps_2024.h5@1.77.0"], + datasets=["populace_us_2024"], years=[], ) materialize.assert_called_once_with( - "gs://policyengine-us-data/enhanced_cps_2024.h5@1.77.0" + "us", "populace_us_2024", data_dir=Path("./data") ) - microsimulation.assert_called_once_with(dataset="/tmp/enhanced_cps_2024.h5") + microsimulation.assert_called_once_with(dataset="/tmp/populace_us_2024.h5") def test_uk_create_datasets_passes_materialized_source_to_country_package( @@ -129,9 +153,11 @@ def test_uk_create_datasets_passes_materialized_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) - materialize = Mock(return_value="/tmp/enhanced_frs_2023_24.h5") + materialize = Mock( + return_value=_materialized("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") + ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "materialize_dataset_source", materialize) + monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -139,11 +165,11 @@ def test_uk_create_datasets_passes_materialized_source_to_country_package( ) uk_datasets.create_datasets( - datasets=["gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3"], + datasets=["populace_uk_2023"], years=[], ) materialize.assert_called_once_with( - "gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3" + "uk", "populace_uk_2023", data_dir=Path("./data") ) - microsimulation.assert_called_once_with(dataset="/tmp/enhanced_frs_2023_24.h5") + microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") diff --git a/tests/test_release_manifests.py b/tests/test_release_manifests.py index 2768766d..41ab85ac 100644 --- a/tests/test_release_manifests.py +++ b/tests/test_release_manifests.py @@ -16,6 +16,10 @@ from policyengine.core.tax_benefit_model import TaxBenefitModel from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion +from policyengine.provenance.dataset_materialization import ( + MaterializedDataset, + resolve_bundle_dataset_plan, +) from policyengine.provenance.manifest import ( ArtifactPathReference, CountryReleaseManifest, @@ -84,6 +88,30 @@ UK_CERTIFIED_DATASET_URI = ( "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2024_25.h5@1.56.16" ) + + +def _materialized_dataset( + country_id: str, + dataset: str, + path: str, +) -> MaterializedDataset: + plan = resolve_bundle_dataset_plan(country_id, dataset) + return MaterializedDataset( + country_id=plan.country_id, + dataset=plan.dataset, + data_package_name=plan.data_package_name, + repo_id=plan.repo_id, + repo_type=plan.repo_type, + revision=plan.revision, + source_uri=plan.source_uri, + expected_sha256=plan.expected_sha256, + actual_sha256=plan.expected_sha256, + path=Path(path), + cache_hit=False, + build_id=plan.build_id, + ) + + UK_LEGACY_DATA_RELEASE_REVISION = "655dd07e4bb9c777b00dac044949611f1feb824f" UK_LEGACY_FRS_DATASET_URI = ( "hf://policyengine/policyengine-uk-data-private/frs_2023_24.h5" @@ -363,6 +391,7 @@ def test__given_us_ensure_datasets_without_dataset__then_uses_certified_default( datasets=["populace_us_2024"], years=[2026], data_folder="./data", + allow_unmanaged=False, ) def test__given_explicit_uri__then_managed_resolution_requires_opt_in(self): @@ -924,8 +953,12 @@ def test__given_us_managed_microsimulation__then_passes_certified_dataset_and_bu ) with patch.object( us_model, - "materialize_dataset_source", - return_value="/tmp/populace_us_2024.h5", + "materialize_bundle_dataset", + return_value=_materialized_dataset( + "us", + "populace_us_2024", + "/tmp/populace_us_2024.h5", + ), ): microsim = us_model.managed_microsimulation() @@ -941,6 +974,11 @@ def test__given_us_managed_microsimulation__then_passes_certified_dataset_and_bu ) dataset_source = microsim.policyengine_bundle["runtime_dataset_source"] assert dataset_source == "/tmp/populace_us_2024.h5" + assert ( + microsim.policyengine_bundle["runtime_dataset_sha256"] + == get_release_manifest("us").datasets["populace_us_2024"].sha256 + ) + assert microsim.policyengine_bundle["runtime_dataset_repo_type"] == "dataset" def test__given_us_unmanaged_dataset_uri__then_source_is_not_rewritten(self): dataset = "hf://policyengine/policyengine-us-data/cps_2023.h5@1.73.0" @@ -964,7 +1002,7 @@ def test__given_us_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( us_model, - "materialize_dataset_source", + "materialize_unmanaged_dataset_source", return_value="/tmp/cps_2023.h5", ): microsim = us_model.managed_microsimulation( @@ -1033,8 +1071,12 @@ def test__given_uk_managed_dataset_name__then_resolves_within_bundle(self): ) with patch.object( uk_model, - "materialize_dataset_source", - return_value="/tmp/enhanced_frs_2024_25.h5", + "materialize_bundle_dataset", + return_value=_materialized_dataset( + "uk", + "enhanced_frs_2024_25", + "/tmp/enhanced_frs_2024_25.h5", + ), ): microsim = uk_model.managed_microsimulation( dataset="enhanced_frs_2024_25" @@ -1051,6 +1093,11 @@ def test__given_uk_managed_dataset_name__then_resolves_within_bundle(self): ) dataset_source = microsim.policyengine_bundle["runtime_dataset_source"] assert dataset_source == "/tmp/enhanced_frs_2024_25.h5" + assert ( + microsim.policyengine_bundle["runtime_dataset_sha256"] + == get_release_manifest("uk").datasets["enhanced_frs_2024_25"].sha256 + ) + assert microsim.policyengine_bundle["runtime_dataset_repo_type"] == "model" def test__given_uk_unmanaged_dataset_uri__then_source_is_not_rewritten(self): dataset = "hf://policyengine/policyengine-uk-data-private/frs_2022_23.h5@1.40.4" @@ -1074,7 +1121,7 @@ def test__given_uk_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( uk_model, - "materialize_dataset_source", + "materialize_unmanaged_dataset_source", return_value="/tmp/frs_2022_23.h5", ): microsim = uk_model.managed_microsimulation( diff --git a/tests/test_us_long_term_datasets.py b/tests/test_us_long_term_datasets.py index e154cdff..d3e6ebbd 100644 --- a/tests/test_us_long_term_datasets.py +++ b/tests/test_us_long_term_datasets.py @@ -2,6 +2,7 @@ import json from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock import h5py import pandas as pd @@ -9,6 +10,10 @@ from microdf import MicroDataFrame import policyengine.tax_benefit_models.us.datasets as us_datasets_module +from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, + MaterializedDataset, +) from policyengine.tax_benefit_models.us.datasets import ( PolicyEngineUSDataset, USYearData, @@ -148,6 +153,24 @@ def _manifest_with_long_term_sha( ) +def _materialized_long_term(path: Path, dataset_uri: str) -> MaterializedDataset: + actual_sha256 = _sha256(path) + return MaterializedDataset( + country_id="us", + dataset="long_term_cps_2100", + data_package_name="policyengine-us-data", + repo_id="policyengine/policyengine-us-data", + repo_type="model", + revision="abc123", + source_uri=dataset_uri, + expected_sha256=actual_sha256, + actual_sha256=actual_sha256, + path=path, + cache_hit=True, + metadata_path=Path(f"{path}.metadata.json"), + ) + + def test__load_long_term_datasets__loads_h5_and_sidecar_metadata(tmp_path): h5_path = tmp_path / "2075.h5" _write_us_h5(h5_path, 2075) @@ -267,7 +290,6 @@ def test__load_managed_long_term_datasets__loads_bundled_local_mirror( policyengine_us={"version": "1.700.0"}, ) dataset_uri = "hf://policyengine/policyengine-us-data/long_term/2100.h5@abc123" - monkeypatch.setattr( us_datasets_module, "get_release_manifest", @@ -287,13 +309,8 @@ def test__load_managed_long_term_datasets__loads_bundled_local_mirror( ) monkeypatch.setattr( us_datasets_module, - "resolve_managed_dataset_reference", - lambda country_id, dataset: dataset_uri, - ) - monkeypatch.setattr( - us_datasets_module, - "resolve_local_managed_dataset_source", - lambda country_id, uri: str(h5_path), + "materialize_bundle_dataset", + lambda *args, **kwargs: _materialized_long_term(h5_path, dataset_uri), ) datasets = load_managed_long_term_datasets( @@ -307,11 +324,12 @@ def test__load_managed_long_term_datasets__loads_bundled_local_mirror( dataset = datasets["long_term_cps_2100"] assert dataset.filepath == str(h5_path) - assert dataset.metadata["policyengine_bundle"] == { - "managed_by": "policyengine.py", - "runtime_dataset": "long_term_cps_2100", - "runtime_dataset_uri": dataset_uri, - } + bundle = dataset.metadata["policyengine_bundle"] + assert bundle["managed_by"] == "policyengine.py" + assert bundle["runtime_dataset"] == "long_term_cps_2100" + assert bundle["runtime_dataset_uri"] == dataset_uri + assert bundle["runtime_dataset_data_package"] == "policyengine-us-data" + assert bundle["runtime_dataset_sha256"] == _sha256(h5_path) def test__load_managed_long_term_datasets__defaults_to_manifest_model_version( @@ -330,13 +348,8 @@ def test__load_managed_long_term_datasets__defaults_to_manifest_model_version( ) monkeypatch.setattr( us_datasets_module, - "resolve_managed_dataset_reference", - lambda country_id, dataset: dataset_uri, - ) - monkeypatch.setattr( - us_datasets_module, - "resolve_local_managed_dataset_source", - lambda country_id, uri: str(h5_path), + "materialize_bundle_dataset", + lambda *args, **kwargs: _materialized_long_term(h5_path, dataset_uri), ) with pytest.raises(ValueError, match="policyengine_us.version"): @@ -350,7 +363,6 @@ def test__load_managed_long_term_datasets__checks_manifest_sha256( h5_path = tmp_path / "2100.h5" _write_us_h5(h5_path, 2100) _write_metadata(h5_path, 2100, policyengine_us={"version": "1.691.12"}) - dataset_uri = "hf://policyengine/policyengine-us-data/long_term/2100.h5@abc123" monkeypatch.setattr( us_datasets_module, @@ -359,13 +371,8 @@ def test__load_managed_long_term_datasets__checks_manifest_sha256( ) monkeypatch.setattr( us_datasets_module, - "resolve_managed_dataset_reference", - lambda country_id, dataset: dataset_uri, - ) - monkeypatch.setattr( - us_datasets_module, - "resolve_local_managed_dataset_source", - lambda country_id, uri: str(h5_path), + "materialize_bundle_dataset", + Mock(side_effect=DatasetMaterializationError("sha256 mismatch")), ) with pytest.raises(ValueError, match="sha256"): @@ -391,41 +398,48 @@ def test__load_managed_long_term_datasets__checks_metadata_sha256( ) monkeypatch.setattr( us_datasets_module, - "resolve_managed_dataset_reference", - lambda country_id, dataset: dataset_uri, - ) - monkeypatch.setattr( - us_datasets_module, - "resolve_local_managed_dataset_source", - lambda country_id, uri: str(h5_path), + "materialize_bundle_dataset", + lambda *args, **kwargs: _materialized_long_term(h5_path, dataset_uri), ) with pytest.raises(ValueError, match="metadata"): load_managed_long_term_datasets([2100]) -def test__load_managed_long_term_datasets__requires_local_mirror( +def test__load_managed_long_term_datasets__materializes_without_local_mirror( monkeypatch, + tmp_path, ): + h5_path = tmp_path / "2100.h5" + _write_us_h5(h5_path, 2100) + _write_metadata(h5_path, 2100, policyengine_us={"version": "1.691.12"}) dataset_uri = "hf://policyengine/policyengine-us-data/long_term/2100.h5@abc123" + manifest = _manifest_with_long_term_sha(_sha256(h5_path)) monkeypatch.setattr( us_datasets_module, "get_release_manifest", - lambda country_id: _manifest_with_long_term_sha("0" * 64), + lambda country_id: manifest, ) + materialize = Mock(return_value=_materialized_long_term(h5_path, dataset_uri)) monkeypatch.setattr( us_datasets_module, - "resolve_managed_dataset_reference", - lambda country_id, dataset: dataset_uri, + "materialize_bundle_dataset", + materialize, ) - monkeypatch.setattr( - us_datasets_module, - "resolve_local_managed_dataset_source", - lambda country_id, uri: uri, + + datasets = load_managed_long_term_datasets( + [2100], + data_folder=str(tmp_path), + require_runtime_policyengine_us_match=False, ) - with pytest.raises(FileNotFoundError, match="no local mirror exists"): - load_managed_long_term_datasets([2100]) + assert datasets["long_term_cps_2100"].filepath == str(h5_path) + materialize.assert_called_once_with( + "us", + "long_term_cps_2100", + data_dir=tmp_path, + manifest=manifest, + ) def test__load_long_term_datasets__rejects_policyengine_us_version_mismatch( From acab43e61a81c26906db81ffb61e9d8e4a80642e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:45:49 +0400 Subject: [PATCH 05/18] Remove obsolete dataset GCS materialization --- pyproject.toml | 2 - .../provenance/dataset_sources.py | 119 ---- src/policyengine/utils/data/__init__.py | 1 - .../data/caching_google_storage_client.py | 105 --- .../data/version_aware_storage_client.py | 127 ---- src/policyengine/utils/google_cloud_bucket.py | 48 -- tests/fixtures/region_fixtures.py | 12 +- tests/test_dataset_runtime.py | 92 +++ tests/test_dataset_sources.py | 175 ----- tests/test_region.py | 4 +- uv.lock | 610 +++--------------- 11 files changed, 176 insertions(+), 1119 deletions(-) delete mode 100644 src/policyengine/provenance/dataset_sources.py delete mode 100644 src/policyengine/utils/data/__init__.py delete mode 100644 src/policyengine/utils/data/caching_google_storage_client.py delete mode 100644 src/policyengine/utils/data/version_aware_storage_client.py delete mode 100644 src/policyengine/utils/google_cloud_bucket.py create mode 100644 tests/test_dataset_runtime.py delete mode 100644 tests/test_dataset_sources.py diff --git a/pyproject.toml b/pyproject.toml index b5cd9ead..b0affd10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,6 @@ dependencies = [ "requests>=2.31.0", "psutil>=5.9.0", "packaging>=23.0", - "google-cloud-storage>=3.1.0,<4.0.0", - "diskcache>=5.6.3,<6.0.0", ] [project.scripts] diff --git a/src/policyengine/provenance/dataset_sources.py b/src/policyengine/provenance/dataset_sources.py deleted file mode 100644 index dfbdb311..00000000 --- a/src/policyengine/provenance/dataset_sources.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Runtime dataset source materialization.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Optional - -from policyengine.utils.google_cloud_bucket import download_file_from_gcs - - -@dataclass(frozen=True) -class GCSArtifactReference: - bucket: str - path: str - version: Optional[str] = None - - -@dataclass(frozen=True) -class HFArtifactReference: - repo_id: str - path: str - version: Optional[str] = None - - -def _select_version( - uri_version: Optional[str], - requested_version: Optional[str], -) -> Optional[str]: - if ( - uri_version is not None - and requested_version is not None - and uri_version != requested_version - ): - raise ValueError( - "Conflicting dataset versions: " - f"URI requests {uri_version!r} but version is {requested_version!r}" - ) - return uri_version or requested_version - - -def parse_gs_uri(uri: str) -> GCSArtifactReference: - if not uri.startswith("gs://"): - raise ValueError(f"Invalid GCS dataset URI: {uri!r}") - - path_with_bucket, version = ( - uri[5:].rsplit("@", maxsplit=1) if "@" in uri[5:] else (uri[5:], None) - ) - bucket, separator, path = path_with_bucket.partition("/") - if not bucket or not separator or not path: - raise ValueError( - "Invalid GCS dataset URI. Expected format " - f"'gs://bucket/path/to/file[@version]', got {uri!r}." - ) - return GCSArtifactReference(bucket=bucket, path=path, version=version) - - -def parse_hf_uri(uri: str) -> HFArtifactReference: - if not uri.startswith("hf://"): - raise ValueError(f"Invalid Hugging Face dataset URI: {uri!r}") - - path_with_repo, version = ( - uri[5:].rsplit("@", maxsplit=1) if "@" in uri[5:] else (uri[5:], None) - ) - parts = path_with_repo.split("/", maxsplit=2) - if len(parts) != 3 or not all(parts): - raise ValueError( - "Invalid Hugging Face dataset URI. Expected format " - f"'hf://owner/repo/path/to/file[@revision]', got {uri!r}." - ) - return HFArtifactReference( - repo_id=f"{parts[0]}/{parts[1]}", - path=parts[2], - version=version, - ) - - -def materialize_dataset_source( - dataset_source: str, - *, - version: Optional[str] = None, -) -> str: - """Return a local file path for supported remote dataset URIs.""" - - if dataset_source.startswith("gs://"): - reference = parse_gs_uri(dataset_source) - local_path, _ = download_file_from_gcs( - reference.bucket, - reference.path, - version=_select_version(reference.version, version), - ) - return local_path - - if dataset_source.startswith("hf://"): - from policyengine_core.tools.hugging_face import ( - download_huggingface_dataset, - ) - - reference = parse_hf_uri(dataset_source) - try: - return download_huggingface_dataset( - reference.repo_id, - reference.path, - version=_select_version(reference.version, version), - ) - except Exception: - # The core helper assumes a model-type repo; certified data - # releases may live in dataset-type repos (e.g. - # policyengine/populace-us). Retry with the dataset repo type - # before surfacing the original failure. - from huggingface_hub import hf_hub_download - - return hf_hub_download( - repo_id=reference.repo_id, - repo_type="dataset", - filename=reference.path, - revision=_select_version(reference.version, version), - ) - - return dataset_source diff --git a/src/policyengine/utils/data/__init__.py b/src/policyengine/utils/data/__init__.py deleted file mode 100644 index 97419b4e..00000000 --- a/src/policyengine/utils/data/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Dataset download helpers.""" diff --git a/src/policyengine/utils/data/caching_google_storage_client.py b/src/policyengine/utils/data/caching_google_storage_client.py deleted file mode 100644 index 08dfcb9b..00000000 --- a/src/policyengine/utils/data/caching_google_storage_client.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Disk-cached Google Cloud Storage downloads.""" - -from __future__ import annotations - -import logging -import os -import tempfile -from contextlib import AbstractContextManager -from pathlib import Path -from typing import Optional - -import diskcache - -from .version_aware_storage_client import VersionAwareStorageClient - -logger = logging.getLogger(__name__) - - -def _atomic_write(target: Path, content: bytes) -> None: - target.parent.mkdir(parents=True, exist_ok=True) - temp_path = None - try: - with tempfile.NamedTemporaryFile( - dir=target.parent, - delete=False, - ) as temp_file: - temp_path = Path(temp_file.name) - temp_file.write(content) - os.replace(temp_path, target) - finally: - if temp_path is not None and temp_path.exists(): - temp_path.unlink() - - -class CachingGoogleStorageClient(AbstractContextManager): - """Download GCS objects through a CRC-keyed disk cache.""" - - def __init__(self) -> None: - self.client = VersionAwareStorageClient() - self.cache = diskcache.Cache() - - @staticmethod - def _data_key(bucket: str, key: str, version: Optional[str] = None) -> str: - return f"{bucket}.{key}.{version}.data" - - @staticmethod - def _crc_key(bucket: str, key: str, version: Optional[str] = None) -> str: - return f"{bucket}.{key}.{version}.crc" - - def download( - self, - bucket: str, - key: str, - target: Path, - version: Optional[str] = None, - return_version: bool = False, - ) -> Optional[str]: - if version is None: - version = self.client.latest_metadata_version(bucket, key) - logger.warning( - "No version specified for %s/%s; using latest metadata version %s", - bucket, - key, - version, - ) - - self.sync(bucket, key, version) - data = self.cache.get(self._data_key(bucket, key, version)) - if isinstance(data, bytes): - _atomic_write(target, data) - return version if return_version else None - - raise TypeError( - f"Expected cached data for {bucket}/{key}@{version} to be bytes" - ) - - def sync( - self, - bucket: str, - key: str, - version: Optional[str] = None, - ) -> None: - crc = self.client.crc32c(bucket, key, version=version) - if crc is None: - raise FileNotFoundError(f"Unable to find gs://{bucket}/{key}") - - data_key = self._data_key(bucket, key, version) - crc_key = self._crc_key(bucket, key, version) - if self.cache.get(crc_key, default=None) == crc: - return - - content, downloaded_crc = self.client.download(bucket, key, version=version) - with self.cache as cache: - cache.set(data_key, content) - cache.set(crc_key, downloaded_crc) - - def clear(self) -> None: - self.cache.clear() - - def __enter__(self) -> CachingGoogleStorageClient: - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.clear() - return None diff --git a/src/policyengine/utils/data/version_aware_storage_client.py b/src/policyengine/utils/data/version_aware_storage_client.py deleted file mode 100644 index cc886271..00000000 --- a/src/policyengine/utils/data/version_aware_storage_client.py +++ /dev/null @@ -1,127 +0,0 @@ -"""GCS client helpers for generation and metadata-versioned objects.""" - -from __future__ import annotations - -import logging -from typing import Optional - -from google.cloud.storage import Blob, Bucket, Client - -logger = logging.getLogger(__name__) - - -class VersionAwareStorageClient: - """Resolve GCS objects by generation, metadata version, or latest object.""" - - def __init__(self) -> None: - self.client = Client() - - def get_blob( - self, - bucket_name: str, - key: str, - version: Optional[str] = None, - ) -> Blob: - bucket = self.client.bucket(bucket_name) - - if version is None: - logger.debug( - "No version specified for %s/%s, using latest", - bucket_name, - key, - ) - return bucket.blob(key) - - if version.isdigit(): - try: - blob = bucket.blob(key, generation=int(version)) - blob.reload() - logger.info( - "Found %s/%s with generation %s", - bucket_name, - key, - version, - ) - return blob - except Exception as exc: - logger.debug( - "Generation lookup failed for %s/%s@%s: %s", - bucket_name, - key, - version, - exc, - ) - - return self._get_blob_by_metadata_version(bucket, key, version) - - def _get_blob_by_metadata_version( - self, - bucket: Bucket, - key: str, - version: str, - ) -> Blob: - logger.debug( - "Searching for %s/%s with metadata version %s", - bucket.name, - key, - version, - ) - matching_blobs = [ - blob - for blob in bucket.list_blobs(prefix=key, versions=True) - if blob.name == key - and blob.metadata is not None - and blob.metadata.get("version") == version - ] - - if not matching_blobs: - raise ValueError( - f"No blob found with version {version!r} for {bucket.name}/{key}" - ) - - newest_blob = max(matching_blobs, key=lambda blob: blob.generation) - logger.info( - "Found %s/%s with metadata version %s and generation %s", - bucket.name, - key, - version, - newest_blob.generation, - ) - return newest_blob - - def crc32c( - self, - bucket_name: str, - key: str, - version: Optional[str] = None, - ) -> Optional[str]: - blob = self.get_blob(bucket_name, key, version) - blob.reload() - return blob.crc32c - - def download( - self, - bucket_name: str, - key: str, - version: Optional[str] = None, - ) -> tuple[bytes, Optional[str]]: - blob = self.get_blob(bucket_name, key, version) - content = blob.download_as_bytes() - return content, blob.crc32c - - def latest_metadata_version( - self, - bucket_name: str, - key: str, - ) -> Optional[str]: - blob = self.client.get_bucket(bucket_name).get_blob(key) - if blob is None: - logger.warning("No blob found for %s/%s", bucket_name, key) - return None - if blob.metadata is None: - logger.warning("No metadata found for %s/%s", bucket_name, key) - return None - version = blob.metadata.get("version") - if version is None: - logger.warning("No metadata version found for %s/%s", bucket_name, key) - return version diff --git a/src/policyengine/utils/google_cloud_bucket.py b/src/policyengine/utils/google_cloud_bucket.py deleted file mode 100644 index c6da587a..00000000 --- a/src/policyengine/utils/google_cloud_bucket.py +++ /dev/null @@ -1,48 +0,0 @@ -"""High-level dataset downloads from Google Cloud Storage.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Optional, Tuple - -from policyengine.utils.data.caching_google_storage_client import ( - CachingGoogleStorageClient, -) - -DATASETS_DIR = Path(".datasets") - -_caching_client: Optional[CachingGoogleStorageClient] = None - - -def _get_client() -> CachingGoogleStorageClient: - global _caching_client - if _caching_client is None: - _caching_client = CachingGoogleStorageClient() - return _caching_client - - -def _clear_client() -> None: - """Reset the singleton client. Intended for tests.""" - - global _caching_client - _caching_client = None - - -def download_file_from_gcs( - bucket_name: str, - gcs_key: str, - version: Optional[str] = None, -) -> Tuple[str, Optional[str]]: - """Download a GCS object into `.datasets`, preserving object path.""" - - local_path = DATASETS_DIR / gcs_key - local_path.parent.mkdir(parents=True, exist_ok=True) - - resolved_version = _get_client().download( - bucket_name, - gcs_key, - local_path, - version=version, - return_version=True, - ) - return str(local_path), resolved_version diff --git a/tests/fixtures/region_fixtures.py b/tests/fixtures/region_fixtures.py index 3dc8a639..383330d0 100644 --- a/tests/fixtures/region_fixtures.py +++ b/tests/fixtures/region_fixtures.py @@ -9,7 +9,9 @@ def create_national_region( country_code: str = "us", label: str = "United States", - dataset_path: str = "gs://policyengine-us-data/enhanced_cps_2024.h5", + dataset_path: str = ( + "hf://policyengine/populace-us/populace_us_2024.h5@certified-release" + ), ) -> Region: """Create a national region.""" return Region( @@ -24,7 +26,7 @@ def create_state_region( state_code: str, state_name: str, parent_code: str = "us", - bucket: str = "gs://policyengine-us-data", + repository: str = "hf://policyengine/policyengine-us-data", ) -> Region: """Create a state region with dedicated dataset.""" return Region( @@ -32,7 +34,7 @@ def create_state_region( label=state_name, region_type="state", parent_code=parent_code, - dataset_path=f"{bucket}/states/{state_code}.h5", + dataset_path=f"{repository}/states/{state_code}.h5@certified-release", state_code=state_code, state_name=state_name, ) @@ -99,7 +101,9 @@ def create_sample_us_registry() -> RegionRegistry: label="California", region_type="state", parent_code="us", - dataset_path="gs://policyengine-us-data/states/CA.h5", + dataset_path=( + "hf://policyengine/policyengine-us-data/states/CA.h5@certified-release" + ), state_code="CA", state_name="California", ) diff --git a/tests/test_dataset_runtime.py b/tests/test_dataset_runtime.py new file mode 100644 index 00000000..11f82b56 --- /dev/null +++ b/tests/test_dataset_runtime.py @@ -0,0 +1,92 @@ +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +from policyengine.provenance.dataset_materialization import ( + MaterializedDataset, + resolve_bundle_dataset_plan, +) + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _load_module_from_path(module_name: str, path: Path): + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDataset: + plan = resolve_bundle_dataset_plan(country_id, dataset) + return MaterializedDataset( + country_id=plan.country_id, + dataset=plan.dataset, + data_package_name=plan.data_package_name, + repo_id=plan.repo_id, + repo_type=plan.repo_type, + revision=plan.revision, + source_uri=plan.source_uri, + expected_sha256=plan.expected_sha256, + actual_sha256=plan.expected_sha256, + path=Path(path), + cache_hit=False, + build_id=plan.build_id, + ) + + +def test_us_create_datasets_passes_verified_bundle_source_to_country_package( + monkeypatch, +): + us_datasets = _load_module_from_path( + "_test_policyengine_us_datasets", + REPO_ROOT / "src/policyengine/tax_benefit_models/us/datasets.py", + ) + materialize = Mock( + return_value=_materialized("us", "populace_us_2024", "/tmp/populace_us_2024.h5") + ) + microsimulation = Mock() + monkeypatch.setattr(us_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setitem( + sys.modules, + "policyengine_us", + SimpleNamespace(Microsimulation=microsimulation), + ) + + us_datasets.create_datasets(datasets=["populace_us_2024"], years=[]) + + materialize.assert_called_once_with( + "us", "populace_us_2024", data_dir=Path("./data") + ) + microsimulation.assert_called_once_with(dataset="/tmp/populace_us_2024.h5") + + +def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( + monkeypatch, +): + uk_datasets = _load_module_from_path( + "_test_policyengine_uk_datasets", + REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", + ) + materialize = Mock( + return_value=_materialized("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") + ) + microsimulation = Mock() + monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setitem( + sys.modules, + "policyengine_uk", + SimpleNamespace(Microsimulation=microsimulation), + ) + + uk_datasets.create_datasets(datasets=["populace_uk_2023"], years=[]) + + materialize.assert_called_once_with( + "uk", "populace_uk_2023", data_dir=Path("./data") + ) + microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") diff --git a/tests/test_dataset_sources.py b/tests/test_dataset_sources.py deleted file mode 100644 index 4b7ff960..00000000 --- a/tests/test_dataset_sources.py +++ /dev/null @@ -1,175 +0,0 @@ -import importlib.util -import sys -from pathlib import Path -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from policyengine.provenance import dataset_sources -from policyengine.provenance.dataset_materialization import ( - MaterializedDataset, - resolve_bundle_dataset_plan, -) -from policyengine.provenance.dataset_sources import ( - materialize_dataset_source, - parse_gs_uri, - parse_hf_uri, -) - -REPO_ROOT = Path(__file__).resolve().parents[1] - - -def _load_module_from_path(module_name: str, path: Path): - spec = importlib.util.spec_from_file_location(module_name, path) - assert spec is not None - assert spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - return module - - -def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDataset: - plan = resolve_bundle_dataset_plan(country_id, dataset) - return MaterializedDataset( - country_id=plan.country_id, - dataset=plan.dataset, - data_package_name=plan.data_package_name, - repo_id=plan.repo_id, - repo_type=plan.repo_type, - revision=plan.revision, - source_uri=plan.source_uri, - expected_sha256=plan.expected_sha256, - actual_sha256=plan.expected_sha256, - path=Path(path), - cache_hit=False, - build_id=plan.build_id, - ) - - -def test_parse_gs_uri_extracts_bucket_path_and_version(): - reference = parse_gs_uri("gs://policyengine-us-data/states/CA.h5@1.77.0") - - assert reference.bucket == "policyengine-us-data" - assert reference.path == "states/CA.h5" - assert reference.version == "1.77.0" - - -def test_parse_hf_uri_extracts_repo_path_and_revision(): - reference = parse_hf_uri( - "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.77.0" - ) - - assert reference.repo_id == "policyengine/policyengine-us-data" - assert reference.path == "enhanced_cps_2024.h5" - assert reference.version == "1.77.0" - - -def test_materialize_dataset_source_downloads_gcs_uri(monkeypatch): - download = Mock(return_value=(".datasets/enhanced_cps_2024.h5", "1.77.0")) - monkeypatch.setattr(dataset_sources, "download_file_from_gcs", download) - - result = materialize_dataset_source( - "gs://policyengine-us-data/enhanced_cps_2024.h5@1.77.0" - ) - - assert result == ".datasets/enhanced_cps_2024.h5" - download.assert_called_once_with( - "policyengine-us-data", - "enhanced_cps_2024.h5", - version="1.77.0", - ) - - -def test_materialize_dataset_source_downloads_hf_uri(monkeypatch): - download = Mock(return_value="/tmp/enhanced_cps_2024.h5") - monkeypatch.setattr( - "policyengine_core.tools.hugging_face.download_huggingface_dataset", - download, - ) - - result = materialize_dataset_source( - "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.77.0" - ) - - assert result == "/tmp/enhanced_cps_2024.h5" - download.assert_called_once_with( - "policyengine/policyengine-us-data", - "enhanced_cps_2024.h5", - version="1.77.0", - ) - - -def test_materialize_dataset_source_preserves_local_path(): - assert materialize_dataset_source("/tmp/enhanced_cps_2024.h5") == ( - "/tmp/enhanced_cps_2024.h5" - ) - - -def test_materialize_dataset_source_rejects_conflicting_versions(): - with pytest.raises(ValueError, match="Conflicting dataset versions"): - materialize_dataset_source( - "gs://policyengine-us-data/enhanced_cps_2024.h5@1.77.0", - version="1.78.0", - ) - - -def test_us_create_datasets_passes_materialized_source_to_country_package( - monkeypatch, -): - us_datasets = _load_module_from_path( - "_test_policyengine_us_datasets", - REPO_ROOT / "src/policyengine/tax_benefit_models/us/datasets.py", - ) - - materialize = Mock( - return_value=_materialized("us", "populace_us_2024", "/tmp/populace_us_2024.h5") - ) - microsimulation = Mock() - monkeypatch.setattr(us_datasets, "materialize_bundle_dataset", materialize) - monkeypatch.setitem( - sys.modules, - "policyengine_us", - SimpleNamespace(Microsimulation=microsimulation), - ) - - us_datasets.create_datasets( - datasets=["populace_us_2024"], - years=[], - ) - - materialize.assert_called_once_with( - "us", "populace_us_2024", data_dir=Path("./data") - ) - microsimulation.assert_called_once_with(dataset="/tmp/populace_us_2024.h5") - - -def test_uk_create_datasets_passes_materialized_source_to_country_package( - monkeypatch, -): - uk_datasets = _load_module_from_path( - "_test_policyengine_uk_datasets", - REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", - ) - - materialize = Mock( - return_value=_materialized("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") - ) - microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) - monkeypatch.setitem( - sys.modules, - "policyengine_uk", - SimpleNamespace(Microsimulation=microsimulation), - ) - - uk_datasets.create_datasets( - datasets=["populace_uk_2023"], - years=[], - ) - - materialize.assert_called_once_with( - "uk", "populace_uk_2023", data_dir=Path("./data") - ) - microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") diff --git a/tests/test_region.py b/tests/test_region.py index fa54124a..3275fe01 100644 --- a/tests/test_region.py +++ b/tests/test_region.py @@ -38,7 +38,9 @@ def test__given_dataset_path__then_region_has_dedicated_dataset(self): region = REGION_WITH_DATASET # Then - assert region.dataset_path == "gs://policyengine-us-data/states/CA.h5" + assert region.dataset_path == ( + "hf://policyengine/policyengine-us-data/states/CA.h5@certified-release" + ) assert region.parent_code == "us" assert region.state_code == "CA" assert not region.requires_filter diff --git a/uv.lock b/uv.lock index 4b5b4430..e0a1ebd3 100644 --- a/uv.lock +++ b/uv.lock @@ -129,10 +129,10 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "msgpack", marker = "python_full_version < '3.10'" }, - { name = "ndindex", marker = "python_full_version < '3.10'" }, - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "py-cpuinfo", marker = "python_full_version < '3.10'" }, + { name = "msgpack" }, + { name = "ndindex" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" } }, + { name = "py-cpuinfo" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f7/60/5bc8601f8ffcd5d8787b346898de8a0b454d031c3e158e3bbc312003984e/blosc2-2.5.1.tar.gz", hash = "sha256:47d5df50e7286edf81e629ece35f87f13f55c13c5e8545832188c420c75d1659", size = 4676483, upload-time = "2024-01-25T12:31:31.168Z" } wheels = [ @@ -177,13 +177,13 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "msgpack", marker = "python_full_version >= '3.10'" }, - { name = "ndindex", marker = "python_full_version >= '3.10'" }, - { name = "numexpr", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10' and platform_machine != 'wasm32'" }, - { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "platformdirs", marker = "python_full_version >= '3.10'" }, - { name = "py-cpuinfo", marker = "python_full_version >= '3.10' and platform_machine != 'wasm32'" }, - { name = "requests", marker = "python_full_version >= '3.10'" }, + { name = "msgpack" }, + { name = "ndindex" }, + { name = "numexpr", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'wasm32'" }, + { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" } }, + { name = "platformdirs" }, + { name = "py-cpuinfo", marker = "platform_machine != 'wasm32'" }, + { name = "requests" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1b/ba/c1d8a8303197d8d8d4ec9b83851fcabd2d57dc5d17af2406643dd015544e/blosc2-3.7.2.tar.gz", hash = "sha256:3e80bd0399241829e4a2100bef9d4de042da979514f5df6aa3378981823f1d9b", size = 3804422, upload-time = "2025-08-19T10:21:46.164Z" } wheels = [ @@ -261,7 +261,7 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy'" }, + { name = "pycparser" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ @@ -434,7 +434,7 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b9/2e/0090cbf739cee7d23781ad4b89a9894a41538e4fcf4c31dcdd705b78eb8b/click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a", size = 226593, upload-time = "2024-12-21T18:38:44.339Z" } wheels = [ @@ -452,7 +452,7 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.10' and sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/60/6c/8ca2efa64cf75a977a0d7fac081354553ebe483345c734fb6b6515d96bbc/click-8.2.1.tar.gz", hash = "sha256:27c491cc05d968d271d5a1db13e3b5a184636d9d930f148c50b038f0d0646202", size = 286342, upload-time = "2025-05-20T23:19:49.832Z" } wheels = [ @@ -594,7 +594,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version < '3.10'" }, + { name = "tomli" }, ] [[package]] @@ -718,137 +718,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version >= '3.10' and python_full_version <= '3.11'" }, -] - -[[package]] -name = "cryptography" -version = "47.0.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version <= '3.9'", -] -dependencies = [ - { name = "cffi", marker = "python_full_version <= '3.9' and platform_python_implementation != 'PyPy'" }, - { name = "typing-extensions", marker = "python_full_version <= '3.9'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ef/b2/7ffa7fe8207a8c42147ffe70c3e360b228160c1d85dc3faff16aaa3244c0/cryptography-47.0.0.tar.gz", hash = "sha256:9f8e55fe4e63613a5e1cc5819030f27b97742d720203a087802ce4ce9ceb52bb", size = 830863, upload-time = "2026-04-24T19:54:57.056Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a4/98/40dfe932134bdcae4f6ab5927c87488754bf9eb79297d7e0070b78dd58e9/cryptography-47.0.0-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:160ad728f128972d362e714054f6ba0067cab7fb350c5202a9ae8ae4ce3ef1a0", size = 7912214, upload-time = "2026-04-24T19:53:03.864Z" }, - { url = "https://files.pythonhosted.org/packages/34/c6/2733531243fba725f58611b918056b277692f1033373dcc8bd01af1c05d4/cryptography-47.0.0-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b9a8943e359b7615db1a3ba587994618e094ff3d6fa5a390c73d079ce18b3973", size = 4644617, upload-time = "2026-04-24T19:53:06.909Z" }, - { url = "https://files.pythonhosted.org/packages/00/e3/b27be1a670a9b87f855d211cf0e1174a5d721216b7616bd52d8581d912ed/cryptography-47.0.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f5c15764f261394b22aef6b00252f5195f46f2ca300bec57149474e2538b31f8", size = 4668186, upload-time = "2026-04-24T19:53:09.053Z" }, - { url = "https://files.pythonhosted.org/packages/81/b9/8443cfe5d17d482d348cee7048acf502bb89a51b6382f06240fd290d4ca3/cryptography-47.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9c59ab0e0fa3a180a5a9c59f3a5abe3ef90d474bc56d7fadfbe80359491b615b", size = 4651244, upload-time = "2026-04-24T19:53:11.217Z" }, - { url = "https://files.pythonhosted.org/packages/5d/5e/13ed0cdd0eb88ba159d6dd5ebfece8cb901dbcf1ae5ac4072e28b55d3153/cryptography-47.0.0-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:34b4358b925a5ea3e14384ca781a2c0ef7ac219b57bb9eacc4457078e2b19f92", size = 5252906, upload-time = "2026-04-24T19:53:13.532Z" }, - { url = "https://files.pythonhosted.org/packages/64/16/ed058e1df0f33d440217cd120d41d5dda9dd215a80b8187f68483185af82/cryptography-47.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0024b87d47ae2399165a6bfb20d24888881eeab83ae2566d62467c5ff0030ce7", size = 4701842, upload-time = "2026-04-24T19:53:15.618Z" }, - { url = "https://files.pythonhosted.org/packages/02/e0/3d30986b30fdbd9e969abbdf8ba00ed0618615144341faeb57f395a084fe/cryptography-47.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:1e47422b5557bb82d3fff997e8d92cff4e28b9789576984f08c248d2b3535d93", size = 4289313, upload-time = "2026-04-24T19:53:17.755Z" }, - { url = "https://files.pythonhosted.org/packages/df/fd/32db38e3ad0cb331f0691cb4c7a8a6f176f679124dee746b3af6633db4d9/cryptography-47.0.0-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:6f29f36582e6151d9686235e586dd35bb67491f024767d10b842e520dc6a07ac", size = 4650964, upload-time = "2026-04-24T19:53:20.062Z" }, - { url = "https://files.pythonhosted.org/packages/86/53/5395d944dfd48cb1f67917f533c609c34347185ef15eb4308024c876f274/cryptography-47.0.0-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:a9b761f012a943b7de0e828843c5688d0de94a0578d44d6c85a1bae32f87791f", size = 5207817, upload-time = "2026-04-24T19:53:22.498Z" }, - { url = "https://files.pythonhosted.org/packages/34/4f/e5711b28e1901f7d480a2b1b688b645aa4c77c73f10731ed17e7f7db3f0d/cryptography-47.0.0-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:4e1de79e047e25d6e9f8cea71c86b4a53aced64134f0f003bbcbf3655fd172c8", size = 4701544, upload-time = "2026-04-24T19:53:24.356Z" }, - { url = "https://files.pythonhosted.org/packages/22/22/c8ddc25de3010fc8da447648f5a092c40e7a8fadf01dd6d255d9c0b9373d/cryptography-47.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef6b3634087f18d2155b1e8ce264e5345a753da2c5fa9815e7d41315c90f8318", size = 4783536, upload-time = "2026-04-24T19:53:26.665Z" }, - { url = "https://files.pythonhosted.org/packages/66/b6/d4a68f4ea999c6d89e8498579cba1c5fcba4276284de7773b17e4fa69293/cryptography-47.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:11dbb9f50a0f1bb9757b3d8c27c1101780efb8f0bdecfb12439c22a74d64c001", size = 4926106, upload-time = "2026-04-24T19:53:28.686Z" }, - { url = "https://files.pythonhosted.org/packages/54/ed/5f524db1fade9c013aa618e1c99c6ed05e8ffc9ceee6cda22fed22dda3f4/cryptography-47.0.0-cp311-abi3-win32.whl", hash = "sha256:7fda2f02c9015db3f42bb8a22324a454516ed10a8c29ca6ece6cdbb5efe2a203", size = 3258581, upload-time = "2026-04-24T19:53:31.058Z" }, - { url = "https://files.pythonhosted.org/packages/b2/dc/1b901990b174786569029f67542b3edf72ac068b6c3c8683c17e6a2f5363/cryptography-47.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:f5c3296dab66202f1b18a91fa266be93d6aa0c2806ea3d67762c69f60adc71aa", size = 3775309, upload-time = "2026-04-24T19:53:33.054Z" }, - { url = "https://files.pythonhosted.org/packages/14/88/7aa18ad9c11bc87689affa5ce4368d884b517502d75739d475fc6f4a03c7/cryptography-47.0.0-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:be12cb6a204f77ed968bcefe68086eb061695b540a3dd05edac507a3111b25f0", size = 7904299, upload-time = "2026-04-24T19:53:35.003Z" }, - { url = "https://files.pythonhosted.org/packages/07/55/c18f75724544872f234678fdedc871391722cb34a2aee19faa9f63100bb2/cryptography-47.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2ebd84adf0728c039a3be2700289378e1c164afc6748df1a5ed456767bef9ba7", size = 4631180, upload-time = "2026-04-24T19:53:37.517Z" }, - { url = "https://files.pythonhosted.org/packages/ee/65/31a5cc0eaca99cec5bafffe155d407115d96136bb161e8b49e0ef73f09a7/cryptography-47.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f68d6fbc7fbbcfb0939fea72c3b96a9f9a6edfc0e1b1d29778a2066030418b1", size = 4653529, upload-time = "2026-04-24T19:53:39.775Z" }, - { url = "https://files.pythonhosted.org/packages/e5/bc/641c0519a495f3bfd0421b48d7cd325c4336578523ccd76ea322b6c29c7a/cryptography-47.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:6651d32eff255423503aa276739da98c30f26c40cbeffcc6048e0d54ef704c0c", size = 4638570, upload-time = "2026-04-24T19:53:42.129Z" }, - { url = "https://files.pythonhosted.org/packages/2b/f2/300327b0a47f6dc94dd8b71b57052aefe178bb51745073d73d80604f11ab/cryptography-47.0.0-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:3fb8fa48075fad7193f2e5496135c6a76ac4b2aa5a38433df0a539296b377829", size = 5238019, upload-time = "2026-04-24T19:53:44.577Z" }, - { url = "https://files.pythonhosted.org/packages/e9/5a/5b5cf994391d4bf9d9c7efd4c66aabe4d95227256627f8fea6cff7dfadbd/cryptography-47.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:11438c7518132d95f354fa01a4aa2f806d172a061a7bed18cf18cbdacdb204d7", size = 4686832, upload-time = "2026-04-24T19:53:47.015Z" }, - { url = "https://files.pythonhosted.org/packages/dc/2c/ae950e28fd6475c852fc21a44db3e6b5bcc1261d1e370f2b6e42fa800fef/cryptography-47.0.0-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:8c1a736bbb3288005796c3f7ccb9453360d7fed483b13b9f468aea5171432923", size = 4269301, upload-time = "2026-04-24T19:53:48.97Z" }, - { url = "https://files.pythonhosted.org/packages/67/fb/6a39782e150ffe5cc1b0018cb6ddc48bf7ca62b498d7539ffc8a758e977d/cryptography-47.0.0-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:f1557695e5c2b86e204f6ce9470497848634100787935ab7adc5397c54abd7ab", size = 4638110, upload-time = "2026-04-24T19:53:51.011Z" }, - { url = "https://files.pythonhosted.org/packages/8e/d7/0b3c71090a76e5c203164a47688b697635ece006dcd2499ab3a4dbd3f0bd/cryptography-47.0.0-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:f9a034b642b960767fb343766ae5ba6ad653f2e890ddd82955aef288ffea8736", size = 5194988, upload-time = "2026-04-24T19:53:52.962Z" }, - { url = "https://files.pythonhosted.org/packages/63/33/63a961498a9df51721ab578c5a2622661411fc520e00bd83b0cc64eb20c4/cryptography-47.0.0-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:b1c76fca783aa7698eb21eb14f9c4aa09452248ee54a627d125025a43f83e7a7", size = 4686563, upload-time = "2026-04-24T19:53:55.274Z" }, - { url = "https://files.pythonhosted.org/packages/b7/bf/5ee5b145248f92250de86145d1c1d6edebbd57a7fe7caa4dedb5d4cf06a1/cryptography-47.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:4f7722c97826770bab8ae92959a2e7b20a5e9e9bf4deae68fd86c3ca457bab52", size = 4770094, upload-time = "2026-04-24T19:53:57.753Z" }, - { url = "https://files.pythonhosted.org/packages/92/43/21d220b2da5d517773894dacdcdb5c682c28d3fffce65548cb06e87d5501/cryptography-47.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:09f6d7bf6724f8db8b32f11eccf23efc8e759924bc5603800335cf8859a3ddbd", size = 4913811, upload-time = "2026-04-24T19:54:00.236Z" }, - { url = "https://files.pythonhosted.org/packages/31/98/dc4ad376ac5f1a1a7d4a83f7b0c6f2bcad36b5d2d8f30aeb482d3a7d9582/cryptography-47.0.0-cp314-cp314t-win32.whl", hash = "sha256:6eebcaf0df1d21ce1f90605c9b432dd2c4f4ab665ac29a40d5e3fc68f51b5e63", size = 3237158, upload-time = "2026-04-24T19:54:02.606Z" }, - { url = "https://files.pythonhosted.org/packages/bc/da/97f62d18306b5133468bc3f8cc73a3111e8cdc8cf8d3e69474d6e5fd2d1b/cryptography-47.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:51c9313e90bd1690ec5a75ed047c27c0b8e6c570029712943d6116ef9a90620b", size = 3758706, upload-time = "2026-04-24T19:54:04.433Z" }, - { url = "https://files.pythonhosted.org/packages/e0/34/a4fae8ae7c3bc227460c9ae43f56abf1b911da0ec29e0ebac53bb0a4b6b7/cryptography-47.0.0-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:14432c8a9bcb37009784f9594a62fae211a2ae9543e96c92b2a8e4c3cd5cd0c4", size = 7904072, upload-time = "2026-04-24T19:54:06.411Z" }, - { url = "https://files.pythonhosted.org/packages/01/64/d7b1e54fdb69f22d24a64bb3e88dc718b31c7fb10ef0b9691a3cf7eeea6e/cryptography-47.0.0-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:07efe86201817e7d3c18781ca9770bc0db04e1e48c994be384e4602bc38f8f27", size = 4635767, upload-time = "2026-04-24T19:54:08.519Z" }, - { url = "https://files.pythonhosted.org/packages/8b/7b/cca826391fb2a94efdcdfe4631eb69306ee1cff0b22f664a412c90713877/cryptography-47.0.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b45761c6ec22b7c726d6a829558777e32d0f1c8be7c3f3480f9c912d5ee8a10", size = 4654350, upload-time = "2026-04-24T19:54:10.795Z" }, - { url = "https://files.pythonhosted.org/packages/4c/65/4b57bcc823f42a991627c51c2f68c9fd6eb1393c1756aac876cba2accae2/cryptography-47.0.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:edd4da498015da5b9f26d38d3bfc2e90257bfa9cbed1f6767c282a0025ae649b", size = 4643394, upload-time = "2026-04-24T19:54:13.275Z" }, - { url = "https://files.pythonhosted.org/packages/f4/c4/2c5fbeea70adbbca2bbae865e1d605d6a4a7f8dbd9d33eaf69645087f06c/cryptography-47.0.0-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:9af828c0d5a65c70ec729cd7495a4bf1a67ecb66417b8f02ff125ab8a6326a74", size = 5225777, upload-time = "2026-04-24T19:54:15.18Z" }, - { url = "https://files.pythonhosted.org/packages/7e/b8/ac57107ef32749d2b244e36069bb688792a363aaaa3acc9e3cf84c130315/cryptography-47.0.0-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:256d07c78a04d6b276f5df935a9923275f53bd1522f214447fdf365494e2d515", size = 4688771, upload-time = "2026-04-24T19:54:17.835Z" }, - { url = "https://files.pythonhosted.org/packages/56/fc/9f1de22ff8be99d991f240a46863c52d475404c408886c5a38d2b5c3bb26/cryptography-47.0.0-cp38-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:5d0e362ff51041b0c0d219cc7d6924d7b8996f57ce5712bdcef71eb3c65a59cc", size = 4270753, upload-time = "2026-04-24T19:54:19.963Z" }, - { url = "https://files.pythonhosted.org/packages/00/68/d70c852797aa68e8e48d12e5a87170c43f67bb4a59403627259dd57d15de/cryptography-47.0.0-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:1581aef4219f7ca2849d0250edaa3866212fb74bf5667284f46aa92f9e65c1ca", size = 4642911, upload-time = "2026-04-24T19:54:21.818Z" }, - { url = "https://files.pythonhosted.org/packages/a5/51/661cbee74f594c5d97ff82d34f10d5551c085ca4668645f4606ebd22bd5d/cryptography-47.0.0-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:a49a3eb5341b9503fa3000a9a0db033161db90d47285291f53c2a9d2cd1b7f76", size = 5181411, upload-time = "2026-04-24T19:54:24.376Z" }, - { url = "https://files.pythonhosted.org/packages/94/87/f2b6c374a82cf076cfa1416992ac8e8ec94d79facc37aec87c1a5cb72352/cryptography-47.0.0-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:2207a498b03275d0051589e326b79d4cf59985c99031b05bb292ac52631c37fe", size = 4688262, upload-time = "2026-04-24T19:54:26.946Z" }, - { url = "https://files.pythonhosted.org/packages/14/e2/8b7462f4acf21ec509616f0245018bb197194ab0b65c2ea21a0bdd53c0eb/cryptography-47.0.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7a02675e2fabd0c0fc04c868b8781863cbf1967691543c22f5470500ff840b31", size = 4775506, upload-time = "2026-04-24T19:54:28.926Z" }, - { url = "https://files.pythonhosted.org/packages/70/75/158e494e4c08dc05e039da5bb48553826bd26c23930cf8d3cd5f21fa8921/cryptography-47.0.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:80887c5cbd1774683cb126f0ab4184567f080071d5acf62205acb354b4b753b7", size = 4912060, upload-time = "2026-04-24T19:54:30.869Z" }, - { url = "https://files.pythonhosted.org/packages/06/bd/0a9d3edbf5eadbac926d7b9b3cd0c4be584eeeae4a003d24d9eda4affbbd/cryptography-47.0.0-cp38-abi3-win32.whl", hash = "sha256:ed67ea4e0cfb5faa5bc7ecb6e2b8838f3807a03758eec239d6c21c8769355310", size = 3248487, upload-time = "2026-04-24T19:54:33.494Z" }, - { url = "https://files.pythonhosted.org/packages/60/80/5681af756d0da3a599b7bdb586fac5a1540f1bcefd2717a20e611ddade45/cryptography-47.0.0-cp38-abi3-win_amd64.whl", hash = "sha256:835d2d7f47cdc53b3224e90810fb1d36ca94ea29cc1801fb4c1bc43876735769", size = 3755737, upload-time = "2026-04-24T19:54:35.408Z" }, - { url = "https://files.pythonhosted.org/packages/1b/a0/928c9ce0d120a40a81aa99e3ba383e87337b9ac9ef9f6db02e4d7822424d/cryptography-47.0.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:7f1207974a904e005f762869996cf620e9bf79ecb4622f148550bb48e0eb35a7", size = 3909893, upload-time = "2026-04-24T19:54:38.334Z" }, - { url = "https://files.pythonhosted.org/packages/81/75/d691e284750df5d9569f2b1ce4a00a71e1d79566da83b2b3e5549c84917f/cryptography-47.0.0-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:1a405c08857258c11016777e11c02bacbe7ef596faf259305d282272a3a05cbe", size = 4587867, upload-time = "2026-04-24T19:54:40.619Z" }, - { url = "https://files.pythonhosted.org/packages/07/d6/1b90f1a4e453009730b4545286f0b39bb348d805c11181fc31544e4f9a65/cryptography-47.0.0-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:20fdbe3e38fb67c385d233c89371fa27f9909f6ebca1cecc20c13518dae65475", size = 4627192, upload-time = "2026-04-24T19:54:42.849Z" }, - { url = "https://files.pythonhosted.org/packages/dc/53/cb358a80e9e359529f496870dd08c102aa8a4b5b9f9064f00f0d6ed5b527/cryptography-47.0.0-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:f7db373287273d8af1414cf95dc4118b13ffdc62be521997b0f2b270771fef50", size = 4587486, upload-time = "2026-04-24T19:54:44.908Z" }, - { url = "https://files.pythonhosted.org/packages/8b/57/aaa3d53876467a226f9a7a82fd14dd48058ad2de1948493442dfa16e2ffd/cryptography-47.0.0-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:9fe6b7c64926c765f9dff301f9c1b867febcda5768868ca084e18589113732ab", size = 4626327, upload-time = "2026-04-24T19:54:47.813Z" }, - { url = "https://files.pythonhosted.org/packages/ab/9c/51f28c3550276bcf35660703ba0ab829a90b88be8cd98a71ef23c2413913/cryptography-47.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:cffbba3392df0fa8629bb7f43454ee2925059ee158e23c54620b9063912b86c8", size = 3698916, upload-time = "2026-04-24T19:54:49.782Z" }, -] - -[[package]] -name = "cryptography" -version = "48.0.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", - "python_full_version > '3.9' and python_full_version < '3.10'", -] -dependencies = [ - { name = "cffi", marker = "python_full_version > '3.9' and platform_python_implementation != 'PyPy'" }, - { name = "typing-extensions", marker = "python_full_version > '3.9' and python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9f/a9/db8f313fdcd85d767d4973515e1db101f9c71f95fced83233de224673757/cryptography-48.0.0.tar.gz", hash = "sha256:5c3932f4436d1cccb036cb0eaef46e6e2db91035166f1ad6505c3c9d5a635920", size = 832984, upload-time = "2026-05-04T22:59:38.133Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/3d/01f6dd9190170a5a241e0e98c2d04be3664a9e6f5b9b872cde63aff1c3dd/cryptography-48.0.0-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:0c558d2cdffd8f4bbb30fc7134c74d2ca9a476f830bb053074498fbc86f41ed6", size = 8001587, upload-time = "2026-05-04T22:57:36.803Z" }, - { url = "https://files.pythonhosted.org/packages/b2/6e/e90527eef33f309beb811cf7c982c3aeffcce8e3edb178baa4ca3ae4a6fa/cryptography-48.0.0-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f5333311663ea94f75dd408665686aaf426563556bb5283554a3539177e03b8c", size = 4690433, upload-time = "2026-05-04T22:57:40.373Z" }, - { url = "https://files.pythonhosted.org/packages/90/04/673510ed51ddff56575f306cf1617d80411ee76831ccd3097599140efdfe/cryptography-48.0.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7995ef305d7165c3f11ae07f2517e5a4f1d5c18da1376a0a9ed496336b69e5f3", size = 4710620, upload-time = "2026-05-04T22:57:42.935Z" }, - { url = "https://files.pythonhosted.org/packages/14/d5/e9c4ef932c8d800490c34d8bd589d64a31d5890e27ec9e9ad532be893294/cryptography-48.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:40ba1f85eaa6959837b1d51c9767e230e14612eea4ef110ee8854ada22da1bf5", size = 4696283, upload-time = "2026-05-04T22:57:45.294Z" }, - { url = "https://files.pythonhosted.org/packages/0c/29/174b9dfb60b12d59ecfc6cfa04bc88c21b42a54f01b8aae09bb6e51e4c7f/cryptography-48.0.0-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:369a6348999f94bbd53435c894377b20ab95f25a9065c283570e70150d8abc3c", size = 5296573, upload-time = "2026-05-04T22:57:47.933Z" }, - { url = "https://files.pythonhosted.org/packages/95/38/0d29a6fd7d0d1373f0c0c88a04ba20e359b257753ac497564cd660fc1d55/cryptography-48.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a0e692c683f4df67815a2d258b324e66f4738bd7a96a218c826dce4f4bd05d8f", size = 4743677, upload-time = "2026-05-04T22:57:50.067Z" }, - { url = "https://files.pythonhosted.org/packages/30/be/eef653013d5c63b6a490529e0316f9ac14a37602965d4903efed1399f32b/cryptography-48.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:18349bbc56f4743c8b12dc32e2bccb2cf83ee8b69a3bba74ef8ae857e26b3d25", size = 4330808, upload-time = "2026-05-04T22:57:52.301Z" }, - { url = "https://files.pythonhosted.org/packages/84/9e/500463e87abb7a0a0f9f256ec21123ecde0a7b5541a15e840ea54551fd81/cryptography-48.0.0-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:7e8eac43dfca5c4cccc6dad9a80504436fca53bb9bc3100a2386d730fbe6b602", size = 4695941, upload-time = "2026-05-04T22:57:54.603Z" }, - { url = "https://files.pythonhosted.org/packages/e3/dc/7303087450c2ec9e7fbb750e17c2abfbc658f23cbd0e54009509b7cc4091/cryptography-48.0.0-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:9ccdac7d40688ecb5a3b4a604b8a88c8002e3442d6c60aead1db2a89a041560c", size = 5252579, upload-time = "2026-05-04T22:57:57.207Z" }, - { url = "https://files.pythonhosted.org/packages/d0/c0/7101d3b7215edcdc90c45da544961fd8ed2d6448f77577460fa75a8443f7/cryptography-48.0.0-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:bd72e68b06bb1e96913f97dd4901119bc17f39d4586a5adf2d3e47bc2b9d58b5", size = 4743326, upload-time = "2026-05-04T22:57:59.535Z" }, - { url = "https://files.pythonhosted.org/packages/ac/d8/5b833bad13016f562ab9d063d68199a4bd121d18458e439515601d3357ec/cryptography-48.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:59baa2cb386c4f0b9905bd6eb4c2a79a69a128408fd31d32ca4d7102d4156321", size = 4826672, upload-time = "2026-05-04T22:58:01.996Z" }, - { url = "https://files.pythonhosted.org/packages/98/e1/7074eb8bf3c135558c73fc2bcf0f5633f912e6fb87e868a55c454080ef09/cryptography-48.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9249e3cd978541d665967ac2cb2787fd6a62bddf1e75b3e347a594d7dacf4f74", size = 4972574, upload-time = "2026-05-04T22:58:03.968Z" }, - { url = "https://files.pythonhosted.org/packages/04/70/e5a1b41d325f797f39427aa44ef8baf0be500065ab6d8e10369d850d4a4f/cryptography-48.0.0-cp311-abi3-win32.whl", hash = "sha256:9c459db21422be75e2809370b829a87eb37f74cd785fc4aa9ea1e5f43b47cda4", size = 3294868, upload-time = "2026-05-04T22:58:06.467Z" }, - { url = "https://files.pythonhosted.org/packages/f4/ac/8ac51b4a5fc5932eb7ee5c517ba7dc8cd834f0048962b6b352f00f41ebf9/cryptography-48.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:5b012212e08b8dd5edc78ef54da83dd9892fd9105323b3993eff6bea65dc21d7", size = 3817107, upload-time = "2026-05-04T22:58:08.845Z" }, - { url = "https://files.pythonhosted.org/packages/6b/84/70e3feea9feea87fd7cbe77efb2712ae1e3e6edf10749dc6e95f4e60e455/cryptography-48.0.0-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:3cb07a3ed6431663cd321ea8a000a1314c74211f823e4177fefa2255e057d1ec", size = 7986556, upload-time = "2026-05-04T22:58:11.172Z" }, - { url = "https://files.pythonhosted.org/packages/89/6e/18e07a618bb5442ba10cf4df16e99c071365528aa570dfcb8c02e25a303b/cryptography-48.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8c7378637d7d88016fa6791c159f698b3d3eed28ebf844ac36b9dc04a14dae18", size = 4684776, upload-time = "2026-05-04T22:58:13.712Z" }, - { url = "https://files.pythonhosted.org/packages/be/6a/4ea3b4c6c6759794d5ee2103c304a5076dc4b19ae1f9fe47dba439e159e9/cryptography-48.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc90c0b39b2e3c65ef52c804b72e3c58f8a04ab2a1871272798e5f9572c17d20", size = 4698121, upload-time = "2026-05-04T22:58:16.448Z" }, - { url = "https://files.pythonhosted.org/packages/2f/59/6ff6ad6cae03bb887da2a5860b2c9805f8dac969ef01ce563336c49bd1d1/cryptography-48.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:76341972e1eff8b4bea859f09c0d3e64b96ce931b084f9b9b7db8ef364c30eff", size = 4690042, upload-time = "2026-05-04T22:58:18.544Z" }, - { url = "https://files.pythonhosted.org/packages/ca/b4/fc334ed8cfd705aca282fe4d8f5ae64a8e0f74932e9feecb344610cf6e4d/cryptography-48.0.0-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:55b7718303bf06a5753dcdccf2f3945cf18ad7bffde41b61226e4db31ab89a9c", size = 5282526, upload-time = "2026-05-04T22:58:20.75Z" }, - { url = "https://files.pythonhosted.org/packages/11/08/9f8c5386cc4cd90d8255c7cdd0f5baf459a08502a09de30dc51f553d38dc/cryptography-48.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:a64697c641c7b1b2178e573cbc31c7c6684cd56883a478d75143dbb7118036db", size = 4733116, upload-time = "2026-05-04T22:58:23.627Z" }, - { url = "https://files.pythonhosted.org/packages/b8/77/99307d7574045699f8805aa500fa0fb83422d115b5400a064ddd306d7750/cryptography-48.0.0-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:561215ea3879cb1cbbf272867e2efda62476f240fb58c64de6b393ae19246741", size = 4316030, upload-time = "2026-05-04T22:58:25.581Z" }, - { url = "https://files.pythonhosted.org/packages/fd/36/a608b98337af3cb2aff4818e406649d30572b7031918b04c87d979495348/cryptography-48.0.0-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:ad64688338ed4bc1a6618076ba75fd7194a5f1797ac60b47afe926285adb3166", size = 4689640, upload-time = "2026-05-04T22:58:27.747Z" }, - { url = "https://files.pythonhosted.org/packages/dd/a6/825010a291b4438aecc1f568bc428189fc1175515223632477c07dc0a6df/cryptography-48.0.0-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:906cbf0670286c6e0044156bc7d4af9cbb0ef6db9f73e52c3ec56ba6bdde5336", size = 5237657, upload-time = "2026-05-04T22:58:29.848Z" }, - { url = "https://files.pythonhosted.org/packages/b9/09/4e76a09b4caa29aad535ddc806f5d4c5d01885bd978bd984fbc6ca032cae/cryptography-48.0.0-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:ea8990436d914540a40ab24b6a77c0969695ed52f4a4874c5137ccf7045a7057", size = 4732362, upload-time = "2026-05-04T22:58:32.009Z" }, - { url = "https://files.pythonhosted.org/packages/18/78/444fa04a77d0cb95f417dda20d450e13c56ba8e5220fc892a1658f44f882/cryptography-48.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c18684a7f0cc9a3cb60328f496b8e3372def7c5d2df39ac267878b05565aaaae", size = 4819580, upload-time = "2026-05-04T22:58:34.254Z" }, - { url = "https://files.pythonhosted.org/packages/38/85/ea67067c70a1fd4be2c63d35eeed82658023021affccc7b17705f8527dd2/cryptography-48.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9be5aafa5736574f8f15f262adc81b2a9869e2cfe9014d52a44633905b40d52c", size = 4963283, upload-time = "2026-05-04T22:58:36.376Z" }, - { url = "https://files.pythonhosted.org/packages/75/54/cc6d0f3deac3e81c7f847e8a189a12b6cdd65059b43dad25d4316abd849a/cryptography-48.0.0-cp314-cp314t-win32.whl", hash = "sha256:c17dfe85494deaeddc5ce251aebd1d60bbe6afc8b62071bb0b469431a000124f", size = 3270954, upload-time = "2026-05-04T22:58:38.791Z" }, - { url = "https://files.pythonhosted.org/packages/49/67/cc947e288c0758a4e5473d1dcb743037ab7785541265a969240b8885441a/cryptography-48.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:27241b1dc9962e056062a8eef1991d02c3a24569c95975bd2322a8a52c6e5e12", size = 3797313, upload-time = "2026-05-04T22:58:40.746Z" }, - { url = "https://files.pythonhosted.org/packages/f2/63/61d4a4e1c6b6bab6ce1e213cd36a24c415d90e76d78c5eb8577c5541d2e8/cryptography-48.0.0-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:58d00498e8933e4a194f3076aee1b4a97dfec1a6da444535755822fe5d8b0b86", size = 7983482, upload-time = "2026-05-04T22:58:43.769Z" }, - { url = "https://files.pythonhosted.org/packages/d5/ac/f5b5995b87770c693e2596559ffafe195b4033a57f14a82268a2842953f3/cryptography-48.0.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:614d0949f4790582d2cc25553abd09dd723025f0c0e7c67376a1d77196743d6e", size = 4683266, upload-time = "2026-05-04T22:58:46.064Z" }, - { url = "https://files.pythonhosted.org/packages/ec/c6/8b14f67e18338fbc4adb76f66c001f5c3610b3e2d1837f268f47a347dbbb/cryptography-48.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7ce4bfae76319a532a2dc68f82cc32f5676ee792a983187dac07183690e5c66f", size = 4696228, upload-time = "2026-05-04T22:58:48.22Z" }, - { url = "https://files.pythonhosted.org/packages/ea/73/f808fbae9514bd91b47875b003f13e284c8c6bdfd904b7944e803937eec1/cryptography-48.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:2eb992bbd4661238c5a397594c83f5b4dc2bc5b848c365c8f991b6780efcc5c7", size = 4689097, upload-time = "2026-05-04T22:58:50.9Z" }, - { url = "https://files.pythonhosted.org/packages/93/01/d86632d7d28db8ae83221995752eeb6639ffb374c2d22955648cf8d52797/cryptography-48.0.0-cp39-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:22a5cb272895dce158b2cacdfdc3debd299019659f42947dbdac6f32d68fe832", size = 5283582, upload-time = "2026-05-04T22:58:53.017Z" }, - { url = "https://files.pythonhosted.org/packages/02/e1/50edc7a50334807cc4791fc4a0ce7468b4a1416d9138eab358bfc9a3d70b/cryptography-48.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2b4d59804e8408e2fea7d1fbaf218e5ec984325221db76e6a241a9abd6cdd95c", size = 4730479, upload-time = "2026-05-04T22:58:55.611Z" }, - { url = "https://files.pythonhosted.org/packages/6f/af/99a582b1b1641ff5911ac559beb45097cf79efd4ead4657f578ef1af2d47/cryptography-48.0.0-cp39-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:984a20b0f62a26f48a3396c72e4bc34c66e356d356bf370053066b3b6d54634a", size = 4326481, upload-time = "2026-05-04T22:58:57.607Z" }, - { url = "https://files.pythonhosted.org/packages/90/ee/89aa26a06ef0a7d7611788ffd571a7c50e368cc6a4d5eef8b4884e866edb/cryptography-48.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:5a5ed8fde7a1d09376ca0b40e68cd59c69fe23b1f9768bd5824f54681626032a", size = 4688713, upload-time = "2026-05-04T22:59:00.077Z" }, - { url = "https://files.pythonhosted.org/packages/70/ba/bcb1b0bb7a33d4c7c0c4d4c7874b4a62ae4f56113a5f4baefa362dfb1f0f/cryptography-48.0.0-cp39-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:8cd666227ef7af430aa5914a9910e0ddd703e75f039cef0825cd0da71b6b711a", size = 5238165, upload-time = "2026-05-04T22:59:02.317Z" }, - { url = "https://files.pythonhosted.org/packages/c9/70/ca4003b1ce5ca3dc3186ada51908c8a9b9ff7d5cab83cc0d43ee14ec144f/cryptography-48.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:9071196d81abc88b3516ac8cdfad32e2b66dd4a5393a8e68a961e9161ddc6239", size = 4729947, upload-time = "2026-05-04T22:59:05.255Z" }, - { url = "https://files.pythonhosted.org/packages/44/a0/4ec7cf774207905aef1a8d11c3750d5a1db805eb380ee4e16df317870128/cryptography-48.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1e2d54c8be6152856a36f0882ab231e70f8ec7f14e93cf87db8a2ed056bf160c", size = 4822059, upload-time = "2026-05-04T22:59:07.802Z" }, - { url = "https://files.pythonhosted.org/packages/1e/75/a2e55f99c16fcac7b5d6c1eb19ad8e00799854d6be5ca845f9259eae1681/cryptography-48.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a5da777e32ffed6f85a7b2b3f7c5cbc88c146bfcd0a1d7baf5fcc6c52ee35dd4", size = 4960575, upload-time = "2026-05-04T22:59:09.851Z" }, - { url = "https://files.pythonhosted.org/packages/b8/23/6e6f32143ab5d8b36ca848a502c4bcd477ae75b9e1677e3530d669062578/cryptography-48.0.0-cp39-abi3-win32.whl", hash = "sha256:77a2ccbbe917f6710e05ba9adaa25fb5075620bf3ea6fb751997875aff4ae4bd", size = 3279117, upload-time = "2026-05-04T22:59:12.019Z" }, - { url = "https://files.pythonhosted.org/packages/9d/9a/0fea98a70cf1749d41d738836f6349d97945f7c89433a259a6c2642eefeb/cryptography-48.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:16cd65b9330583e4619939b3a3843eec1e6e789744bb01e7c7e2e62e33c239c8", size = 3792100, upload-time = "2026-05-04T22:59:14.884Z" }, - { url = "https://files.pythonhosted.org/packages/be/d2/024b5e06be9d44cb021fb0e1a03d34d63989cf56a0fe62f3dfbab695b9b4/cryptography-48.0.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:84cf79f0dc8b36ac5da873481716e87aef31fcfa0444f9e1d8b4b2cece142855", size = 3950391, upload-time = "2026-05-04T22:59:17.415Z" }, - { url = "https://files.pythonhosted.org/packages/bc/17/3861e17c56fa0fd37491a14a8673fdb77c57fc5693cafe745ea8b06dba75/cryptography-48.0.0-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:fdfef35d751d510fcef5252703621574364fec16418c4a1e5e1055248401054b", size = 4637126, upload-time = "2026-05-04T22:59:20.197Z" }, - { url = "https://files.pythonhosted.org/packages/f0/0a/7e226dbff530f21480727eb764973a7bff2b912f8e15cd4f129e71b56d1d/cryptography-48.0.0-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:0890f502ddf7d9c6426129c3f49f5c0a39278ed7cd6322c8755ffca6ee675a13", size = 4667270, upload-time = "2026-05-04T22:59:22.647Z" }, - { url = "https://files.pythonhosted.org/packages/3b/f2/5a72274ca9f1b2a8b44a662ee0bf1b435909deb473d6f97bcd035bcdbc71/cryptography-48.0.0-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:ecde28a596bead48b0cfd2a1b4416c3d43074c2d785e3a398d7ec1fc4d0f7fbb", size = 4636797, upload-time = "2026-05-04T22:59:24.912Z" }, - { url = "https://files.pythonhosted.org/packages/b4/e1/48cedb2fe63626e91ded1edad159e2a4fb8b6906c4425eb7749673077ce7/cryptography-48.0.0-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:4defde8685ae324a9eb9d818717e93b4638ef67070ac9bc15b8ca85f63048355", size = 4666800, upload-time = "2026-05-04T22:59:27.474Z" }, - { url = "https://files.pythonhosted.org/packages/a2/ca/7e8365deec19afb2b2c7be7c1c0aa8f99633b54e90c570999acda93260fc/cryptography-48.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:db63bf618e5dea46c07de12e900fe1cdd2541e6dc9dbae772a70b7d4d4765f6a", size = 3739536, upload-time = "2026-05-04T22:59:29.61Z" }, + { name = "tomli", marker = "python_full_version <= '3.11'" }, ] [[package]] @@ -902,15 +772,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] -[[package]] -name = "diskcache" -version = "5.6.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/3f/21/1c1ffc1a039ddcc459db43cc108658f32c57d271d7289a2794e401d0fdb6/diskcache-5.6.3.tar.gz", hash = "sha256:2c3a3fa2743d8535d832ec61c2054a1641f41775aa7c556758a109941e33e4fc", size = 67916, upload-time = "2023-08-31T06:12:00.316Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" }, -] - [[package]] name = "docutils" version = "0.21.2" @@ -943,7 +804,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1002,229 +863,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/34/2b07b72bee02a63241d654f5d8af87a2de977c59638eec41ca356ab915cd/furo-2025.7.19-py3-none-any.whl", hash = "sha256:bdea869822dfd2b494ea84c0973937e35d1575af088b6721a29c7f7878adc9e3", size = 342175, upload-time = "2025-07-19T10:52:02.399Z" }, ] -[[package]] -name = "google-api-core" -version = "2.30.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-auth", version = "2.50.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "google-auth", version = "2.53.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "googleapis-common-protos" }, - { name = "proto-plus", version = "1.27.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "proto-plus", version = "1.28.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "protobuf", version = "6.33.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "protobuf", version = "7.35.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/16/ce/502a57fb0ec752026d24df1280b162294b22a0afb98a326084f9a979138b/google_api_core-2.30.3.tar.gz", hash = "sha256:e601a37f148585319b26db36e219df68c5d07b6382cff2d580e83404e44d641b", size = 177001, upload-time = "2026-04-10T00:41:28.035Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/03/15/e56f351cf6ef1cfea58e6ac226a7318ed1deb2218c4b3cc9bd9e4b786c5a/google_api_core-2.30.3-py3-none-any.whl", hash = "sha256:a85761ba72c444dad5d611c2220633480b2b6be2521eca69cca2dbb3ffd6bfe8", size = 173274, upload-time = "2026-04-09T22:57:16.198Z" }, -] - -[[package]] -name = "google-auth" -version = "2.50.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version > '3.9' and python_full_version < '3.10'", - "python_full_version <= '3.9'", -] -dependencies = [ - { name = "cryptography", version = "47.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version <= '3.9'" }, - { name = "cryptography", version = "48.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version > '3.9' and python_full_version < '3.10'" }, - { name = "pyasn1-modules", marker = "python_full_version < '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5f/18/238d7021d151bdab868f23433817b027dd759135202f4dfce0670d1230ca/google_auth-2.50.0.tar.gz", hash = "sha256:f35eafb191195328e8ce10a7883970877e7aeb49c2bfaa54aa0e394316d353d0", size = 336523, upload-time = "2026-04-30T21:19:29.659Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/37/cf/4880c2137c14280b2f59975cdf12cc442bc0ae1f9ea473a26eaa0c146786/google_auth-2.50.0-py3-none-any.whl", hash = "sha256:04382175e28b94f49694977f0a792688b59a668def1499e9d8de996dc9ce5b15", size = 246495, upload-time = "2026-04-30T21:19:27.664Z" }, -] - -[[package]] -name = "google-auth" -version = "2.53.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", -] -dependencies = [ - { name = "cryptography", version = "48.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "pyasn1-modules", marker = "python_full_version >= '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c6/ad/ff781329bbbdc0974a098d996e89c9e1f7024262f9e3eec442fbb9ad1ac6/google_auth-2.53.0.tar.gz", hash = "sha256:e7e6aa16f6bee7b2b264830fd04f08087a1d5a836df516251a5d15327b246c9c", size = 335844, upload-time = "2026-05-15T20:53:07.928Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/4a/c9/db44165ba7c581268c6d46017ef63339110378305062830104fc7fa144cb/google_auth-2.53.0-py3-none-any.whl", hash = "sha256:6e7449917c599b35126a99ec268ec6880301f2fea41dce198fe8fd83ff642b68", size = 246071, upload-time = "2026-05-15T20:53:05.609Z" }, -] - -[[package]] -name = "google-cloud-core" -version = "2.5.1" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version > '3.9' and python_full_version < '3.10'", - "python_full_version <= '3.9'", -] -dependencies = [ - { name = "google-api-core", marker = "python_full_version < '3.10'" }, - { name = "google-auth", version = "2.50.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/dc/24/6ca08b0a03c7b0c620427503ab00353a4ae806b848b93bcea18b6b76fde6/google_cloud_core-2.5.1.tar.gz", hash = "sha256:3dc94bdec9d05a31d9f355045ed0f369fbc0d8c665076c734f065d729800f811", size = 36078, upload-time = "2026-03-30T22:50:08.057Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/73/d9/5bb050cb32826466aa9b25f79e2ca2879fe66cb76782d4ed798dd7506151/google_cloud_core-2.5.1-py3-none-any.whl", hash = "sha256:ea62cdf502c20e3e14be8a32c05ed02113d7bef454e40ff3fab6fe1ec9f1f4e7", size = 29452, upload-time = "2026-03-30T22:48:31.567Z" }, -] - -[[package]] -name = "google-cloud-core" -version = "2.6.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", -] -dependencies = [ - { name = "google-api-core", marker = "python_full_version >= '3.10'" }, - { name = "google-auth", version = "2.53.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a8/dd/1eef226e470369b26824a505c34482c0b493bc35fe8e0c6b003b5feca21a/google_cloud_core-2.6.0.tar.gz", hash = "sha256:e76149739f90fac1fc6757c09f47eaccb3145b54adbd7759b0f7c4b235f46c83", size = 36001, upload-time = "2026-05-07T08:04:04.124Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/84/4a/98da8930ab109c73d9a5d13782a9ebb81ea8c111f6d534a567b71d23e52b/google_cloud_core-2.6.0-py3-none-any.whl", hash = "sha256:6d63ac8e5eca6d9e4319d0a1e2265fadcd7f1049904378caecfa01cf52dd869e", size = 29390, upload-time = "2026-05-07T08:02:34.672Z" }, -] - -[[package]] -name = "google-cloud-storage" -version = "3.9.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version > '3.9' and python_full_version < '3.10'", - "python_full_version <= '3.9'", -] -dependencies = [ - { name = "google-api-core", marker = "python_full_version < '3.10'" }, - { name = "google-auth", version = "2.50.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "google-cloud-core", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "google-crc32c", marker = "python_full_version < '3.10'" }, - { name = "google-resumable-media", version = "2.8.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "requests", marker = "python_full_version < '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f7/b1/4f0798e88285b50dfc60ed3a7de071def538b358db2da468c2e0deecbb40/google_cloud_storage-3.9.0.tar.gz", hash = "sha256:f2d8ca7db2f652be757e92573b2196e10fbc09649b5c016f8b422ad593c641cc", size = 17298544, upload-time = "2026-02-02T13:36:34.119Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/46/0b/816a6ae3c9fd096937d2e5f9670558908811d57d59ddf69dd4b83b326fd1/google_cloud_storage-3.9.0-py3-none-any.whl", hash = "sha256:2dce75a9e8b3387078cbbdad44757d410ecdb916101f8ba308abf202b6968066", size = 321324, upload-time = "2026-02-02T13:36:32.271Z" }, -] - -[[package]] -name = "google-cloud-storage" -version = "3.10.1" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", -] -dependencies = [ - { name = "google-api-core", marker = "python_full_version >= '3.10'" }, - { name = "google-auth", version = "2.53.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "google-cloud-core", version = "2.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "google-crc32c", marker = "python_full_version >= '3.10'" }, - { name = "google-resumable-media", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, - { name = "requests", marker = "python_full_version >= '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/4c/47/205eb8e9a1739b5345843e5a425775cbdc472cc38e7eda082ba5b8d02450/google_cloud_storage-3.10.1.tar.gz", hash = "sha256:97db9aa4460727982040edd2bd13ff3d5e2260b5331ad22895802da1fc2a5286", size = 17309950, upload-time = "2026-03-23T09:35:23.409Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ad/ff/ca9ab2417fa913d75aae38bf40bf856bb2749a604b2e0f701b37cfcd23cc/google_cloud_storage-3.10.1-py3-none-any.whl", hash = "sha256:a72f656759b7b99bda700f901adcb3425a828d4a29f911bc26b3ea79c5b1217f", size = 324453, upload-time = "2026-03-23T09:35:21.368Z" }, -] - -[[package]] -name = "google-crc32c" -version = "1.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/03/41/4b9c02f99e4c5fb477122cd5437403b552873f014616ac1d19ac8221a58d/google_crc32c-1.8.0.tar.gz", hash = "sha256:a428e25fb7691024de47fecfbff7ff957214da51eddded0da0ae0e0f03a2cf79", size = 14192, upload-time = "2025-12-16T00:35:25.142Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/95/ac/6f7bc93886a823ab545948c2dd48143027b2355ad1944c7cf852b338dc91/google_crc32c-1.8.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:0470b8c3d73b5f4e3300165498e4cf25221c7eb37f1159e221d1825b6df8a7ff", size = 31296, upload-time = "2025-12-16T00:19:07.261Z" }, - { url = "https://files.pythonhosted.org/packages/f7/97/a5accde175dee985311d949cfcb1249dcbb290f5ec83c994ea733311948f/google_crc32c-1.8.0-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:119fcd90c57c89f30040b47c211acee231b25a45d225e3225294386f5d258288", size = 30870, upload-time = "2025-12-16T00:29:17.669Z" }, - { url = "https://files.pythonhosted.org/packages/3d/63/bec827e70b7a0d4094e7476f863c0dbd6b5f0f1f91d9c9b32b76dcdfeb4e/google_crc32c-1.8.0-cp310-cp310-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:6f35aaffc8ccd81ba3162443fabb920e65b1f20ab1952a31b13173a67811467d", size = 33214, upload-time = "2025-12-16T00:40:19.618Z" }, - { url = "https://files.pythonhosted.org/packages/63/bc/11b70614df04c289128d782efc084b9035ef8466b3d0a8757c1b6f5cf7ac/google_crc32c-1.8.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:864abafe7d6e2c4c66395c1eb0fe12dc891879769b52a3d56499612ca93b6092", size = 33589, upload-time = "2025-12-16T00:40:20.7Z" }, - { url = "https://files.pythonhosted.org/packages/3e/00/a08a4bc24f1261cc5b0f47312d8aebfbe4b53c2e6307f1b595605eed246b/google_crc32c-1.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:db3fe8eaf0612fc8b20fa21a5f25bd785bc3cd5be69f8f3412b0ac2ffd49e733", size = 34437, upload-time = "2025-12-16T00:35:19.437Z" }, - { url = "https://files.pythonhosted.org/packages/5d/ef/21ccfaab3d5078d41efe8612e0ed0bfc9ce22475de074162a91a25f7980d/google_crc32c-1.8.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:014a7e68d623e9a4222d663931febc3033c5c7c9730785727de2a81f87d5bab8", size = 31298, upload-time = "2025-12-16T00:20:32.241Z" }, - { url = "https://files.pythonhosted.org/packages/c5/b8/f8413d3f4b676136e965e764ceedec904fe38ae8de0cdc52a12d8eb1096e/google_crc32c-1.8.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:86cfc00fe45a0ac7359e5214a1704e51a99e757d0272554874f419f79838c5f7", size = 30872, upload-time = "2025-12-16T00:33:58.785Z" }, - { url = "https://files.pythonhosted.org/packages/f6/fd/33aa4ec62b290477181c55bb1c9302c9698c58c0ce9a6ab4874abc8b0d60/google_crc32c-1.8.0-cp311-cp311-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:19b40d637a54cb71e0829179f6cb41835f0fbd9e8eb60552152a8b52c36cbe15", size = 33243, upload-time = "2025-12-16T00:40:21.46Z" }, - { url = "https://files.pythonhosted.org/packages/71/03/4820b3bd99c9653d1a5210cb32f9ba4da9681619b4d35b6a052432df4773/google_crc32c-1.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:17446feb05abddc187e5441a45971b8394ea4c1b6efd88ab0af393fd9e0a156a", size = 33608, upload-time = "2025-12-16T00:40:22.204Z" }, - { url = "https://files.pythonhosted.org/packages/7c/43/acf61476a11437bf9733fb2f70599b1ced11ec7ed9ea760fdd9a77d0c619/google_crc32c-1.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:71734788a88f551fbd6a97be9668a0020698e07b2bf5b3aa26a36c10cdfb27b2", size = 34439, upload-time = "2025-12-16T00:35:20.458Z" }, - { url = "https://files.pythonhosted.org/packages/e9/5f/7307325b1198b59324c0fa9807cafb551afb65e831699f2ce211ad5c8240/google_crc32c-1.8.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:4b8286b659c1335172e39563ab0a768b8015e88e08329fa5321f774275fc3113", size = 31300, upload-time = "2025-12-16T00:21:56.723Z" }, - { url = "https://files.pythonhosted.org/packages/21/8e/58c0d5d86e2220e6a37befe7e6a94dd2f6006044b1a33edf1ff6d9f7e319/google_crc32c-1.8.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:2a3dc3318507de089c5384cc74d54318401410f82aa65b2d9cdde9d297aca7cb", size = 30867, upload-time = "2025-12-16T00:38:31.302Z" }, - { url = "https://files.pythonhosted.org/packages/ce/a9/a780cc66f86335a6019f557a8aaca8fbb970728f0efd2430d15ff1beae0e/google_crc32c-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:14f87e04d613dfa218d6135e81b78272c3b904e2a7053b841481b38a7d901411", size = 33364, upload-time = "2025-12-16T00:40:22.96Z" }, - { url = "https://files.pythonhosted.org/packages/21/3f/3457ea803db0198c9aaca2dd373750972ce28a26f00544b6b85088811939/google_crc32c-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb5c869c2923d56cb0c8e6bcdd73c009c36ae39b652dbe46a05eb4ef0ad01454", size = 33740, upload-time = "2025-12-16T00:40:23.96Z" }, - { url = "https://files.pythonhosted.org/packages/df/c0/87c2073e0c72515bb8733d4eef7b21548e8d189f094b5dad20b0ecaf64f6/google_crc32c-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:3cc0c8912038065eafa603b238abf252e204accab2a704c63b9e14837a854962", size = 34437, upload-time = "2025-12-16T00:35:21.395Z" }, - { url = "https://files.pythonhosted.org/packages/d1/db/000f15b41724589b0e7bc24bc7a8967898d8d3bc8caf64c513d91ef1f6c0/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:3ebb04528e83b2634857f43f9bb8ef5b2bbe7f10f140daeb01b58f972d04736b", size = 31297, upload-time = "2025-12-16T00:23:20.709Z" }, - { url = "https://files.pythonhosted.org/packages/d7/0d/8ebed0c39c53a7e838e2a486da8abb0e52de135f1b376ae2f0b160eb4c1a/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:450dc98429d3e33ed2926fc99ee81001928d63460f8538f21a5d6060912a8e27", size = 30867, upload-time = "2025-12-16T00:43:14.628Z" }, - { url = "https://files.pythonhosted.org/packages/ce/42/b468aec74a0354b34c8cbf748db20d6e350a68a2b0912e128cabee49806c/google_crc32c-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3b9776774b24ba76831609ffbabce8cdf6fa2bd5e9df37b594221c7e333a81fa", size = 33344, upload-time = "2025-12-16T00:40:24.742Z" }, - { url = "https://files.pythonhosted.org/packages/1c/e8/b33784d6fc77fb5062a8a7854e43e1e618b87d5ddf610a88025e4de6226e/google_crc32c-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:89c17d53d75562edfff86679244830599ee0a48efc216200691de8b02ab6b2b8", size = 33694, upload-time = "2025-12-16T00:40:25.505Z" }, - { url = "https://files.pythonhosted.org/packages/92/b1/d3cbd4d988afb3d8e4db94ca953df429ed6db7282ed0e700d25e6c7bfc8d/google_crc32c-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:57a50a9035b75643996fbf224d6661e386c7162d1dfdab9bc4ca790947d1007f", size = 34435, upload-time = "2025-12-16T00:35:22.107Z" }, - { url = "https://files.pythonhosted.org/packages/21/88/8ecf3c2b864a490b9e7010c84fd203ec8cf3b280651106a3a74dd1b0ca72/google_crc32c-1.8.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:e6584b12cb06796d285d09e33f63309a09368b9d806a551d8036a4207ea43697", size = 31301, upload-time = "2025-12-16T00:24:48.527Z" }, - { url = "https://files.pythonhosted.org/packages/36/c6/f7ff6c11f5ca215d9f43d3629163727a272eabc356e5c9b2853df2bfe965/google_crc32c-1.8.0-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:f4b51844ef67d6cf2e9425983274da75f18b1597bb2c998e1c0a0e8d46f8f651", size = 30868, upload-time = "2025-12-16T00:48:12.163Z" }, - { url = "https://files.pythonhosted.org/packages/56/15/c25671c7aad70f8179d858c55a6ae8404902abe0cdcf32a29d581792b491/google_crc32c-1.8.0-cp314-cp314-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b0d1a7afc6e8e4635564ba8aa5c0548e3173e41b6384d7711a9123165f582de2", size = 33381, upload-time = "2025-12-16T00:40:26.268Z" }, - { url = "https://files.pythonhosted.org/packages/42/fa/f50f51260d7b0ef5d4898af122d8a7ec5a84e2984f676f746445f783705f/google_crc32c-1.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8b3f68782f3cbd1bce027e48768293072813469af6a61a86f6bb4977a4380f21", size = 33734, upload-time = "2025-12-16T00:40:27.028Z" }, - { url = "https://files.pythonhosted.org/packages/08/a5/7b059810934a09fb3ccb657e0843813c1fee1183d3bc2c8041800374aa2c/google_crc32c-1.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:d511b3153e7011a27ab6ee6bb3a5404a55b994dc1a7322c0b87b29606d9790e2", size = 34878, upload-time = "2025-12-16T00:35:23.142Z" }, - { url = "https://files.pythonhosted.org/packages/42/c5/4c4cde2e7e54d9cde5c3d131f54a609eb3a77e60a04ec348051f61071fc2/google_crc32c-1.8.0-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:ba6aba18daf4d36ad4412feede6221414692f44d17e5428bdd81ad3fc1eee5dc", size = 31291, upload-time = "2025-12-16T00:17:45.878Z" }, - { url = "https://files.pythonhosted.org/packages/31/1d/abae5a7ca05c07dc7f129b32f7f8cce5314172fd2300c7aec305427a637e/google_crc32c-1.8.0-cp39-cp39-macosx_12_0_x86_64.whl", hash = "sha256:87b0072c4ecc9505cfa16ee734b00cd7721d20a0f595be4d40d3d21b41f65ae2", size = 30862, upload-time = "2025-12-16T00:25:03.867Z" }, - { url = "https://files.pythonhosted.org/packages/1e/c4/7032f0e87ee0b0f65669ac8a1022beabd80afe5da69f4bbf49eb7fea9c40/google_crc32c-1.8.0-cp39-cp39-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3d488e98b18809f5e322978d4506373599c0c13e6c5ad13e53bb44758e18d215", size = 33063, upload-time = "2025-12-16T00:40:27.789Z" }, - { url = "https://files.pythonhosted.org/packages/cc/fc/831d92dd02bc145523590db3927a73300f5121a34b56c2696e4305411b67/google_crc32c-1.8.0-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:01f126a5cfddc378290de52095e2c7052be2ba7656a9f0caf4bcd1bfb1833f8a", size = 33434, upload-time = "2025-12-16T00:40:28.555Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ef/74accbe6e6892c3bcbe5a7ed8d650a23b3042ddcc0b301896c22e6733bea/google_crc32c-1.8.0-cp39-cp39-win_amd64.whl", hash = "sha256:61f58b28e0b21fcb249a8247ad0db2e64114e201e2e9b4200af020f3b6242c9f", size = 34432, upload-time = "2025-12-16T00:35:24.136Z" }, - { url = "https://files.pythonhosted.org/packages/52/c5/c171e4d8c44fec1422d801a6d2e5d7ddabd733eeda505c79730ee9607f07/google_crc32c-1.8.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:87fa445064e7db928226b2e6f0d5304ab4cd0339e664a4e9a25029f384d9bb93", size = 28615, upload-time = "2025-12-16T00:40:29.298Z" }, - { url = "https://files.pythonhosted.org/packages/9c/97/7d75fe37a7a6ed171a2cf17117177e7aab7e6e0d115858741b41e9dd4254/google_crc32c-1.8.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f639065ea2042d5c034bf258a9f085eaa7af0cd250667c0635a3118e8f92c69c", size = 28800, upload-time = "2025-12-16T00:40:30.322Z" }, -] - -[[package]] -name = "google-resumable-media" -version = "2.8.2" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version > '3.9' and python_full_version < '3.10'", - "python_full_version <= '3.9'", -] -dependencies = [ - { name = "google-crc32c", marker = "python_full_version < '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3f/d1/b1ea14b93b6b78f57fc580125de44e9f593ab88dd2460f1a8a8d18f74754/google_resumable_media-2.8.2.tar.gz", hash = "sha256:f3354a182ebd193ae3f42e3ef95e6c9b10f128320de23ac7637236713b1acd70", size = 2164510, upload-time = "2026-03-30T23:34:25.369Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/f8/50bfaf4658431ff9de45c5c3935af7ab01157a4903c603cd0eee6e78e087/google_resumable_media-2.8.2-py3-none-any.whl", hash = "sha256:82b6d8ccd11765268cdd2a2123f417ec806b8eef3000a9a38dfe3033da5fb220", size = 81511, upload-time = "2026-03-30T23:34:09.671Z" }, -] - -[[package]] -name = "google-resumable-media" -version = "2.9.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", -] -dependencies = [ - { name = "google-crc32c", marker = "python_full_version >= '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/00/4b/0b235beccc310d0a48adbc7246b719d173cca6c88c572dfa4b090e39143c/google_resumable_media-2.9.0.tar.gz", hash = "sha256:f7cfb224846a9dd444d125115dfbe8ef02a2b893e78f087762fe716a255a734b", size = 2164534, upload-time = "2026-05-07T08:04:44.236Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/07/73/3518e63deb1667c5409a4579e28daf5e84479a87a72c547e0487f7883dcd/google_resumable_media-2.9.0-py3-none-any.whl", hash = "sha256:c8901e88e389af8bed64d9696c74d8bad961865eb2236e13e0bfca9bb0a65ca3", size = 81507, upload-time = "2026-05-07T08:03:23.809Z" }, -] - -[[package]] -name = "googleapis-common-protos" -version = "1.75.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "protobuf", version = "6.33.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "protobuf", version = "7.35.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b5/c8/f439cffde755cffa462bfbb156278fa6f9d09119719af9814b858fd4f81f/googleapis_common_protos-1.75.0.tar.gz", hash = "sha256:53a062ff3c32552fbd62c11fe23768b78e4ddf0494d5e5fd97d3f4689c75fbbd", size = 151035, upload-time = "2026-05-07T08:04:49.423Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/c8/e2645aa8ed02fd4c7a2f59d68783b65b1f3cbdfe39a6308e156509d1fee8/googleapis_common_protos-1.75.0-py3-none-any.whl", hash = "sha256:961ed60399c457ceb0ee8f285a84c870aabc9c6a832b9d37bb281b5bebde43ed", size = 300631, upload-time = "2026-05-07T08:03:30.345Z" }, -] - [[package]] name = "greenlet" version = "3.2.4" @@ -1301,7 +939,7 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/5d/57/dfb3c5c3f1bf5f5ef2e59a22dec4ff1f3d7408b55bfcefcfb0ea69ef21c6/h5py-3.14.0.tar.gz", hash = "sha256:2372116b2e0d5d3e5e705b7f663f7c8d96fa79a4052d250484ef91d24d6a08f4", size = 424323, upload-time = "2025-06-06T14:06:15.01Z" } wheels = [ @@ -1343,7 +981,7 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/db/33/acd0ce6863b6c0d7735007df01815403f5589a21ff8c2e1ee2587a38f548/h5py-3.16.0.tar.gz", hash = "sha256:a0dbaad796840ccaa67a4c144a0d0c8080073c34c76d5a6941d6818678ef2738", size = 446526, upload-time = "2026-03-06T13:49:08.07Z" } wheels = [ @@ -1465,7 +1103,7 @@ name = "importlib-resources" version = "6.5.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "zipp", marker = "python_full_version < '3.10'" }, + { name = "zipp" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cf/8c/f834fbf984f691b4f7ff60f50b514cc3de5cc08abfc3295564dd89c5e2e7/importlib_resources-6.5.2.tar.gz", hash = "sha256:185f87adef5bcc288449d98fb4fba07cea78bc036455dd44c5fc4a2fe78fed2c", size = 44693, upload-time = "2025-01-03T18:51:56.698Z" } wheels = [ @@ -1515,17 +1153,17 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version < '3.10'" }, - { name = "exceptiongroup", marker = "python_full_version < '3.10'" }, - { name = "jedi", marker = "python_full_version < '3.10'" }, - { name = "matplotlib-inline", marker = "python_full_version < '3.10'" }, - { name = "pexpect", marker = "python_full_version < '3.10' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version < '3.10'" }, - { name = "pygments", marker = "python_full_version < '3.10'" }, - { name = "stack-data", marker = "python_full_version < '3.10'" }, - { name = "traitlets", marker = "python_full_version < '3.10'" }, - { name = "typing-extensions", marker = "python_full_version < '3.10'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "exceptiongroup" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b1/b9/3ba6c45a6df813c09a48bac313c22ff83efa26cbb55011218d925a46e2ad/ipython-8.18.1.tar.gz", hash = "sha256:ca6f079bb33457c66e233e4580ebfc4128855b4cf6370dddd73842a9563e8a27", size = 5486330, upload-time = "2023-11-27T09:58:34.596Z" } wheels = [ @@ -1543,17 +1181,17 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.10' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version >= '3.10'" }, - { name = "exceptiongroup", marker = "python_full_version == '3.10.*'" }, - { name = "jedi", marker = "python_full_version >= '3.10'" }, - { name = "matplotlib-inline", marker = "python_full_version >= '3.10'" }, - { name = "pexpect", marker = "python_full_version >= '3.10' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version >= '3.10'" }, - { name = "pygments", marker = "python_full_version >= '3.10'" }, - { name = "stack-data", marker = "python_full_version >= '3.10'" }, - { name = "traitlets", marker = "python_full_version >= '3.10'" }, - { name = "typing-extensions", marker = "python_full_version >= '3.10' and python_full_version < '3.12'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/85/31/10ac88f3357fc276dc8a64e8880c82e80e7459326ae1d0a211b40abf6665/ipython-8.37.0.tar.gz", hash = "sha256:ca815841e1a41a1e6b73a0b08f3038af9b2252564d01fc405356d34033012216", size = 5606088, upload-time = "2025-05-31T16:39:09.613Z" } wheels = [ @@ -2028,7 +1666,7 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "markdown-it-py", marker = "python_full_version < '3.10'" }, + { name = "markdown-it-py" }, ] sdist = { url = "https://files.pythonhosted.org/packages/19/03/a2ecab526543b152300717cf232bb4bb8605b6edb946c845016fa9c9c9fd/mdit_py_plugins-0.4.2.tar.gz", hash = "sha256:5f2cd1fdb606ddf152d37ec30e46101a60512bc0e5fa1a7002c36647b09e26b5", size = 43542, upload-time = "2024-09-09T20:27:49.564Z" } wheels = [ @@ -2046,7 +1684,7 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "markdown-it-py", marker = "python_full_version >= '3.10'" }, + { name = "markdown-it-py" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" } wheels = [ @@ -2143,11 +1781,11 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "librt", marker = "python_full_version < '3.10' and platform_python_implementation != 'PyPy'" }, - { name = "mypy-extensions", marker = "python_full_version < '3.10'" }, - { name = "pathspec", marker = "python_full_version < '3.10'" }, - { name = "tomli", marker = "python_full_version < '3.10'" }, - { name = "typing-extensions", marker = "python_full_version < '3.10'" }, + { name = "librt", marker = "platform_python_implementation != 'PyPy'" }, + { name = "mypy-extensions" }, + { name = "pathspec" }, + { name = "tomli" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f5/db/4efed9504bc01309ab9c2da7e352cc223569f05478012b5d9ece38fd44d2/mypy-1.19.1.tar.gz", hash = "sha256:19d88bb05303fe63f71dd2c6270daca27cb9401c4ca8255fe50d1d920e0eb9ba", size = 3582404, upload-time = "2025-12-15T05:03:48.42Z" } wheels = [ @@ -2201,11 +1839,11 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "librt", marker = "python_full_version >= '3.10' and platform_python_implementation != 'PyPy'" }, - { name = "mypy-extensions", marker = "python_full_version >= '3.10'" }, - { name = "pathspec", marker = "python_full_version >= '3.10'" }, - { name = "tomli", marker = "python_full_version == '3.10.*'" }, - { name = "typing-extensions", marker = "python_full_version >= '3.10'" }, + { name = "librt", marker = "platform_python_implementation != 'PyPy'" }, + { name = "mypy-extensions" }, + { name = "pathspec" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f8/5c/b0089fe7fef0a994ae5ee07029ced0526082c6cfaaa4c10d40a10e33b097/mypy-1.20.0.tar.gz", hash = "sha256:eb96c84efcc33f0b5e0e04beacf00129dd963b67226b01c00b9dfc8affb464c3", size = 3815028, upload-time = "2026-03-31T16:55:14.959Z" } wheels = [ @@ -2459,7 +2097,7 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/21/67/c7415cf04ebe418193cfd6595ae03e3a64d76dac7b9c010098b39cc7992e/numexpr-2.10.2.tar.gz", hash = "sha256:b0aff6b48ebc99d2f54f27b5f73a58cb92fde650aeff1b397c71c8788b4fff1a", size = 106787, upload-time = "2024-11-23T13:34:23.127Z" } wheels = [ @@ -2511,7 +2149,7 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/d2/8f/2cc977e91adbfbcdb6b49fdb9147e1d1c7566eb2c0c1e737e9a47020b5ca/numexpr-2.11.0.tar.gz", hash = "sha256:75b2c01a4eda2e7c357bc67a3f5c3dd76506c15b5fd4dc42845ef2e182181bad", size = 108960, upload-time = "2025-06-09T11:05:56.79Z" } wheels = [ @@ -2820,12 +2458,9 @@ wheels = [ [[package]] name = "policyengine" -version = "5.0.2" +version = "5.0.4" source = { editable = "." } dependencies = [ - { name = "diskcache" }, - { name = "google-cloud-storage", version = "3.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "google-cloud-storage", version = "3.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, { name = "h5py", version = "3.14.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "h5py", version = "3.16.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, { name = "jsonschema" }, @@ -2883,9 +2518,7 @@ us = [ requires-dist = [ { name = "autodoc-pydantic", marker = "extra == 'dev'" }, { name = "build", marker = "extra == 'dev'" }, - { name = "diskcache", specifier = ">=5.6.3,<6.0.0" }, { name = "furo", marker = "extra == 'dev'" }, - { name = "google-cloud-storage", specifier = ">=3.1.0,<4.0.0" }, { name = "h5py", specifier = ">=3.0.0" }, { name = "itables", marker = "extra == 'dev'" }, { name = "jsonschema", specifier = ">=4.0.0" }, @@ -2998,82 +2631,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, ] -[[package]] -name = "proto-plus" -version = "1.27.2" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version > '3.9' and python_full_version < '3.10'", - "python_full_version <= '3.9'", -] -dependencies = [ - { name = "protobuf", version = "6.33.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/81/0d/94dfe80193e79d55258345901acd2917523d56e8381bc4dee7fd38e3868a/proto_plus-1.27.2.tar.gz", hash = "sha256:b2adde53adadf75737c44d3dcb0104fde65250dfc83ad59168b4aa3e574b6a24", size = 57204, upload-time = "2026-03-26T22:18:57.174Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/84/f3/1fba73eeffafc998a25d59703b63f8be4fe8a5cb12eaff7386a0ba0f7125/proto_plus-1.27.2-py3-none-any.whl", hash = "sha256:6432f75893d3b9e70b9c412f1d2f03f65b11fb164b793d14ae2ca01821d22718", size = 50450, upload-time = "2026-03-26T22:13:42.927Z" }, -] - -[[package]] -name = "proto-plus" -version = "1.28.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", -] -dependencies = [ - { name = "protobuf", version = "7.35.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c9/56/e647b0c675392d2da368da7b6f158f7368b18542fd6f7d7400a2f39de000/proto_plus-1.28.0.tar.gz", hash = "sha256:38e5696342835b08fc116f30a25665b29531cda9d5d5643e9b81fc312385abd9", size = 57221, upload-time = "2026-05-07T08:04:50.811Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/20/b122d4626976acb81132036d2ad1bb35a1a8775fceb837ec30964622516a/proto_plus-1.28.0-py3-none-any.whl", hash = "sha256:a630604310899e73c59ec302e5765c058d412b2f090b9c79c8822589f14955b8", size = 50410, upload-time = "2026-05-07T08:03:31.962Z" }, -] - -[[package]] -name = "protobuf" -version = "6.33.6" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version > '3.9' and python_full_version < '3.10'", - "python_full_version <= '3.9'", -] -sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, - { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, - { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, - { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, - { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, - { url = "https://files.pythonhosted.org/packages/0c/bd/88a687e9147329fc7e6c26a058fc52214c47190688a496bb283000a4d2a3/protobuf-6.33.6-cp39-cp39-win32.whl", hash = "sha256:bd56799fb262994b2c2faa1799693c95cc2e22c62f56fb43af311cae45d26f0e", size = 425861, upload-time = "2026-03-18T19:04:57.064Z" }, - { url = "https://files.pythonhosted.org/packages/84/d6/fab384eea064bfc3b273183e4e09bb3a3cf4ec83876b3828c09fcacbb651/protobuf-6.33.6-cp39-cp39-win_amd64.whl", hash = "sha256:f443a394af5ed23672bc6c486be138628fbe5c651ccbc536873d7da23d1868cf", size = 437109, upload-time = "2026-03-18T19:04:58.713Z" }, - { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, -] - -[[package]] -name = "protobuf" -version = "7.35.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13'", - "python_full_version == '3.12.*'", - "python_full_version == '3.11.*'", - "python_full_version == '3.10.*'", -] -sdist = { url = "https://files.pythonhosted.org/packages/60/fd/5b1491d9e4b586d621c54f4c36b888714164b6875f8d6afa3f9072906a51/protobuf-7.35.0.tar.gz", hash = "sha256:a2efd84605f41e559f1881b0912b44099d0a2ac9bf46b3474823f10fb393b0e6", size = 458677, upload-time = "2026-05-19T23:02:29.197Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/83/ee/93d06e358a4aa32280b00e722d3ea0a1f25fc3cc5778d80581c9cca2c10e/protobuf-7.35.0-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:66be6c513931c794fa92c080ffee41671390da3d79da219cf9c0c0907f035dda", size = 433225, upload-time = "2026-05-19T23:02:19.884Z" }, - { url = "https://files.pythonhosted.org/packages/8b/39/1c76c2da93f3c507e958e0aecee2391cc44d4625de6c728bbc555195b5a8/protobuf-7.35.0-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:fcbe42a4ac09d3ec9c987ddfcd956afd0b15f1ff613bd8371bde9405ffd5c8e5", size = 328847, upload-time = "2026-05-19T23:02:22.3Z" }, - { url = "https://files.pythonhosted.org/packages/91/1a/39f7ce90a238c1a987a4d81ec26379e02ca0aff367de68e4a1fa474215b9/protobuf-7.35.0-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:4cbf5cc286130e06a6c9bbefac442431173906dfcc979712183d4adcc01b37ee", size = 344030, upload-time = "2026-05-19T23:02:23.591Z" }, - { url = "https://files.pythonhosted.org/packages/70/5b/6baf9008817964454055ff3fe65f1de0b5f1e26c80c82f7fb108b7cd4ea3/protobuf-7.35.0-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:6c0f98f10c8a05ea30f8993dfef2de093d27b490fdae78bb60c8343795d55011", size = 327130, upload-time = "2026-05-19T23:02:24.637Z" }, - { url = "https://files.pythonhosted.org/packages/8e/e5/e46adb0badc388bfb84877a5f9f026aff63f60e611016cf64dbe77e05446/protobuf-7.35.0-cp310-abi3-win32.whl", hash = "sha256:4c4617b83ade0e279d1d2bfe04025a1adb87f9ed657de038620dc0ff959357f6", size = 428946, upload-time = "2026-05-19T23:02:25.741Z" }, - { url = "https://files.pythonhosted.org/packages/a7/ab/547fbd9e16d879dd13c167478f8ae0a83a428008ca07a5e06acdc23ad473/protobuf-7.35.0-cp310-abi3-win_amd64.whl", hash = "sha256:f05bcadf9a2a6b8dda047007075135fb7d08c73d9177aabc067e1be46881a201", size = 439996, upload-time = "2026-05-19T23:02:26.808Z" }, - { url = "https://files.pythonhosted.org/packages/b8/ef/50433d346c56657a70d27f156c7b349ac59a068b01de4eb796e747eecc43/protobuf-7.35.0-py3-none-any.whl", hash = "sha256:c13f325cf242bad135c350629eeb5d54b24228eb472fb3e2e9ebbd4c5dc20ca0", size = 171659, upload-time = "2026-05-19T23:02:27.842Z" }, -] - [[package]] name = "psutil" version = "6.1.1" @@ -3116,27 +2673,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, ] -[[package]] -name = "pyasn1" -version = "0.6.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5c/5f/6583902b6f79b399c9c40674ac384fd9cd77805f9e6205075f828ef11fb2/pyasn1-0.6.3.tar.gz", hash = "sha256:697a8ecd6d98891189184ca1fa05d1bb00e2f84b5977c481452050549c8a72cf", size = 148685, upload-time = "2026-03-17T01:06:53.382Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl", hash = "sha256:a80184d120f0864a52a073acc6fc642847d0be408e7c7252f31390c0f4eadcde", size = 83997, upload-time = "2026-03-17T01:06:52.036Z" }, -] - -[[package]] -name = "pyasn1-modules" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyasn1" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, -] - [[package]] name = "pybtex" version = "0.25.1" @@ -4222,11 +3758,11 @@ resolution-markers = [ "python_full_version <= '3.9'", ] dependencies = [ - { name = "blosc2", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "numexpr", version = "2.10.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, - { name = "packaging", marker = "python_full_version < '3.10'" }, - { name = "py-cpuinfo", marker = "python_full_version < '3.10'" }, + { name = "blosc2", version = "2.5.1", source = { registry = "https://pypi.org/simple" } }, + { name = "numexpr", version = "2.10.2", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" } }, + { name = "packaging" }, + { name = "py-cpuinfo" }, ] sdist = { url = "https://files.pythonhosted.org/packages/31/83/8a13be8338219c3fe0aa7357d1ec4edb27bc346e0f224df7212892b243b5/tables-3.9.2.tar.gz", hash = "sha256:d470263c2e50c4b7c8635a0d99ac1ff2f9e704c24d71e5fa33c4529e7d0ad9c3", size = 4683437, upload-time = "2023-11-27T11:53:17.229Z" } wheels = [ @@ -4256,12 +3792,12 @@ resolution-markers = [ "python_full_version == '3.10.*'", ] dependencies = [ - { name = "blosc2", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, - { name = "numexpr", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, - { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, - { name = "packaging", marker = "python_full_version == '3.10.*'" }, - { name = "py-cpuinfo", marker = "python_full_version == '3.10.*'" }, - { name = "typing-extensions", marker = "python_full_version == '3.10.*'" }, + { name = "blosc2", version = "3.7.2", source = { registry = "https://pypi.org/simple" } }, + { name = "numexpr", version = "2.11.0", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" } }, + { name = "packaging" }, + { name = "py-cpuinfo" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0d/5d/96708a84e9fcd29d1f684d56d4c38a23d29b1c934599a072a49f27ccfa71/tables-3.10.1.tar.gz", hash = "sha256:4aa07ac734b9c037baeaf44aec64ec902ad247f57811b59f30c4e31d31f126cf", size = 4762413, upload-time = "2024-08-17T09:57:47.127Z" } wheels = [ @@ -4296,11 +3832,11 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "blosc2", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "numexpr", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "packaging", marker = "python_full_version >= '3.11'" }, - { name = "py-cpuinfo", marker = "python_full_version >= '3.11'" }, + { name = "blosc2", version = "3.7.2", source = { registry = "https://pypi.org/simple" } }, + { name = "numexpr", version = "2.11.0", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.1.3", source = { registry = "https://pypi.org/simple" } }, + { name = "packaging" }, + { name = "py-cpuinfo" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/a3/d213ebe7376d48055bd55a29cd9f99061afa0dcece608f94a5025d797b0a/tables-3.11.1.tar.gz", hash = "sha256:78abcf413091bc7c1e4e8c10fbbb438d1ac0b5a87436c5b972c3e8253871b6fb", size = 4790533, upload-time = "2026-03-01T11:43:36.036Z" } wheels = [ From ecf229877f8dd9172d44e9a93b9182aa5427a37f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:47:50 +0400 Subject: [PATCH 06/18] Document bundle-based dataset loading --- changelog.d/bundle-dataset-gcs.removed.md | 2 + .../bundle-dataset-materialization.changed.md | 5 ++ docs/bundles.md | 28 ++++++++++ docs/data-publishing-design.md | 11 ++-- docs/microsim.md | 54 +++++++++++-------- docs/release-bundles.md | 15 ++++++ 6 files changed, 91 insertions(+), 24 deletions(-) create mode 100644 changelog.d/bundle-dataset-gcs.removed.md create mode 100644 changelog.d/bundle-dataset-materialization.changed.md diff --git a/changelog.d/bundle-dataset-gcs.removed.md b/changelog.d/bundle-dataset-gcs.removed.md new file mode 100644 index 00000000..09f0ca73 --- /dev/null +++ b/changelog.d/bundle-dataset-gcs.removed.md @@ -0,0 +1,2 @@ +Removed the legacy dataset-specific GCS downloader and its direct dependencies. +The separate UK geography asset implementation is unchanged. diff --git a/changelog.d/bundle-dataset-materialization.changed.md b/changelog.d/bundle-dataset-materialization.changed.md new file mode 100644 index 00000000..c69157f2 --- /dev/null +++ b/changelog.d/bundle-dataset-materialization.changed.md @@ -0,0 +1,5 @@ +PolicyEngine.py now materializes managed datasets from the exact Hugging Face +repository type, immutable revision, and SHA-256 certified by its release +bundle. Bundle installation and US and UK calculation entry points share this +implementation, while explicitly unmanaged local and Hugging Face sources +remain opt-in. diff --git a/docs/bundles.md b/docs/bundles.md index e3099709..1c6f8bc6 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -34,12 +34,40 @@ datasets into `./data`, moves replaced dataset files into `./data/.policyengine-bundle-receipt.json` receipt that records the target Python. +Dataset pre-download uses the same materialization function as US and UK +calculations. For every managed artifact, PolicyEngine.py reads the source data +package name, Hugging Face repository type, immutable revision, and SHA-256 from +the bundle. It reuses an existing file only when its hash matches, downloads and +verifies replacements before moving the old file into the backup directory, and +records the verified result in the receipt. + The bundle manifest can certify additional regional datasets, such as US state datasets. Those artifacts are part of the citable bundle manifest, but `policyengine bundle install` does not eagerly download every regional file. Runtime callers should use the manifest's regional dataset URI when a regional simulation needs one. +To materialize a default or named artifact without installing the complete +package scaffold: + +```python +from policyengine.provenance import materialize_bundle_dataset + +result = materialize_bundle_dataset("us", "populace_us_2024") +print(result.path) +print(result.actual_sha256) +``` + +`materialize_bundle_dataset` returns a Pydantic model containing the selected +source package, repository type, revision, expected and actual hashes, local +path, and cache status. `policyengine-*-data` and `populace-data` artifacts are +selected by their bundle package names. Callers do not infer repository type +from the repository name. + +Managed datasets are downloaded from the Hugging Face artifact specified in the +bundle. GCS dataset URIs are unsupported. The separate UK geography lookup files +retain their existing storage implementation. + Country-specific and package-only installs are supported: ```bash diff --git a/docs/data-publishing-design.md b/docs/data-publishing-design.md index 9bdb33fb..3e8c40da 100644 --- a/docs/data-publishing-design.md +++ b/docs/data-publishing-design.md @@ -162,17 +162,22 @@ concrete `artifact_sha256` pin in the country release manifest. After that, the release manifest is what papers cite; the storage channel is just the cache. -## Consumer resolver (what `pe.py` changes) +## Consumer resolver (historical proposal) -Minimal. The existing `pe.us.ensure_datasets` takes a URI today: +PolicyEngine.py now resolves managed datasets by logical name through its +certified release bundle: ```python pe.us.ensure_datasets( - datasets=["hf://policyengine/populace-us/populace_us_2024.h5@"], + datasets=["populace_us_2024"], years=[2026], ) ``` +Direct Hugging Face references require `allow_unmanaged=True`; GCS dataset +references are unsupported. The `pe-data://` examples below remain an +unimplemented design proposal rather than a description of current behavior. + Under the substrate, the URI scheme gains a new prefix: ```python diff --git a/docs/microsim.md b/docs/microsim.md index 34fb376b..4f108737 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -40,7 +40,15 @@ datasets = pe.us.ensure_datasets( dataset = datasets["populace_us_2024_2026"] ``` -The default US dataset is **Populace US 2024** — a Populace-built dataset calibrated to IRS, CMS, SNAP, Census, and other administrative totals. The UK default is **Populace UK 2023** — a Populace-built Family Resources Survey dataset calibrated to UK administrative targets. +The default US dataset is **Populace US 2024** — a Populace-built dataset +calibrated to IRS, CMS, SNAP, Census, and other administrative totals. The +current UK certified default is **Enhanced FRS 2024–25**, supplied by +`policyengine-uk-data`. **Populace UK 2023** remains available as a named, +non-default bundle dataset. + +PolicyEngine.py obtains the repository type, immutable revision, and SHA-256 +from the installed release bundle. A cached file or local data-repository mirror +is reused only after hash verification. List datasets already known to the country: @@ -97,7 +105,7 @@ ca = Simulation( UK population data uses licensed Family Resources Survey inputs. The default UK release bundle points to the private -`policyengine/populace-uk-private` Hugging Face dataset repository. Set +`policyengine/policyengine-uk-data-private` Hugging Face repository. Set `HUGGING_FACE_TOKEN` to a token from a Hugging Face account with access: ```bash @@ -113,11 +121,11 @@ import policyengine as pe from policyengine.core import Simulation datasets = pe.uk.ensure_datasets( - datasets=["populace_uk_2023"], + datasets=["enhanced_frs_2024_25"], years=[2026], data_folder="./data", ) -dataset = datasets["populace_uk_2023_2026"] +dataset = datasets["enhanced_frs_2024_25_2026"] simulation = Simulation( dataset=dataset, @@ -126,28 +134,25 @@ simulation = Simulation( simulation.run() ``` -To download the raw h5 artifact directly from Hugging Face, use -`huggingface_hub` and specify `repo_type="dataset"`: +To materialize the raw certified artifact without creating uprated yearly +datasets, use PolicyEngine.py's bundle API: ```python -import os -from huggingface_hub import hf_hub_download - -path = hf_hub_download( - repo_id="policyengine/populace-uk-private", - filename="populace_uk_2023.h5", - repo_type="dataset", - token=os.environ["HUGGING_FACE_TOKEN"], +from policyengine.provenance import materialize_bundle_dataset + +result = materialize_bundle_dataset( + "uk", + "enhanced_frs_2024_25", ) -print(path) +print(result.path) +print(result.actual_sha256) ``` -The repository URL is -. A 404 from -the website or `RepositoryNotFoundError` from the Hub API usually means the -browser or token is not authenticated as an account with access, or that the -Hub call omitted `repo_type="dataset"`. +The bundle API uses the repository type recorded in the bundle, so callers do +not need repository-specific download logic. Authentication or authorization +failures are reported directly and do not cause a retry against another +repository type. ## Simulations @@ -219,6 +224,7 @@ Smaller custom H5 datasets can be passed explicitly for testing: datasets = pe.us.ensure_datasets( datasets=["/path/to/smoke_test_populace_us_2024.h5"], years=[2026], + allow_unmanaged=True, ) ``` @@ -235,7 +241,13 @@ sim = managed_microsimulation() # `sim` is a policyengine_us.Microsimulation — use its API directly ``` -Pass `allow_unmanaged=True` with a custom `dataset=` to opt out of the release bundle. +Pass `allow_unmanaged=True` with a custom `dataset=` to opt out of the release +bundle. Explicit local paths and Hugging Face URIs remain supported in this +mode. GCS dataset URIs are not supported. + +For managed simulations, `sim.policyengine_bundle` records the actual source +package, repository type, revision, expected and actual SHA-256, local path, and +whether an already verified file was reused. ## Pinned model versions diff --git a/docs/release-bundles.md b/docs/release-bundles.md index cc9c5126..b9bd780f 100644 --- a/docs/release-bundles.md +++ b/docs/release-bundles.md @@ -44,6 +44,11 @@ Python. Existing dataset files with the same filename are moved to `./data/.policyengine-bundle-backups//`. +The command invokes the same bundle materializer used by calculations. The +materializer selects its internal strategy from the artifact's data package +name, uses the manifest's exact Hugging Face repository type and immutable +revision, and verifies the certified SHA-256 before exposing the local file. + Regional datasets may also be certified in the bundle manifest. They are not eagerly downloaded by `policyengine bundle install`; callers should materialize the certified regional URI from the manifest when they run a regional @@ -153,8 +158,10 @@ sibling `dataset_overlays.{country}` map: "dataset_overlays": { "us": { "populace_us_2024_acs_local": { + "data_package_name": "populace-data", "path": "populace_us_2024_acs_local.h5", "repo_id": "policyengine/populace-us", + "repo_type": "dataset", "revision": "populace-us-2024-buildo-acs-local-...", "sha256": "..." } @@ -170,6 +177,11 @@ default resolution is untouched. Because certification only rewrites `data_releases`, overlays survive re-certification without any manual re-add step. +Cross-package overlays must declare `data_package_name` and `repo_type` +explicitly. Ordinary artifacts inherit these values from the country release's +primary `data_package`. This prevents runtime code from guessing how a +repository should be addressed. + Earlier releases (policyengine 4.15.x–4.16.x) were certified through the `PolicyEngine/policyengine-bundles` archive flow; those bundles remain the historical record of their certifications. @@ -601,6 +613,9 @@ The target implementation in `policyengine.py` should add: - explicit runtime bundle metadata on simulations, APIs, and app responses - checksum-backed dataset resolution from the certified bundle manifest +The checksum-backed runtime resolution described above is now implemented. +Managed materialization is Hugging Face-only; GCS is not a dataset source. + ## Why not let `policyengine.py` build all country data directly? Because that would centralise the wrong concerns: From 3e2b36eeb2d7898c3b9be89a3a356c818daabafb Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:58:23 +0400 Subject: [PATCH 07/18] Align dataset defaults and provenance metadata --- .../data/bundle/uk.trace.tro.jsonld | 4 +- .../data/bundle/us.trace.tro.jsonld | 4 +- .../provenance/dataset_materialization.py | 1 + .../tax_benefit_models/uk/datasets.py | 19 +++--- tests/test_dataset_materialization.py | 8 +++ tests/test_dataset_runtime.py | 63 +++++++++++++++++++ 6 files changed, 86 insertions(+), 13 deletions(-) diff --git a/src/policyengine/data/bundle/uk.trace.tro.jsonld b/src/policyengine/data/bundle/uk.trace.tro.jsonld index 5fd41cb6..da05bf47 100644 --- a/src/policyengine/data/bundle/uk.trace.tro.jsonld +++ b/src/policyengine/data/bundle/uk.trace.tro.jsonld @@ -75,7 +75,7 @@ "@type": "trov:ResearchArtifact", "schema:name": "policyengine.py bundle manifest for uk", "trov:mimeType": "application/json", - "trov:sha256": "7906f50dcef3adb586ef4a6e3011caf0d3b4096c53c8843c41551e96f929537a" + "trov:sha256": "b2f6dc37f9597ae0932c3b3b926f17cfd1cb622727bebe7a998f94c84b30adde" }, { "@id": "composition/1/artifact/data_release_manifest", @@ -102,7 +102,7 @@ "trov:hasFingerprint": { "@id": "composition/1/fingerprint", "@type": "trov:CompositionFingerprint", - "trov:sha256": "a7da1e779a5933e4be99c8b1f6a1b72972faa6f904d11059f64ef343807b2c38" + "trov:sha256": "b3eaab80f13125c431d83da8fb4015814145513481fa6703f96502d9da787fa0" } }, "trov:hasPerformance": { diff --git a/src/policyengine/data/bundle/us.trace.tro.jsonld b/src/policyengine/data/bundle/us.trace.tro.jsonld index b79d80c4..6e2ee6f8 100644 --- a/src/policyengine/data/bundle/us.trace.tro.jsonld +++ b/src/policyengine/data/bundle/us.trace.tro.jsonld @@ -75,7 +75,7 @@ "@type": "trov:ResearchArtifact", "schema:name": "policyengine.py bundle manifest for us", "trov:mimeType": "application/json", - "trov:sha256": "7906f50dcef3adb586ef4a6e3011caf0d3b4096c53c8843c41551e96f929537a" + "trov:sha256": "b2f6dc37f9597ae0932c3b3b926f17cfd1cb622727bebe7a998f94c84b30adde" }, { "@id": "composition/1/artifact/data_release_manifest", @@ -102,7 +102,7 @@ "trov:hasFingerprint": { "@id": "composition/1/fingerprint", "@type": "trov:CompositionFingerprint", - "trov:sha256": "1cdd1b5d66894fdbafbbc270cc2056859ad02c944673c6476318087bf0be61dc" + "trov:sha256": "747a5d2b33a4daa2104af3fa832c89bc2754ac88c06c03c0a444030d6fb93c37" } }, "trov:hasPerformance": { diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 02a351d7..695f64cf 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -184,6 +184,7 @@ def resolve_bundle_dataset_plan( build_id=( country_manifest.certified_data_artifact.build_id if country_manifest.certified_data_artifact is not None + and reference.data_package_name is None else None ), metadata_expected_sha256=reference.metadata_sha256, diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index 086c96a2..2fcb2d7d 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -118,13 +118,13 @@ def __repr__(self) -> str: def create_datasets( - datasets: list[str] = [ - "populace_uk_2023", - ], + datasets: Optional[list[str]] = None, years: list[int] = [2026, 2027, 2028, 2029, 2030], data_folder: str = "./data", allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUKDataset]: + if datasets is None: + datasets = [get_release_manifest("uk").default_dataset] result = {} for dataset in datasets: manifest = get_release_manifest("uk") @@ -218,12 +218,12 @@ def create_datasets( def load_datasets( - datasets: list[str] = [ - "populace_uk_2023", - ], + datasets: Optional[list[str]] = None, years: list[int] = [2026, 2027, 2028, 2029, 2030], data_folder: str = "./data", ) -> dict[str, PolicyEngineUKDataset]: + if datasets is None: + datasets = [get_release_manifest("uk").default_dataset] result = {} for dataset in datasets: resolved_dataset = resolve_dataset_reference("uk", dataset) @@ -245,9 +245,7 @@ def load_datasets( def ensure_datasets( - datasets: list[str] = [ - "populace_uk_2023", - ], + datasets: Optional[list[str]] = None, years: list[int] = [2026, 2027, 2028, 2029, 2030], data_folder: str = "./data", allow_unmanaged: bool = False, @@ -262,6 +260,9 @@ def ensure_datasets( Returns: Dictionary mapping dataset keys to PolicyEngineUKDataset objects """ + if datasets is None: + datasets = [get_release_manifest("uk").default_dataset] + # Check if all dataset files exist all_exist = True for dataset in datasets: diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 633bf537..1bcf0946 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -28,6 +28,12 @@ def _manifest() -> CountryReleaseManifest: "repo_type": "model", }, "default_dataset": "enhanced_frs_2024_25", + "certified_data_artifact": { + "dataset": "enhanced_frs_2024_25", + "uri": "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2024_25.h5@uk-release", + "sha256": "a" * 64, + "build_id": "policyengine-uk-data-test-build", + }, "datasets": { "enhanced_frs_2024_25": { "path": "enhanced_frs_2024_25.h5", @@ -55,6 +61,7 @@ def test_resolve_bundle_dataset_plan_inherits_primary_package(tmp_path): assert plan.repo_type == "model" assert plan.revision == "uk-release" assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" + assert plan.build_id == "policyengine-uk-data-test-build" def test_resolve_bundle_dataset_plan_uses_cross_package_overlay(tmp_path): @@ -69,6 +76,7 @@ def test_resolve_bundle_dataset_plan_uses_cross_package_overlay(tmp_path): assert plan.repo_id == "policyengine/populace-uk-private" assert plan.repo_type == "model" assert plan.revision == "populace-release" + assert plan.build_id is None def test_bundle_dataset_models_round_trip_json(): diff --git a/tests/test_dataset_runtime.py b/tests/test_dataset_runtime.py index 11f82b56..865de7b3 100644 --- a/tests/test_dataset_runtime.py +++ b/tests/test_dataset_runtime.py @@ -90,3 +90,66 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( "uk", "populace_uk_2023", data_dir=Path("./data") ) microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") + + +def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): + uk_datasets = _load_module_from_path( + "_test_policyengine_uk_default_create_datasets", + REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", + ) + materialize = Mock( + return_value=_materialized( + "uk", + "enhanced_frs_2024_25", + "/tmp/enhanced_frs_2024_25.h5", + ) + ) + microsimulation = Mock() + monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setitem( + sys.modules, + "policyengine_uk", + SimpleNamespace(Microsimulation=microsimulation), + ) + + uk_datasets.create_datasets(years=[]) + + materialize.assert_called_once_with( + "uk", "enhanced_frs_2024_25", data_dir=Path("./data") + ) + microsimulation.assert_called_once_with(dataset="/tmp/enhanced_frs_2024_25.h5") + + +def test_uk_load_datasets_defaults_to_certified_bundle_dataset(monkeypatch): + uk_datasets = _load_module_from_path( + "_test_policyengine_uk_default_load_datasets", + REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", + ) + resolve = Mock( + return_value=( + "hf://policyengine/policyengine-uk-data-private/" + "enhanced_frs_2024_25.h5@1.56.16" + ) + ) + monkeypatch.setattr(uk_datasets, "resolve_dataset_reference", resolve) + + assert uk_datasets.load_datasets(years=[]) == {} + + resolve.assert_called_once_with("uk", "enhanced_frs_2024_25") + + +def test_uk_ensure_datasets_defaults_to_certified_bundle_dataset(monkeypatch): + uk_datasets = _load_module_from_path( + "_test_policyengine_uk_default_ensure_datasets", + REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", + ) + load = Mock(return_value={}) + monkeypatch.setattr(uk_datasets, "load_datasets", load) + + assert uk_datasets.ensure_datasets(years=[]) == {} + + load.assert_called_once_with( + datasets=["enhanced_frs_2024_25"], + years=[], + data_folder="./data", + ) From 9185a41977f24a01d23195864f77d3742d327e7e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:08:04 +0400 Subject: [PATCH 08/18] Exclude cached bytecode from distributions --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index b0affd10..1f956b53 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,9 @@ where = ["src"] [tool.setuptools.package-data] "policyengine" = ["**/*"] +[tool.setuptools.exclude-package-data] +"policyengine" = ["**/__pycache__/*", "**/*.pyc", "**/*.pyo"] + [tool.pytest.ini_options] addopts = "-v" testpaths = [ From 419e46087271fda73e6fb2698e9c041ca5b59f18 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 28 Aug 2026 20:40:10 +0400 Subject: [PATCH 09/18] Reference issue in changelog fragments --- .../{bundle-dataset-materialization.changed.md => 502.changed.md} | 0 changelog.d/{bundle-dataset-gcs.removed.md => 502.removed.md} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename changelog.d/{bundle-dataset-materialization.changed.md => 502.changed.md} (100%) rename changelog.d/{bundle-dataset-gcs.removed.md => 502.removed.md} (100%) diff --git a/changelog.d/bundle-dataset-materialization.changed.md b/changelog.d/502.changed.md similarity index 100% rename from changelog.d/bundle-dataset-materialization.changed.md rename to changelog.d/502.changed.md diff --git a/changelog.d/bundle-dataset-gcs.removed.md b/changelog.d/502.removed.md similarity index 100% rename from changelog.d/bundle-dataset-gcs.removed.md rename to changelog.d/502.removed.md From 70770d269c521b0a43665ace94a4bd6b26ce6dcf Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 28 Aug 2026 21:09:50 +0400 Subject: [PATCH 10/18] Simplify dataset package strategy dispatch --- .../provenance/dataset_materialization.py | 63 +++---------------- tests/test_dataset_materialization.py | 16 ++++- 2 files changed, 23 insertions(+), 56 deletions(-) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 695f64cf..0b95c0c1 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -70,63 +70,16 @@ class MaterializedDataset(BaseModel): metadata_path: Optional[Path] = None -class _DatasetPackageStrategy: - """Package-specific behavior for a bundle dataset source.""" - - def validate(self, plan: BundleDatasetPlan) -> None: - raise NotImplementedError - - def materialize( - self, - plan: BundleDatasetPlan, - *, - session=requests, - ) -> MaterializedDataset: - raise NotImplementedError - - -class _CountryDataPackageStrategy(_DatasetPackageStrategy): - def validate(self, plan: BundleDatasetPlan) -> None: - package_name = plan.data_package_name - if not ( - package_name.startswith("policyengine-") and package_name.endswith("-data") - ): - raise DatasetMaterializationError( - f"Unsupported country data package: {package_name!r}." - ) - - def materialize( - self, - plan: BundleDatasetPlan, - *, - session=requests, - ) -> MaterializedDataset: - return _materialize_country_data_package(plan, session=session) - - -class _PopulaceDataPackageStrategy(_DatasetPackageStrategy): - def validate(self, plan: BundleDatasetPlan) -> None: - if plan.data_package_name != "populace-data": - raise DatasetMaterializationError( - f"Unsupported Populace data package: {plan.data_package_name!r}." - ) +_DatasetPackageType = Literal["country", "populace"] - def materialize( - self, - plan: BundleDatasetPlan, - *, - session=requests, - ) -> MaterializedDataset: - return _materialize_populace_data_package(plan, session=session) - -def _dataset_package_strategy(data_package_name: str) -> _DatasetPackageStrategy: +def _dataset_package_type(data_package_name: str) -> _DatasetPackageType: if data_package_name == "populace-data": - return _PopulaceDataPackageStrategy() + return "populace" if data_package_name.startswith("policyengine-") and data_package_name.endswith( "-data" ): - return _CountryDataPackageStrategy() + return "country" raise DatasetMaterializationError( "Unsupported bundle data package " f"{data_package_name!r}; expected 'populace-data' or " @@ -199,7 +152,7 @@ def resolve_bundle_dataset_plan( else None ), ) - _dataset_package_strategy(data_package_name).validate(plan) + _dataset_package_type(data_package_name) return plan @@ -249,8 +202,10 @@ def materialize_bundle_dataset( metadata_actual_sha256=metadata_actual_sha256, metadata_path=metadata_path, ) - strategy = _dataset_package_strategy(plan.data_package_name) - return strategy.materialize(plan, session=session) + package_type = _dataset_package_type(plan.data_package_name) + if package_type == "populace": + return _materialize_populace_data_package(plan, session=session) + return _materialize_country_data_package(plan, session=session) def _materialize_country_data_package( diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 1bcf0946..0b4074b3 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -7,7 +7,7 @@ BundleDatasetPlan, DatasetMaterializationError, MaterializedDataset, - _dataset_package_strategy, + _dataset_package_type, materialize_bundle_dataset, materialize_unmanaged_dataset_source, resolve_bundle_dataset_plan, @@ -92,9 +92,21 @@ def test_bundle_dataset_models_round_trip_json(): assert MaterializedDataset.model_validate_json(result.model_dump_json()) == result +@pytest.mark.parametrize( + ("package_name", "expected_type"), + [ + ("policyengine-us-data", "country"), + ("policyengine-uk-data", "country"), + ("populace-data", "populace"), + ], +) +def test_dataset_package_type(package_name, expected_type): + assert _dataset_package_type(package_name) == expected_type + + def test_unknown_data_package_is_rejected(): with pytest.raises(DatasetMaterializationError, match="Unsupported bundle"): - _dataset_package_strategy("unknown-data") + _dataset_package_type("unknown-data") def _sha256(payload: bytes) -> str: From 96e1695b94f887fea04af1d560ffdf530a783082 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 28 Aug 2026 21:42:02 +0400 Subject: [PATCH 11/18] Simplify bundle dataset downloads --- docs/bundles.md | 8 +- docs/release-bundles.md | 10 +- src/policyengine/bundle.py | 43 +- src/policyengine/provenance/__init__.py | 9 - .../provenance/dataset_materialization.py | 489 ++++++++---------- src/policyengine/provenance/manifest.py | 6 +- .../tax_benefit_models/uk/datasets.py | 30 +- .../tax_benefit_models/uk/model.py | 74 +-- .../tax_benefit_models/us/datasets.py | 79 +-- .../tax_benefit_models/us/model.py | 74 +-- tests/test_bundle.py | 16 +- tests/test_dataset_materialization.py | 155 +++--- tests/test_dataset_runtime.py | 44 +- tests/test_release_manifests.py | 29 +- tests/test_us_long_term_datasets.py | 10 +- 15 files changed, 411 insertions(+), 665 deletions(-) diff --git a/docs/bundles.md b/docs/bundles.md index 1c6f8bc6..e3c9b6da 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -29,16 +29,14 @@ When run from `uvx` or `pipx`, the installer creates or reuses `./.venv`. Inside an existing virtualenv or conda environment, it installs into that active environment. The installer then installs the exact bundled package scaffold with pip, downloads certified default US and UK -datasets into `./data`, moves replaced dataset files into -`./data/.policyengine-bundle-backups//`, and writes a -`./data/.policyengine-bundle-receipt.json` receipt that records the target -Python. +datasets into `./data`, and writes a `./data/.policyengine-bundle-receipt.json` +receipt that records the target Python. Dataset pre-download uses the same materialization function as US and UK calculations. For every managed artifact, PolicyEngine.py reads the source data package name, Hugging Face repository type, immutable revision, and SHA-256 from the bundle. It reuses an existing file only when its hash matches, downloads and -verifies replacements before moving the old file into the backup directory, and +verifies a replacement before atomically replacing an invalid local file, and records the verified result in the receipt. The bundle manifest can certify additional regional datasets, such as US state diff --git a/docs/release-bundles.md b/docs/release-bundles.md index b9bd780f..e7e2469b 100644 --- a/docs/release-bundles.md +++ b/docs/release-bundles.md @@ -41,13 +41,13 @@ environment. It installs the bundled Python packages with pip, downloads the certified default US and UK datasets into `./data`, and writes a `./data/.policyengine-bundle-receipt.json` receipt that records the target Python. -Existing dataset files with the same filename are moved to -`./data/.policyengine-bundle-backups//`. +An existing file is reused when its SHA-256 matches the manifest. Otherwise, a +verified download atomically replaces it. The command invokes the same bundle materializer used by calculations. The -materializer selects its internal strategy from the artifact's data package -name, uses the manifest's exact Hugging Face repository type and immutable -revision, and verifies the certified SHA-256 before exposing the local file. +materializer uses the manifest's exact Hugging Face repository type, immutable +revision, and certified SHA-256. The data package name is retained only as +provenance metadata; it does not select a download implementation. Regional datasets may also be certified in the bundle manifest. They are not eagerly downloaded by `policyengine bundle install`; callers should materialize diff --git a/src/policyengine/bundle.py b/src/policyengine/bundle.py index 02f9dec6..6dd63091 100644 --- a/src/policyengine/bundle.py +++ b/src/policyengine/bundle.py @@ -22,13 +22,12 @@ import requests from policyengine.provenance.dataset_materialization import ( - BACKUP_DIR_NAME, - BundleDatasetPlan, DatasetMaterializationError, MaterializedDataset, + _materialize_resolved_dataset, + _resolve_bundle_dataset, + _ResolvedBundleDataset, _sha256_file, - materialize_bundle_dataset, - resolve_bundle_dataset_plan, ) from policyengine.provenance.manifest import CountryReleaseManifest @@ -270,7 +269,7 @@ def install_package_scaffold( def _confirm_dataset_install( - plans: Sequence[BundleDatasetPlan], + plans: Sequence[_ResolvedBundleDataset], *, data_dir: Path, yes: bool, @@ -281,10 +280,7 @@ def _confirm_dataset_install( "This will download certified PolicyEngine datasets for " f"{countries} into {data_dir}." ) - print( - "Existing matching dataset files will be moved to " - f"{data_dir / BACKUP_DIR_NAME}//." - ) + print("Existing files with the certified content will be reused.") if yes or dry_run: return answer = input("Continue? [y/N] ").strip().lower() @@ -293,7 +289,7 @@ def _confirm_dataset_install( def _receipt_dataset( - plan: BundleDatasetPlan, + plan: _ResolvedBundleDataset, release: Mapping[str, Any], *, materialized: Optional[MaterializedDataset] = None, @@ -301,15 +297,15 @@ def _receipt_dataset( receipt = { "country": plan.country_id, "dataset": plan.dataset, - "version": release.get("version") or plan.build_id, + "version": release.get("version") or release.get("build_id"), "uri": plan.source_uri, "path": str(plan.destination), "release_manifest_uri": release.get("release_manifest_uri"), "data_package_name": plan.data_package_name, "repo_type": plan.repo_type, } - if plan.build_id: - receipt["build_id"] = plan.build_id + if release.get("build_id"): + receipt["build_id"] = release["build_id"] receipt["expected_sha256"] = plan.expected_sha256 if materialized is not None: receipt["installed_sha256"] = materialized.actual_sha256 @@ -321,7 +317,7 @@ def _selected_dataset_plans( countries: Sequence[str], *, data_dir: Path, -) -> list[tuple[BundleDatasetPlan, CountryReleaseManifest, Mapping[str, Any]]]: +) -> list[tuple[_ResolvedBundleDataset, Mapping[str, Any]]]: releases = manifest.get("data_releases") if not isinstance(releases, Mapping): raise BundleError("Bundle manifest does not contain data releases.") @@ -335,14 +331,14 @@ def _selected_dataset_plans( ) try: country_manifest = CountryReleaseManifest.model_validate(release) - plan = resolve_bundle_dataset_plan( + plan = _resolve_bundle_dataset( country, data_dir=data_dir, manifest=country_manifest, ) except (ValueError, DatasetMaterializationError) as exc: raise BundleError(str(exc)) from exc - selected.append((plan, country_manifest, release)) + selected.append((plan, release)) return selected @@ -416,18 +412,13 @@ def install_bundle( yes=yes, dry_run=dry_run, ) - for plan, country_manifest, release in dataset_entries: + for plan, release in dataset_entries: if dry_run: print(f"download {plan.source_uri} -> {plan.destination}") installed_datasets.append(_receipt_dataset(plan, release)) continue try: - materialized = materialize_bundle_dataset( - plan.country_id, - plan.dataset, - data_dir=data_dir, - manifest=country_manifest, - ) + materialized = _materialize_resolved_dataset(plan) except DatasetMaterializationError as exc: raise BundleError(str(exc)) from exc installed_datasets.append( @@ -638,7 +629,7 @@ def _dataset_checks( if isinstance(dataset, Mapping) and dataset.get("country"): receipt_datasets[str(dataset["country"])] = dataset checks = [] - for plan, _, release in _selected_dataset_plans( + for plan, release in _selected_dataset_plans( manifest, countries, data_dir=data_dir ): receipt_dataset = receipt_datasets.get(plan.country_id) @@ -647,11 +638,11 @@ def _dataset_checks( def _dataset_check( - plan: BundleDatasetPlan, + plan: _ResolvedBundleDataset, release: Mapping[str, Any], receipt_dataset: Optional[Mapping[str, Any]], ) -> dict[str, Any]: - expected_version = release.get("version") or plan.build_id + expected_version = release.get("version") or release.get("build_id") check: dict[str, Any] = { "country": plan.country_id, "dataset": plan.dataset, diff --git a/src/policyengine/provenance/__init__.py b/src/policyengine/provenance/__init__.py index 47010765..af781ea1 100644 --- a/src/policyengine/provenance/__init__.py +++ b/src/policyengine/provenance/__init__.py @@ -24,9 +24,6 @@ from .certification import ( certify_data_release as certify_data_release, ) -from .dataset_materialization import ( - BundleDatasetPlan as BundleDatasetPlan, -) from .dataset_materialization import ( DatasetMaterializationError as DatasetMaterializationError, ) @@ -36,12 +33,6 @@ from .dataset_materialization import ( materialize_bundle_dataset as materialize_bundle_dataset, ) -from .dataset_materialization import ( - materialize_unmanaged_dataset_source as materialize_unmanaged_dataset_source, -) -from .dataset_materialization import ( - resolve_bundle_dataset_plan as resolve_bundle_dataset_plan, -) from .manifest import ( CertifiedDataArtifact as CertifiedDataArtifact, ) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 0b95c0c1..ef99bc10 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -1,14 +1,13 @@ -"""Resolve and materialize datasets certified by a PolicyEngine.py bundle.""" +"""Make bundle-certified datasets available as verified local files.""" from __future__ import annotations import hashlib import os -import shutil import tempfile -from datetime import datetime, timezone +from dataclasses import dataclass from pathlib import Path -from typing import Literal, Optional, Union, cast +from typing import Literal, Optional, Union from urllib.parse import quote import requests @@ -18,142 +17,91 @@ CountryReleaseManifest, _artifact_revision, build_hf_uri, + dataset_logical_name, get_release_manifest, resolve_local_managed_dataset_source, + resolve_managed_dataset_reference, ) DEFAULT_DATA_DIR = Path("./data") -BACKUP_DIR_NAME = ".policyengine-bundle-backups" DOWNLOAD_TIMEOUT_SECONDS = 60 class DatasetMaterializationError(ValueError): - """Raised when a dataset cannot be resolved or materialized safely.""" + """Raised when a dataset cannot be made available safely.""" -class BundleDatasetPlan(BaseModel): - """Exact bundle metadata needed to materialize one managed dataset.""" +class MaterializedDataset(BaseModel): + """Verified local representation of one bundle-managed dataset.""" - country_id: str - dataset: str data_package_name: str - repo_id: str repo_type: Literal["model", "dataset"] - path: str revision: str - expected_sha256: str source_uri: str - destination: Path - build_id: Optional[str] = None - metadata_expected_sha256: Optional[str] = None - metadata_source_uri: Optional[str] = None - metadata_destination: Optional[Path] = None + expected_sha256: str + actual_sha256: str + path: Path + cache_hit: bool + metadata_path: Optional[Path] = None -class MaterializedDataset(BaseModel): - """Verified local representation of one bundle-managed dataset.""" - +@dataclass(frozen=True) +class _ResolvedBundleDataset: country_id: str dataset: str data_package_name: str repo_id: str repo_type: Literal["model", "dataset"] + path: str revision: str - source_uri: str expected_sha256: str - actual_sha256: str - path: Path - cache_hit: bool - build_id: Optional[str] = None + destination: Path metadata_expected_sha256: Optional[str] = None - metadata_actual_sha256: Optional[str] = None - metadata_path: Optional[Path] = None + @property + def source_uri(self) -> str: + return build_hf_uri(self.repo_id, self.path, self.revision) -_DatasetPackageType = Literal["country", "populace"] + @property + def metadata_destination(self) -> Path: + return Path(f"{self.destination}.metadata.json") -def _dataset_package_type(data_package_name: str) -> _DatasetPackageType: - if data_package_name == "populace-data": - return "populace" - if data_package_name.startswith("policyengine-") and data_package_name.endswith( - "-data" - ): - return "country" - raise DatasetMaterializationError( - "Unsupported bundle data package " - f"{data_package_name!r}; expected 'populace-data' or " - "'policyengine--data'." - ) - - -def resolve_bundle_dataset_plan( +def _resolve_bundle_dataset( country_id: str, dataset: Optional[str] = None, *, data_dir: Path = DEFAULT_DATA_DIR, manifest: Optional[CountryReleaseManifest] = None, -) -> BundleDatasetPlan: - """Resolve one logical dataset to its exact bundle-certified source.""" - +) -> _ResolvedBundleDataset: country_manifest = manifest or get_release_manifest(country_id) dataset_name = dataset or country_manifest.default_dataset reference = country_manifest.datasets.get(dataset_name) if reference is None: raise DatasetMaterializationError( f"Unknown managed dataset {dataset_name!r} for country " - f"{country_id!r}. Known datasets: " - f"{sorted(country_manifest.datasets)}" + f"{country_id!r}. Known datasets: {sorted(country_manifest.datasets)}" ) - - data_package_name = ( - reference.data_package_name or country_manifest.data_package.name - ) - repo_id = reference.repo_id or country_manifest.data_package.repo_id - repo_type = reference.repo_type or country_manifest.data_package.repo_type - revision = reference.revision or _artifact_revision(country_manifest.data_package) - if repo_type not in {"model", "dataset"}: - raise DatasetMaterializationError( - f"Dataset {dataset_name!r} has unsupported Hugging Face repository " - f"type {repo_type!r}." - ) - validated_repo_type = cast(Literal["model", "dataset"], repo_type) if not reference.sha256: raise DatasetMaterializationError( f"Managed dataset {dataset_name!r} is missing a certified sha256." ) - plan = BundleDatasetPlan( + return _ResolvedBundleDataset( country_id=country_id, dataset=dataset_name, - data_package_name=data_package_name, - repo_id=repo_id, - repo_type=validated_repo_type, + data_package_name=( + reference.data_package_name or country_manifest.data_package.name + ), + repo_id=reference.repo_id or country_manifest.data_package.repo_id, + repo_type=reference.repo_type or country_manifest.data_package.repo_type, path=reference.path, - revision=revision, + revision=reference.revision + or _artifact_revision(country_manifest.data_package), expected_sha256=reference.sha256, - source_uri=build_hf_uri(repo_id, reference.path, revision), destination=data_dir / Path(reference.path).name, - build_id=( - country_manifest.certified_data_artifact.build_id - if country_manifest.certified_data_artifact is not None - and reference.data_package_name is None - else None - ), metadata_expected_sha256=reference.metadata_sha256, - metadata_source_uri=( - build_hf_uri(repo_id, f"{reference.path}.metadata.json", revision) - if reference.metadata_sha256 - else None - ), - metadata_destination=( - data_dir / f"{Path(reference.path).name}.metadata.json" - if reference.metadata_sha256 - else None - ), ) - _dataset_package_type(data_package_name) - return plan def materialize_bundle_dataset( @@ -161,212 +109,193 @@ def materialize_bundle_dataset( dataset: Optional[str] = None, *, data_dir: Path = DEFAULT_DATA_DIR, - manifest: Optional[CountryReleaseManifest] = None, - allow_local_mirror: bool = True, - session=requests, ) -> MaterializedDataset: - """Download and verify one dataset certified by the release bundle.""" + """Return a verified local copy of a dataset from the installed bundle.""" - plan = resolve_bundle_dataset_plan( - country_id, - dataset, - data_dir=data_dir, - manifest=manifest, - ) - local_source = resolve_local_managed_dataset_source( - country_id, - plan.source_uri, - allow_local_mirror=allow_local_mirror, + return _materialize_resolved_dataset( + _resolve_bundle_dataset(country_id, dataset, data_dir=data_dir) ) - if local_source != plan.source_uri: - local_path = Path(local_source).expanduser() - if local_path.is_file(): - actual_sha256 = _sha256_file(local_path) - if actual_sha256 == plan.expected_sha256: - metadata_path = Path(f"{local_path}.metadata.json") - if plan.metadata_expected_sha256 is None: - return _materialized_result( - plan, - actual_sha256=actual_sha256, - cache_hit=True, - path=local_path, - ) - if metadata_path.is_file(): - metadata_actual_sha256 = _sha256_file(metadata_path) - if metadata_actual_sha256 == plan.metadata_expected_sha256: - return _materialized_result( - plan, - actual_sha256=actual_sha256, - cache_hit=True, - path=local_path, - metadata_actual_sha256=metadata_actual_sha256, - metadata_path=metadata_path, - ) - package_type = _dataset_package_type(plan.data_package_name) - if package_type == "populace": - return _materialize_populace_data_package(plan, session=session) - return _materialize_country_data_package(plan, session=session) - - -def _materialize_country_data_package( - plan: BundleDatasetPlan, - *, - session=requests, -) -> MaterializedDataset: - return _materialize_managed_hf_dataset(plan, session=session) - - -def _materialize_populace_data_package( - plan: BundleDatasetPlan, - *, - session=requests, -) -> MaterializedDataset: - return _materialize_managed_hf_dataset(plan, session=session) -def _materialize_managed_hf_dataset( - plan: BundleDatasetPlan, +def _materialize_resolved_dataset( + resolved: _ResolvedBundleDataset, *, session=requests, ) -> MaterializedDataset: - destination = plan.destination - if destination.is_file(): - actual_sha256 = _sha256_file(destination) - if actual_sha256 == plan.expected_sha256: - metadata_actual_sha256, metadata_path = _materialize_metadata( - plan, - session=session, - ) - return _materialized_result( - plan, - actual_sha256=actual_sha256, - cache_hit=True, - metadata_actual_sha256=metadata_actual_sha256, - metadata_path=metadata_path, - ) - - url = _hf_download_url( - repo_id=plan.repo_id, - repo_type=plan.repo_type, - path=plan.path, - revision=plan.revision, + local_source = resolve_local_managed_dataset_source( + resolved.country_id, + resolved.source_uri, ) - downloaded = _download_to_temp( - url, - destination=destination, - source_description=(f"{plan.country_id.upper()} dataset {plan.dataset!r}"), - session=session, + local_path = Path(local_source).expanduser() + local_sha256 = ( + _matching_sha256(local_path, resolved.expected_sha256) + if local_source != resolved.source_uri + else None ) - try: - actual_sha256 = _sha256_file(downloaded) - if actual_sha256 != plan.expected_sha256: - raise DatasetMaterializationError( - f"Downloaded {plan.country_id.upper()} dataset {plan.dataset!r} " - f"has sha256 {actual_sha256}, expected {plan.expected_sha256}." - ) - _backup_existing(destination) - destination.parent.mkdir(parents=True, exist_ok=True) - os.replace(downloaded, destination) - finally: - downloaded.unlink(missing_ok=True) - metadata_actual_sha256, metadata_path = _materialize_metadata( - plan, + if local_sha256 is not None: + path = local_path + actual_sha256 = local_sha256 + cache_hit = True + else: + path, actual_sha256, cache_hit = _materialize_verified_file( + url=_hf_download_url( + repo_id=resolved.repo_id, + repo_type=resolved.repo_type, + path=resolved.path, + revision=resolved.revision, + ), + destination=resolved.destination, + expected_sha256=resolved.expected_sha256, + description=(f"{resolved.country_id.upper()} dataset {resolved.dataset!r}"), + session=session, + ) + + metadata_path = _materialize_metadata( + resolved, + dataset_path=path, session=session, ) - return _materialized_result( - plan, + return MaterializedDataset( + data_package_name=resolved.data_package_name, + repo_type=resolved.repo_type, + revision=resolved.revision, + source_uri=resolved.source_uri, + expected_sha256=resolved.expected_sha256, actual_sha256=actual_sha256, - cache_hit=False, - metadata_actual_sha256=metadata_actual_sha256, + path=path, + cache_hit=cache_hit, metadata_path=metadata_path, ) -def _materialized_result( - plan: BundleDatasetPlan, +def _materialize_metadata( + resolved: _ResolvedBundleDataset, *, - actual_sha256: str, - cache_hit: bool, - path: Optional[Path] = None, - metadata_actual_sha256: Optional[str] = None, - metadata_path: Optional[Path] = None, -) -> MaterializedDataset: - return MaterializedDataset( - country_id=plan.country_id, - dataset=plan.dataset, - data_package_name=plan.data_package_name, - repo_id=plan.repo_id, - repo_type=plan.repo_type, - revision=plan.revision, - source_uri=plan.source_uri, - expected_sha256=plan.expected_sha256, - actual_sha256=actual_sha256, - path=path or plan.destination, - cache_hit=cache_hit, - build_id=plan.build_id, - metadata_expected_sha256=plan.metadata_expected_sha256, - metadata_actual_sha256=metadata_actual_sha256, - metadata_path=metadata_path, + dataset_path: Path, + session=requests, +) -> Optional[Path]: + expected_sha256 = resolved.metadata_expected_sha256 + if expected_sha256 is None: + return None + + local_path = Path(f"{dataset_path}.metadata.json") + if _matching_sha256(local_path, expected_sha256) is not None: + return local_path + + path, _, _ = _materialize_verified_file( + url=_hf_download_url( + repo_id=resolved.repo_id, + repo_type=resolved.repo_type, + path=f"{resolved.path}.metadata.json", + revision=resolved.revision, + ), + destination=resolved.metadata_destination, + expected_sha256=expected_sha256, + description=f"metadata for {resolved.dataset!r}", + session=session, ) + return path -def _materialize_metadata( - plan: BundleDatasetPlan, +def _materialize_verified_file( *, + url: str, + destination: Path, + expected_sha256: str, + description: str, session=requests, -) -> tuple[Optional[str], Optional[Path]]: - if plan.metadata_expected_sha256 is None: - return None, None - if plan.metadata_destination is None: - raise DatasetMaterializationError( - f"Managed dataset {plan.dataset!r} has a metadata hash but no " - "metadata destination." - ) +) -> tuple[Path, str, bool]: + existing_sha256 = _matching_sha256(destination, expected_sha256) + if existing_sha256 is not None: + return destination, existing_sha256, True - destination = plan.metadata_destination - if destination.is_file(): - actual_sha256 = _sha256_file(destination) - if actual_sha256 == plan.metadata_expected_sha256: - return actual_sha256, destination - - url = _hf_download_url( - repo_id=plan.repo_id, - repo_type=plan.repo_type, - path=f"{plan.path}.metadata.json", - revision=plan.revision, - ) downloaded = _download_to_temp( url, destination=destination, - source_description=f"metadata for {plan.dataset!r}", + description=description, session=session, ) try: actual_sha256 = _sha256_file(downloaded) - if actual_sha256 != plan.metadata_expected_sha256: + if actual_sha256 != expected_sha256: raise DatasetMaterializationError( - f"Downloaded metadata for dataset {plan.dataset!r} has sha256 " - f"{actual_sha256}, expected {plan.metadata_expected_sha256}." + f"Downloaded {description} has sha256 {actual_sha256}, " + f"expected {expected_sha256}." ) - _backup_existing(destination) os.replace(downloaded, destination) finally: downloaded.unlink(missing_ok=True) - return actual_sha256, destination + return destination, actual_sha256, False -class _UnmanagedHFReference(BaseModel): - repo_id: str - path: str - revision: str +def _materialize_dataset_request( + country_id: str, + dataset: Optional[str], + *, + allow_unmanaged: bool, + data_dir: Path = DEFAULT_DATA_DIR, +) -> tuple[str, str, Optional[MaterializedDataset]]: + manifest = get_release_manifest(country_id) + managed_dataset = None + if dataset is None: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + elif dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + + if managed_dataset is not None: + materialized = _materialize_resolved_dataset( + _resolve_bundle_dataset( + country_id, + managed_dataset, + data_dir=data_dir, + manifest=manifest, + ) + ) + return materialized.source_uri, str(materialized.path), materialized + source_uri = resolve_managed_dataset_reference( + country_id, + dataset, + allow_unmanaged=allow_unmanaged, + ) + return ( + source_uri, + _materialize_unmanaged_dataset_source(source_uri, data_dir=data_dir), + None, + ) -class _DatasetNotFoundError(DatasetMaterializationError): - pass + +def _runtime_dataset_provenance( + source_uri: str, + local_path: str, + materialized: Optional[MaterializedDataset], + *, + logical_name: Optional[str] = None, +) -> dict[str, object]: + provenance: dict[str, object] = { + "managed_by": "policyengine.py", + "runtime_dataset": logical_name or dataset_logical_name(source_uri), + "runtime_dataset_uri": source_uri, + "runtime_dataset_source": local_path, + } + if materialized is not None: + provenance.update( + { + "runtime_dataset_data_package": materialized.data_package_name, + "runtime_dataset_repo_type": materialized.repo_type, + "runtime_dataset_revision": materialized.revision, + "runtime_dataset_expected_sha256": materialized.expected_sha256, + "runtime_dataset_sha256": materialized.actual_sha256, + "runtime_dataset_cache_hit": materialized.cache_hit, + } + ) + return provenance -def materialize_unmanaged_dataset_source( +def _materialize_unmanaged_dataset_source( dataset_source: Union[str, Path], *, version: Optional[str] = None, @@ -377,11 +306,6 @@ def materialize_unmanaged_dataset_source( """Return a local path for an explicitly unmanaged local or HF source.""" source = str(dataset_source) - if source.startswith("gs://"): - raise DatasetMaterializationError( - "GCS dataset sources are no longer supported. Publish the dataset " - "on Hugging Face and reference that artifact instead." - ) if not source.startswith("hf://"): if "://" in source: raise DatasetMaterializationError( @@ -389,8 +313,8 @@ def materialize_unmanaged_dataset_source( ) return source - reference = _parse_unmanaged_hf_reference(source, version=version) - destination = data_dir / Path(reference.path).name + repo_id, path, revision = _parse_unmanaged_hf_reference(source, version=version) + destination = data_dir / Path(path).name repo_types: list[Literal["model", "dataset"]] = ( [repo_type] if repo_type is not None else ["model", "dataset"] ) @@ -398,21 +322,23 @@ def materialize_unmanaged_dataset_source( try: downloaded = _download_to_temp( _hf_download_url( - repo_id=reference.repo_id, + repo_id=repo_id, repo_type=candidate_repo_type, - path=reference.path, - revision=reference.revision, + path=path, + revision=revision, ), destination=destination, - source_description=f"unmanaged dataset {source!r}", + description=f"unmanaged dataset {source!r}", session=session, ) except _DatasetNotFoundError: if index + 1 < len(repo_types): continue raise - destination.parent.mkdir(parents=True, exist_ok=True) - os.replace(downloaded, destination) + try: + os.replace(downloaded, destination) + finally: + downloaded.unlink(missing_ok=True) return str(destination) raise DatasetMaterializationError(f"Could not materialize dataset {source!r}.") @@ -422,7 +348,7 @@ def _parse_unmanaged_hf_reference( uri: str, *, version: Optional[str], -) -> _UnmanagedHFReference: +) -> tuple[str, str, str]: path_with_repo, uri_revision = ( uri[5:].rsplit("@", maxsplit=1) if "@" in uri[5:] else (uri[5:], None) ) @@ -437,11 +363,11 @@ def _parse_unmanaged_hf_reference( "Invalid Hugging Face dataset URI. Expected format " f"'hf://owner/repo/path/to/file[@revision]', got {uri!r}." ) - return _UnmanagedHFReference( - repo_id=f"{parts[0]}/{parts[1]}", - path=parts[2], - revision=uri_revision or version or "main", - ) + return f"{parts[0]}/{parts[1]}", parts[2], uri_revision or version or "main" + + +class _DatasetNotFoundError(DatasetMaterializationError): + pass def _hf_download_url( @@ -462,14 +388,13 @@ def _download_to_temp( url: str, *, destination: Path, - source_description: str, + description: str, session=requests, ) -> Path: destination.parent.mkdir(parents=True, exist_ok=True) - suffix = destination.suffix or ".download" file_descriptor, temp_name = tempfile.mkstemp( prefix=".policyengine-download-", - suffix=suffix, + suffix=destination.suffix or ".download", dir=destination.parent, ) os.close(file_descriptor) @@ -483,14 +408,12 @@ def _download_to_temp( ) as response: if response.status_code in {401, 403}: raise DatasetMaterializationError( - f"Could not download {source_description}: Hugging Face " - "rejected the configured credentials. Set HUGGING_FACE_TOKEN " - "to a token with access to the certified repository." + f"Could not download {description}: Hugging Face rejected " + "the configured credentials. Set HUGGING_FACE_TOKEN to a " + "token with access to the repository." ) if response.status_code == 404: - raise _DatasetNotFoundError( - f"Could not find {source_description} at {url}." - ) + raise _DatasetNotFoundError(f"Could not find {description} at {url}.") response.raise_for_status() with temp_path.open("wb") as output: for chunk in response.iter_content(chunk_size=1024 * 1024): @@ -511,13 +434,11 @@ def _hugging_face_auth_headers() -> dict[str, str]: return {"Authorization": f"Bearer {token}"} if token else {} -def _backup_existing(path: Path) -> None: - if not path.exists(): - return - timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") - backup_dir = path.parent / BACKUP_DIR_NAME / timestamp - backup_dir.mkdir(parents=True, exist_ok=True) - shutil.move(str(path), str(backup_dir / path.name)) +def _matching_sha256(path: Path, expected_sha256: str) -> Optional[str]: + if not path.is_file(): + return None + actual_sha256 = _sha256_file(path) + return actual_sha256 if actual_sha256 == expected_sha256 else None def _sha256_file(path: Path) -> str: diff --git a/src/policyengine/provenance/manifest.py b/src/policyengine/provenance/manifest.py index b2795c0a..e18c719b 100644 --- a/src/policyengine/provenance/manifest.py +++ b/src/policyengine/provenance/manifest.py @@ -5,7 +5,7 @@ from importlib import import_module from importlib.resources import files from pathlib import Path -from typing import Optional +from typing import Literal, Optional import requests from pydantic import BaseModel, Field @@ -31,7 +31,7 @@ class PackageVersion(BaseModel): class DataPackageVersion(PackageVersion): repo_id: str - repo_type: str = "model" + repo_type: Literal["model", "dataset"] = "model" release_manifest_path: str = "release_manifest.json" release_manifest_revision: Optional[str] = None @@ -67,7 +67,7 @@ class ArtifactPathReference(BaseModel): # Set when the artifact lives outside the data package's repo (inherited # datasets keep their original repo + revision pins). repo_id: Optional[str] = None - repo_type: Optional[str] = None + repo_type: Optional[Literal["model", "dataset"]] = None class ArtifactPathTemplate(BaseModel): diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index 2fcb2d7d..b61185e0 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -7,14 +7,12 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( - materialize_bundle_dataset, - materialize_unmanaged_dataset_source, + _materialize_dataset_request, ) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, - resolve_managed_dataset_reference, ) @@ -127,26 +125,12 @@ def create_datasets( datasets = [get_release_manifest("uk").default_dataset] result = {} for dataset in datasets: - manifest = get_release_manifest("uk") - managed_dataset = dataset if dataset in manifest.datasets else None - if managed_dataset is not None: - materialized = materialize_bundle_dataset( - "uk", - managed_dataset, - data_dir=Path(data_folder), - ) - resolved_dataset = materialized.source_uri - runtime_dataset = str(materialized.path) - else: - resolved_dataset = resolve_managed_dataset_reference( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - ) - runtime_dataset = materialize_unmanaged_dataset_source( - resolved_dataset, - data_dir=Path(data_folder), - ) + resolved_dataset, runtime_dataset, _ = _materialize_dataset_request( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + data_dir=Path(data_folder), + ) dataset_stem = dataset_logical_name(resolved_dataset) from policyengine_uk import Microsimulation diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index 12c8b93f..c1ed4003 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -1,19 +1,13 @@ import datetime -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Optional import pandas as pd from microdf import MicroDataFrame from policyengine.core import TaxBenefitModel from policyengine.provenance.dataset_materialization import ( - MaterializedDataset, - materialize_bundle_dataset, - materialize_unmanaged_dataset_source, -) -from policyengine.provenance.manifest import ( - dataset_logical_name, - get_release_manifest, - resolve_managed_dataset_reference, + _materialize_dataset_request, + _runtime_dataset_provenance, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion from policyengine.tax_benefit_models.common.model_version import ( @@ -268,31 +262,6 @@ def run(self, simulation: "Simulation") -> "Simulation": ) -def _managed_release_bundle( - dataset_uri: str, - dataset_source: Optional[str] = None, - materialized: Optional[MaterializedDataset] = None, -) -> dict[str, Any]: - bundle: dict[str, Any] = dict(uk_latest.release_bundle) - bundle["runtime_dataset"] = dataset_logical_name(dataset_uri) - bundle["runtime_dataset_uri"] = dataset_uri - if dataset_source: - bundle["runtime_dataset_source"] = dataset_source - if materialized is not None: - bundle.update( - { - "runtime_dataset_data_package": materialized.data_package_name, - "runtime_dataset_repo_type": materialized.repo_type, - "runtime_dataset_revision": materialized.revision, - "runtime_dataset_expected_sha256": materialized.expected_sha256, - "runtime_dataset_sha256": materialized.actual_sha256, - "runtime_dataset_cache_hit": materialized.cache_hit, - } - ) - bundle["managed_by"] = "policyengine.py" - return bundle - - def managed_microsimulation( *, dataset: Optional[str] = None, @@ -314,27 +283,11 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - manifest = get_release_manifest("uk") - managed_dataset = None - if dataset is None: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - elif dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - - materialized = None - if managed_dataset is not None: - materialized = materialize_bundle_dataset("uk", managed_dataset) - dataset_uri = materialized.source_uri - runtime_dataset_source = str(materialized.path) - else: - dataset_uri = resolve_managed_dataset_reference( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - ) - runtime_dataset_source = materialize_unmanaged_dataset_source(dataset_uri) + dataset_uri, runtime_dataset_source, materialized = _materialize_dataset_request( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + ) runtime_dataset = runtime_dataset_source if isinstance(runtime_dataset_source, str) and "://" not in runtime_dataset_source: from policyengine_uk.data.dataset_schema import ( @@ -347,10 +300,13 @@ def managed_microsimulation( elif UKSingleYearDataset.validate_file_path(runtime_dataset_source, False): runtime_dataset = UKSingleYearDataset(runtime_dataset_source) microsim = Microsimulation(dataset=runtime_dataset, **kwargs) - microsim.policyengine_bundle = _managed_release_bundle( - dataset_uri, - runtime_dataset_source, - materialized, + microsim.policyengine_bundle = dict(uk_latest.release_bundle) + microsim.policyengine_bundle.update( + _runtime_dataset_provenance( + dataset_uri, + runtime_dataset_source, + materialized, + ) ) return microsim diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index e3691868..806d6a54 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -14,14 +14,14 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( MaterializedDataset, + _materialize_dataset_request, + _runtime_dataset_provenance, materialize_bundle_dataset, - materialize_unmanaged_dataset_source, ) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, - resolve_managed_dataset_reference, ) @@ -299,26 +299,12 @@ def create_datasets( datasets = datasets or [get_release_manifest("us").default_dataset] result = {} for dataset in datasets: - manifest = get_release_manifest("us") - managed_dataset = dataset if dataset in manifest.datasets else None - if managed_dataset is not None: - materialized = materialize_bundle_dataset( - "us", - managed_dataset, - data_dir=Path(data_folder), - ) - resolved_dataset = materialized.source_uri - runtime_dataset = str(materialized.path) - else: - resolved_dataset = resolve_managed_dataset_reference( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) - runtime_dataset = materialize_unmanaged_dataset_source( - resolved_dataset, - data_dir=Path(data_folder), - ) + resolved_dataset, runtime_dataset, _ = _materialize_dataset_request( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + data_dir=Path(data_folder), + ) dataset_stem = dataset_logical_name(resolved_dataset) sim = Microsimulation(dataset=runtime_dataset) @@ -531,9 +517,12 @@ def _metadata_path_for_h5(path: Path) -> Path: def _load_dataset_metadata( - path: Path, require_metadata: bool + path: Path, + require_metadata: bool, + *, + metadata_path: Optional[Path] = None, ) -> tuple[dict, Optional[Path]]: - metadata_path = _metadata_path_for_h5(path) + metadata_path = metadata_path or _metadata_path_for_h5(path) if not metadata_path.exists(): if require_metadata: raise FileNotFoundError( @@ -872,23 +861,13 @@ def _build_long_term_dataset( if dataset_uri is not None: dataset.metadata.setdefault("policyengine_bundle", {}) dataset.metadata["policyengine_bundle"].update( - { - "managed_by": "policyengine.py", - "runtime_dataset": _long_term_dataset_key(dataset_name, year), - "runtime_dataset_uri": dataset_uri, - } - ) - if materialized is not None: - dataset.metadata["policyengine_bundle"].update( - { - "runtime_dataset_data_package": (materialized.data_package_name), - "runtime_dataset_repo_type": materialized.repo_type, - "runtime_dataset_revision": materialized.revision, - "runtime_dataset_expected_sha256": (materialized.expected_sha256), - "runtime_dataset_sha256": materialized.actual_sha256, - "runtime_dataset_cache_hit": materialized.cache_hit, - } + _runtime_dataset_provenance( + dataset_uri, + str(path), + materialized, + logical_name=_long_term_dataset_key(dataset_name, year), ) + ) return dataset @@ -1107,24 +1086,14 @@ def load_managed_long_term_datasets( "us", key, data_dir=Path(data_folder), - manifest=manifest, ) dataset_uri = materialized.source_uri path = materialized.path - metadata, metadata_path = _load_dataset_metadata(path, require_metadata) - if path_reference.metadata_sha256: - if metadata_path is None: - raise FileNotFoundError( - f"Managed long-term dataset {key!r} at {path} is missing " - "metadata sidecar required by the bundled manifest." - ) - metadata_sha256 = _sha256_file(metadata_path) - if metadata_sha256 != path_reference.metadata_sha256: - raise ValueError( - f"Managed long-term dataset {key!r} metadata at " - f"{metadata_path} has sha256 {metadata_sha256}, expected " - f"{path_reference.metadata_sha256}." - ) + metadata, metadata_path = _load_dataset_metadata( + path, + require_metadata, + metadata_path=materialized.metadata_path, + ) _validate_loaded_long_term_metadata( metadata=metadata, metadata_path=metadata_path, diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 402396c3..2890e363 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -1,19 +1,13 @@ import datetime -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Optional import pandas as pd from microdf import MicroDataFrame from policyengine.core import TaxBenefitModel from policyengine.provenance.dataset_materialization import ( - MaterializedDataset, - materialize_bundle_dataset, - materialize_unmanaged_dataset_source, -) -from policyengine.provenance.manifest import ( - dataset_logical_name, - get_release_manifest, - resolve_managed_dataset_reference, + _materialize_dataset_request, + _runtime_dataset_provenance, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion from policyengine.tax_benefit_models.common.model_version import ( @@ -411,31 +405,6 @@ def _build_simulation_from_dataset(self, microsim, dataset, system): microsim.set_input(column, dataset.year, df[column].values) -def _managed_release_bundle( - dataset_uri: str, - dataset_source: Optional[str] = None, - materialized: Optional[MaterializedDataset] = None, -) -> dict[str, Any]: - bundle: dict[str, Any] = dict(us_latest.release_bundle) - bundle["runtime_dataset"] = dataset_logical_name(dataset_uri) - bundle["runtime_dataset_uri"] = dataset_uri - if dataset_source: - bundle["runtime_dataset_source"] = dataset_source - if materialized is not None: - bundle.update( - { - "runtime_dataset_data_package": materialized.data_package_name, - "runtime_dataset_repo_type": materialized.repo_type, - "runtime_dataset_revision": materialized.revision, - "runtime_dataset_expected_sha256": materialized.expected_sha256, - "runtime_dataset_sha256": materialized.actual_sha256, - "runtime_dataset_cache_hit": materialized.cache_hit, - } - ) - bundle["managed_by"] = "policyengine.py" - return bundle - - def managed_microsimulation( *, dataset: Optional[str] = None, @@ -457,32 +426,19 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - manifest = get_release_manifest("us") - managed_dataset = None - if dataset is None: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - elif dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - - materialized = None - if managed_dataset is not None: - materialized = materialize_bundle_dataset("us", managed_dataset) - dataset_uri = materialized.source_uri - runtime_dataset_source = str(materialized.path) - else: - dataset_uri = resolve_managed_dataset_reference( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) - runtime_dataset_source = materialize_unmanaged_dataset_source(dataset_uri) + dataset_uri, runtime_dataset_source, materialized = _materialize_dataset_request( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + ) microsim = Microsimulation(dataset=runtime_dataset_source, **kwargs) - microsim.policyengine_bundle = _managed_release_bundle( - dataset_uri, - runtime_dataset_source, - materialized, + microsim.policyengine_bundle = dict(us_latest.release_bundle) + microsim.policyengine_bundle.update( + _runtime_dataset_provenance( + dataset_uri, + runtime_dataset_source, + materialized, + ) ) return microsim diff --git a/tests/test_bundle.py b/tests/test_bundle.py index f6806c68..a28efaaa 100644 --- a/tests/test_bundle.py +++ b/tests/test_bundle.py @@ -73,7 +73,7 @@ def test_selected_dataset_plan_uses_certified_release_metadata(tmp_path): bundle.get_current_bundle(), ["uk"], data_dir=tmp_path ) - plan, _, release = entries[0] + plan, release = entries[0] assert plan.country_id == "uk" assert plan.data_package_name == "policyengine-uk-data" assert plan.repo_type == "model" @@ -151,21 +151,12 @@ def test_install_bundle_materializes_defaults_and_records_receipt( bundle, "install_package_scaffold", lambda *args, **kwargs: None ) - def fake_materialize(country_id, dataset, *, data_dir, manifest): - plan = bundle.resolve_bundle_dataset_plan( - country_id, - dataset, - data_dir=data_dir, - manifest=manifest, - ) + def fake_materialize(plan): calls.append(plan) plan.destination.parent.mkdir(parents=True, exist_ok=True) plan.destination.write_bytes(b"materialized") return MaterializedDataset( - country_id=plan.country_id, - dataset=plan.dataset, data_package_name=plan.data_package_name, - repo_id=plan.repo_id, repo_type=plan.repo_type, revision=plan.revision, source_uri=plan.source_uri, @@ -173,10 +164,9 @@ def fake_materialize(country_id, dataset, *, data_dir, manifest): actual_sha256=plan.expected_sha256, path=plan.destination, cache_hit=False, - build_id=plan.build_id, ) - monkeypatch.setattr(bundle, "materialize_bundle_dataset", fake_materialize) + monkeypatch.setattr(bundle, "_materialize_resolved_dataset", fake_materialize) result = bundle.install_bundle( python=sys.executable, diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 0b4074b3..7a66e8ea 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -4,13 +4,11 @@ import pytest from policyengine.provenance.dataset_materialization import ( - BundleDatasetPlan, DatasetMaterializationError, MaterializedDataset, - _dataset_package_type, - materialize_bundle_dataset, - materialize_unmanaged_dataset_source, - resolve_bundle_dataset_plan, + _materialize_resolved_dataset, + _materialize_unmanaged_dataset_source, + _resolve_bundle_dataset, ) from policyengine.provenance.manifest import CountryReleaseManifest @@ -53,19 +51,18 @@ def _manifest() -> CountryReleaseManifest: ) -def test_resolve_bundle_dataset_plan_inherits_primary_package(tmp_path): - plan = resolve_bundle_dataset_plan("uk", data_dir=tmp_path, manifest=_manifest()) +def test_resolve_bundle_dataset_inherits_primary_package(tmp_path): + plan = _resolve_bundle_dataset("uk", data_dir=tmp_path, manifest=_manifest()) assert plan.data_package_name == "policyengine-uk-data" assert plan.repo_id == "policyengine/policyengine-uk-data-private" assert plan.repo_type == "model" assert plan.revision == "uk-release" assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" - assert plan.build_id == "policyengine-uk-data-test-build" -def test_resolve_bundle_dataset_plan_uses_cross_package_overlay(tmp_path): - plan = resolve_bundle_dataset_plan( +def test_resolve_bundle_dataset_uses_cross_package_overlay(tmp_path): + plan = _resolve_bundle_dataset( "uk", "populace_uk_2023", data_dir=tmp_path, @@ -76,15 +73,16 @@ def test_resolve_bundle_dataset_plan_uses_cross_package_overlay(tmp_path): assert plan.repo_id == "policyengine/populace-uk-private" assert plan.repo_type == "model" assert plan.revision == "populace-release" - assert plan.build_id is None -def test_bundle_dataset_models_round_trip_json(): - plan = resolve_bundle_dataset_plan("uk", manifest=_manifest()) - assert BundleDatasetPlan.model_validate_json(plan.model_dump_json()) == plan - +def test_materialized_dataset_round_trips_json(): + plan = _resolve_bundle_dataset("uk", manifest=_manifest()) result = MaterializedDataset( - **plan.model_dump(exclude={"path", "destination"}), + data_package_name=plan.data_package_name, + repo_type=plan.repo_type, + revision=plan.revision, + source_uri=plan.source_uri, + expected_sha256=plan.expected_sha256, actual_sha256=plan.expected_sha256, path=Path("data/enhanced_frs_2024_25.h5"), cache_hit=True, @@ -92,23 +90,6 @@ def test_bundle_dataset_models_round_trip_json(): assert MaterializedDataset.model_validate_json(result.model_dump_json()) == result -@pytest.mark.parametrize( - ("package_name", "expected_type"), - [ - ("policyengine-us-data", "country"), - ("policyengine-uk-data", "country"), - ("populace-data", "populace"), - ], -) -def test_dataset_package_type(package_name, expected_type): - assert _dataset_package_type(package_name) == expected_type - - -def test_unknown_data_package_is_rejected(): - with pytest.raises(DatasetMaterializationError, match="Unsupported bundle"): - _dataset_package_type("unknown-data") - - def _sha256(payload: bytes) -> str: return hashlib.sha256(payload).hexdigest() @@ -155,13 +136,20 @@ def _manifest_with_hash( return manifest +def _materialize(manifest, tmp_path, session): + resolved = _resolve_bundle_dataset( + "uk", + data_dir=tmp_path, + manifest=manifest, + ) + return _materialize_resolved_dataset(resolved, session=session) + + def test_materialize_country_data_package_uses_model_repo_url(tmp_path): session = _Session(_Response(b"country-data")) manifest = _manifest_with_hash(_sha256(b"country-data")) - result = materialize_bundle_dataset( - "uk", data_dir=tmp_path, manifest=manifest, session=session - ) + result = _materialize(manifest, tmp_path, session) assert result.path.read_bytes() == b"country-data" assert result.cache_hit is False @@ -178,9 +166,7 @@ def test_materialize_populace_package_uses_dataset_repo_url(tmp_path): repo_type="dataset", ) - result = materialize_bundle_dataset( - "uk", data_dir=tmp_path, manifest=manifest, session=session - ) + result = _materialize(manifest, tmp_path, session) assert result.path.read_bytes() == b"populace-data" assert session.calls[0][0].startswith( @@ -196,13 +182,10 @@ def test_materialize_downloads_and_verifies_metadata_sidecar(tmp_path): reference.metadata_sha256 = _sha256(metadata_payload) session = _Session(_Response(dataset_payload), _Response(metadata_payload)) - result = materialize_bundle_dataset( - "uk", data_dir=tmp_path, manifest=manifest, session=session - ) + result = _materialize(manifest, tmp_path, session) assert result.metadata_path == (tmp_path / "enhanced_frs_2024_25.h5.metadata.json") assert result.metadata_path.read_bytes() == metadata_payload - assert result.metadata_actual_sha256 == reference.metadata_sha256 assert session.calls[1][0].endswith("/enhanced_frs_2024_25.h5.metadata.json") @@ -213,9 +196,7 @@ def test_materialize_reuses_only_hash_verified_cache(tmp_path): destination.write_bytes(payload) session = _Session() - result = materialize_bundle_dataset( - "uk", data_dir=tmp_path, manifest=manifest, session=session - ) + result = _materialize(manifest, tmp_path, session) assert result.cache_hit is True assert session.calls == [] @@ -233,15 +214,39 @@ def test_materialize_reuses_hash_verified_local_mirror(monkeypatch, tmp_path): ) session = _Session() - result = materialize_bundle_dataset( - "uk", data_dir=tmp_path, manifest=manifest, session=session - ) + result = _materialize(manifest, tmp_path, session) assert result.path == mirror assert result.cache_hit is True assert session.calls == [] +def test_materialize_reuses_local_mirror_and_downloads_missing_metadata( + monkeypatch, tmp_path +): + payload = b"certified-local-mirror" + metadata_payload = b'{"year": 2025}' + manifest = _manifest_with_hash(_sha256(payload)) + manifest.datasets[manifest.default_dataset].metadata_sha256 = _sha256( + metadata_payload + ) + mirror = tmp_path / "mirror" / "enhanced_frs_2024_25.h5" + mirror.parent.mkdir() + mirror.write_bytes(payload) + monkeypatch.setattr( + "policyengine.provenance.dataset_materialization.resolve_local_managed_dataset_source", + lambda *args, **kwargs: str(mirror), + ) + session = _Session(_Response(metadata_payload)) + + result = _materialize(manifest, tmp_path, session) + + assert result.path == mirror + assert result.metadata_path == tmp_path / "enhanced_frs_2024_25.h5.metadata.json" + assert result.metadata_path.read_bytes() == metadata_payload + assert len(session.calls) == 1 + + def test_materialize_ignores_mismatched_local_mirror(monkeypatch, tmp_path): payload = b"certified-download" manifest = _manifest_with_hash(_sha256(payload)) @@ -254,34 +259,22 @@ def test_materialize_ignores_mismatched_local_mirror(monkeypatch, tmp_path): ) session = _Session(_Response(payload)) - result = materialize_bundle_dataset( - "uk", data_dir=tmp_path, manifest=manifest, session=session - ) + result = _materialize(manifest, tmp_path, session) assert result.path == tmp_path / "enhanced_frs_2024_25.h5" assert result.path.read_bytes() == payload assert mirror.read_bytes() == b"wrong" -def test_materialize_replaces_and_backs_up_mismatched_cache(tmp_path): +def test_materialize_replaces_mismatched_cache(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) destination = tmp_path / "enhanced_frs_2024_25.h5" destination.write_bytes(b"old") - materialize_bundle_dataset( - "uk", - data_dir=tmp_path, - manifest=manifest, - session=_Session(_Response(payload)), - ) + _materialize(manifest, tmp_path, _Session(_Response(payload))) assert destination.read_bytes() == payload - backups = list( - (tmp_path / ".policyengine-bundle-backups").glob("*/enhanced_frs_2024_25.h5") - ) - assert len(backups) == 1 - assert backups[0].read_bytes() == b"old" def test_hash_failure_does_not_replace_existing_cache(tmp_path): @@ -290,15 +283,9 @@ def test_hash_failure_does_not_replace_existing_cache(tmp_path): destination.write_bytes(b"old") with pytest.raises(DatasetMaterializationError, match="sha256"): - materialize_bundle_dataset( - "uk", - data_dir=tmp_path, - manifest=manifest, - session=_Session(_Response(b"wrong")), - ) + _materialize(manifest, tmp_path, _Session(_Response(b"wrong"))) assert destination.read_bytes() == b"old" - assert not (tmp_path / ".policyengine-bundle-backups").exists() def test_materialize_passes_hugging_face_token(monkeypatch, tmp_path): @@ -306,12 +293,7 @@ def test_materialize_passes_hugging_face_token(monkeypatch, tmp_path): payload = b"certified" session = _Session(_Response(payload)) - materialize_bundle_dataset( - "uk", - data_dir=tmp_path, - manifest=_manifest_with_hash(_sha256(payload)), - session=session, - ) + _materialize(_manifest_with_hash(_sha256(payload)), tmp_path, session) assert session.calls[0][1]["headers"] == {"Authorization": "Bearer secret-token"} @@ -321,11 +303,10 @@ def test_managed_auth_failure_does_not_retry_repo_type(tmp_path, status_code): session = _Session(_Response(status_code=status_code)) with pytest.raises(DatasetMaterializationError, match="credentials"): - materialize_bundle_dataset( - "uk", - data_dir=tmp_path, - manifest=_manifest_with_hash(_sha256(b"certified")), - session=session, + _materialize( + _manifest_with_hash(_sha256(b"certified")), + tmp_path, + session, ) assert len(session.calls) == 1 @@ -334,7 +315,7 @@ def test_managed_auth_failure_does_not_retry_repo_type(tmp_path, status_code): def test_unmanaged_hf_retries_dataset_repo_only_after_not_found(tmp_path): session = _Session(_Response(status_code=404), _Response(b"dataset")) - result = materialize_unmanaged_dataset_source( + result = _materialize_unmanaged_dataset_source( "hf://policyengine/example/data.h5@release", data_dir=tmp_path, session=session, @@ -349,7 +330,7 @@ def test_unmanaged_auth_failure_does_not_retry(tmp_path): session = _Session(_Response(status_code=403)) with pytest.raises(DatasetMaterializationError, match="credentials"): - materialize_unmanaged_dataset_source( + _materialize_unmanaged_dataset_source( "hf://policyengine/example/data.h5@release", data_dir=tmp_path, session=session, @@ -359,9 +340,9 @@ def test_unmanaged_auth_failure_does_not_retry(tmp_path): def test_unmanaged_local_path_is_preserved(): - assert materialize_unmanaged_dataset_source("/tmp/custom.h5") == ("/tmp/custom.h5") + assert _materialize_unmanaged_dataset_source("/tmp/custom.h5") == ("/tmp/custom.h5") def test_unmanaged_gcs_source_is_rejected(): - with pytest.raises(DatasetMaterializationError, match="no longer supported"): - materialize_unmanaged_dataset_source("gs://bucket/data.h5@release") + with pytest.raises(DatasetMaterializationError, match="Unsupported unmanaged"): + _materialize_unmanaged_dataset_source("gs://bucket/data.h5@release") diff --git a/tests/test_dataset_runtime.py b/tests/test_dataset_runtime.py index 865de7b3..9c420884 100644 --- a/tests/test_dataset_runtime.py +++ b/tests/test_dataset_runtime.py @@ -6,7 +6,7 @@ from policyengine.provenance.dataset_materialization import ( MaterializedDataset, - resolve_bundle_dataset_plan, + _resolve_bundle_dataset, ) REPO_ROOT = Path(__file__).resolve().parents[1] @@ -23,12 +23,9 @@ def _load_module_from_path(module_name: str, path: Path): def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDataset: - plan = resolve_bundle_dataset_plan(country_id, dataset) + plan = _resolve_bundle_dataset(country_id, dataset) return MaterializedDataset( - country_id=plan.country_id, - dataset=plan.dataset, data_package_name=plan.data_package_name, - repo_id=plan.repo_id, repo_type=plan.repo_type, revision=plan.revision, source_uri=plan.source_uri, @@ -36,10 +33,14 @@ def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDatas actual_sha256=plan.expected_sha256, path=Path(path), cache_hit=False, - build_id=plan.build_id, ) +def _materialized_request(country_id: str, dataset: str, path: str): + materialized = _materialized(country_id, dataset, path) + return materialized.source_uri, str(materialized.path), materialized + + def test_us_create_datasets_passes_verified_bundle_source_to_country_package( monkeypatch, ): @@ -48,10 +49,12 @@ def test_us_create_datasets_passes_verified_bundle_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/us/datasets.py", ) materialize = Mock( - return_value=_materialized("us", "populace_us_2024", "/tmp/populace_us_2024.h5") + return_value=_materialized_request( + "us", "populace_us_2024", "/tmp/populace_us_2024.h5" + ) ) microsimulation = Mock() - monkeypatch.setattr(us_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setattr(us_datasets, "_materialize_dataset_request", materialize) monkeypatch.setitem( sys.modules, "policyengine_us", @@ -61,7 +64,10 @@ def test_us_create_datasets_passes_verified_bundle_source_to_country_package( us_datasets.create_datasets(datasets=["populace_us_2024"], years=[]) materialize.assert_called_once_with( - "us", "populace_us_2024", data_dir=Path("./data") + "us", + "populace_us_2024", + allow_unmanaged=False, + data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/populace_us_2024.h5") @@ -74,10 +80,12 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) materialize = Mock( - return_value=_materialized("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") + return_value=_materialized_request( + "uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5" + ) ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setattr(uk_datasets, "_materialize_dataset_request", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -87,7 +95,10 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( uk_datasets.create_datasets(datasets=["populace_uk_2023"], years=[]) materialize.assert_called_once_with( - "uk", "populace_uk_2023", data_dir=Path("./data") + "uk", + "populace_uk_2023", + allow_unmanaged=False, + data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") @@ -98,14 +109,14 @@ def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) materialize = Mock( - return_value=_materialized( + return_value=_materialized_request( "uk", "enhanced_frs_2024_25", "/tmp/enhanced_frs_2024_25.h5", ) ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setattr(uk_datasets, "_materialize_dataset_request", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -115,7 +126,10 @@ def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): uk_datasets.create_datasets(years=[]) materialize.assert_called_once_with( - "uk", "enhanced_frs_2024_25", data_dir=Path("./data") + "uk", + "enhanced_frs_2024_25", + allow_unmanaged=False, + data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/enhanced_frs_2024_25.h5") diff --git a/tests/test_release_manifests.py b/tests/test_release_manifests.py index 41ab85ac..996c5ecc 100644 --- a/tests/test_release_manifests.py +++ b/tests/test_release_manifests.py @@ -18,7 +18,7 @@ from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion from policyengine.provenance.dataset_materialization import ( MaterializedDataset, - resolve_bundle_dataset_plan, + _resolve_bundle_dataset, ) from policyengine.provenance.manifest import ( ArtifactPathReference, @@ -95,12 +95,9 @@ def _materialized_dataset( dataset: str, path: str, ) -> MaterializedDataset: - plan = resolve_bundle_dataset_plan(country_id, dataset) + plan = _resolve_bundle_dataset(country_id, dataset) return MaterializedDataset( - country_id=plan.country_id, - dataset=plan.dataset, data_package_name=plan.data_package_name, - repo_id=plan.repo_id, repo_type=plan.repo_type, revision=plan.revision, source_uri=plan.source_uri, @@ -108,10 +105,14 @@ def _materialized_dataset( actual_sha256=plan.expected_sha256, path=Path(path), cache_hit=False, - build_id=plan.build_id, ) +def _materialized_dataset_request(country_id: str, dataset: str, path: str): + materialized = _materialized_dataset(country_id, dataset, path) + return materialized.source_uri, str(materialized.path), materialized + + UK_LEGACY_DATA_RELEASE_REVISION = "655dd07e4bb9c777b00dac044949611f1feb824f" UK_LEGACY_FRS_DATASET_URI = ( "hf://policyengine/policyengine-uk-data-private/frs_2023_24.h5" @@ -953,8 +954,8 @@ def test__given_us_managed_microsimulation__then_passes_certified_dataset_and_bu ) with patch.object( us_model, - "materialize_bundle_dataset", - return_value=_materialized_dataset( + "_materialize_dataset_request", + return_value=_materialized_dataset_request( "us", "populace_us_2024", "/tmp/populace_us_2024.h5", @@ -1002,8 +1003,8 @@ def test__given_us_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( us_model, - "materialize_unmanaged_dataset_source", - return_value="/tmp/cps_2023.h5", + "_materialize_dataset_request", + return_value=(dataset, "/tmp/cps_2023.h5", None), ): microsim = us_model.managed_microsimulation( dataset=dataset, @@ -1071,8 +1072,8 @@ def test__given_uk_managed_dataset_name__then_resolves_within_bundle(self): ) with patch.object( uk_model, - "materialize_bundle_dataset", - return_value=_materialized_dataset( + "_materialize_dataset_request", + return_value=_materialized_dataset_request( "uk", "enhanced_frs_2024_25", "/tmp/enhanced_frs_2024_25.h5", @@ -1121,8 +1122,8 @@ def test__given_uk_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( uk_model, - "materialize_unmanaged_dataset_source", - return_value="/tmp/frs_2022_23.h5", + "_materialize_dataset_request", + return_value=(dataset, "/tmp/frs_2022_23.h5", None), ): microsim = uk_model.managed_microsimulation( dataset=dataset, diff --git a/tests/test_us_long_term_datasets.py b/tests/test_us_long_term_datasets.py index d3e6ebbd..27d259c3 100644 --- a/tests/test_us_long_term_datasets.py +++ b/tests/test_us_long_term_datasets.py @@ -156,10 +156,7 @@ def _manifest_with_long_term_sha( def _materialized_long_term(path: Path, dataset_uri: str) -> MaterializedDataset: actual_sha256 = _sha256(path) return MaterializedDataset( - country_id="us", - dataset="long_term_cps_2100", data_package_name="policyengine-us-data", - repo_id="policyengine/policyengine-us-data", repo_type="model", revision="abc123", source_uri=dataset_uri, @@ -340,7 +337,6 @@ def test__load_managed_long_term_datasets__defaults_to_manifest_model_version( _write_us_h5(h5_path, 2100) _write_metadata(h5_path, 2100, policyengine_us={"version": "1.691.10"}) dataset_uri = "hf://policyengine/policyengine-us-data/long_term/2100.h5@abc123" - monkeypatch.setattr( us_datasets_module, "get_release_manifest", @@ -379,14 +375,13 @@ def test__load_managed_long_term_datasets__checks_manifest_sha256( load_managed_long_term_datasets([2100]) -def test__load_managed_long_term_datasets__checks_metadata_sha256( +def test__load_managed_long_term_datasets__propagates_metadata_hash_failure( monkeypatch, tmp_path, ): h5_path = tmp_path / "2100.h5" _write_us_h5(h5_path, 2100) _write_metadata(h5_path, 2100, policyengine_us={"version": "1.691.12"}) - dataset_uri = "hf://policyengine/policyengine-us-data/long_term/2100.h5@abc123" monkeypatch.setattr( us_datasets_module, @@ -399,7 +394,7 @@ def test__load_managed_long_term_datasets__checks_metadata_sha256( monkeypatch.setattr( us_datasets_module, "materialize_bundle_dataset", - lambda *args, **kwargs: _materialized_long_term(h5_path, dataset_uri), + Mock(side_effect=DatasetMaterializationError("metadata sha256 mismatch")), ) with pytest.raises(ValueError, match="metadata"): @@ -438,7 +433,6 @@ def test__load_managed_long_term_datasets__materializes_without_local_mirror( "us", "long_term_cps_2100", data_dir=tmp_path, - manifest=manifest, ) From 9ae33aba890108e0f3e1caa6dfdcd227fb9aa0da Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:43:02 +0400 Subject: [PATCH 12/18] Reduce bundle dataset download layers --- changelog.d/502.changed.md | 7 +- docs/bundles.md | 12 +- docs/microsim.md | 5 +- src/policyengine/bundle.py | 28 +- src/policyengine/core/run_record.py | 12 +- src/policyengine/provenance/__init__.py | 3 - .../provenance/dataset_materialization.py | 432 ++++-------------- src/policyengine/provenance/manifest.py | 90 +--- .../common/dataset_source.py | 98 ++++ .../common/model_version.py | 29 ++ .../tax_benefit_models/uk/datasets.py | 45 +- .../tax_benefit_models/uk/model.py | 47 +- .../tax_benefit_models/us/datasets.py | 69 ++- .../tax_benefit_models/us/model.py | 47 +- src/policyengine/utils/hashing.py | 15 + tests/test_bundle.py | 10 +- tests/test_dataset_materialization.py | 192 ++------ tests/test_dataset_runtime.py | 28 +- tests/test_dataset_source.py | 76 +++ tests/test_release_manifests.py | 48 +- tests/test_us_long_term_datasets.py | 4 +- 21 files changed, 577 insertions(+), 720 deletions(-) create mode 100644 src/policyengine/tax_benefit_models/common/dataset_source.py create mode 100644 src/policyengine/utils/hashing.py create mode 100644 tests/test_dataset_source.py diff --git a/changelog.d/502.changed.md b/changelog.d/502.changed.md index c69157f2..48f94c18 100644 --- a/changelog.d/502.changed.md +++ b/changelog.d/502.changed.md @@ -1,5 +1,4 @@ -PolicyEngine.py now materializes managed datasets from the exact Hugging Face -repository type, immutable revision, and SHA-256 certified by its release +PolicyEngine.py now reuses or downloads managed datasets from the exact Hugging +Face repository type, immutable revision, and SHA-256 recorded in its release bundle. Bundle installation and US and UK calculation entry points share this -implementation, while explicitly unmanaged local and Hugging Face sources -remain opt-in. +implementation, while explicit local paths and Hugging Face URIs remain opt-in. diff --git a/docs/bundles.md b/docs/bundles.md index e3c9b6da..cb0c0cf1 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -53,14 +53,14 @@ from policyengine.provenance import materialize_bundle_dataset result = materialize_bundle_dataset("us", "populace_us_2024") print(result.path) -print(result.actual_sha256) +print(result.sha256) ``` -`materialize_bundle_dataset` returns a Pydantic model containing the selected -source package, repository type, revision, expected and actual hashes, local -path, and cache status. `policyengine-*-data` and `populace-data` artifacts are -selected by their bundle package names. Callers do not infer repository type -from the repository name. +`materialize_bundle_dataset` returns the selected source package, repository +type, revision, verified SHA-256, local path, and optional metadata path. +`policyengine-*-data` and `populace-data` artifacts use the repository type +recorded in the bundle. Callers do not infer repository type from the repository +name. Managed datasets are downloaded from the Hugging Face artifact specified in the bundle. GCS dataset URIs are unsupported. The separate UK geography lookup files diff --git a/docs/microsim.md b/docs/microsim.md index 4f108737..0a4a520c 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -146,7 +146,7 @@ result = materialize_bundle_dataset( ) print(result.path) -print(result.actual_sha256) +print(result.sha256) ``` The bundle API uses the repository type recorded in the bundle, so callers do @@ -246,8 +246,7 @@ bundle. Explicit local paths and Hugging Face URIs remain supported in this mode. GCS dataset URIs are not supported. For managed simulations, `sim.policyengine_bundle` records the actual source -package, repository type, revision, expected and actual SHA-256, local path, and -whether an already verified file was reused. +package, repository type, revision, verified SHA-256, and local path. ## Pinned model versions diff --git a/src/policyengine/bundle.py b/src/policyengine/bundle.py index 6dd63091..c4518673 100644 --- a/src/policyengine/bundle.py +++ b/src/policyengine/bundle.py @@ -24,12 +24,12 @@ from policyengine.provenance.dataset_materialization import ( DatasetMaterializationError, MaterializedDataset, - _materialize_resolved_dataset, + _BundleDatasetSpec, _resolve_bundle_dataset, - _ResolvedBundleDataset, - _sha256_file, + _reuse_or_download_bundle_files, ) from policyengine.provenance.manifest import CountryReleaseManifest +from policyengine.utils.hashing import sha256_file BUNDLE_MANIFEST_RESOURCE = ("data", "bundle", "manifest.json") BUNDLE_HISTORY_RESOURCE = ("data", "bundles") @@ -269,7 +269,7 @@ def install_package_scaffold( def _confirm_dataset_install( - plans: Sequence[_ResolvedBundleDataset], + plans: Sequence[_BundleDatasetSpec], *, data_dir: Path, yes: bool, @@ -289,7 +289,7 @@ def _confirm_dataset_install( def _receipt_dataset( - plan: _ResolvedBundleDataset, + plan: _BundleDatasetSpec, release: Mapping[str, Any], *, materialized: Optional[MaterializedDataset] = None, @@ -306,9 +306,9 @@ def _receipt_dataset( } if release.get("build_id"): receipt["build_id"] = release["build_id"] - receipt["expected_sha256"] = plan.expected_sha256 + receipt["expected_sha256"] = plan.sha256 if materialized is not None: - receipt["installed_sha256"] = materialized.actual_sha256 + receipt["installed_sha256"] = materialized.sha256 return receipt @@ -317,7 +317,7 @@ def _selected_dataset_plans( countries: Sequence[str], *, data_dir: Path, -) -> list[tuple[_ResolvedBundleDataset, Mapping[str, Any]]]: +) -> list[tuple[_BundleDatasetSpec, Mapping[str, Any]]]: releases = manifest.get("data_releases") if not isinstance(releases, Mapping): raise BundleError("Bundle manifest does not contain data releases.") @@ -418,7 +418,7 @@ def install_bundle( installed_datasets.append(_receipt_dataset(plan, release)) continue try: - materialized = _materialize_resolved_dataset(plan) + materialized = _reuse_or_download_bundle_files(plan) except DatasetMaterializationError as exc: raise BundleError(str(exc)) from exc installed_datasets.append( @@ -638,7 +638,7 @@ def _dataset_checks( def _dataset_check( - plan: _ResolvedBundleDataset, + plan: _BundleDatasetSpec, release: Mapping[str, Any], receipt_dataset: Optional[Mapping[str, Any]], ) -> dict[str, Any]: @@ -648,7 +648,7 @@ def _dataset_check( "dataset": plan.dataset, "expected_version": expected_version, "expected_path": str(plan.destination), - "expected_sha256": plan.expected_sha256, + "expected_sha256": plan.sha256, } if receipt_dataset is None: check["status"] = "missing_receipt" @@ -661,11 +661,9 @@ def _dataset_check( if not path.exists(): check["status"] = "missing_file" return check - actual_sha256 = _sha256_file(path) + actual_sha256 = sha256_file(path) check["installed_version"] = receipt_dataset.get("version") check["installed_sha256"] = actual_sha256 check["path"] = str(path) - check["status"] = ( - "ok" if actual_sha256 == plan.expected_sha256 else "sha256_mismatch" - ) + check["status"] = "ok" if actual_sha256 == plan.sha256 else "sha256_mismatch" return check diff --git a/src/policyengine/core/run_record.py b/src/policyengine/core/run_record.py index 13f46bf8..10118077 100644 --- a/src/policyengine/core/run_record.py +++ b/src/policyengine/core/run_record.py @@ -20,7 +20,6 @@ from __future__ import annotations -import hashlib from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, Mapping, Optional, Union @@ -31,6 +30,7 @@ extract_bundle_tro_reference, serialize_trace_tro, ) +from policyengine.utils.hashing import sha256_file if TYPE_CHECKING: from .dynamic import Dynamic @@ -51,14 +51,6 @@ class SimulationRunRecord: tro: dict = field(default_factory=dict) -def _sha256_file(path: Union[str, Path]) -> str: - digest = hashlib.sha256() - with open(path, "rb") as handle: - for chunk in iter(lambda: handle.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - def reform_specification( reform: Optional[Union[Policy, Dynamic]], ) -> Optional[dict[str, Any]]: @@ -119,7 +111,7 @@ def _dataset_reference(dataset: Any) -> dict[str, Any]: return { "name": dataset.name, "file": filepath.name, - "sha256": _sha256_file(filepath), + "sha256": sha256_file(filepath), "year": dataset.year, } diff --git a/src/policyengine/provenance/__init__.py b/src/policyengine/provenance/__init__.py index af781ea1..001cb18d 100644 --- a/src/policyengine/provenance/__init__.py +++ b/src/policyengine/provenance/__init__.py @@ -81,9 +81,6 @@ from .manifest import ( resolve_dataset_reference as resolve_dataset_reference, ) -from .manifest import ( - resolve_local_managed_dataset_source as resolve_local_managed_dataset_source, -) from .manifest import ( resolve_managed_dataset_reference as resolve_managed_dataset_reference, ) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index ef99bc10..a845f8b5 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -2,25 +2,23 @@ from __future__ import annotations -import hashlib import os import tempfile from dataclasses import dataclass from pathlib import Path -from typing import Literal, Optional, Union -from urllib.parse import quote +from typing import Literal, Optional import requests -from pydantic import BaseModel + +from policyengine.utils.hashing import sha256_file from .manifest import ( CountryReleaseManifest, _artifact_revision, build_hf_uri, - dataset_logical_name, get_release_manifest, - resolve_local_managed_dataset_source, - resolve_managed_dataset_reference, + https_dataset_uri, + hugging_face_auth_headers, ) DEFAULT_DATA_DIR = Path("./data") @@ -28,25 +26,26 @@ class DatasetMaterializationError(ValueError): - """Raised when a dataset cannot be made available safely.""" + """Raised when a bundle dataset cannot be made available safely.""" -class MaterializedDataset(BaseModel): - """Verified local representation of one bundle-managed dataset.""" +@dataclass(frozen=True) +class MaterializedDataset: + """Local file and provenance values for a verified bundle dataset.""" data_package_name: str repo_type: Literal["model", "dataset"] revision: str source_uri: str - expected_sha256: str - actual_sha256: str + sha256: str path: Path - cache_hit: bool metadata_path: Optional[Path] = None @dataclass(frozen=True) -class _ResolvedBundleDataset: +class _BundleDatasetSpec: + """Manifest values required to inspect or download one bundle dataset.""" + country_id: str dataset: str data_package_name: str @@ -54,9 +53,9 @@ class _ResolvedBundleDataset: repo_type: Literal["model", "dataset"] path: str revision: str - expected_sha256: str + sha256: str destination: Path - metadata_expected_sha256: Optional[str] = None + metadata_sha256: Optional[str] = None @property def source_uri(self) -> str: @@ -73,7 +72,7 @@ def _resolve_bundle_dataset( *, data_dir: Path = DEFAULT_DATA_DIR, manifest: Optional[CountryReleaseManifest] = None, -) -> _ResolvedBundleDataset: +) -> _BundleDatasetSpec: country_manifest = manifest or get_release_manifest(country_id) dataset_name = dataset or country_manifest.default_dataset reference = country_manifest.datasets.get(dataset_name) @@ -87,7 +86,7 @@ def _resolve_bundle_dataset( f"Managed dataset {dataset_name!r} is missing a certified sha256." ) - return _ResolvedBundleDataset( + return _BundleDatasetSpec( country_id=country_id, dataset=dataset_name, data_package_name=( @@ -98,9 +97,9 @@ def _resolve_bundle_dataset( path=reference.path, revision=reference.revision or _artifact_revision(country_manifest.data_package), - expected_sha256=reference.sha256, + sha256=reference.sha256, destination=data_dir / Path(reference.path).name, - metadata_expected_sha256=reference.metadata_sha256, + metadata_sha256=reference.metadata_sha256, ) @@ -112,338 +111,93 @@ def materialize_bundle_dataset( ) -> MaterializedDataset: """Return a verified local copy of a dataset from the installed bundle.""" - return _materialize_resolved_dataset( + return _reuse_or_download_bundle_files( _resolve_bundle_dataset(country_id, dataset, data_dir=data_dir) ) -def _materialize_resolved_dataset( - resolved: _ResolvedBundleDataset, - *, - session=requests, +def _reuse_or_download_bundle_files( + dataset: _BundleDatasetSpec, ) -> MaterializedDataset: - local_source = resolve_local_managed_dataset_source( - resolved.country_id, - resolved.source_uri, - ) - local_path = Path(local_source).expanduser() - local_sha256 = ( - _matching_sha256(local_path, resolved.expected_sha256) - if local_source != resolved.source_uri - else None - ) - - if local_sha256 is not None: - path = local_path - actual_sha256 = local_sha256 - cache_hit = True - else: - path, actual_sha256, cache_hit = _materialize_verified_file( - url=_hf_download_url( - repo_id=resolved.repo_id, - repo_type=resolved.repo_type, - path=resolved.path, - revision=resolved.revision, - ), - destination=resolved.destination, - expected_sha256=resolved.expected_sha256, - description=(f"{resolved.country_id.upper()} dataset {resolved.dataset!r}"), - session=session, + files = [ + ( + dataset.path, + dataset.destination, + dataset.sha256, + f"{dataset.country_id.upper()} dataset {dataset.dataset!r}", ) - - metadata_path = _materialize_metadata( - resolved, - dataset_path=path, - session=session, - ) - return MaterializedDataset( - data_package_name=resolved.data_package_name, - repo_type=resolved.repo_type, - revision=resolved.revision, - source_uri=resolved.source_uri, - expected_sha256=resolved.expected_sha256, - actual_sha256=actual_sha256, - path=path, - cache_hit=cache_hit, - metadata_path=metadata_path, - ) - - -def _materialize_metadata( - resolved: _ResolvedBundleDataset, - *, - dataset_path: Path, - session=requests, -) -> Optional[Path]: - expected_sha256 = resolved.metadata_expected_sha256 - if expected_sha256 is None: - return None - - local_path = Path(f"{dataset_path}.metadata.json") - if _matching_sha256(local_path, expected_sha256) is not None: - return local_path - - path, _, _ = _materialize_verified_file( - url=_hf_download_url( - repo_id=resolved.repo_id, - repo_type=resolved.repo_type, - path=f"{resolved.path}.metadata.json", - revision=resolved.revision, - ), - destination=resolved.metadata_destination, - expected_sha256=expected_sha256, - description=f"metadata for {resolved.dataset!r}", - session=session, - ) - return path - - -def _materialize_verified_file( - *, - url: str, - destination: Path, - expected_sha256: str, - description: str, - session=requests, -) -> tuple[Path, str, bool]: - existing_sha256 = _matching_sha256(destination, expected_sha256) - if existing_sha256 is not None: - return destination, existing_sha256, True - - downloaded = _download_to_temp( - url, - destination=destination, - description=description, - session=session, - ) - try: - actual_sha256 = _sha256_file(downloaded) - if actual_sha256 != expected_sha256: - raise DatasetMaterializationError( - f"Downloaded {description} has sha256 {actual_sha256}, " - f"expected {expected_sha256}." - ) - os.replace(downloaded, destination) - finally: - downloaded.unlink(missing_ok=True) - return destination, actual_sha256, False - - -def _materialize_dataset_request( - country_id: str, - dataset: Optional[str], - *, - allow_unmanaged: bool, - data_dir: Path = DEFAULT_DATA_DIR, -) -> tuple[str, str, Optional[MaterializedDataset]]: - manifest = get_release_manifest(country_id) - managed_dataset = None - if dataset is None: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - elif dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - - if managed_dataset is not None: - materialized = _materialize_resolved_dataset( - _resolve_bundle_dataset( - country_id, - managed_dataset, - data_dir=data_dir, - manifest=manifest, + ] + if dataset.metadata_sha256 is not None: + files.append( + ( + f"{dataset.path}.metadata.json", + dataset.metadata_destination, + dataset.metadata_sha256, + f"metadata for {dataset.dataset!r}", ) ) - return materialized.source_uri, str(materialized.path), materialized - - source_uri = resolve_managed_dataset_reference( - country_id, - dataset, - allow_unmanaged=allow_unmanaged, - ) - return ( - source_uri, - _materialize_unmanaged_dataset_source(source_uri, data_dir=data_dir), - None, - ) + for repository_path, destination, expected_sha256, description in files: + if destination.is_file() and sha256_file(destination) == expected_sha256: + continue -def _runtime_dataset_provenance( - source_uri: str, - local_path: str, - materialized: Optional[MaterializedDataset], - *, - logical_name: Optional[str] = None, -) -> dict[str, object]: - provenance: dict[str, object] = { - "managed_by": "policyengine.py", - "runtime_dataset": logical_name or dataset_logical_name(source_uri), - "runtime_dataset_uri": source_uri, - "runtime_dataset_source": local_path, - } - if materialized is not None: - provenance.update( - { - "runtime_dataset_data_package": materialized.data_package_name, - "runtime_dataset_repo_type": materialized.repo_type, - "runtime_dataset_revision": materialized.revision, - "runtime_dataset_expected_sha256": materialized.expected_sha256, - "runtime_dataset_sha256": materialized.actual_sha256, - "runtime_dataset_cache_hit": materialized.cache_hit, - } + destination.parent.mkdir(parents=True, exist_ok=True) + file_descriptor, temp_name = tempfile.mkstemp( + prefix=".policyengine-download-", + suffix=destination.suffix or ".download", + dir=destination.parent, ) - return provenance - - -def _materialize_unmanaged_dataset_source( - dataset_source: Union[str, Path], - *, - version: Optional[str] = None, - data_dir: Path = DEFAULT_DATA_DIR, - repo_type: Optional[Literal["model", "dataset"]] = None, - session=requests, -) -> str: - """Return a local path for an explicitly unmanaged local or HF source.""" - - source = str(dataset_source) - if not source.startswith("hf://"): - if "://" in source: - raise DatasetMaterializationError( - f"Unsupported unmanaged dataset URI: {source!r}." - ) - return source - - repo_id, path, revision = _parse_unmanaged_hf_reference(source, version=version) - destination = data_dir / Path(path).name - repo_types: list[Literal["model", "dataset"]] = ( - [repo_type] if repo_type is not None else ["model", "dataset"] - ) - for index, candidate_repo_type in enumerate(repo_types): - try: - downloaded = _download_to_temp( - _hf_download_url( - repo_id=repo_id, - repo_type=candidate_repo_type, - path=path, - revision=revision, - ), - destination=destination, - description=f"unmanaged dataset {source!r}", - session=session, - ) - except _DatasetNotFoundError: - if index + 1 < len(repo_types): - continue - raise - try: - os.replace(downloaded, destination) - finally: - downloaded.unlink(missing_ok=True) - return str(destination) - - raise DatasetMaterializationError(f"Could not materialize dataset {source!r}.") - - -def _parse_unmanaged_hf_reference( - uri: str, - *, - version: Optional[str], -) -> tuple[str, str, str]: - path_with_repo, uri_revision = ( - uri[5:].rsplit("@", maxsplit=1) if "@" in uri[5:] else (uri[5:], None) - ) - if uri_revision is not None and version is not None and uri_revision != version: - raise DatasetMaterializationError( - "Conflicting dataset versions: " - f"URI requests {uri_revision!r} but version is {version!r}." - ) - parts = path_with_repo.split("/", maxsplit=2) - if len(parts) != 3 or not all(parts): - raise DatasetMaterializationError( - "Invalid Hugging Face dataset URI. Expected format " - f"'hf://owner/repo/path/to/file[@revision]', got {uri!r}." + os.close(file_descriptor) + temporary_path = Path(temp_name) + url = https_dataset_uri( + dataset.repo_id, + repository_path, + dataset.revision, + repo_type=dataset.repo_type, ) - return f"{parts[0]}/{parts[1]}", parts[2], uri_revision or version or "main" - - -class _DatasetNotFoundError(DatasetMaterializationError): - pass - - -def _hf_download_url( - *, - repo_id: str, - repo_type: Literal["model", "dataset"], - path: str, - revision: str, -) -> str: - prefix = "datasets/" if repo_type == "dataset" else "" - return ( - f"https://huggingface.co/{prefix}{repo_id}/resolve/" - f"{quote(revision, safe='')}/{quote(path)}" - ) - - -def _download_to_temp( - url: str, - *, - destination: Path, - description: str, - session=requests, -) -> Path: - destination.parent.mkdir(parents=True, exist_ok=True) - file_descriptor, temp_name = tempfile.mkstemp( - prefix=".policyengine-download-", - suffix=destination.suffix or ".download", - dir=destination.parent, - ) - os.close(file_descriptor) - temp_path = Path(temp_name) - try: - with session.get( - url, - headers=_hugging_face_auth_headers(), - stream=True, - timeout=DOWNLOAD_TIMEOUT_SECONDS, - ) as response: - if response.status_code in {401, 403}: + try: + with requests.get( + url, + headers=hugging_face_auth_headers(), + stream=True, + timeout=DOWNLOAD_TIMEOUT_SECONDS, + ) as response: + if response.status_code in {401, 403}: + raise DatasetMaterializationError( + f"Could not download {description}: Hugging Face rejected " + "the configured credentials. Set HUGGING_FACE_TOKEN to a " + "token with access to the repository." + ) + if response.status_code == 404: + raise DatasetMaterializationError( + f"Could not find {description} at {url}." + ) + response.raise_for_status() + with temporary_path.open("wb") as output: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + output.write(chunk) + + actual_sha256 = sha256_file(temporary_path) + if actual_sha256 != expected_sha256: raise DatasetMaterializationError( - f"Could not download {description}: Hugging Face rejected " - "the configured credentials. Set HUGGING_FACE_TOKEN to a " - "token with access to the repository." + f"Downloaded {description} has sha256 {actual_sha256}, " + f"expected {expected_sha256}." ) - if response.status_code == 404: - raise _DatasetNotFoundError(f"Could not find {description} at {url}.") - response.raise_for_status() - with temp_path.open("wb") as output: - for chunk in response.iter_content(chunk_size=1024 * 1024): - if chunk: - output.write(chunk) - except Exception: - temp_path.unlink(missing_ok=True) - raise - return temp_path - + os.replace(temporary_path, destination) + finally: + temporary_path.unlink(missing_ok=True) -def _hugging_face_auth_headers() -> dict[str, str]: - token = ( - os.environ.get("HUGGING_FACE_TOKEN") - or os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_HUB_TOKEN") + return MaterializedDataset( + data_package_name=dataset.data_package_name, + repo_type=dataset.repo_type, + revision=dataset.revision, + source_uri=dataset.source_uri, + sha256=dataset.sha256, + path=dataset.destination, + metadata_path=( + dataset.metadata_destination + if dataset.metadata_sha256 is not None + else None + ), ) - return {"Authorization": f"Bearer {token}"} if token else {} - - -def _matching_sha256(path: Path, expected_sha256: str) -> Optional[str]: - if not path.is_file(): - return None - actual_sha256 = _sha256_file(path) - return actual_sha256 if actual_sha256 == expected_sha256 else None - - -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() diff --git a/src/policyengine/provenance/manifest.py b/src/policyengine/provenance/manifest.py index e18c719b..99ca72be 100644 --- a/src/policyengine/provenance/manifest.py +++ b/src/policyengine/provenance/manifest.py @@ -2,20 +2,16 @@ import json import os from functools import lru_cache -from importlib import import_module from importlib.resources import files from pathlib import Path from typing import Literal, Optional +from urllib.parse import quote import requests from pydantic import BaseModel, Field HF_REQUEST_TIMEOUT_SECONDS = 30 PYPI_REQUEST_TIMEOUT_SECONDS = 30 -LOCAL_DATA_REPO_HINTS = { - "us": ("policyengine_us", "policyengine-us-data", "policyengine_us_data"), - "uk": ("policyengine_uk", "policyengine-uk-data", "policyengine_uk_data"), -} class DataReleaseManifestUnavailableError(ValueError): @@ -203,7 +199,21 @@ def https_dataset_uri( ) -> str: """Return a dereferenceable HTTPS URI for a Hugging Face dataset artifact.""" prefix = "datasets/" if repo_type == "dataset" else "" - return f"https://huggingface.co/{prefix}{repo_id}/resolve/{revision}/{path_in_repo}" + return ( + f"https://huggingface.co/{prefix}{repo_id}/resolve/" + f"{quote(revision, safe='')}/{quote(path_in_repo)}" + ) + + +def hugging_face_auth_headers() -> dict[str, str]: + """Return authentication headers for Hugging Face requests.""" + + token = ( + os.environ.get("HUGGING_FACE_TOKEN") + or os.environ.get("HF_TOKEN") + or os.environ.get("HUGGINGFACE_HUB_TOKEN") + ) + return {"Authorization": f"Bearer {token}"} if token else {} def _artifact_revision(data_package: "DataPackageVersion") -> str: @@ -325,15 +335,10 @@ def get_release_manifest(country_id: str) -> CountryReleaseManifest: def get_data_release_manifest(country_id: str) -> DataReleaseManifest: country_manifest = get_release_manifest(country_id) - headers = {} - token = os.environ.get("HUGGING_FACE_TOKEN") - if token: - headers["Authorization"] = f"Bearer {token}" - try: response = requests.get( https_release_manifest_uri(country_manifest.data_package), - headers=headers, + headers=hugging_face_auth_headers(), timeout=HF_REQUEST_TIMEOUT_SECONDS, ) except requests.RequestException as exc: @@ -598,67 +603,6 @@ def resolve_managed_dataset_reference( ) -def resolve_local_managed_dataset_source( - country_id: str, - dataset_uri: str, - *, - allow_local_mirror: bool = True, -) -> str: - """Resolve a local mirror of a managed dataset when available. - - This preserves the bundled dataset URI for provenance while allowing local - development environments with sibling data-repo checkouts to load the - exact certified artifact from disk rather than re-downloading it. - """ - - if not allow_local_mirror or not dataset_uri.startswith("hf://"): - return dataset_uri - - local_hint = LOCAL_DATA_REPO_HINTS.get(country_id) - if local_hint is None: - return dataset_uri - - path_without_revision = dataset_uri[5:].rsplit("@", 1)[0] - parts = path_without_revision.split("/", 2) - if len(parts) != 3: - return dataset_uri - _, _, path_in_repo = parts - - model_module_name, data_repo_name, data_package_name = local_hint - explicit_repo_roots = [] - country_env = f"POLICYENGINE_{country_id.upper()}_DATA_REPO" - for env_name in (country_env, "POLICYENGINE_LOCAL_DATA_REPO_ROOT"): - env_value = os.environ.get(env_name) - if env_value: - explicit_repo_roots.extend( - [ - Path(env_value).expanduser(), - Path(env_value).expanduser() / data_repo_name, - ] - ) - - for candidate_repo_root in explicit_repo_roots: - local_path = candidate_repo_root / data_package_name / "storage" / path_in_repo - if local_path.exists(): - return str(local_path) - - try: - model_module = import_module(model_module_name) - except ImportError: - return dataset_uri - - repo_root = Path(model_module.__file__).resolve().parents[1] - local_path = ( - repo_root.with_name(data_repo_name) - / data_package_name - / "storage" - / path_in_repo - ) - if local_path.exists(): - return str(local_path) - return dataset_uri - - def dataset_logical_name(dataset: str) -> str: return Path(dataset.rsplit("@", 1)[0]).stem diff --git a/src/policyengine/tax_benefit_models/common/dataset_source.py b/src/policyengine/tax_benefit_models/common/dataset_source.py new file mode 100644 index 00000000..b1511a0d --- /dev/null +++ b/src/policyengine/tax_benefit_models/common/dataset_source.py @@ -0,0 +1,98 @@ +"""Downloads for explicit Hugging Face dataset inputs.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + +import requests + +from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, +) +from policyengine.provenance.manifest import ( + https_dataset_uri, + hugging_face_auth_headers, +) + +DEFAULT_DATA_DIR = Path("./data") +DOWNLOAD_TIMEOUT_SECONDS = 60 + + +def download_hf_dataset( + dataset_uri: str, + *, + data_dir: Path = DEFAULT_DATA_DIR, +) -> str: + """Download an explicit Hugging Face dataset URI and return its local path.""" + + if not dataset_uri.startswith("hf://"): + raise DatasetMaterializationError( + f"Expected an hf:// dataset URI, got {dataset_uri!r}." + ) + + path_with_repo, revision = ( + dataset_uri[5:].rsplit("@", maxsplit=1) + if "@" in dataset_uri[5:] + else (dataset_uri[5:], "main") + ) + parts = path_with_repo.split("/", maxsplit=2) + if len(parts) != 3 or not all(parts): + raise DatasetMaterializationError( + "Invalid Hugging Face dataset URI. Expected format " + f"'hf://owner/repo/path/to/file[@revision]', got {dataset_uri!r}." + ) + + repo_id = f"{parts[0]}/{parts[1]}" + repository_path = parts[2] + destination = data_dir / Path(repository_path).name + destination.parent.mkdir(parents=True, exist_ok=True) + + for repo_type in ("model", "dataset"): + file_descriptor, temp_name = tempfile.mkstemp( + prefix=".policyengine-download-", + suffix=destination.suffix or ".download", + dir=destination.parent, + ) + os.close(file_descriptor) + temporary_path = Path(temp_name) + url = https_dataset_uri( + repo_id, + repository_path, + revision, + repo_type=repo_type, + ) + try: + with requests.get( + url, + headers=hugging_face_auth_headers(), + stream=True, + timeout=DOWNLOAD_TIMEOUT_SECONDS, + ) as response: + if response.status_code in {401, 403}: + raise DatasetMaterializationError( + "Could not download explicit dataset " + f"{dataset_uri!r}: Hugging Face rejected the configured " + "credentials. Set HUGGING_FACE_TOKEN to a token with " + "access to the repository." + ) + if response.status_code == 404: + if repo_type == "model": + continue + raise DatasetMaterializationError( + f"Could not find explicit dataset {dataset_uri!r}." + ) + response.raise_for_status() + with temporary_path.open("wb") as output: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + output.write(chunk) + os.replace(temporary_path, destination) + return str(destination) + finally: + temporary_path.unlink(missing_ok=True) + + raise DatasetMaterializationError( + f"Could not download explicit dataset {dataset_uri!r}." + ) diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index e8c3b696..57ac573b 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -35,6 +35,7 @@ ) from policyengine.provenance.manifest import ( certify_data_release_compatibility, + dataset_logical_name, get_release_manifest, ) from policyengine.utils.entity_utils import build_entity_relationships @@ -45,6 +46,7 @@ if TYPE_CHECKING: from policyengine.core.simulation import Simulation + from policyengine.provenance.dataset_materialization import MaterializedDataset def output_dataset_filepath(simulation: Simulation) -> Path: @@ -61,6 +63,33 @@ def output_dataset_filepath(simulation: Simulation) -> Path: return parent / (simulation.id + ".h5") +def build_runtime_dataset_provenance( + source_uri: str, + local_path: str, + materialized: Optional[MaterializedDataset], + *, + logical_name: Optional[str] = None, +) -> dict[str, object]: + """Return the dataset fields recorded with a managed simulation.""" + + provenance: dict[str, object] = { + "managed_by": "policyengine.py", + "runtime_dataset": logical_name or dataset_logical_name(source_uri), + "runtime_dataset_uri": source_uri, + "runtime_dataset_source": local_path, + } + if materialized is not None: + provenance.update( + { + "runtime_dataset_data_package": materialized.data_package_name, + "runtime_dataset_repo_type": materialized.repo_type, + "runtime_dataset_revision": materialized.revision, + "runtime_dataset_sha256": materialized.sha256, + } + ) + return provenance + + class MicrosimulationModelVersion(TaxBenefitModelVersion): """Shared init / save / load logic for country microsim model versions. diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index b61185e0..bcc0d0df 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -7,12 +7,17 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( - _materialize_dataset_request, + DatasetMaterializationError, + materialize_bundle_dataset, ) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, + resolve_managed_dataset_reference, +) +from policyengine.tax_benefit_models.common.dataset_source import ( + download_hf_dataset, ) @@ -125,12 +130,38 @@ def create_datasets( datasets = [get_release_manifest("uk").default_dataset] result = {} for dataset in datasets: - resolved_dataset, runtime_dataset, _ = _materialize_dataset_request( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - data_dir=Path(data_folder), - ) + manifest = get_release_manifest("uk") + if dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + else: + managed_dataset = None + if managed_dataset is not None: + materialized = materialize_bundle_dataset( + "uk", + managed_dataset, + data_dir=Path(data_folder), + ) + resolved_dataset = materialized.source_uri + runtime_dataset = str(materialized.path) + else: + resolved_dataset = resolve_managed_dataset_reference( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + ) + if resolved_dataset.startswith("hf://"): + runtime_dataset = download_hf_dataset( + resolved_dataset, + data_dir=Path(data_folder), + ) + elif "://" in resolved_dataset: + raise DatasetMaterializationError( + f"Unsupported explicit dataset URI: {resolved_dataset!r}." + ) + else: + runtime_dataset = resolved_dataset dataset_stem = dataset_logical_name(resolved_dataset) from policyengine_uk import Microsimulation diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index c1ed4003..a513c849 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -6,10 +6,20 @@ from policyengine.core import TaxBenefitModel from policyengine.provenance.dataset_materialization import ( - _materialize_dataset_request, - _runtime_dataset_provenance, + DatasetMaterializationError, + materialize_bundle_dataset, +) +from policyengine.provenance.manifest import ( + get_release_manifest, + resolve_managed_dataset_reference, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion +from policyengine.tax_benefit_models.common.dataset_source import ( + download_hf_dataset, +) +from policyengine.tax_benefit_models.common.model_version import ( + build_runtime_dataset_provenance, +) from policyengine.tax_benefit_models.common.model_version import ( output_dataset_filepath as _output_dataset_filepath, ) @@ -283,11 +293,32 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - dataset_uri, runtime_dataset_source, materialized = _materialize_dataset_request( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - ) + manifest = get_release_manifest("uk") + if dataset is None or dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + else: + managed_dataset = None + materialized = None + if managed_dataset is not None: + materialized = materialize_bundle_dataset("uk", managed_dataset) + dataset_uri = materialized.source_uri + runtime_dataset_source = str(materialized.path) + else: + dataset_uri = resolve_managed_dataset_reference( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + ) + if dataset_uri.startswith("hf://"): + runtime_dataset_source = download_hf_dataset(dataset_uri) + elif "://" in dataset_uri: + raise DatasetMaterializationError( + f"Unsupported explicit dataset URI: {dataset_uri!r}." + ) + else: + runtime_dataset_source = dataset_uri runtime_dataset = runtime_dataset_source if isinstance(runtime_dataset_source, str) and "://" not in runtime_dataset_source: from policyengine_uk.data.dataset_schema import ( @@ -302,7 +333,7 @@ def managed_microsimulation( microsim = Microsimulation(dataset=runtime_dataset, **kwargs) microsim.policyengine_bundle = dict(uk_latest.release_bundle) microsim.policyengine_bundle.update( - _runtime_dataset_provenance( + build_runtime_dataset_provenance( dataset_uri, runtime_dataset_source, materialized, diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 806d6a54..abdb3aec 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -13,16 +13,23 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, MaterializedDataset, - _materialize_dataset_request, - _runtime_dataset_provenance, materialize_bundle_dataset, ) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, + resolve_managed_dataset_reference, ) +from policyengine.tax_benefit_models.common.dataset_source import ( + download_hf_dataset, +) +from policyengine.tax_benefit_models.common.model_version import ( + build_runtime_dataset_provenance, +) +from policyengine.utils.hashing import sha256_file class USYearData(YearData): @@ -299,12 +306,38 @@ def create_datasets( datasets = datasets or [get_release_manifest("us").default_dataset] result = {} for dataset in datasets: - resolved_dataset, runtime_dataset, _ = _materialize_dataset_request( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - data_dir=Path(data_folder), - ) + manifest = get_release_manifest("us") + if dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + else: + managed_dataset = None + if managed_dataset is not None: + materialized = materialize_bundle_dataset( + "us", + managed_dataset, + data_dir=Path(data_folder), + ) + resolved_dataset = materialized.source_uri + runtime_dataset = str(materialized.path) + else: + resolved_dataset = resolve_managed_dataset_reference( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + ) + if resolved_dataset.startswith("hf://"): + runtime_dataset = download_hf_dataset( + resolved_dataset, + data_dir=Path(data_folder), + ) + elif "://" in resolved_dataset: + raise DatasetMaterializationError( + f"Unsupported explicit dataset URI: {resolved_dataset!r}." + ) + else: + runtime_dataset = resolved_dataset dataset_stem = dataset_logical_name(resolved_dataset) sim = Microsimulation(dataset=runtime_dataset) @@ -583,7 +616,7 @@ def _runtime_policyengine_us_metadata() -> dict[str, Any]: result["direct_url"] = {} package_file = _runtime_policyengine_us_package_file() if package_file is not None: - result["package_file_sha256"] = _sha256_file(package_file) + result["package_file_sha256"] = sha256_file(package_file) result["package_tree_sha256"] = _sha256_directory(package_file.parent) return result @@ -861,7 +894,7 @@ def _build_long_term_dataset( if dataset_uri is not None: dataset.metadata.setdefault("policyengine_bundle", {}) dataset.metadata["policyengine_bundle"].update( - _runtime_dataset_provenance( + build_runtime_dataset_provenance( dataset_uri, str(path), materialized, @@ -927,14 +960,6 @@ def _validate_loaded_long_term_metadata( ) -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - def load_long_term_datasets( years: list[int], data_folder: str = "./projected_datasets", @@ -1051,10 +1076,10 @@ def load_managed_long_term_datasets( """Load bundled long-term US datasets from the managed release manifest. Each requested year must have a logical dataset entry named - ``{dataset_name}_{year}`` in the bundled US manifest. A verified local mirror - is reused when available; otherwise the exact bundle-certified Hugging Face - artifact and its certified metadata sidecar are materialized into - ``data_folder``. + ``{dataset_name}_{year}`` in the bundled US manifest. A verified file in + ``data_folder`` is reused when available; otherwise the exact + bundle-certified Hugging Face artifact and its certified metadata sidecar + are downloaded there. """ manifest = get_release_manifest("us") diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 2890e363..43331b72 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -6,10 +6,20 @@ from policyengine.core import TaxBenefitModel from policyengine.provenance.dataset_materialization import ( - _materialize_dataset_request, - _runtime_dataset_provenance, + DatasetMaterializationError, + materialize_bundle_dataset, +) +from policyengine.provenance.manifest import ( + get_release_manifest, + resolve_managed_dataset_reference, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion +from policyengine.tax_benefit_models.common.dataset_source import ( + download_hf_dataset, +) +from policyengine.tax_benefit_models.common.model_version import ( + build_runtime_dataset_provenance, +) from policyengine.tax_benefit_models.common.model_version import ( output_dataset_filepath as _output_dataset_filepath, ) @@ -426,15 +436,36 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - dataset_uri, runtime_dataset_source, materialized = _materialize_dataset_request( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) + manifest = get_release_manifest("us") + if dataset is None or dataset == manifest.default_dataset_uri: + managed_dataset = manifest.default_dataset + elif dataset in manifest.datasets: + managed_dataset = dataset + else: + managed_dataset = None + materialized = None + if managed_dataset is not None: + materialized = materialize_bundle_dataset("us", managed_dataset) + dataset_uri = materialized.source_uri + runtime_dataset_source = str(materialized.path) + else: + dataset_uri = resolve_managed_dataset_reference( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + ) + if dataset_uri.startswith("hf://"): + runtime_dataset_source = download_hf_dataset(dataset_uri) + elif "://" in dataset_uri: + raise DatasetMaterializationError( + f"Unsupported explicit dataset URI: {dataset_uri!r}." + ) + else: + runtime_dataset_source = dataset_uri microsim = Microsimulation(dataset=runtime_dataset_source, **kwargs) microsim.policyengine_bundle = dict(us_latest.release_bundle) microsim.policyengine_bundle.update( - _runtime_dataset_provenance( + build_runtime_dataset_provenance( dataset_uri, runtime_dataset_source, materialized, diff --git a/src/policyengine/utils/hashing.py b/src/policyengine/utils/hashing.py new file mode 100644 index 00000000..2030d908 --- /dev/null +++ b/src/policyengine/utils/hashing.py @@ -0,0 +1,15 @@ +"""File hashing utilities.""" + +import hashlib +from pathlib import Path +from typing import Union + + +def sha256_file(path: Union[str, Path]) -> str: + """Return the SHA-256 digest of a file.""" + + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/tests/test_bundle.py b/tests/test_bundle.py index a28efaaa..70d514db 100644 --- a/tests/test_bundle.py +++ b/tests/test_bundle.py @@ -78,7 +78,7 @@ def test_selected_dataset_plan_uses_certified_release_metadata(tmp_path): assert plan.data_package_name == "policyengine-uk-data" assert plan.repo_type == "model" assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" - assert plan.expected_sha256 == release["datasets"][plan.dataset]["sha256"] + assert plan.sha256 == release["datasets"][plan.dataset]["sha256"] def test_install_bundle_package_only_uses_explicit_python(monkeypatch, tmp_path): @@ -160,13 +160,11 @@ def fake_materialize(plan): repo_type=plan.repo_type, revision=plan.revision, source_uri=plan.source_uri, - expected_sha256=plan.expected_sha256, - actual_sha256=plan.expected_sha256, + sha256=plan.sha256, path=plan.destination, - cache_hit=False, ) - monkeypatch.setattr(bundle, "_materialize_resolved_dataset", fake_materialize) + monkeypatch.setattr(bundle, "_reuse_or_download_bundle_files", fake_materialize) result = bundle.install_bundle( python=sys.executable, @@ -177,7 +175,7 @@ def fake_materialize(plan): assert [plan.country_id for plan in calls] == ["uk"] assert result["datasets"][0]["data_package_name"] == "policyengine-uk-data" - assert result["datasets"][0]["installed_sha256"] == calls[0].expected_sha256 + assert result["datasets"][0]["installed_sha256"] == calls[0].sha256 receipt = bundle.read_receipt(tmp_path) assert receipt is not None assert receipt["datasets"] == result["datasets"] diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 7a66e8ea..52fbf4e9 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -1,14 +1,13 @@ import hashlib -from pathlib import Path +from unittest.mock import patch import pytest +import policyengine.provenance.dataset_materialization as dataset_materialization from policyengine.provenance.dataset_materialization import ( DatasetMaterializationError, - MaterializedDataset, - _materialize_resolved_dataset, - _materialize_unmanaged_dataset_source, _resolve_bundle_dataset, + _reuse_or_download_bundle_files, ) from policyengine.provenance.manifest import CountryReleaseManifest @@ -52,42 +51,27 @@ def _manifest() -> CountryReleaseManifest: def test_resolve_bundle_dataset_inherits_primary_package(tmp_path): - plan = _resolve_bundle_dataset("uk", data_dir=tmp_path, manifest=_manifest()) + dataset = _resolve_bundle_dataset("uk", data_dir=tmp_path, manifest=_manifest()) - assert plan.data_package_name == "policyengine-uk-data" - assert plan.repo_id == "policyengine/policyengine-uk-data-private" - assert plan.repo_type == "model" - assert plan.revision == "uk-release" - assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" + assert dataset.data_package_name == "policyengine-uk-data" + assert dataset.repo_id == "policyengine/policyengine-uk-data-private" + assert dataset.repo_type == "model" + assert dataset.revision == "uk-release" + assert dataset.destination == tmp_path / "enhanced_frs_2024_25.h5" def test_resolve_bundle_dataset_uses_cross_package_overlay(tmp_path): - plan = _resolve_bundle_dataset( + dataset = _resolve_bundle_dataset( "uk", "populace_uk_2023", data_dir=tmp_path, manifest=_manifest(), ) - assert plan.data_package_name == "populace-data" - assert plan.repo_id == "policyengine/populace-uk-private" - assert plan.repo_type == "model" - assert plan.revision == "populace-release" - - -def test_materialized_dataset_round_trips_json(): - plan = _resolve_bundle_dataset("uk", manifest=_manifest()) - result = MaterializedDataset( - data_package_name=plan.data_package_name, - repo_type=plan.repo_type, - revision=plan.revision, - source_uri=plan.source_uri, - expected_sha256=plan.expected_sha256, - actual_sha256=plan.expected_sha256, - path=Path("data/enhanced_frs_2024_25.h5"), - cache_hit=True, - ) - assert MaterializedDataset.model_validate_json(result.model_dump_json()) == result + assert dataset.data_package_name == "populace-data" + assert dataset.repo_id == "policyengine/populace-uk-private" + assert dataset.repo_type == "model" + assert dataset.revision == "populace-release" def _sha256(payload: bytes) -> str: @@ -136,29 +120,30 @@ def _manifest_with_hash( return manifest -def _materialize(manifest, tmp_path, session): - resolved = _resolve_bundle_dataset( +def _download(manifest, tmp_path, session): + dataset = _resolve_bundle_dataset( "uk", data_dir=tmp_path, manifest=manifest, ) - return _materialize_resolved_dataset(resolved, session=session) + with patch.object(dataset_materialization.requests, "get", side_effect=session.get): + return _reuse_or_download_bundle_files(dataset) -def test_materialize_country_data_package_uses_model_repo_url(tmp_path): +def test_country_data_package_uses_model_repository_url(tmp_path): session = _Session(_Response(b"country-data")) manifest = _manifest_with_hash(_sha256(b"country-data")) - result = _materialize(manifest, tmp_path, session) + result = _download(manifest, tmp_path, session) assert result.path.read_bytes() == b"country-data" - assert result.cache_hit is False + assert result.sha256 == _sha256(b"country-data") assert session.calls[0][0].startswith( "https://huggingface.co/policyengine/policyengine-uk-data-private/resolve/" ) -def test_materialize_populace_package_uses_dataset_repo_url(tmp_path): +def test_populace_package_uses_dataset_repository_url(tmp_path): session = _Session(_Response(b"populace-data")) manifest = _manifest_with_hash( _sha256(b"populace-data"), @@ -166,7 +151,7 @@ def test_materialize_populace_package_uses_dataset_repo_url(tmp_path): repo_type="dataset", ) - result = _materialize(manifest, tmp_path, session) + result = _download(manifest, tmp_path, session) assert result.path.read_bytes() == b"populace-data" assert session.calls[0][0].startswith( @@ -174,175 +159,76 @@ def test_materialize_populace_package_uses_dataset_repo_url(tmp_path): ) -def test_materialize_downloads_and_verifies_metadata_sidecar(tmp_path): +def test_downloads_and_verifies_metadata_sidecar(tmp_path): dataset_payload = b"long-term-data" metadata_payload = b'{"year": 2100}' manifest = _manifest_with_hash(_sha256(dataset_payload)) - reference = manifest.datasets[manifest.default_dataset] - reference.metadata_sha256 = _sha256(metadata_payload) + manifest.datasets[manifest.default_dataset].metadata_sha256 = _sha256( + metadata_payload + ) session = _Session(_Response(dataset_payload), _Response(metadata_payload)) - result = _materialize(manifest, tmp_path, session) + result = _download(manifest, tmp_path, session) assert result.metadata_path == (tmp_path / "enhanced_frs_2024_25.h5.metadata.json") assert result.metadata_path.read_bytes() == metadata_payload assert session.calls[1][0].endswith("/enhanced_frs_2024_25.h5.metadata.json") -def test_materialize_reuses_only_hash_verified_cache(tmp_path): +def test_reuses_destination_when_hash_matches(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) destination = tmp_path / "enhanced_frs_2024_25.h5" destination.write_bytes(payload) session = _Session() - result = _materialize(manifest, tmp_path, session) - - assert result.cache_hit is True - assert session.calls == [] - - -def test_materialize_reuses_hash_verified_local_mirror(monkeypatch, tmp_path): - payload = b"certified-local-mirror" - manifest = _manifest_with_hash(_sha256(payload)) - mirror = tmp_path / "mirror" / "enhanced_frs_2024_25.h5" - mirror.parent.mkdir() - mirror.write_bytes(payload) - monkeypatch.setattr( - "policyengine.provenance.dataset_materialization.resolve_local_managed_dataset_source", - lambda *args, **kwargs: str(mirror), - ) - session = _Session() - - result = _materialize(manifest, tmp_path, session) + result = _download(manifest, tmp_path, session) - assert result.path == mirror - assert result.cache_hit is True + assert result.path == destination assert session.calls == [] -def test_materialize_reuses_local_mirror_and_downloads_missing_metadata( - monkeypatch, tmp_path -): - payload = b"certified-local-mirror" - metadata_payload = b'{"year": 2025}' - manifest = _manifest_with_hash(_sha256(payload)) - manifest.datasets[manifest.default_dataset].metadata_sha256 = _sha256( - metadata_payload - ) - mirror = tmp_path / "mirror" / "enhanced_frs_2024_25.h5" - mirror.parent.mkdir() - mirror.write_bytes(payload) - monkeypatch.setattr( - "policyengine.provenance.dataset_materialization.resolve_local_managed_dataset_source", - lambda *args, **kwargs: str(mirror), - ) - session = _Session(_Response(metadata_payload)) - - result = _materialize(manifest, tmp_path, session) - - assert result.path == mirror - assert result.metadata_path == tmp_path / "enhanced_frs_2024_25.h5.metadata.json" - assert result.metadata_path.read_bytes() == metadata_payload - assert len(session.calls) == 1 - - -def test_materialize_ignores_mismatched_local_mirror(monkeypatch, tmp_path): - payload = b"certified-download" - manifest = _manifest_with_hash(_sha256(payload)) - mirror = tmp_path / "mirror" / "enhanced_frs_2024_25.h5" - mirror.parent.mkdir() - mirror.write_bytes(b"wrong") - monkeypatch.setattr( - "policyengine.provenance.dataset_materialization.resolve_local_managed_dataset_source", - lambda *args, **kwargs: str(mirror), - ) - session = _Session(_Response(payload)) - - result = _materialize(manifest, tmp_path, session) - - assert result.path == tmp_path / "enhanced_frs_2024_25.h5" - assert result.path.read_bytes() == payload - assert mirror.read_bytes() == b"wrong" - - -def test_materialize_replaces_mismatched_cache(tmp_path): +def test_replaces_destination_when_hash_does_not_match(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) destination = tmp_path / "enhanced_frs_2024_25.h5" destination.write_bytes(b"old") - _materialize(manifest, tmp_path, _Session(_Response(payload))) + _download(manifest, tmp_path, _Session(_Response(payload))) assert destination.read_bytes() == payload -def test_hash_failure_does_not_replace_existing_cache(tmp_path): +def test_hash_failure_preserves_existing_destination(tmp_path): manifest = _manifest_with_hash(_sha256(b"expected")) destination = tmp_path / "enhanced_frs_2024_25.h5" destination.write_bytes(b"old") with pytest.raises(DatasetMaterializationError, match="sha256"): - _materialize(manifest, tmp_path, _Session(_Response(b"wrong"))) + _download(manifest, tmp_path, _Session(_Response(b"wrong"))) assert destination.read_bytes() == b"old" -def test_materialize_passes_hugging_face_token(monkeypatch, tmp_path): +def test_download_passes_hugging_face_token(monkeypatch, tmp_path): monkeypatch.setenv("HUGGING_FACE_TOKEN", "secret-token") payload = b"certified" session = _Session(_Response(payload)) - _materialize(_manifest_with_hash(_sha256(payload)), tmp_path, session) + _download(_manifest_with_hash(_sha256(payload)), tmp_path, session) assert session.calls[0][1]["headers"] == {"Authorization": "Bearer secret-token"} @pytest.mark.parametrize("status_code", [401, 403]) -def test_managed_auth_failure_does_not_retry_repo_type(tmp_path, status_code): +def test_authentication_failure_does_not_retry_repository_type(tmp_path, status_code): session = _Session(_Response(status_code=status_code)) with pytest.raises(DatasetMaterializationError, match="credentials"): - _materialize( + _download( _manifest_with_hash(_sha256(b"certified")), tmp_path, session, ) assert len(session.calls) == 1 - - -def test_unmanaged_hf_retries_dataset_repo_only_after_not_found(tmp_path): - session = _Session(_Response(status_code=404), _Response(b"dataset")) - - result = _materialize_unmanaged_dataset_source( - "hf://policyengine/example/data.h5@release", - data_dir=tmp_path, - session=session, - ) - - assert Path(result).read_bytes() == b"dataset" - assert "/policyengine/example/" in session.calls[0][0] - assert "/datasets/policyengine/example/" in session.calls[1][0] - - -def test_unmanaged_auth_failure_does_not_retry(tmp_path): - session = _Session(_Response(status_code=403)) - - with pytest.raises(DatasetMaterializationError, match="credentials"): - _materialize_unmanaged_dataset_source( - "hf://policyengine/example/data.h5@release", - data_dir=tmp_path, - session=session, - ) - - assert len(session.calls) == 1 - - -def test_unmanaged_local_path_is_preserved(): - assert _materialize_unmanaged_dataset_source("/tmp/custom.h5") == ("/tmp/custom.h5") - - -def test_unmanaged_gcs_source_is_rejected(): - with pytest.raises(DatasetMaterializationError, match="Unsupported unmanaged"): - _materialize_unmanaged_dataset_source("gs://bucket/data.h5@release") diff --git a/tests/test_dataset_runtime.py b/tests/test_dataset_runtime.py index 9c420884..da9b59ee 100644 --- a/tests/test_dataset_runtime.py +++ b/tests/test_dataset_runtime.py @@ -29,18 +29,11 @@ def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDatas repo_type=plan.repo_type, revision=plan.revision, source_uri=plan.source_uri, - expected_sha256=plan.expected_sha256, - actual_sha256=plan.expected_sha256, + sha256=plan.sha256, path=Path(path), - cache_hit=False, ) -def _materialized_request(country_id: str, dataset: str, path: str): - materialized = _materialized(country_id, dataset, path) - return materialized.source_uri, str(materialized.path), materialized - - def test_us_create_datasets_passes_verified_bundle_source_to_country_package( monkeypatch, ): @@ -49,12 +42,10 @@ def test_us_create_datasets_passes_verified_bundle_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/us/datasets.py", ) materialize = Mock( - return_value=_materialized_request( - "us", "populace_us_2024", "/tmp/populace_us_2024.h5" - ) + return_value=_materialized("us", "populace_us_2024", "/tmp/populace_us_2024.h5") ) microsimulation = Mock() - monkeypatch.setattr(us_datasets, "_materialize_dataset_request", materialize) + monkeypatch.setattr(us_datasets, "materialize_bundle_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_us", @@ -66,7 +57,6 @@ def test_us_create_datasets_passes_verified_bundle_source_to_country_package( materialize.assert_called_once_with( "us", "populace_us_2024", - allow_unmanaged=False, data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/populace_us_2024.h5") @@ -80,12 +70,10 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) materialize = Mock( - return_value=_materialized_request( - "uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5" - ) + return_value=_materialized("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "_materialize_dataset_request", materialize) + monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -97,7 +85,6 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( materialize.assert_called_once_with( "uk", "populace_uk_2023", - allow_unmanaged=False, data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") @@ -109,14 +96,14 @@ def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) materialize = Mock( - return_value=_materialized_request( + return_value=_materialized( "uk", "enhanced_frs_2024_25", "/tmp/enhanced_frs_2024_25.h5", ) ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "_materialize_dataset_request", materialize) + monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -128,7 +115,6 @@ def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): materialize.assert_called_once_with( "uk", "enhanced_frs_2024_25", - allow_unmanaged=False, data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/enhanced_frs_2024_25.h5") diff --git a/tests/test_dataset_source.py b/tests/test_dataset_source.py new file mode 100644 index 00000000..ccc0300b --- /dev/null +++ b/tests/test_dataset_source.py @@ -0,0 +1,76 @@ +from pathlib import Path +from unittest.mock import patch + +import pytest + +import policyengine.tax_benefit_models.common.dataset_source as dataset_source +from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, +) +from policyengine.tax_benefit_models.common.dataset_source import download_hf_dataset + + +class _Response: + def __init__(self, payload: bytes = b"dataset", status_code: int = 200): + self.payload = payload + self.status_code = status_code + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + return False + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + def iter_content(self, chunk_size): + yield self.payload + + +class _Session: + def __init__(self, *responses: _Response): + self.responses = list(responses) + self.calls = [] + + def get(self, url, **kwargs): + self.calls.append((url, kwargs)) + return self.responses.pop(0) + + +def _download(uri, tmp_path, session): + with patch.object(dataset_source.requests, "get", side_effect=session.get): + return download_hf_dataset(uri, data_dir=tmp_path) + + +def test_explicit_hf_download_retries_dataset_repository_after_model_404(tmp_path): + session = _Session(_Response(status_code=404), _Response(b"dataset")) + + result = _download( + "hf://policyengine/example/data.h5@release", + tmp_path, + session, + ) + + assert Path(result).read_bytes() == b"dataset" + assert "/policyengine/example/" in session.calls[0][0] + assert "/datasets/policyengine/example/" in session.calls[1][0] + + +def test_explicit_hf_authentication_failure_does_not_retry(tmp_path): + session = _Session(_Response(status_code=403)) + + with pytest.raises(DatasetMaterializationError, match="credentials"): + _download( + "hf://policyengine/example/data.h5@release", + tmp_path, + session, + ) + + assert len(session.calls) == 1 + + +def test_explicit_hf_download_rejects_non_hf_uri(tmp_path): + with pytest.raises(DatasetMaterializationError, match="Expected an hf://"): + download_hf_dataset("gs://bucket/data.h5@release", data_dir=tmp_path) diff --git a/tests/test_release_manifests.py b/tests/test_release_manifests.py index 996c5ecc..46678720 100644 --- a/tests/test_release_manifests.py +++ b/tests/test_release_manifests.py @@ -33,7 +33,6 @@ https_release_manifest_uri, resolve_dataset_reference, resolve_default_datasets, - resolve_local_managed_dataset_source, resolve_managed_dataset_reference, resolve_region_dataset_path, ) @@ -101,18 +100,11 @@ def _materialized_dataset( repo_type=plan.repo_type, revision=plan.revision, source_uri=plan.source_uri, - expected_sha256=plan.expected_sha256, - actual_sha256=plan.expected_sha256, + sha256=plan.sha256, path=Path(path), - cache_hit=False, ) -def _materialized_dataset_request(country_id: str, dataset: str, path: str): - materialized = _materialized_dataset(country_id, dataset, path) - return materialized.source_uri, str(materialized.path), materialized - - UK_LEGACY_DATA_RELEASE_REVISION = "655dd07e4bb9c777b00dac044949611f1feb824f" UK_LEGACY_FRS_DATASET_URI = ( "hf://policyengine/policyengine-uk-data-private/frs_2023_24.h5" @@ -454,28 +446,6 @@ def test__given_versioned_dataset_url__then_logical_name_drops_version(self): assert dataset_logical_name(dataset) == "enhanced_cps_2024" - def test__given_explicit_local_data_repo__then_resolves_local_mirror( - self, monkeypatch, tmp_path - ): - local_dataset = ( - tmp_path - / "policyengine-us-data" - / "policyengine_us_data" - / "storage" - / "long_term" - / "2100.h5" - ) - local_dataset.parent.mkdir(parents=True) - local_dataset.write_text("", encoding="utf-8") - monkeypatch.setenv("POLICYENGINE_LOCAL_DATA_REPO_ROOT", str(tmp_path)) - - resolved = resolve_local_managed_dataset_source( - "us", - "hf://policyengine/policyengine-us-data/long_term/2100.h5@candidate", - ) - - assert resolved == str(local_dataset) - def test__given_country__then_can_fetch_data_release_manifest(self): get_data_release_manifest.cache_clear() payload = { @@ -954,8 +924,8 @@ def test__given_us_managed_microsimulation__then_passes_certified_dataset_and_bu ) with patch.object( us_model, - "_materialize_dataset_request", - return_value=_materialized_dataset_request( + "materialize_bundle_dataset", + return_value=_materialized_dataset( "us", "populace_us_2024", "/tmp/populace_us_2024.h5", @@ -1003,8 +973,8 @@ def test__given_us_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( us_model, - "_materialize_dataset_request", - return_value=(dataset, "/tmp/cps_2023.h5", None), + "download_hf_dataset", + return_value="/tmp/cps_2023.h5", ): microsim = us_model.managed_microsimulation( dataset=dataset, @@ -1072,8 +1042,8 @@ def test__given_uk_managed_dataset_name__then_resolves_within_bundle(self): ) with patch.object( uk_model, - "_materialize_dataset_request", - return_value=_materialized_dataset_request( + "materialize_bundle_dataset", + return_value=_materialized_dataset( "uk", "enhanced_frs_2024_25", "/tmp/enhanced_frs_2024_25.h5", @@ -1122,8 +1092,8 @@ def test__given_uk_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( uk_model, - "_materialize_dataset_request", - return_value=(dataset, "/tmp/frs_2022_23.h5", None), + "download_hf_dataset", + return_value="/tmp/frs_2022_23.h5", ): microsim = uk_model.managed_microsimulation( dataset=dataset, diff --git a/tests/test_us_long_term_datasets.py b/tests/test_us_long_term_datasets.py index 27d259c3..d688f2ce 100644 --- a/tests/test_us_long_term_datasets.py +++ b/tests/test_us_long_term_datasets.py @@ -160,10 +160,8 @@ def _materialized_long_term(path: Path, dataset_uri: str) -> MaterializedDataset repo_type="model", revision="abc123", source_uri=dataset_uri, - expected_sha256=actual_sha256, - actual_sha256=actual_sha256, + sha256=actual_sha256, path=path, - cache_hit=True, metadata_path=Path(f"{path}.metadata.json"), ) From 593c19eac4314e179aead505026a57f5a2801914 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:49:04 +0400 Subject: [PATCH 13/18] Rename bundle dataset tests --- tests/test_us_long_term_datasets.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_us_long_term_datasets.py b/tests/test_us_long_term_datasets.py index d688f2ce..60846221 100644 --- a/tests/test_us_long_term_datasets.py +++ b/tests/test_us_long_term_datasets.py @@ -273,7 +273,7 @@ def test__load_long_term_datasets__rejects_support_contract_mismatch(tmp_path): ) -def test__load_managed_long_term_datasets__loads_bundled_local_mirror( +def test__load_managed_long_term_datasets__loads_verified_bundle_file( monkeypatch, tmp_path, ): @@ -399,7 +399,7 @@ def test__load_managed_long_term_datasets__propagates_metadata_hash_failure( load_managed_long_term_datasets([2100]) -def test__load_managed_long_term_datasets__materializes_without_local_mirror( +def test__load_managed_long_term_datasets__requests_file_in_data_folder( monkeypatch, tmp_path, ): From 33d60b925fbc3795561bfdcdb7fc7ced0fe71c33 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:13:27 +0400 Subject: [PATCH 14/18] Unify dataset source selection --- changelog.d/502.changed.md | 9 +- docs/bundles.md | 12 +- docs/microsim.md | 6 +- src/policyengine/provenance/__init__.py | 5 +- .../provenance/dataset_materialization.py | 150 +++++++++++++++++- .../common/dataset_source.py | 98 ------------ .../tax_benefit_models/uk/datasets.py | 54 ++----- .../tax_benefit_models/uk/model.py | 59 ++----- .../tax_benefit_models/us/datasets.py | 62 ++------ .../tax_benefit_models/us/model.py | 49 ++---- tests/test_dataset_runtime.py | 27 +++- tests/test_dataset_source.py | 63 +++++++- tests/test_release_manifests.py | 36 ++++- tests/test_us_long_term_datasets.py | 20 ++- 14 files changed, 334 insertions(+), 316 deletions(-) delete mode 100644 src/policyengine/tax_benefit_models/common/dataset_source.py diff --git a/changelog.d/502.changed.md b/changelog.d/502.changed.md index 48f94c18..34d8b4ad 100644 --- a/changelog.d/502.changed.md +++ b/changelog.d/502.changed.md @@ -1,4 +1,5 @@ -PolicyEngine.py now reuses or downloads managed datasets from the exact Hugging -Face repository type, immutable revision, and SHA-256 recorded in its release -bundle. Bundle installation and US and UK calculation entry points share this -implementation, while explicit local paths and Hugging Face URIs remain opt-in. +PolicyEngine.py now exposes one `materialize_dataset` entry point for +bundle-managed datasets, explicit Hugging Face URIs, and explicit local paths. +Managed datasets are reused or downloaded from the exact repository type, +immutable revision, and SHA-256 recorded in the release bundle; the other two +inputs remain opt-in. diff --git a/docs/bundles.md b/docs/bundles.md index cb0c0cf1..7457f757 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -49,15 +49,17 @@ To materialize a default or named artifact without installing the complete package scaffold: ```python -from policyengine.provenance import materialize_bundle_dataset +from policyengine.provenance import materialize_dataset -result = materialize_bundle_dataset("us", "populace_us_2024") +result = materialize_dataset("us", "populace_us_2024") print(result.path) -print(result.sha256) +print(result.bundle_dataset.sha256) ``` -`materialize_bundle_dataset` returns the selected source package, repository -type, revision, verified SHA-256, local path, and optional metadata path. +`materialize_dataset` returns the selected source URI and local path. For a +bundle-managed input, `bundle_dataset` also contains the selected source +package, repository type, revision, verified SHA-256, and optional metadata +path. `policyengine-*-data` and `populace-data` artifacts use the repository type recorded in the bundle. Callers do not infer repository type from the repository name. diff --git a/docs/microsim.md b/docs/microsim.md index 0a4a520c..f3c2e1da 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -138,15 +138,15 @@ To materialize the raw certified artifact without creating uprated yearly datasets, use PolicyEngine.py's bundle API: ```python -from policyengine.provenance import materialize_bundle_dataset +from policyengine.provenance import materialize_dataset -result = materialize_bundle_dataset( +result = materialize_dataset( "uk", "enhanced_frs_2024_25", ) print(result.path) -print(result.sha256) +print(result.bundle_dataset.sha256) ``` The bundle API uses the repository type recorded in the bundle, so callers do diff --git a/src/policyengine/provenance/__init__.py b/src/policyengine/provenance/__init__.py index 001cb18d..fe1f8046 100644 --- a/src/policyengine/provenance/__init__.py +++ b/src/policyengine/provenance/__init__.py @@ -27,11 +27,14 @@ from .dataset_materialization import ( DatasetMaterializationError as DatasetMaterializationError, ) +from .dataset_materialization import ( + DatasetSource as DatasetSource, +) from .dataset_materialization import ( MaterializedDataset as MaterializedDataset, ) from .dataset_materialization import ( - materialize_bundle_dataset as materialize_bundle_dataset, + materialize_dataset as materialize_dataset, ) from .manifest import ( CertifiedDataArtifact as CertifiedDataArtifact, diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index a845f8b5..eb0a2830 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -16,9 +16,11 @@ CountryReleaseManifest, _artifact_revision, build_hf_uri, + dataset_logical_name, get_release_manifest, https_dataset_uri, hugging_face_auth_headers, + resolve_managed_dataset_reference, ) DEFAULT_DATA_DIR = Path("./data") @@ -26,7 +28,7 @@ class DatasetMaterializationError(ValueError): - """Raised when a bundle dataset cannot be made available safely.""" + """Raised when a requested dataset cannot be made available safely.""" @dataclass(frozen=True) @@ -42,6 +44,19 @@ class MaterializedDataset: metadata_path: Optional[Path] = None +@dataclass(frozen=True) +class DatasetSource: + """Dataset source selected for a country-package calculation.""" + + source_uri: str + path: str + bundle_dataset: Optional[MaterializedDataset] = None + + @property + def name(self) -> str: + return dataset_logical_name(self.source_uri) + + @dataclass(frozen=True) class _BundleDatasetSpec: """Manifest values required to inspect or download one bundle dataset.""" @@ -103,19 +118,142 @@ def _resolve_bundle_dataset( ) -def materialize_bundle_dataset( +def materialize_dataset( country_id: str, dataset: Optional[str] = None, *, + allow_unmanaged: bool = False, data_dir: Path = DEFAULT_DATA_DIR, -) -> MaterializedDataset: - """Return a verified local copy of a dataset from the installed bundle.""" +) -> DatasetSource: + """Select a dataset source and return the local file used for calculation.""" + + manifest = get_release_manifest(country_id) + if dataset is None or dataset == manifest.default_dataset_uri: + return _use_bundle_dataset( + country_id, + manifest.default_dataset, + data_dir=data_dir, + manifest=manifest, + ) + if dataset in manifest.datasets: + return _use_bundle_dataset( + country_id, + dataset, + data_dir=data_dir, + manifest=manifest, + ) + + source_uri = resolve_managed_dataset_reference( + country_id, + dataset, + allow_unmanaged=allow_unmanaged, + ) + if source_uri.startswith("hf://"): + return _download_hugging_face_dataset(source_uri, data_dir=data_dir) + if "://" in source_uri: + raise DatasetMaterializationError( + f"Unsupported explicit dataset URI: {source_uri!r}." + ) + return _use_local_dataset(source_uri) + + +def _use_bundle_dataset( + country_id: str, + dataset: str, + *, + data_dir: Path, + manifest: CountryReleaseManifest, +) -> DatasetSource: + bundle_dataset = _reuse_or_download_bundle_files( + _resolve_bundle_dataset( + country_id, + dataset, + data_dir=data_dir, + manifest=manifest, + ) + ) + return DatasetSource( + source_uri=bundle_dataset.source_uri, + path=str(bundle_dataset.path), + bundle_dataset=bundle_dataset, + ) + + +def _download_hugging_face_dataset( + dataset_uri: str, + *, + data_dir: Path, +) -> DatasetSource: + path_with_repo, revision = ( + dataset_uri[5:].rsplit("@", maxsplit=1) + if "@" in dataset_uri[5:] + else (dataset_uri[5:], "main") + ) + parts = path_with_repo.split("/", maxsplit=2) + if len(parts) != 3 or not all(parts): + raise DatasetMaterializationError( + "Invalid Hugging Face dataset URI. Expected format " + f"'hf://owner/repo/path/to/file[@revision]', got {dataset_uri!r}." + ) + + repo_id = f"{parts[0]}/{parts[1]}" + repository_path = parts[2] + destination = data_dir / Path(repository_path).name + destination.parent.mkdir(parents=True, exist_ok=True) + + for repo_type in ("model", "dataset"): + file_descriptor, temp_name = tempfile.mkstemp( + prefix=".policyengine-download-", + suffix=destination.suffix or ".download", + dir=destination.parent, + ) + os.close(file_descriptor) + temporary_path = Path(temp_name) + url = https_dataset_uri( + repo_id, + repository_path, + revision, + repo_type=repo_type, + ) + try: + with requests.get( + url, + headers=hugging_face_auth_headers(), + stream=True, + timeout=DOWNLOAD_TIMEOUT_SECONDS, + ) as response: + if response.status_code in {401, 403}: + raise DatasetMaterializationError( + "Could not download explicit dataset " + f"{dataset_uri!r}: Hugging Face rejected the configured " + "credentials. Set HUGGING_FACE_TOKEN to a token with " + "access to the repository." + ) + if response.status_code == 404: + if repo_type == "model": + continue + raise DatasetMaterializationError( + f"Could not find explicit dataset {dataset_uri!r}." + ) + response.raise_for_status() + with temporary_path.open("wb") as output: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + output.write(chunk) + os.replace(temporary_path, destination) + return DatasetSource(source_uri=dataset_uri, path=str(destination)) + finally: + temporary_path.unlink(missing_ok=True) - return _reuse_or_download_bundle_files( - _resolve_bundle_dataset(country_id, dataset, data_dir=data_dir) + raise DatasetMaterializationError( + f"Could not download explicit dataset {dataset_uri!r}." ) +def _use_local_dataset(dataset_path: str) -> DatasetSource: + return DatasetSource(source_uri=dataset_path, path=dataset_path) + + def _reuse_or_download_bundle_files( dataset: _BundleDatasetSpec, ) -> MaterializedDataset: diff --git a/src/policyengine/tax_benefit_models/common/dataset_source.py b/src/policyengine/tax_benefit_models/common/dataset_source.py deleted file mode 100644 index b1511a0d..00000000 --- a/src/policyengine/tax_benefit_models/common/dataset_source.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Downloads for explicit Hugging Face dataset inputs.""" - -from __future__ import annotations - -import os -import tempfile -from pathlib import Path - -import requests - -from policyengine.provenance.dataset_materialization import ( - DatasetMaterializationError, -) -from policyengine.provenance.manifest import ( - https_dataset_uri, - hugging_face_auth_headers, -) - -DEFAULT_DATA_DIR = Path("./data") -DOWNLOAD_TIMEOUT_SECONDS = 60 - - -def download_hf_dataset( - dataset_uri: str, - *, - data_dir: Path = DEFAULT_DATA_DIR, -) -> str: - """Download an explicit Hugging Face dataset URI and return its local path.""" - - if not dataset_uri.startswith("hf://"): - raise DatasetMaterializationError( - f"Expected an hf:// dataset URI, got {dataset_uri!r}." - ) - - path_with_repo, revision = ( - dataset_uri[5:].rsplit("@", maxsplit=1) - if "@" in dataset_uri[5:] - else (dataset_uri[5:], "main") - ) - parts = path_with_repo.split("/", maxsplit=2) - if len(parts) != 3 or not all(parts): - raise DatasetMaterializationError( - "Invalid Hugging Face dataset URI. Expected format " - f"'hf://owner/repo/path/to/file[@revision]', got {dataset_uri!r}." - ) - - repo_id = f"{parts[0]}/{parts[1]}" - repository_path = parts[2] - destination = data_dir / Path(repository_path).name - destination.parent.mkdir(parents=True, exist_ok=True) - - for repo_type in ("model", "dataset"): - file_descriptor, temp_name = tempfile.mkstemp( - prefix=".policyengine-download-", - suffix=destination.suffix or ".download", - dir=destination.parent, - ) - os.close(file_descriptor) - temporary_path = Path(temp_name) - url = https_dataset_uri( - repo_id, - repository_path, - revision, - repo_type=repo_type, - ) - try: - with requests.get( - url, - headers=hugging_face_auth_headers(), - stream=True, - timeout=DOWNLOAD_TIMEOUT_SECONDS, - ) as response: - if response.status_code in {401, 403}: - raise DatasetMaterializationError( - "Could not download explicit dataset " - f"{dataset_uri!r}: Hugging Face rejected the configured " - "credentials. Set HUGGING_FACE_TOKEN to a token with " - "access to the repository." - ) - if response.status_code == 404: - if repo_type == "model": - continue - raise DatasetMaterializationError( - f"Could not find explicit dataset {dataset_uri!r}." - ) - response.raise_for_status() - with temporary_path.open("wb") as output: - for chunk in response.iter_content(chunk_size=1024 * 1024): - if chunk: - output.write(chunk) - os.replace(temporary_path, destination) - return str(destination) - finally: - temporary_path.unlink(missing_ok=True) - - raise DatasetMaterializationError( - f"Could not download explicit dataset {dataset_uri!r}." - ) diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index bcc0d0df..105dfdd0 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -7,17 +7,12 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( - DatasetMaterializationError, - materialize_bundle_dataset, + materialize_dataset, ) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, - resolve_managed_dataset_reference, -) -from policyengine.tax_benefit_models.common.dataset_source import ( - download_hf_dataset, ) @@ -126,46 +121,19 @@ def create_datasets( data_folder: str = "./data", allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUKDataset]: - if datasets is None: - datasets = [get_release_manifest("uk").default_dataset] + dataset_requests: list[Optional[str]] = [None] if datasets is None else datasets result = {} - for dataset in datasets: - manifest = get_release_manifest("uk") - if dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - else: - managed_dataset = None - if managed_dataset is not None: - materialized = materialize_bundle_dataset( - "uk", - managed_dataset, - data_dir=Path(data_folder), - ) - resolved_dataset = materialized.source_uri - runtime_dataset = str(materialized.path) - else: - resolved_dataset = resolve_managed_dataset_reference( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - ) - if resolved_dataset.startswith("hf://"): - runtime_dataset = download_hf_dataset( - resolved_dataset, - data_dir=Path(data_folder), - ) - elif "://" in resolved_dataset: - raise DatasetMaterializationError( - f"Unsupported explicit dataset URI: {resolved_dataset!r}." - ) - else: - runtime_dataset = resolved_dataset - dataset_stem = dataset_logical_name(resolved_dataset) + for dataset in dataset_requests: + source = materialize_dataset( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + data_dir=Path(data_folder), + ) + dataset_stem = source.name from policyengine_uk import Microsimulation - sim = Microsimulation(dataset=runtime_dataset) + sim = Microsimulation(dataset=source.path) for year in years: year_dataset = sim.dataset[year] diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index a513c849..b0e22ad1 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -6,17 +6,9 @@ from policyengine.core import TaxBenefitModel from policyengine.provenance.dataset_materialization import ( - DatasetMaterializationError, - materialize_bundle_dataset, -) -from policyengine.provenance.manifest import ( - get_release_manifest, - resolve_managed_dataset_reference, + materialize_dataset, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion -from policyengine.tax_benefit_models.common.dataset_source import ( - download_hf_dataset, -) from policyengine.tax_benefit_models.common.model_version import ( build_runtime_dataset_provenance, ) @@ -293,50 +285,29 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - manifest = get_release_manifest("uk") - if dataset is None or dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - else: - managed_dataset = None - materialized = None - if managed_dataset is not None: - materialized = materialize_bundle_dataset("uk", managed_dataset) - dataset_uri = materialized.source_uri - runtime_dataset_source = str(materialized.path) - else: - dataset_uri = resolve_managed_dataset_reference( - "uk", - dataset, - allow_unmanaged=allow_unmanaged, - ) - if dataset_uri.startswith("hf://"): - runtime_dataset_source = download_hf_dataset(dataset_uri) - elif "://" in dataset_uri: - raise DatasetMaterializationError( - f"Unsupported explicit dataset URI: {dataset_uri!r}." - ) - else: - runtime_dataset_source = dataset_uri - runtime_dataset = runtime_dataset_source - if isinstance(runtime_dataset_source, str) and "://" not in runtime_dataset_source: + source = materialize_dataset( + "uk", + dataset, + allow_unmanaged=allow_unmanaged, + ) + runtime_dataset = source.path + if "://" not in source.path: from policyengine_uk.data.dataset_schema import ( UKMultiYearDataset, UKSingleYearDataset, ) - if UKMultiYearDataset.validate_file_path(runtime_dataset_source, False): - runtime_dataset = UKMultiYearDataset(runtime_dataset_source) - elif UKSingleYearDataset.validate_file_path(runtime_dataset_source, False): - runtime_dataset = UKSingleYearDataset(runtime_dataset_source) + if UKMultiYearDataset.validate_file_path(source.path, False): + runtime_dataset = UKMultiYearDataset(source.path) + elif UKSingleYearDataset.validate_file_path(source.path, False): + runtime_dataset = UKSingleYearDataset(source.path) microsim = Microsimulation(dataset=runtime_dataset, **kwargs) microsim.policyengine_bundle = dict(uk_latest.release_bundle) microsim.policyengine_bundle.update( build_runtime_dataset_provenance( - dataset_uri, - runtime_dataset_source, - materialized, + source.source_uri, + source.path, + source.bundle_dataset, ) ) return microsim diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index abdb3aec..4c65c9b0 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -13,18 +13,13 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( - DatasetMaterializationError, MaterializedDataset, - materialize_bundle_dataset, + materialize_dataset, ) from policyengine.provenance.manifest import ( dataset_logical_name, get_release_manifest, resolve_dataset_reference, - resolve_managed_dataset_reference, -) -from policyengine.tax_benefit_models.common.dataset_source import ( - download_hf_dataset, ) from policyengine.tax_benefit_models.common.model_version import ( build_runtime_dataset_provenance, @@ -303,43 +298,17 @@ def create_datasets( """ from policyengine_us import Microsimulation - datasets = datasets or [get_release_manifest("us").default_dataset] + dataset_requests: list[Optional[str]] = datasets or [None] result = {} - for dataset in datasets: - manifest = get_release_manifest("us") - if dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - else: - managed_dataset = None - if managed_dataset is not None: - materialized = materialize_bundle_dataset( - "us", - managed_dataset, - data_dir=Path(data_folder), - ) - resolved_dataset = materialized.source_uri - runtime_dataset = str(materialized.path) - else: - resolved_dataset = resolve_managed_dataset_reference( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) - if resolved_dataset.startswith("hf://"): - runtime_dataset = download_hf_dataset( - resolved_dataset, - data_dir=Path(data_folder), - ) - elif "://" in resolved_dataset: - raise DatasetMaterializationError( - f"Unsupported explicit dataset URI: {resolved_dataset!r}." - ) - else: - runtime_dataset = resolved_dataset - dataset_stem = dataset_logical_name(resolved_dataset) - sim = Microsimulation(dataset=runtime_dataset) + for dataset in dataset_requests: + source = materialize_dataset( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + data_dir=Path(data_folder), + ) + dataset_stem = source.name + sim = Microsimulation(dataset=source.path) for year in years: # Get all input variables from the simulation @@ -1107,13 +1076,16 @@ def load_managed_long_term_datasets( f"Managed long-term dataset {key!r} is missing a sha256 in " "the bundled US release manifest." ) - materialized = materialize_bundle_dataset( + source = materialize_dataset( "us", key, data_dir=Path(data_folder), ) - dataset_uri = materialized.source_uri - path = materialized.path + materialized = source.bundle_dataset + if materialized is None: + raise ValueError(f"Bundle dataset {key!r} was not bundle-managed.") + dataset_uri = source.source_uri + path = Path(source.path) metadata, metadata_path = _load_dataset_metadata( path, require_metadata, diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 43331b72..c517855f 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -6,17 +6,9 @@ from policyengine.core import TaxBenefitModel from policyengine.provenance.dataset_materialization import ( - DatasetMaterializationError, - materialize_bundle_dataset, -) -from policyengine.provenance.manifest import ( - get_release_manifest, - resolve_managed_dataset_reference, + materialize_dataset, ) from policyengine.tax_benefit_models.common import MicrosimulationModelVersion -from policyengine.tax_benefit_models.common.dataset_source import ( - download_hf_dataset, -) from policyengine.tax_benefit_models.common.model_version import ( build_runtime_dataset_provenance, ) @@ -436,39 +428,18 @@ def managed_microsimulation( "**kwargs, so policyengine.py can enforce the release bundle." ) - manifest = get_release_manifest("us") - if dataset is None or dataset == manifest.default_dataset_uri: - managed_dataset = manifest.default_dataset - elif dataset in manifest.datasets: - managed_dataset = dataset - else: - managed_dataset = None - materialized = None - if managed_dataset is not None: - materialized = materialize_bundle_dataset("us", managed_dataset) - dataset_uri = materialized.source_uri - runtime_dataset_source = str(materialized.path) - else: - dataset_uri = resolve_managed_dataset_reference( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) - if dataset_uri.startswith("hf://"): - runtime_dataset_source = download_hf_dataset(dataset_uri) - elif "://" in dataset_uri: - raise DatasetMaterializationError( - f"Unsupported explicit dataset URI: {dataset_uri!r}." - ) - else: - runtime_dataset_source = dataset_uri - microsim = Microsimulation(dataset=runtime_dataset_source, **kwargs) + source = materialize_dataset( + "us", + dataset, + allow_unmanaged=allow_unmanaged, + ) + microsim = Microsimulation(dataset=source.path, **kwargs) microsim.policyengine_bundle = dict(us_latest.release_bundle) microsim.policyengine_bundle.update( build_runtime_dataset_provenance( - dataset_uri, - runtime_dataset_source, - materialized, + source.source_uri, + source.path, + source.bundle_dataset, ) ) return microsim diff --git a/tests/test_dataset_runtime.py b/tests/test_dataset_runtime.py index da9b59ee..9d97db0e 100644 --- a/tests/test_dataset_runtime.py +++ b/tests/test_dataset_runtime.py @@ -5,6 +5,7 @@ from unittest.mock import Mock from policyengine.provenance.dataset_materialization import ( + DatasetSource, MaterializedDataset, _resolve_bundle_dataset, ) @@ -34,6 +35,15 @@ def _materialized(country_id: str, dataset: str, path: str) -> MaterializedDatas ) +def _source(country_id: str, dataset: str, path: str) -> DatasetSource: + bundle_dataset = _materialized(country_id, dataset, path) + return DatasetSource( + source_uri=bundle_dataset.source_uri, + path=str(bundle_dataset.path), + bundle_dataset=bundle_dataset, + ) + + def test_us_create_datasets_passes_verified_bundle_source_to_country_package( monkeypatch, ): @@ -42,10 +52,10 @@ def test_us_create_datasets_passes_verified_bundle_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/us/datasets.py", ) materialize = Mock( - return_value=_materialized("us", "populace_us_2024", "/tmp/populace_us_2024.h5") + return_value=_source("us", "populace_us_2024", "/tmp/populace_us_2024.h5") ) microsimulation = Mock() - monkeypatch.setattr(us_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setattr(us_datasets, "materialize_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_us", @@ -57,6 +67,7 @@ def test_us_create_datasets_passes_verified_bundle_source_to_country_package( materialize.assert_called_once_with( "us", "populace_us_2024", + allow_unmanaged=False, data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/populace_us_2024.h5") @@ -70,10 +81,10 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) materialize = Mock( - return_value=_materialized("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") + return_value=_source("uk", "populace_uk_2023", "/tmp/populace_uk_2023.h5") ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setattr(uk_datasets, "materialize_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -85,6 +96,7 @@ def test_uk_create_datasets_passes_verified_bundle_source_to_country_package( materialize.assert_called_once_with( "uk", "populace_uk_2023", + allow_unmanaged=False, data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/populace_uk_2023.h5") @@ -96,14 +108,14 @@ def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) materialize = Mock( - return_value=_materialized( + return_value=_source( "uk", "enhanced_frs_2024_25", "/tmp/enhanced_frs_2024_25.h5", ) ) microsimulation = Mock() - monkeypatch.setattr(uk_datasets, "materialize_bundle_dataset", materialize) + monkeypatch.setattr(uk_datasets, "materialize_dataset", materialize) monkeypatch.setitem( sys.modules, "policyengine_uk", @@ -114,7 +126,8 @@ def test_uk_create_datasets_defaults_to_certified_bundle_dataset(monkeypatch): materialize.assert_called_once_with( "uk", - "enhanced_frs_2024_25", + None, + allow_unmanaged=False, data_dir=Path("./data"), ) microsimulation.assert_called_once_with(dataset="/tmp/enhanced_frs_2024_25.h5") diff --git a/tests/test_dataset_source.py b/tests/test_dataset_source.py index ccc0300b..2b325647 100644 --- a/tests/test_dataset_source.py +++ b/tests/test_dataset_source.py @@ -3,11 +3,12 @@ import pytest -import policyengine.tax_benefit_models.common.dataset_source as dataset_source +import policyengine.provenance.dataset_materialization as dataset_source from policyengine.provenance.dataset_materialization import ( DatasetMaterializationError, + DatasetSource, + materialize_dataset, ) -from policyengine.tax_benefit_models.common.dataset_source import download_hf_dataset class _Response: @@ -41,7 +42,12 @@ def get(self, url, **kwargs): def _download(uri, tmp_path, session): with patch.object(dataset_source.requests, "get", side_effect=session.get): - return download_hf_dataset(uri, data_dir=tmp_path) + return materialize_dataset( + "us", + uri, + allow_unmanaged=True, + data_dir=tmp_path, + ) def test_explicit_hf_download_retries_dataset_repository_after_model_404(tmp_path): @@ -53,7 +59,9 @@ def test_explicit_hf_download_retries_dataset_repository_after_model_404(tmp_pat session, ) - assert Path(result).read_bytes() == b"dataset" + assert Path(result.path).read_bytes() == b"dataset" + assert result.source_uri == "hf://policyengine/example/data.h5@release" + assert result.bundle_dataset is None assert "/policyengine/example/" in session.calls[0][0] assert "/datasets/policyengine/example/" in session.calls[1][0] @@ -72,5 +80,48 @@ def test_explicit_hf_authentication_failure_does_not_retry(tmp_path): def test_explicit_hf_download_rejects_non_hf_uri(tmp_path): - with pytest.raises(DatasetMaterializationError, match="Expected an hf://"): - download_hf_dataset("gs://bucket/data.h5@release", data_dir=tmp_path) + with pytest.raises(DatasetMaterializationError, match="Unsupported explicit"): + materialize_dataset( + "us", + "gs://bucket/data.h5@release", + allow_unmanaged=True, + data_dir=tmp_path, + ) + + +def test_bundle_dataset_uses_bundle_strategy(tmp_path): + expected = DatasetSource( + source_uri="hf://policyengine/populace-us/populace_us_2024.h5@release", + path=str(tmp_path / "populace_us_2024.h5"), + ) + + with patch.object( + dataset_source, + "_use_bundle_dataset", + return_value=expected, + ) as use_bundle: + result = materialize_dataset( + "us", + "populace_us_2024", + data_dir=tmp_path, + ) + + assert result is expected + use_bundle.assert_called_once() + + +def test_explicit_local_path_uses_local_strategy(tmp_path): + local_path = tmp_path / "custom.h5" + local_path.touch() + + result = materialize_dataset( + "us", + str(local_path), + allow_unmanaged=True, + data_dir=tmp_path, + ) + + assert result == DatasetSource( + source_uri=str(local_path), + path=str(local_path), + ) diff --git a/tests/test_release_manifests.py b/tests/test_release_manifests.py index 46678720..24bb05be 100644 --- a/tests/test_release_manifests.py +++ b/tests/test_release_manifests.py @@ -17,6 +17,7 @@ from policyengine.core.tax_benefit_model import TaxBenefitModel from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion from policyengine.provenance.dataset_materialization import ( + DatasetSource, MaterializedDataset, _resolve_bundle_dataset, ) @@ -105,6 +106,19 @@ def _materialized_dataset( ) +def _dataset_source( + country_id: str, + dataset: str, + path: str, +) -> DatasetSource: + bundle_dataset = _materialized_dataset(country_id, dataset, path) + return DatasetSource( + source_uri=bundle_dataset.source_uri, + path=str(bundle_dataset.path), + bundle_dataset=bundle_dataset, + ) + + UK_LEGACY_DATA_RELEASE_REVISION = "655dd07e4bb9c777b00dac044949611f1feb824f" UK_LEGACY_FRS_DATASET_URI = ( "hf://policyengine/policyengine-uk-data-private/frs_2023_24.h5" @@ -924,8 +938,8 @@ def test__given_us_managed_microsimulation__then_passes_certified_dataset_and_bu ) with patch.object( us_model, - "materialize_bundle_dataset", - return_value=_materialized_dataset( + "materialize_dataset", + return_value=_dataset_source( "us", "populace_us_2024", "/tmp/populace_us_2024.h5", @@ -973,8 +987,11 @@ def test__given_us_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( us_model, - "download_hf_dataset", - return_value="/tmp/cps_2023.h5", + "materialize_dataset", + return_value=DatasetSource( + source_uri=dataset, + path="/tmp/cps_2023.h5", + ), ): microsim = us_model.managed_microsimulation( dataset=dataset, @@ -1042,8 +1059,8 @@ def test__given_uk_managed_dataset_name__then_resolves_within_bundle(self): ) with patch.object( uk_model, - "materialize_bundle_dataset", - return_value=_materialized_dataset( + "materialize_dataset", + return_value=_dataset_source( "uk", "enhanced_frs_2024_25", "/tmp/enhanced_frs_2024_25.h5", @@ -1092,8 +1109,11 @@ def test__given_uk_unmanaged_dataset_uri__then_source_is_not_rewritten(self): ) with patch.object( uk_model, - "download_hf_dataset", - return_value="/tmp/frs_2022_23.h5", + "materialize_dataset", + return_value=DatasetSource( + source_uri=dataset, + path="/tmp/frs_2022_23.h5", + ), ): microsim = uk_model.managed_microsimulation( dataset=dataset, diff --git a/tests/test_us_long_term_datasets.py b/tests/test_us_long_term_datasets.py index 60846221..db523784 100644 --- a/tests/test_us_long_term_datasets.py +++ b/tests/test_us_long_term_datasets.py @@ -12,6 +12,7 @@ import policyengine.tax_benefit_models.us.datasets as us_datasets_module from policyengine.provenance.dataset_materialization import ( DatasetMaterializationError, + DatasetSource, MaterializedDataset, ) from policyengine.tax_benefit_models.us.datasets import ( @@ -153,9 +154,9 @@ def _manifest_with_long_term_sha( ) -def _materialized_long_term(path: Path, dataset_uri: str) -> MaterializedDataset: +def _materialized_long_term(path: Path, dataset_uri: str) -> DatasetSource: actual_sha256 = _sha256(path) - return MaterializedDataset( + bundle_dataset = MaterializedDataset( data_package_name="policyengine-us-data", repo_type="model", revision="abc123", @@ -164,6 +165,11 @@ def _materialized_long_term(path: Path, dataset_uri: str) -> MaterializedDataset path=path, metadata_path=Path(f"{path}.metadata.json"), ) + return DatasetSource( + source_uri=dataset_uri, + path=str(path), + bundle_dataset=bundle_dataset, + ) def test__load_long_term_datasets__loads_h5_and_sidecar_metadata(tmp_path): @@ -304,7 +310,7 @@ def test__load_managed_long_term_datasets__loads_verified_bundle_file( ) monkeypatch.setattr( us_datasets_module, - "materialize_bundle_dataset", + "materialize_dataset", lambda *args, **kwargs: _materialized_long_term(h5_path, dataset_uri), ) @@ -342,7 +348,7 @@ def test__load_managed_long_term_datasets__defaults_to_manifest_model_version( ) monkeypatch.setattr( us_datasets_module, - "materialize_bundle_dataset", + "materialize_dataset", lambda *args, **kwargs: _materialized_long_term(h5_path, dataset_uri), ) @@ -365,7 +371,7 @@ def test__load_managed_long_term_datasets__checks_manifest_sha256( ) monkeypatch.setattr( us_datasets_module, - "materialize_bundle_dataset", + "materialize_dataset", Mock(side_effect=DatasetMaterializationError("sha256 mismatch")), ) @@ -391,7 +397,7 @@ def test__load_managed_long_term_datasets__propagates_metadata_hash_failure( ) monkeypatch.setattr( us_datasets_module, - "materialize_bundle_dataset", + "materialize_dataset", Mock(side_effect=DatasetMaterializationError("metadata sha256 mismatch")), ) @@ -416,7 +422,7 @@ def test__load_managed_long_term_datasets__requests_file_in_data_folder( materialize = Mock(return_value=_materialized_long_term(h5_path, dataset_uri)) monkeypatch.setattr( us_datasets_module, - "materialize_bundle_dataset", + "materialize_dataset", materialize, ) From 7ddacaf63f3cf7507b6f2decc384d21649ad366e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:22:30 +0400 Subject: [PATCH 15/18] Remove redundant bundle source check --- .../tax_benefit_models/us/datasets.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 4c65c9b0..7032ea61 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -521,10 +521,8 @@ def _metadata_path_for_h5(path: Path) -> Path: def _load_dataset_metadata( path: Path, require_metadata: bool, - *, - metadata_path: Optional[Path] = None, ) -> tuple[dict, Optional[Path]]: - metadata_path = metadata_path or _metadata_path_for_h5(path) + metadata_path = _metadata_path_for_h5(path) if not metadata_path.exists(): if require_metadata: raise FileNotFoundError( @@ -1081,16 +1079,9 @@ def load_managed_long_term_datasets( key, data_dir=Path(data_folder), ) - materialized = source.bundle_dataset - if materialized is None: - raise ValueError(f"Bundle dataset {key!r} was not bundle-managed.") dataset_uri = source.source_uri path = Path(source.path) - metadata, metadata_path = _load_dataset_metadata( - path, - require_metadata, - metadata_path=materialized.metadata_path, - ) + metadata, metadata_path = _load_dataset_metadata(path, require_metadata) _validate_loaded_long_term_metadata( metadata=metadata, metadata_path=metadata_path, @@ -1133,7 +1124,7 @@ def load_managed_long_term_datasets( metadata=metadata, metadata_path=metadata_path, dataset_uri=dataset_uri, - materialized=materialized, + materialized=source.bundle_dataset, ) return result From eba89f19fda9dec7740b41524bd93bdb5af89cad Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:50:09 +0400 Subject: [PATCH 16/18] Correct UK Populace repository type --- src/policyengine/data/bundle/manifest.json | 2 +- .../data/bundle/uk.trace.tro.jsonld | 4 +-- .../data/bundle/us.trace.tro.jsonld | 4 +-- tests/test_dataset_materialization.py | 26 ++++++++++++++++--- 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/src/policyengine/data/bundle/manifest.json b/src/policyengine/data/bundle/manifest.json index 98c4d1f0..8ace212b 100644 --- a/src/policyengine/data/bundle/manifest.json +++ b/src/policyengine/data/bundle/manifest.json @@ -220,7 +220,7 @@ "data_package_name": "populace-data", "path": "populace_uk_2023.h5", "repo_id": "policyengine/populace-uk-private", - "repo_type": "model", + "repo_type": "dataset", "revision": "populace-uk-2023-dd68c73-4aa4b14-20260619T023711Z", "sha256": "f17306ccb2aad7ff0130be3589b560afb2e2a12a943570911cd0c77f07934833" } diff --git a/src/policyengine/data/bundle/uk.trace.tro.jsonld b/src/policyengine/data/bundle/uk.trace.tro.jsonld index da05bf47..fffb3692 100644 --- a/src/policyengine/data/bundle/uk.trace.tro.jsonld +++ b/src/policyengine/data/bundle/uk.trace.tro.jsonld @@ -75,7 +75,7 @@ "@type": "trov:ResearchArtifact", "schema:name": "policyengine.py bundle manifest for uk", "trov:mimeType": "application/json", - "trov:sha256": "b2f6dc37f9597ae0932c3b3b926f17cfd1cb622727bebe7a998f94c84b30adde" + "trov:sha256": "d11510105984e5f681b35efb2fa5c58cc6a0cf730ab97e3e1a240b2b49eea862" }, { "@id": "composition/1/artifact/data_release_manifest", @@ -102,7 +102,7 @@ "trov:hasFingerprint": { "@id": "composition/1/fingerprint", "@type": "trov:CompositionFingerprint", - "trov:sha256": "b3eaab80f13125c431d83da8fb4015814145513481fa6703f96502d9da787fa0" + "trov:sha256": "dfdc28587fe4b5181ab6dc803659c6f2222a4c27a3f3e9cc04a3f00ad3b6d5d5" } }, "trov:hasPerformance": { diff --git a/src/policyengine/data/bundle/us.trace.tro.jsonld b/src/policyengine/data/bundle/us.trace.tro.jsonld index 6e2ee6f8..15aa4ccc 100644 --- a/src/policyengine/data/bundle/us.trace.tro.jsonld +++ b/src/policyengine/data/bundle/us.trace.tro.jsonld @@ -75,7 +75,7 @@ "@type": "trov:ResearchArtifact", "schema:name": "policyengine.py bundle manifest for us", "trov:mimeType": "application/json", - "trov:sha256": "b2f6dc37f9597ae0932c3b3b926f17cfd1cb622727bebe7a998f94c84b30adde" + "trov:sha256": "d11510105984e5f681b35efb2fa5c58cc6a0cf730ab97e3e1a240b2b49eea862" }, { "@id": "composition/1/artifact/data_release_manifest", @@ -102,7 +102,7 @@ "trov:hasFingerprint": { "@id": "composition/1/fingerprint", "@type": "trov:CompositionFingerprint", - "trov:sha256": "747a5d2b33a4daa2104af3fa832c89bc2754ac88c06c03c0a444030d6fb93c37" + "trov:sha256": "5b724a715ccfdfc1cb76f0ac110414956099037f8d8b628c90f86913d4762842" } }, "trov:hasPerformance": { diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 52fbf4e9..256e004d 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -9,7 +9,10 @@ _resolve_bundle_dataset, _reuse_or_download_bundle_files, ) -from policyengine.provenance.manifest import CountryReleaseManifest +from policyengine.provenance.manifest import ( + CountryReleaseManifest, + https_dataset_uri, +) def _manifest() -> CountryReleaseManifest: @@ -41,7 +44,7 @@ def _manifest() -> CountryReleaseManifest: "data_package_name": "populace-data", "path": "populace_uk_2023.h5", "repo_id": "policyengine/populace-uk-private", - "repo_type": "model", + "repo_type": "dataset", "revision": "populace-release", "sha256": "b" * 64, }, @@ -70,10 +73,27 @@ def test_resolve_bundle_dataset_uses_cross_package_overlay(tmp_path): assert dataset.data_package_name == "populace-data" assert dataset.repo_id == "policyengine/populace-uk-private" - assert dataset.repo_type == "model" + assert dataset.repo_type == "dataset" assert dataset.revision == "populace-release" +def test_bundled_uk_populace_dataset_uses_dataset_repository_url(): + dataset = _resolve_bundle_dataset("uk", "populace_uk_2023") + + url = https_dataset_uri( + dataset.repo_id, + dataset.path, + dataset.revision, + repo_type=dataset.repo_type, + ) + + assert dataset.data_package_name == "populace-data" + assert dataset.repo_type == "dataset" + assert url.startswith( + "https://huggingface.co/datasets/policyengine/populace-uk-private/" + ) + + def _sha256(payload: bytes) -> str: return hashlib.sha256(payload).hexdigest() From fdb4f0c72e2cda9864c0b1181f42b551342efc77 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:57:03 +0400 Subject: [PATCH 17/18] Use Pydantic dataset result models --- .../provenance/dataset_materialization.py | 11 +++++++---- tests/test_dataset_source.py | 17 +++++++++++++++-- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index eb0a2830..35824357 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -9,6 +9,7 @@ from typing import Literal, Optional import requests +from pydantic import BaseModel, ConfigDict from policyengine.utils.hashing import sha256_file @@ -31,10 +32,11 @@ class DatasetMaterializationError(ValueError): """Raised when a requested dataset cannot be made available safely.""" -@dataclass(frozen=True) -class MaterializedDataset: +class MaterializedDataset(BaseModel): """Local file and provenance values for a verified bundle dataset.""" + model_config = ConfigDict(frozen=True) + data_package_name: str repo_type: Literal["model", "dataset"] revision: str @@ -44,10 +46,11 @@ class MaterializedDataset: metadata_path: Optional[Path] = None -@dataclass(frozen=True) -class DatasetSource: +class DatasetSource(BaseModel): """Dataset source selected for a country-package calculation.""" + model_config = ConfigDict(frozen=True) + source_uri: str path: str bundle_dataset: Optional[MaterializedDataset] = None diff --git a/tests/test_dataset_source.py b/tests/test_dataset_source.py index 2b325647..58932565 100644 --- a/tests/test_dataset_source.py +++ b/tests/test_dataset_source.py @@ -7,6 +7,7 @@ from policyengine.provenance.dataset_materialization import ( DatasetMaterializationError, DatasetSource, + MaterializedDataset, materialize_dataset, ) @@ -64,6 +65,7 @@ def test_explicit_hf_download_retries_dataset_repository_after_model_404(tmp_pat assert result.bundle_dataset is None assert "/policyengine/example/" in session.calls[0][0] assert "/datasets/policyengine/example/" in session.calls[1][0] + assert DatasetSource.model_validate_json(result.model_dump_json()) == result def test_explicit_hf_authentication_failure_does_not_retry(tmp_path): @@ -90,9 +92,18 @@ def test_explicit_hf_download_rejects_non_hf_uri(tmp_path): def test_bundle_dataset_uses_bundle_strategy(tmp_path): + bundle_dataset = MaterializedDataset( + data_package_name="populace-data", + repo_type="dataset", + revision="release", + source_uri=("hf://policyengine/populace-us/populace_us_2024.h5@release"), + sha256="a" * 64, + path=tmp_path / "populace_us_2024.h5", + ) expected = DatasetSource( - source_uri="hf://policyengine/populace-us/populace_us_2024.h5@release", - path=str(tmp_path / "populace_us_2024.h5"), + source_uri=bundle_dataset.source_uri, + path=str(bundle_dataset.path), + bundle_dataset=bundle_dataset, ) with patch.object( @@ -108,6 +119,7 @@ def test_bundle_dataset_uses_bundle_strategy(tmp_path): assert result is expected use_bundle.assert_called_once() + assert DatasetSource.model_validate_json(result.model_dump_json()) == result def test_explicit_local_path_uses_local_strategy(tmp_path): @@ -125,3 +137,4 @@ def test_explicit_local_path_uses_local_strategy(tmp_path): source_uri=str(local_path), path=str(local_path), ) + assert DatasetSource.model_validate_json(result.model_dump_json()) == result From d79a6109e82b220e9cbef44499bf8c75d7423a94 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:03:18 +0400 Subject: [PATCH 18/18] Clarify dataset materialization documentation --- changelog.d/502.removed.md | 8 ++++++-- docs/bundles.md | 12 ++++++------ docs/microsim.md | 4 ++-- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/changelog.d/502.removed.md b/changelog.d/502.removed.md index 09f0ca73..abfc9892 100644 --- a/changelog.d/502.removed.md +++ b/changelog.d/502.removed.md @@ -1,2 +1,6 @@ -Removed the legacy dataset-specific GCS downloader and its direct dependencies. -The separate UK geography asset implementation is unchanged. +Removed the legacy dataset-specific GCS downloader and its direct dependencies, +the public `resolve_local_managed_dataset_source` export, and the former +`policyengine.provenance.dataset_sources.materialize_dataset_source` function. +Use `policyengine.provenance.materialize_dataset` for bundle-managed datasets, +explicit Hugging Face references, and explicit local files. The separate UK +geography asset implementation is unchanged. diff --git a/docs/bundles.md b/docs/bundles.md index 7457f757..0b8c37c9 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -32,12 +32,12 @@ exact bundled package scaffold with pip, downloads certified default US and UK datasets into `./data`, and writes a `./data/.policyengine-bundle-receipt.json` receipt that records the target Python. -Dataset pre-download uses the same materialization function as US and UK -calculations. For every managed artifact, PolicyEngine.py reads the source data -package name, Hugging Face repository type, immutable revision, and SHA-256 from -the bundle. It reuses an existing file only when its hash matches, downloads and -verifies a replacement before atomically replacing an invalid local file, and -records the verified result in the receipt. +Dataset pre-download and US and UK calculations share the same verified-download +implementation. For every managed artifact, PolicyEngine.py reads the source +data package name, Hugging Face repository type, immutable revision, and SHA-256 +from the bundle. It reuses an existing file only when its hash matches, downloads +and verifies a replacement before atomically replacing an invalid local file, +and records the verified result in the receipt. The bundle manifest can certify additional regional datasets, such as US state datasets. Those artifacts are part of the citable bundle manifest, but diff --git a/docs/microsim.md b/docs/microsim.md index f3c2e1da..0e3d73e1 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -47,8 +47,8 @@ current UK certified default is **Enhanced FRS 2024–25**, supplied by non-default bundle dataset. PolicyEngine.py obtains the repository type, immutable revision, and SHA-256 -from the installed release bundle. A cached file or local data-repository mirror -is reused only after hash verification. +from the installed release bundle. An existing file in the configured data +directory is reused only after hash verification. List datasets already known to the country: