From e85da77980193762d25c4db72eb9cf7f7d84fb39 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 06:33:38 -0400 Subject: [PATCH 1/9] Stamp US year files with their source dataset and package versions A year file is named by its dataset's logical name, which drops the revision, so year files cut from one release of a dataset were reused after the bundle re-pinned the same name to another revision or sha256. Nothing recorded the policyengine-us or policyengine-core version that calculated the values either. The stamp (revision 2) now also records the source identity (the bundle URI and verified sha256, an explicit URI, or a local file's resolved path, size and mtime) and the installed package versions. ensure_datasets and load_datasets compare it with what the request resolves to now, computed from the manifest, the URI or one stat(), and read only the stamp. Fixes #591 Co-Authored-By: Claude Opus 5.5 --- changelog.d/591.fixed.md | 1 + .../provenance/dataset_materialization.py | 117 ++++++++++++- .../tax_benefit_models/us/datasets.py | 79 ++++++--- .../tax_benefit_models/us/legacy_inputs.py | 162 +++++++++++++++--- 4 files changed, 303 insertions(+), 56 deletions(-) create mode 100644 changelog.d/591.fixed.md diff --git a/changelog.d/591.fixed.md b/changelog.d/591.fixed.md new file mode 100644 index 00000000..200418f3 --- /dev/null +++ b/changelog.d/591.fixed.md @@ -0,0 +1 @@ +US year files now record the source dataset they were cut from and the policyengine-core and policyengine-us versions that cut them. `ensure_datasets` creates a year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source or another version is installed. A year file cut from a bundle dataset is tied to the revision and sha256 the bundle pins, so re-pinning a dataset under the same name (such as the ACS-local replacement) no longer reuses year files cut from the old release. A local source is identified by its path, size and modification time, and an explicit Hugging Face URI by the URI. Deciding still reads only each year file's stamp, never the source. Year files stamped by earlier releases are created again once. diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 35824357..d9b73e14 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -2,11 +2,12 @@ from __future__ import annotations +import datetime import os import tempfile from dataclasses import dataclass from pathlib import Path -from typing import Literal, Optional +from typing import Any, Literal, Optional import requests from pydantic import BaseModel, ConfigDict @@ -257,6 +258,120 @@ def _use_local_dataset(dataset_path: str) -> DatasetSource: return DatasetSource(source_uri=dataset_path, path=dataset_path) +def dataset_source_identity(source: DatasetSource) -> dict[str, Any]: + """Return the identity of a dataset source, for files derived from it. + + A file derived from a source (such as a US year file) records this, so a + later request for the same logical dataset can tell whether the source + it would use now is the one the file was derived from. The identity has + one of three shapes: + + - a bundle dataset: ``{"uri": ..., "sha256": ...}``, the + ``hf://repo/path@revision`` URI the bundle manifest pins and the sha256 + that the materialized file was verified against; + - an explicit Hugging Face URI (``allow_unmanaged=True``): ``{"uri": ...}``, + the URI as given. Nothing verified its contents, so a URI that does not + pin a revision, or pins a branch, names whatever the branch holds; + - a local file (``allow_unmanaged=True``): ``{"path": ..., "size": ..., + "mtime_ns": ...}``, its resolved path, size in bytes and modification + time in nanoseconds, read with one ``stat()``. + + None of them reads the source's contents. See + :func:`requested_dataset_source_identity` for the identity a request + gives without materializing it. + """ + if source.bundle_dataset is not None: + return { + "uri": source.bundle_dataset.source_uri, + "sha256": source.bundle_dataset.sha256, + } + if "://" in source.source_uri: + return {"uri": source.source_uri} + return _local_file_identity(source.path) + + +def requested_dataset_source_identity( + country_id: str, + dataset: Optional[str] = None, +) -> dict[str, Any]: + """Return the identity of the source ``materialize_dataset`` would use. + + This is what :func:`dataset_source_identity` gives for + ``materialize_dataset(country_id, dataset, allow_unmanaged=True)``, + computed without downloading anything or reading the source's contents: + the bundle manifest entry for a bundle dataset, the URI for an explicit + Hugging Face URI, and one ``stat()`` for a local file. A name that is + neither a bundle dataset nor a local file is looked up in the data + release manifest, as ``materialize_dataset`` does. + + It does not enforce bundle selection: an unmanaged request gives its + identity whether or not ``allow_unmanaged`` would be set. + """ + manifest = get_release_manifest(country_id) + if dataset is None or dataset == manifest.default_dataset_uri: + dataset = manifest.default_dataset + if dataset in manifest.datasets: + bundle_dataset = _resolve_bundle_dataset(country_id, dataset, manifest=manifest) + return {"uri": bundle_dataset.source_uri, "sha256": bundle_dataset.sha256} + source_uri = resolve_managed_dataset_reference( + country_id, dataset, allow_unmanaged=True + ) + if "://" in source_uri: + return {"uri": source_uri} + return _local_file_identity(source_uri) + + +def _local_file_identity(path: str | Path) -> dict[str, Any]: + resolved = Path(path).expanduser().resolve() + status = resolved.stat() + return { + "path": str(resolved), + "size": status.st_size, + "mtime_ns": status.st_mtime_ns, + } + + +def is_dataset_source_identity(value: Any) -> bool: + """Return whether ``value`` has a shape :func:`dataset_source_identity` gives.""" + if not isinstance(value, dict): + return False + if set(value) == {"uri", "sha256"}: + return isinstance(value["uri"], str) and isinstance(value["sha256"], str) + if set(value) == {"uri"}: + return isinstance(value["uri"], str) + if set(value) == {"path", "size", "mtime_ns"}: + # ``type(...) is int``: JSON ``true`` loads as ``True``, which equals 1. + return ( + isinstance(value["path"], str) + and type(value["size"]) is int + and type(value["mtime_ns"]) is int + ) + return False + + +def describe_dataset_source_identity(identity: dict[str, Any]) -> str: + """Describe a well-formed source identity in words (see + :func:`is_dataset_source_identity`).""" + if "uri" in identity: + if "sha256" in identity: + return f"{identity['uri']} (sha256 {identity['sha256']})" + return identity["uri"] + return ( + f"{identity['path']} ({identity['size']} bytes, modified " + f"{_describe_mtime_ns(identity['mtime_ns'])})" + ) + + +def _describe_mtime_ns(mtime_ns: int) -> str: + """Format a modification time as UTC ISO 8601, to the nanosecond.""" + seconds, nanoseconds = divmod(mtime_ns, 1_000_000_000) + try: + moment = datetime.datetime.fromtimestamp(seconds, datetime.UTC) + except (OverflowError, OSError, ValueError): + return f"at {mtime_ns} ns since the epoch" + return f"{moment:%Y-%m-%dT%H:%M:%S}.{nanoseconds:09d}Z" + + def _reuse_or_download_bundle_files( dataset: _BundleDatasetSpec, ) -> MaterializedDataset: diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index a5fbfb9b..26a3f36c 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -15,7 +15,9 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( MaterializedDataset, + dataset_source_identity, materialize_dataset, + requested_dataset_source_identity, ) from policyengine.provenance.manifest import ( dataset_logical_name, @@ -369,6 +371,18 @@ def _load_policyengine_core_h5(path: Path, year: int) -> USYearData: ) +def _year_file_source(dataset: str) -> tuple[str, dict[str, Any]]: + """Return a requested dataset's year-file stem and its source identity. + + The stem names the year files (``{stem}_year_{year}.h5``) and drops any + revision, so files cut from every release of a logical dataset share + names; the identity tells them apart. Neither downloads nor reads the + source (see ``requested_dataset_source_identity``). + """ + stem = dataset_logical_name(resolve_dataset_reference("us", dataset)) + return stem, requested_dataset_source_identity("us", dataset) + + def create_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2024, 2025, 2026, 2027, 2028], @@ -378,10 +392,13 @@ def create_datasets( """Create PolicyEngineUSDataset instances from logical dataset names or URLs. Each year file is written to ``{data_folder}/{dataset_stem}_year_{year}.h5`` - and stamped last with the register of renamed stored inputs it was cut - under and the year-file format revision (see - ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which ``ensure_datasets`` - and ``load_datasets`` check. + and stamped last with the identity of the source dataset it was cut from, + the installed policyengine-core and policyengine-us versions, the register + of renamed stored inputs it was cut under and the year-file format + revision (see ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which + ``ensure_datasets`` and ``load_datasets`` check. A local source is + identified before it is loaded, so one that changes while the files are + cut leaves them out of date. Args: datasets: List of logical dataset names or HuggingFace dataset URLs @@ -405,9 +422,11 @@ def create_datasets( data_dir=Path(data_folder), ) dataset_stem = source.name + # What each year file is cut from and under: the source, identified + # before it is loaded, the package versions and the register the + # mapping below uses. + stamp = year_file_stamp(dataset_source_identity(source)) sim = Microsimulation(dataset=source.path, spm=resolve_spm_selection()) - # The register the mapping below uses, stamped into each year file. - stamp = year_file_stamp() # Map stored inputs the engine has since renamed before extracting, # so each year's file stores the live input rather than dropping it. legacy_input_renames = apply_legacy_input_renames_to_microsimulation(sim) @@ -586,13 +605,16 @@ def load_datasets( ) -> dict[str, PolicyEngineUSDataset]: """Load PolicyEngineUSDataset instances from saved HDF5 files. - A year file is refused unless ``create_datasets`` stamped it with this - release's register of renamed stored inputs and year-file format revision - (see ``legacy_inputs.year_file_stamp_problem``). A file cut before a - register entry was added may have lost that input, as files that - policyengine.py 6.1.1 and earlier cut from the certified default lost the - WIC take-up draw. The error says why, and names the renames the file may - lack when its register is the reason. + A year file is refused unless ``create_datasets`` stamped it as cut from + the source the dataset resolves to now, with the installed + policyengine-core and policyengine-us versions, this release's register of + renamed stored inputs and its year-file format revision (see + ``legacy_inputs.year_file_stamp_problem``). A file cut before a register + entry was added may have lost that input, as files that policyengine.py + 6.1.1 and earlier cut from the certified default lost the WIC take-up + draw. A file cut from an earlier release of a dataset the bundle has since + re-pinned holds that release's data. The error says why, naming the + renames the file may lack, the two sources or the two package versions. Args: datasets: List of HuggingFace dataset paths (used to derive file names) @@ -605,12 +627,11 @@ def load_datasets( datasets = datasets or [get_release_manifest("us").default_dataset] result = {} for dataset in datasets: - resolved_dataset = resolve_dataset_reference("us", dataset) - dataset_stem = dataset_logical_name(resolved_dataset) + dataset_stem, source = _year_file_source(dataset) for year in years: filepath = f"{data_folder}/{dataset_stem}_year_{year}.h5" if Path(filepath).exists(): - problem = year_file_stamp_problem(filepath) + problem = year_file_stamp_problem(filepath, source) if problem is not None: raise ValueError( f"US year file {filepath} is out of date: {problem}. " @@ -1263,12 +1284,17 @@ def ensure_datasets( """Ensure datasets exist, loading if available or creating if not. Existing year files are reused only when ``create_datasets`` stamped them - with this release's register of renamed stored inputs and year-file format - revision (see ``legacy_inputs.year_file_stamp_problem``). Only each file's - stamp and renames record are read to decide, never its tables or the - source dataset. If any year file is missing or out of date, every - year file is created again from the source dataset, and each out-of-date - file is logged with the reason. + as cut from the source each dataset resolves to now, with the installed + policyengine-core and policyengine-us versions, this release's register of + renamed stored inputs and its year-file format revision (see + ``legacy_inputs.year_file_stamp_problem``). Only each file's stamp and + renames record are read to decide, never its tables. The source is never + downloaded or read: a bundle dataset is identified by its manifest entry, + an explicit Hugging Face URI by the URI and a local file by one + ``stat()`` (see ``dataset_materialization.dataset_source_identity``). If + any year file is missing or out of date, every year file is created again + from the source dataset, and each out-of-date file is logged with the + reason. Args: datasets: List of HuggingFace dataset paths @@ -1281,18 +1307,17 @@ def ensure_datasets( datasets = datasets or [get_release_manifest("us").default_dataset] # Check that every year file exists and is current. A year file cut - # under another register of renamed stored inputs may have lost one, so - # it is created again rather than reused. + # from another source, with other package versions or under another + # register of renamed stored inputs is created again rather than reused. all_current = True for dataset in datasets: - resolved_dataset = resolve_dataset_reference("us", dataset) - dataset_stem = dataset_logical_name(resolved_dataset) + dataset_stem, source = _year_file_source(dataset) for year in years: filepath = Path(f"{data_folder}/{dataset_stem}_year_{year}.h5") if not filepath.exists(): all_current = False continue - problem = year_file_stamp_problem(filepath) + problem = year_file_stamp_problem(filepath, source) if problem is not None: all_current = False logger.warning( diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 38fb0249..42696432 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -25,12 +25,18 @@ import inspect import json from collections.abc import Iterable, Iterator, Mapping +from importlib import metadata as importlib_metadata from pathlib import Path from typing import Any import numpy as np import pandas as pd +from policyengine.provenance.dataset_materialization import ( + describe_dataset_source_identity, + is_dataset_source_identity, +) + #: Stored input columns that policyengine-us has since renamed, mapped to the #: live input that replaced each one. #: @@ -74,27 +80,50 @@ #: reused. RENAMES_H5_DATASET = "policyengine_legacy_input_renames" -#: H5 dataset stamping a ``create_datasets`` year file, as UTF-8 JSON -#: ``{"register": {legacy: live, ...}, "revision": int}``: the whole register -#: the file was cut under and its year-file format revision (see -#: :func:`year_file_stamp`). ``create_datasets`` writes it after the tables and -#: the renames record, so a year file whose write stopped part-way has none. -#: ``ensure_datasets`` cuts a year file again, and ``load_datasets`` refuses -#: it, unless its stamp equals this release's (see +#: H5 dataset stamping a ``create_datasets`` year file, as UTF-8 JSON (see +#: :func:`year_file_stamp`):: +#: +#: {"packages": {"policyengine-core": str, "policyengine-us": str}, +#: "register": {legacy: live, ...}, +#: "revision": int, +#: "source": {...}} +#: +#: It holds the versions of the packages that calculated the file's values, +#: the whole register the file was cut under, its year-file format revision +#: and the identity of the dataset it was cut from. ``create_datasets`` writes +#: it after the tables and the renames record, so a year file whose write +#: stopped part-way has none. ``ensure_datasets`` cuts a year file again, and +#: ``load_datasets`` refuses it, unless its stamp equals the one a file cut +#: now from the requested dataset would carry (see #: :func:`year_file_stamp_problem`). ``PolicyEngineUSDataset.save()`` neither #: writes nor keeps it, so saving a year file again drops the stamp. YEAR_FILE_STAMP_H5_DATASET = "policyengine_year_file_stamp" #: Format revision of ``create_datasets`` year files, stored in their stamp. -#: Raise it when a change to year-file extraction means that files cut before -#: it must be cut again, for a reason the register does not capture (for -#: example, policyengine.py 6.1.1 and earlier exported only the engine's input -#: variables, so a stored column the engine had renamed was dropped). +#: Raise it when the stamp changes shape, or when a change to year-file +#: extraction means that files cut before it must be cut again, for a reason +#: the rest of the stamp does not capture (for example, policyengine.py 6.1.1 +#: and earlier exported only the engine's input variables, so a stored column +#: the engine had renamed was dropped). #: #: Removing the temporary ACS WIC assumption (:func:`_temporary_acs_wic_values`) #: is such a change: year files cut while it applies store assumed claiming #: as ``takes_up_wic_if_eligible`` for ACS people with no stored decision. -YEAR_FILE_STAMP_REVISION = 1 +#: Re-pinning ``populace_us_2024_acs_local`` to its replacement changes its +#: source identity, so year files cut from the bundle's old pin are cut again +#: without a new revision. Year files cut from the old dataset requested by +#: its explicit URI are reused unless the revision is raised. +#: +#: Revision 2 added the source dataset's identity and the package versions, +#: so year files stamped with revision 1, which record neither, are cut again +#: once. +YEAR_FILE_STAMP_REVISION = 2 + +#: Distributions whose installed versions a year file's stamp records. +#: ``create_datasets`` stores each input as policyengine-us calculates it for +#: the year, and policyengine-core uprates an input from its stored year by a +#: policyengine-us parameter, so either can change a year file's values. +YEAR_FILE_PACKAGES = ("policyengine-core", "policyengine-us") #: ``{year: {entity key: stored table}}``, one table per entity per dataset #: year, in the dataset's own row order. @@ -144,16 +173,42 @@ def write_renames_record(path: str | Path, record: Mapping[str, str]) -> None: _write_json_dataset(path, RENAMES_H5_DATASET, dict(record)) -def year_file_stamp() -> dict[str, Any]: - """Return the stamp for a year file cut now. +def installed_package_versions() -> dict[str, str | None]: + """Return the installed version of each of :data:`YEAR_FILE_PACKAGES`. - It holds the whole register, not only the entries an engine needs mapped, - so checking a stamp never imports policyengine-us. The register is read - when this is called. + ``None`` stands for a package that is not installed. The versions come + from the installed distributions' metadata, so no package is imported. """ + versions: dict[str, str | None] = {} + for name in YEAR_FILE_PACKAGES: + try: + versions[name] = importlib_metadata.version(name) + except importlib_metadata.PackageNotFoundError: + versions[name] = None + return versions + + +def year_file_stamp( + source: Mapping[str, Any], + packages: Mapping[str, str | None] | None = None, +) -> dict[str, Any]: + """Return the stamp for a year file cut now from ``source``. + + ``source`` is the source dataset's identity (see + ``policyengine.provenance.dataset_materialization.dataset_source_identity``). + ``packages`` defaults to :func:`installed_package_versions`. + + The stamp holds the whole register, not only the entries an engine needs + mapped, so checking a stamp never imports policyengine-us. The register + and the package versions are read when this is called. + """ + if packages is None: + packages = installed_package_versions() return { + "packages": dict(sorted(packages.items())), "register": dict(sorted(LEGACY_INPUT_RENAMES.items())), "revision": YEAR_FILE_STAMP_REVISION, + "source": dict(source), } @@ -178,13 +233,21 @@ def write_year_file_stamp(path: str | Path, stamp: Mapping[str, Any]) -> None: _write_json_dataset(path, YEAR_FILE_STAMP_H5_DATASET, dict(stamp)) -def year_file_stamp_problem(path: str | Path) -> str | None: +def year_file_stamp_problem( + path: str | Path, + source: Mapping[str, Any], + packages: Mapping[str, str | None] | None = None, +) -> str | None: """Return why a ``create_datasets`` year file is out of date, or ``None``. - A year file is current only when both of these hold: + ``source`` is the identity of the dataset the year file is requested for + (see ``policyengine.provenance.dataset_materialization + .requested_dataset_source_identity``), and ``packages`` defaults to + :func:`installed_package_versions`. A year file is current only when both + of these hold: - - its stamp equals :func:`year_file_stamp`: the same format revision and - the same register; + - its stamp equals ``year_file_stamp(source, packages)``: the same format + revision, register, source identity and package versions; - it stores the renames record, listing only renames in that register. The register must match exactly. A file cut under a register that lacks @@ -195,6 +258,11 @@ def year_file_stamp_problem(path: str | Path) -> str | None: this release's engine does not define, and a run skips any column the engine does not define. + The source must match exactly too: a file cut from another revision or + file of a dataset that keeps its logical name (and so the year file's + name) holds the other data. So must the package versions, which + calculated every value in the file. + Only the stamp and the record are read, never the tables or the source dataset. A file that cannot be read as HDF5 is out of date. An error that says nothing about the file's contents (no permission, a directory, a file @@ -203,12 +271,14 @@ def year_file_stamp_problem(path: str | Path) -> str | None: Returns: ``None`` for a current year file; otherwise a clause, starting "it" - or "its", that says why the file is out of date and names the - renames it may have lost. + or "its", that says why the file is out of date. It names the renames + the file may have lost, the source it was cut from and the source + requested, and the package versions it was cut with and those + installed, wherever they differ. """ import h5py - current = year_file_stamp() + current = year_file_stamp(source, packages) register = current["register"] try: with h5py.File(path, "r") as file: @@ -247,13 +317,13 @@ def year_file_stamp_problem(path: str | Path) -> str | None: # ``True``, which equals 1. if not ( isinstance(stamp, dict) - and set(stamp) == {"register", "revision"} + and set(stamp) == {"packages", "register", "revision", "source"} and type(revision) is int and _is_str_map(stamp["register"]) + and _is_version_map(stamp["packages"]) + and is_dataset_source_identity(stamp["source"]) ): - return ( - "its stamp of the register of renamed inputs it was cut under is malformed" - ) + return "its stamp of what it was cut from and under is malformed" if record is _ABSENT: return "it has no record of the renamed stored inputs mapped when it was cut" if not _is_str_map(record): @@ -287,6 +357,34 @@ def year_file_stamp_problem(path: str | Path) -> str | None: "it was cut under renames this version of policyengine.py does not " f"register: {_describe(unknown)}" ) + if stamp["source"] != current["source"]: + reasons.append( + "it was cut from " + f"{describe_dataset_source_identity(stamp['source'])}, and the " + "requested source is " + f"{describe_dataset_source_identity(current['source'])}" + ) + stamped_packages, packages = stamp["packages"], current["packages"] + for name in sorted(stamped_packages.keys() | packages.keys()): + if name not in packages: + reasons.append( + f"it records the {name} version it was cut with, which this " + "version of policyengine.py does not check" + ) + continue + if name in stamped_packages and stamped_packages[name] == packages[name]: + continue + cut_with = ( + f"{name} {stamped_packages[name]}" + if stamped_packages.get(name) is not None + else f"no recorded {name} version" + ) + installed = ( + f"{name} {packages[name]} is installed" + if packages[name] is not None + else f"{name} is not installed" + ) + reasons.append(f"it was cut with {cut_with}, and {installed}") return "; ".join(reasons) @@ -331,6 +429,14 @@ def _is_str_map(value: Any) -> bool: ) +def _is_version_map(value: Any) -> bool: + """Whether ``value`` maps package names to versions, ``None`` for absent.""" + return isinstance(value, dict) and all( + isinstance(key, str) and (item is None or isinstance(item, str)) + for key, item in value.items() + ) + + def _describe(renames: Mapping[str, str]) -> str: return ", ".join(f"{legacy} -> {live}" for legacy, live in sorted(renames.items())) From 64f5eb38161662014dadab4eb8b29d50477ca354 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 06:45:39 -0400 Subject: [PATCH 2/9] Test and document year-file source and package identity Unit and Hypothesis tests extend the stamp's invariants (exactness, totality, named differences, ensure agrees with load, decides from the stamps alone) to the source identity and package versions. A new test file checks that identifying a request agrees with materializing it on every route, never reads the source, and changes exactly when the bundle re-pins an entry or a local file's size or mtime changes. Integration tests cover an edited local source, a re-pinned bundle dataset and a new policyengine-us version over the real engine. Refs #591 Co-Authored-By: Claude Opus 5.5 --- changelog.d/591.fixed.md | 2 +- docs/microsim.md | 97 ++- tests/test_dataset_source_identity.py | 533 +++++++++++++++ tests/test_us_legacy_inputs_integration.py | 229 ++++++- tests/test_us_year_file_stamp.py | 755 +++++++++++++++++---- 5 files changed, 1445 insertions(+), 171 deletions(-) create mode 100644 tests/test_dataset_source_identity.py diff --git a/changelog.d/591.fixed.md b/changelog.d/591.fixed.md index 200418f3..ee3602e9 100644 --- a/changelog.d/591.fixed.md +++ b/changelog.d/591.fixed.md @@ -1 +1 @@ -US year files now record the source dataset they were cut from and the policyengine-core and policyengine-us versions that cut them. `ensure_datasets` creates a year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source or another version is installed. A year file cut from a bundle dataset is tied to the revision and sha256 the bundle pins, so re-pinning a dataset under the same name (such as the ACS-local replacement) no longer reuses year files cut from the old release. A local source is identified by its path, size and modification time, and an explicit Hugging Face URI by the URI. Deciding still reads only each year file's stamp, never the source. Year files stamped by earlier releases are created again once. +US year files now record the source dataset they were cut from and the policyengine-core and policyengine-us versions that cut them. `ensure_datasets` creates a year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source or another version is installed. A year file cut from a bundle dataset is tied to the revision and sha256 the bundle pins, so re-pinning a dataset under the same name (such as the ACS-local replacement) no longer reuses year files cut from the old release. A local source is identified by its path, size and modification time, and an explicit Hugging Face URI by the URI. Deciding still reads only each year file's stamp and renames record, and never downloads or reads the source. diff --git a/docs/microsim.md b/docs/microsim.md index 1e63432a..1fc2986f 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -409,17 +409,55 @@ under. ### Year files -`create_datasets` stamps each year file, after writing everything else, with -the register it was cut under and a year-file format revision -(`legacy_inputs.YEAR_FILE_STAMP_H5_DATASET`, UTF-8 JSON -`{"register": {old: new, ...}, "revision": 1}`). `ensure_datasets` reuses a -year file only when its stamp equals this release's exactly and its renames -record lists only renames in the stamped register. Otherwise it logs why and -creates the year files again; `load_datasets` refuses the file with the same -reason, which names the renames the file may lack when its register is the -cause. To decide, both read only the stamp and the record, never the tables -or the source dataset. `legacy_inputs.year_file_stamp_problem(path)` returns -the reason, or `None` for a current year file. +`create_datasets` names each US year file after its dataset's logical name, +`{data_folder}/{name}_year_{year}.h5`, which drops the revision. Files cut +from every release of a dataset therefore share names. After writing +everything else, it stamps each file with what the file was cut from and under +(`legacy_inputs.YEAR_FILE_STAMP_H5_DATASET`, UTF-8 JSON): + +```json +{ + "packages": {"policyengine-core": "3.33.0", "policyengine-us": "2.2.1"}, + "register": {"would_claim_wic": "takes_up_wic_if_eligible"}, + "revision": 2, + "source": { + "sha256": "6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84", + "uri": "hf://policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260915" + } +} +``` + +- `packages`: the installed policyengine-core and policyengine-us versions. + The file stores each input as policyengine-us calculates it for that year, + and policyengine-core uprates an input from its stored year by a + policyengine-us parameter, so either version can change the file's values; +- `register`: the whole `LEGACY_INPUT_RENAMES` it was cut under; +- `revision`: the year-file format revision; +- `source`: the identity of the dataset it was cut from + (`dataset_materialization.dataset_source_identity`), in one of three shapes: + - a bundle dataset: `{"uri", "sha256"}`, the `hf://…@revision` URI the + bundle pins and the sha256 the local copy was verified against; + - an explicit Hugging Face URI (`allow_unmanaged=True`): `{"uri"}`, the URI + as given; + - a local file (`allow_unmanaged=True`): `{"path", "size", "mtime_ns"}`, + its resolved path, size in bytes and modification time. + +`ensure_datasets` reuses a year file only when its stamp equals the one a file +cut now from the requested dataset would carry, and its renames record lists +only renames in the stamped register. Otherwise it logs why and creates the +year files again. `load_datasets` refuses the file with the same reason, which +names the renames the file may lack, the source it was cut from and the +source requested, or the package versions it was cut with and those +installed. `legacy_inputs.year_file_stamp_problem(path, source)` returns the +reason, or `None` for a current year file; `source` is +`dataset_materialization.requested_dataset_source_identity("us", dataset)`. + +Deciding stays cheap. Both functions read only each year file's stamp and +record, never its tables. They identify the requested source without +downloading or reading it: a bundle dataset by its bundle manifest entry, an +explicit URI by the URI, and a local file by one `stat()`. Package versions +come from the installed distributions' metadata, so policyengine-us is not +imported. A year file is out of date when: @@ -434,16 +472,37 @@ A year file is out of date when: - its register has an entry this release lacks, or maps an old name to another new name. It may store a new name this release's engine does not define, and a run skips any column the engine does not define; +- it was cut from another source. The bundle may have re-pinned the dataset + to another revision or sha256. The replacement of + `populace_us_2024_acs_local` (PolicyEngine/microcosm#1154) is one such + re-pin. A local file may have been rewritten, or another file with the + same name may have been requested; +- it was cut with another policyengine-core or policyengine-us version; - it has another format revision, or its stamp or record cannot be read. -So any change to `LEGACY_INPUT_RENAMES` makes every year file out of date, -and two policyengine.py versions with different registers that share a data -folder cut each other's year files again. `ensure_datasets` raises, rather -than cutting the files again, when a year file cannot be opened for reasons -unrelated to its contents: no permission, a directory at its path, or a lock -held by another process. A year file opened directly, as -`PolicyEngineUSDataset(filepath=...)`, is not checked, and saving it again with -`save()` drops its stamp. +So any change to `LEGACY_INPUT_RENAMES`, a re-pinned dataset or a new package +version makes the affected year files out of date. Two environments with +different registers or package versions that share a data folder cut each +other's year files again. + +The source check has limits that follow from not reading the source: + +- a local file is compared by path, size and modification time, the quick + check `rsync` makes by default. Touching or copying the file cuts the + year files again. A rewrite that keeps the size and the modification time + (restored with `touch -r`, or landing within the file system's timestamp + resolution) is not detected; +- an explicit Hugging Face URI is compared as a string. Pin it to a commit + (`hf://owner/repo/file.h5@`): a URI without a revision, or with a + branch, names whatever the branch holds when the files are cut; +- an editable or source install of policyengine-us keeps its version string + while its code changes, so its year files are not cut again. + +`ensure_datasets` raises, rather than cutting the files again, when a year +file cannot be opened for reasons unrelated to its contents: no permission, a +directory at its path, or a lock held by another process. A year file opened +directly, as `PolicyEngineUSDataset(filepath=...)`, is not checked, and saving +it again with `save()` drops its stamp. ## Pinned model versions diff --git a/tests/test_dataset_source_identity.py b/tests/test_dataset_source_identity.py new file mode 100644 index 00000000..d887f87c --- /dev/null +++ b/tests/test_dataset_source_identity.py @@ -0,0 +1,533 @@ +"""Tests for the identity of a dataset source. + +A file derived from a dataset (such as a US year file) records +``dataset_source_identity(materialize_dataset(...))``, and a later request +compares it with ``requested_dataset_source_identity(...)``, which must not +download or read the source. These tests state the invariants that makes +safe, over a fixture bundle manifest whose datasets are tiny local files: + +- **Agreement (differential).** For every way a dataset can be requested + (the default, the default's URI, a bundle name, an explicit Hugging Face + URI, a data release manifest name, and a local path given relative, + absolute, with ``~`` or through a symlink), the identity of the request + equals the identity of what ``materialize_dataset`` makes available. +- **Cheapness.** Identifying a request never downloads, opens, hashes or + loads the source. +- **Sensitivity, bundle datasets.** A request's identity changes exactly + when the bundle re-pins its entry (repository, path, revision or sha256), + and not when any other entry changes. +- **Sensitivity, local files.** A local file's identity changes exactly when + its resolved path, size or modification time changes. Bytes rewritten + under the same size and modification time keep it: the intended blind spot + of a ``stat()`` check, the quick check ``rsync`` makes by default. +- **Shapes.** Every identity produced has a shape + ``is_dataset_source_identity`` accepts, and describing one never raises. +""" + +from __future__ import annotations + +import builtins +import hashlib +import io +import os +from pathlib import Path +from unittest.mock import patch + +import h5py +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import policyengine.provenance.dataset_materialization as dataset_materialization +import policyengine.provenance.manifest as manifest_module +from policyengine.provenance.dataset_materialization import ( + dataset_source_identity, + describe_dataset_source_identity, + is_dataset_source_identity, + materialize_dataset, + requested_dataset_source_identity, +) +from policyengine.provenance.manifest import ( + CountryReleaseManifest, + DataReleaseManifest, + build_hf_uri, +) + +DEFAULT_BYTES = b"default dataset" +OTHER_BYTES = b"other dataset, from another repository" +EXPLICIT_BYTES = b"explicit dataset" + + +def _sha256(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def _manifest(**other_entry) -> CountryReleaseManifest: + """A fixture US bundle: a default dataset and one from another repo.""" + other = { + "path": "nested/tiny_other.h5", + "repo_id": "fixture/other-repo", + "repo_type": "dataset", + "revision": "other-r9", + "sha256": _sha256(OTHER_BYTES), + **other_entry, + } + return CountryReleaseManifest.model_validate( + { + "country_id": "us", + "policyengine_version": "0.0.0", + "model_package": {"name": "policyengine-us", "version": "2.2.1"}, + "data_package": { + "name": "fixture-data", + "version": "0.1.0", + "repo_id": "fixture/data", + "repo_type": "model", + "release_manifest_revision": "fixture-release", + }, + "default_dataset": "tiny_default", + "certified_data_artifact": { + "dataset": "tiny_default", + "uri": build_hf_uri("fixture/data", "tiny_default.h5", "r1"), + "sha256": _sha256(DEFAULT_BYTES), + }, + "datasets": { + "tiny_default": { + "path": "tiny_default.h5", + "revision": "r1", + "sha256": _sha256(DEFAULT_BYTES), + }, + "tiny_other": {key: value for key, value in other.items() if value}, + }, + } + ) + + +DATA_RELEASE = DataReleaseManifest.model_validate( + { + "schema_version": 1, + "data_package": {"name": "fixture-data", "version": "0.1.0"}, + "artifacts": { + "release_only": { + "kind": "dataset", + "path": "release_only.h5", + "repo_id": "fixture/data", + "revision": "release-r2", + } + }, + } +) + + +class _Response: + def __init__(self, payload: bytes): + self.payload = payload + self.status_code = 200 + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def raise_for_status(self): + pass + + def iter_content(self, chunk_size): + yield self.payload + + +def _serve(payload: bytes): + calls = [] + + def get(url, **kwargs): + calls.append(url) + return _Response(payload) + + return get, calls + + +def _refuse(*args, **kwargs): + raise AssertionError("identifying a request downloaded or read its source") + + +def _with_manifest(manifest): + """Patch every module that looks up the bundled US manifest.""" + return ( + patch.object(manifest_module, "get_release_manifest", lambda _: manifest), + patch.object( + dataset_materialization, "get_release_manifest", lambda _: manifest + ), + patch.object( + manifest_module, "get_data_release_manifest", lambda _: DATA_RELEASE + ), + ) + + +@pytest.fixture +def bundle(tmp_path): + """The fixture bundle, with its datasets already on disk, verified.""" + manifest = _manifest() + patchers = _with_manifest(manifest) + for patcher in patchers: + patcher.start() + data_dir = tmp_path / "data" + data_dir.mkdir() + (data_dir / "tiny_default.h5").write_bytes(DEFAULT_BYTES) + (data_dir / "tiny_other.h5").write_bytes(OTHER_BYTES) + try: + yield manifest, data_dir + finally: + for patcher in patchers: + patcher.stop() + + +def _local_file(tmp_path) -> Path: + path = tmp_path / "sources" / "tiny_local.h5" + path.parent.mkdir(exist_ok=True) + path.write_bytes(b"local dataset") + return path + + +#: Each way to request a dataset: (request, bytes served for a download). +#: ``{tmp}`` and ``{home}`` stand for the test's folders. +ROUTES = { + "default": (None, None), + "default URI": (build_hf_uri("fixture/data", "tiny_default.h5", "r1"), None), + "bundle name": ("tiny_default", None), + "bundle name in another repository": ("tiny_other", None), + "explicit URI": ("hf://fixture/explicit/explicit.h5@pinned", EXPLICIT_BYTES), + "explicit URI without a revision": ( + "hf://fixture/explicit/explicit.h5", + EXPLICIT_BYTES, + ), + "explicit URI of a bundle dataset": ( + build_hf_uri("fixture/other-repo", "nested/tiny_other.h5", "other-r9"), + OTHER_BYTES, + ), + "data release manifest name": ("release_only", EXPLICIT_BYTES), + "absolute local path": ("{tmp}/sources/tiny_local.h5", None), + "relative local path": ("sources/tiny_local.h5", None), + "local path under ~": ("~/sources/tiny_local.h5", None), + "symlinked local path": ("{tmp}/link.h5", None), +} + + +@pytest.mark.parametrize("route", sorted(ROUTES)) +def test_a_request_is_identified_as_what_it_materializes( + bundle, tmp_path, monkeypatch, route +): + """Agreement: identifying a request without materializing it gives the + identity of what materializing it makes available.""" + _, data_dir = bundle + _local_file(tmp_path) + (tmp_path / "link.h5").symlink_to(tmp_path / "sources" / "tiny_local.h5") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path)) + dataset, served = ROUTES[route] + if dataset is not None: + dataset = dataset.format(tmp=tmp_path) + + with ( + patch.object(dataset_materialization.requests, "get", _refuse), + patch.object(dataset_materialization, "sha256_file", _refuse), + ): + requested = requested_dataset_source_identity("us", dataset) + + get, calls = _serve(served or b"") + with patch.object( + dataset_materialization.requests, "get", get if served else _refuse + ): + source = materialize_dataset( + "us", dataset, allow_unmanaged=True, data_dir=data_dir + ) + + assert dataset_source_identity(source) == requested + assert is_dataset_source_identity(requested) + assert bool(calls) == bool(served) + + +def test_the_identities_of_each_route(bundle, tmp_path, monkeypatch): + _local_file(tmp_path) + monkeypatch.chdir(tmp_path) + local = tmp_path.resolve() / "sources" / "tiny_local.h5" + status = local.stat() + + assert requested_dataset_source_identity("us") == { + "uri": "hf://fixture/data/tiny_default.h5@r1", + "sha256": _sha256(DEFAULT_BYTES), + } + assert requested_dataset_source_identity("us", "tiny_other") == { + "uri": "hf://fixture/other-repo/nested/tiny_other.h5@other-r9", + "sha256": _sha256(OTHER_BYTES), + } + # An explicit URI is not verified, even when it names a bundle dataset. + assert requested_dataset_source_identity( + "us", "hf://fixture/other-repo/nested/tiny_other.h5@other-r9" + ) == {"uri": "hf://fixture/other-repo/nested/tiny_other.h5@other-r9"} + assert requested_dataset_source_identity("us", "release_only") == { + "uri": "hf://fixture/data/release_only.h5@release-r2" + } + assert requested_dataset_source_identity("us", "sources/tiny_local.h5") == { + "path": str(local), + "size": status.st_size, + "mtime_ns": status.st_mtime_ns, + } + + +def test_a_bundle_entry_without_a_revision_uses_the_release_revision(): + """It inherits the data package's repository and release revision, as + ``_resolve_bundle_dataset`` does.""" + manifest = _manifest(revision=None, repo_id=None, repo_type=None) + patchers = _with_manifest(manifest) + with patchers[0], patchers[1], patchers[2]: + assert requested_dataset_source_identity("us", "tiny_other") == { + "uri": "hf://fixture/data/nested/tiny_other.h5@fixture-release", + "sha256": _sha256(OTHER_BYTES), + } + + +def test_a_bundle_entry_without_a_sha256_is_refused_as_materialize_does(tmp_path): + manifest = _manifest() + manifest.datasets["tiny_other"].sha256 = None + patchers = _with_manifest(manifest) + with patchers[0], patchers[1], patchers[2]: + with pytest.raises( + dataset_materialization.DatasetMaterializationError, match="sha256" + ): + requested_dataset_source_identity("us", "tiny_other") + with pytest.raises( + dataset_materialization.DatasetMaterializationError, match="sha256" + ): + materialize_dataset("us", "tiny_other", data_dir=tmp_path) + + +@pytest.mark.parametrize( + "route", + [ + "default", + "bundle name", + "explicit URI", + "absolute local path", + "relative local path", + ], +) +def test_identifying_a_request_never_reads_the_source( + bundle, tmp_path, monkeypatch, route +): + """Cheapness: no download, no open, no hash, no H5 read of the source.""" + _, data_dir = bundle + source = _local_file(tmp_path).resolve() + monkeypatch.chdir(tmp_path) + dataset, _ = ROUTES[route] + if dataset is not None: + dataset = dataset.format(tmp=tmp_path) + real_open = builtins.open + guarded = {source, (data_dir / "tiny_default.h5").resolve()} + + def open_anything_but_a_source(file, *args, **kwargs): + if isinstance(file, (str, os.PathLike)) and Path(file).resolve() in guarded: + _refuse() + return real_open(file, *args, **kwargs) + + with ( + patch.object(builtins, "open", open_anything_but_a_source), + patch.object(io, "open", open_anything_but_a_source), + patch.object(h5py, "File", _refuse), + patch.object(dataset_materialization, "sha256_file", _refuse), + patch.object(dataset_materialization.requests, "get", _refuse), + ): + identity = requested_dataset_source_identity("us", dataset) + + assert is_dataset_source_identity(identity) + + +# --- Sensitivity --------------------------------------------------------- + +REPOSITORIES = st.sampled_from([None, "fixture/other-repo", "fixture/third-repo"]) +PATHS = st.sampled_from(["nested/tiny_other.h5", "tiny_other.h5"]) +REVISIONS = st.sampled_from([None, "other-r9", "other-r10"]) +SHA256S = st.sampled_from([_sha256(OTHER_BYTES), _sha256(b"replacement")]) + + +@st.composite +def entries(draw): + return { + "repo_id": draw(REPOSITORIES), + "path": draw(PATHS), + "revision": draw(REVISIONS), + "sha256": draw(SHA256S), + } + + +def _identity_under(manifest, dataset): + patchers = _with_manifest(manifest) + with patchers[0], patchers[1], patchers[2]: + return requested_dataset_source_identity("us", dataset) + + +@settings(max_examples=150, deadline=None) +@given(first=entries(), second=entries()) +def test_a_bundle_dataset_is_re_identified_exactly_when_re_pinned(first, second): + """Re-pinning the requested entry (its repository, path, revision or + sha256) changes its identity, and only that does.""" + + def pinned(entry): + return ( + entry["repo_id"] or "fixture/data", + entry["path"], + entry["revision"] or "fixture-release", + entry["sha256"], + ) + + first_identity = _identity_under(_manifest(**first), "tiny_other") + second_identity = _identity_under(_manifest(**second), "tiny_other") + + assert (first_identity == second_identity) == (pinned(first) == pinned(second)) + # The default dataset's identity does not move with another entry. + assert _identity_under(_manifest(**first), None) == _identity_under( + _manifest(**second), None + ) + + +SECOND = 1_000_000_000 +#: Whole seconds, so a file system that keeps only seconds still keeps them. +MTIMES = st.integers(1, 2**31).map(lambda seconds: seconds * SECOND) +CONTENTS = st.binary(min_size=0, max_size=4) + + +@settings(max_examples=150, deadline=None) +@given( + first=CONTENTS, + second=CONTENTS, + first_mtime=MTIMES, + second_mtime=MTIMES, +) +def test_a_local_file_is_re_identified_when_its_size_or_mtime_changes( + tmp_path_factory, first, second, first_mtime, second_mtime +): + """Sensitivity: identities differ exactly when size or mtime does. Equal + size and mtime over other bytes give equal identities: intended.""" + path = tmp_path_factory.mktemp("local") / "tiny_local.h5" + + path.write_bytes(first) + os.utime(path, ns=(first_mtime, first_mtime)) + first_identity = requested_dataset_source_identity("us", str(path)) + path.write_bytes(second) + os.utime(path, ns=(second_mtime, second_mtime)) + second_identity = requested_dataset_source_identity("us", str(path)) + + assert (first_identity == second_identity) == ( + len(first) == len(second) and first_mtime == second_mtime + ) + assert first_identity["mtime_ns"] == first_mtime + assert first_identity["size"] == len(first) + + +def test_a_local_file_elsewhere_is_another_source(tmp_path): + """Two files with one name, so one year-file name, are told apart.""" + first, second = tmp_path / "a" / "tiny.h5", tmp_path / "b" / "tiny.h5" + for path in (first, second): + path.parent.mkdir() + path.write_bytes(b"same bytes") + os.utime(path, ns=(SECOND, SECOND)) + + assert requested_dataset_source_identity( + "us", str(first) + ) != requested_dataset_source_identity("us", str(second)) + + +def test_re_pointing_a_symlink_re_identifies_the_source(tmp_path): + first, second, link = tmp_path / "v1.h5", tmp_path / "v2.h5", tmp_path / "x.h5" + for path in (first, second): + path.write_bytes(b"same bytes") + os.utime(path, ns=(SECOND, SECOND)) + link.symlink_to(first) + before = requested_dataset_source_identity("us", str(link)) + link.unlink() + link.symlink_to(second) + + assert requested_dataset_source_identity("us", str(link)) != before + assert before["path"] == str(first.resolve()) + + +# --- Shapes --------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value", + [ + {"uri": "hf://a/b/c.h5@r", "sha256": "a" * 64}, + {"uri": "hf://a/b/c.h5"}, + {"path": "/a/c.h5", "size": 0, "mtime_ns": -1}, + ], +) +def test_the_three_shapes_are_identities(value): + assert is_dataset_source_identity(value) + + +@pytest.mark.parametrize( + "value", + [ + None, + "hf://a/b/c.h5", + [], + {}, + {"uri": None}, + {"uri": "hf://a/b/c.h5", "sha256": None}, + {"uri": "hf://a/b/c.h5", "revision": "r"}, + {"path": "/a/c.h5", "size": True, "mtime_ns": 1}, + {"path": "/a/c.h5", "size": 1, "mtime_ns": 1.0}, + {"path": "/a/c.h5", "size": 1}, + {"path": 1, "size": 1, "mtime_ns": 1}, + {"path": "/a/c.h5", "size": 1, "mtime_ns": 1, "sha256": "a"}, + ], +) +def test_other_values_are_not_identities(value): + assert not is_dataset_source_identity(value) + + +IDENTITIES = st.one_of( + st.builds(lambda uri: {"uri": uri}, st.text()), + st.builds(lambda uri, sha: {"uri": uri, "sha256": sha}, st.text(), st.text()), + st.builds( + lambda path, size, mtime_ns: {"path": path, "size": size, "mtime_ns": mtime_ns}, + st.text(), + st.integers(), + st.integers(), + ), +) + + +@settings(max_examples=300, deadline=None) +@given(identity=IDENTITIES) +def test_describing_an_identity_is_total(identity): + described = describe_dataset_source_identity(identity) + + assert isinstance(described, str) + assert identity.get("uri", identity.get("path")) in described + + +@pytest.mark.parametrize( + "identity, described", + [ + ({"uri": "hf://a/b/c.h5@r"}, "hf://a/b/c.h5@r"), + ( + {"uri": "hf://a/b/c.h5@r", "sha256": "f" * 64}, + f"hf://a/b/c.h5@r (sha256 {'f' * 64})", + ), + ( + {"path": "/a/c.h5", "size": 3, "mtime_ns": 1_760_000_000_123_456_789}, + "/a/c.h5 (3 bytes, modified 2025-10-09T08:53:20.123456789Z)", + ), + ( + {"path": "/a/c.h5", "size": 3, "mtime_ns": -1}, + "/a/c.h5 (3 bytes, modified 1969-12-31T23:59:59.999999999Z)", + ), + ( + {"path": "/a/c.h5", "size": 3, "mtime_ns": 10**30}, + f"/a/c.h5 (3 bytes, modified at {10**30} ns since the epoch)", + ), + ], +) +def test_an_identity_is_described_in_words(identity, described): + assert describe_dataset_source_identity(identity) == described diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 182d39af..589b3369 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -13,6 +13,10 @@ from __future__ import annotations +import os +import re +from importlib import metadata as importlib_metadata +from pathlib import Path from unittest.mock import patch import h5py @@ -27,6 +31,16 @@ pytest.importorskip("spm_calculator.policyengine_adapter") import policyengine as pe # noqa: E402 +import policyengine.provenance.dataset_materialization as dataset_materialization # noqa: E402 +import policyengine.provenance.manifest as manifest_module # noqa: E402 +from policyengine.provenance.dataset_materialization import ( # noqa: E402 + describe_dataset_source_identity, + requested_dataset_source_identity, +) +from policyengine.provenance.manifest import ( # noqa: E402 + ArtifactPathReference, + get_release_manifest, +) from policyengine.tax_benefit_models.us import datasets as us_datasets # noqa: E402 from policyengine.tax_benefit_models.us import legacy_inputs # noqa: E402 from policyengine.tax_benefit_models.us.datasets import ( # noqa: E402 @@ -49,6 +63,7 @@ from policyengine.tax_benefit_models.us.model import ( # noqa: E402 PolicyEngineUSLatest, ) +from policyengine.utils.hashing import sha256_file # noqa: E402 YEAR = 2024 LEGACY = "would_claim_wic" @@ -104,10 +119,16 @@ def _in_memory_dataset(directory, frames=None) -> PolicyEngineUSDataset: ) -def _write_entity_tables(path) -> str: - """Write the household in the certified Populace default's layout.""" +def _write_entity_tables(path, draw=None) -> str: + """Write the household in the certified Populace default's layout. + + ``draw`` replaces the stored draws (``DRAW``), as a new release would. + """ + frames = _frames() + if draw is not None: + frames["person"][LEGACY] = draw with pd.HDFStore(path, mode="w") as store: - for entity, frame in _frames().items(): + for entity, frame in frames.items(): store.put(entity, frame, format="table", data_columns=True) store.put("_time_period", pd.Series([YEAR]), format="table") return str(path) @@ -556,8 +577,20 @@ def _create_year_file(directory): return source, dataset +def _stamp_for(source): + """The stamp a year file cut now from ``source`` carries.""" + return year_file_stamp(requested_dataset_source_identity("us", source)) + + +def _problem(path, source): + """Why a year file is out of date for a request of ``source``.""" + return year_file_stamp_problem( + path, requested_dataset_source_identity("us", source) + ) + + def test_create_datasets_extracts_the_draw_under_the_live_name(tmp_path): - _, dataset = _create_year_file(tmp_path) + source, dataset = _create_year_file(tmp_path) person = pd.DataFrame(dataset.data.person) assert person[LIVE].tolist() == DRAW @@ -565,9 +598,20 @@ def test_create_datasets_extracts_the_draw_under_the_live_name(tmp_path): assert dataset.metadata["legacy_input_renames"] == RENAME # The year file keeps the record, so reloading it keeps the chain. assert read_renames_record(dataset.filepath) == RENAME - # It is stamped with the register it was cut under, so it is current. - assert read_year_file_stamp(dataset.filepath) == year_file_stamp() - assert year_file_stamp_problem(dataset.filepath) is None + # It is stamped with what it was cut from and under, so it is current. + stamp = read_year_file_stamp(dataset.filepath) + assert stamp == _stamp_for(source) + status = os.stat(source) + assert stamp["source"] == { + "path": str(Path(source).resolve()), + "size": status.st_size, + "mtime_ns": status.st_mtime_ns, + } + assert stamp["packages"] == { + name: importlib_metadata.version(name) + for name in ("policyengine-core", "policyengine-us") + } + assert _problem(dataset.filepath, source) is None reloaded = PolicyEngineUSDataset( name=dataset.name, description=dataset.description, @@ -675,13 +719,13 @@ def test_a_year_file_cut_before_a_register_entry_is_regenerated(tmp_path): with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {}): source, dataset = _create_year_file(tmp_path) # Under the register it was cut with, the file is current. - assert year_file_stamp_problem(dataset.filepath) is None + assert _problem(dataset.filepath, source) is None path = dataset.filepath person = pd.DataFrame(dataset.data.person) assert LIVE not in person.columns assert LEGACY not in person.columns assert read_renames_record(path) == {} - assert read_year_file_stamp(path) == {"register": {}, "revision": 1} + assert read_year_file_stamp(path) == {**_stamp_for(source), "register": {}} # The draw is lost from it: a run gives the infant WIC. assert _person_outputs(_run(dataset))["wic"].iloc[INFANT] > 0 @@ -690,7 +734,7 @@ def test_a_year_file_cut_before_a_register_entry_is_regenerated(tmp_path): "it was cut before these renames were registered, so it may have lost " f"them: {LEGACY} -> {LIVE}" ) - assert year_file_stamp_problem(path) == reason + assert _problem(path, source) == reason with pytest.raises(ValueError, match=f"is out of date: {reason}"): load_datasets( datasets=[source], years=[YEAR], data_folder=str(tmp_path / "data") @@ -699,7 +743,7 @@ def test_a_year_file_cut_before_a_register_entry_is_regenerated(tmp_path): regenerated = _ensure(source, tmp_path / "data") assert regenerated.filepath == path - assert read_year_file_stamp(path) == year_file_stamp() + assert read_year_file_stamp(path) == _stamp_for(source) assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW loaded = _assert_reused(source, tmp_path / "data") assert _person_outputs(_run(loaded))["wic"].iloc[INFANT] == 0 @@ -716,7 +760,7 @@ def test_a_future_register_entry_regenerates_year_files(tmp_path): data_folder = tmp_path / "data" with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {**RENAME, **future}): - reason = year_file_stamp_problem(dataset.filepath) + reason = _problem(dataset.filepath, source) assert reason == ( "it was cut before these renames were registered, so it may have " "lost them: fictitious_old_input -> fictitious_new_input" @@ -736,8 +780,8 @@ def spy(*args, **kwargs): assert created == [[YEAR]] assert read_year_file_stamp(dataset.filepath) == { + **_stamp_for(source), "register": {**RENAME, **future}, - "revision": 1, } # Only the entry the engine needs was applied, and the data is as # before. @@ -746,8 +790,8 @@ def spy(*args, **kwargs): _assert_reused(source, data_folder) # Back under this release's register, the newer file is out of date too. - assert "does not register: fictitious_old_input" in year_file_stamp_problem( - dataset.filepath + assert "does not register: fictitious_old_input" in _problem( + dataset.filepath, source ) @@ -765,7 +809,7 @@ def test_ensure_datasets_regenerates_a_year_file_with_only_the_renames_record( regenerated = _ensure(source, data_folder) - assert read_year_file_stamp(regenerated.filepath) == year_file_stamp() + assert read_year_file_stamp(regenerated.filepath) == _stamp_for(source) assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW _assert_reused(source, data_folder) @@ -803,7 +847,7 @@ def recording_stamp(path, value): ("stamp", paths[1]), ] for path in paths: - assert year_file_stamp_problem(path) is None + assert _problem(path, source) is None def test_a_year_file_whose_stamp_was_never_written_is_out_of_date(tmp_path): @@ -825,9 +869,156 @@ def stop(path, value): # The tables and the record were written, but without the stamp the file # is cut again. assert read_renames_record(path) == RENAME - assert year_file_stamp_problem(path).startswith("it has no stamp") + assert _problem(path, source).startswith("it has no stamp") assert _ensure(source, tmp_path / "data").filepath == str(path) - assert year_file_stamp_problem(path) is None + assert _problem(path, source) is None + + +# --- The year-file stamp: the source and packages a year file was cut from + + +def test_editing_a_local_source_cuts_its_year_files_again(tmp_path): + """A local dataset rewritten in place keeps its name, and so its year + files' names, but not its identity.""" + source, dataset = _create_year_file(tmp_path) + data_folder = tmp_path / "data" + cut_from = requested_dataset_source_identity("us", source) + assert read_year_file_stamp(dataset.filepath)["source"] == cut_from + + # The infant's stored draw becomes True. The rewrite is given a later + # modification time of its own, since a file system can keep coarser + # times than a quick rewrite in a test needs. + _write_entity_tables(source, draw=[False, True, True]) + later = cut_from["mtime_ns"] + 1_000_000_000 + os.utime(source, ns=(later, later)) + now = requested_dataset_source_identity("us", source) + assert now["path"] == cut_from["path"] and now != cut_from + + reason = ( + f"it was cut from {describe_dataset_source_identity(cut_from)}, and the " + f"requested source is {describe_dataset_source_identity(now)}" + ) + assert _problem(dataset.filepath, source) == reason + with pytest.raises(ValueError, match=re.escape(reason)): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + # The year file still holds the old draw, which a run would use. + assert pd.DataFrame(dataset.data.person)[LIVE].tolist() == DRAW + + regenerated = _ensure(source, data_folder) + + assert regenerated.filepath == dataset.filepath + assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == [False, True, True] + assert read_year_file_stamp(dataset.filepath)["source"] == now + loaded = _assert_reused(source, data_folder) + assert _person_outputs(_run(loaded))["wic"].iloc[INFANT] > 0 + + +#: A logical dataset name that a fixture bundle pins, as the certified bundle +#: pins ``populace_us_2024_acs_local``. +BUNDLE_NAME = "tiny_bundle" +BUNDLE_URI = f"hf://fixture/tiny-bundle/{BUNDLE_NAME}.h5" + + +def _pinned(revision, sha256): + """Patch the bundled US manifest to pin ``BUNDLE_NAME`` as given, and + refuse every download: a verified copy is already in the data folder.""" + manifest = get_release_manifest("us").model_copy(deep=True) + manifest.datasets[BUNDLE_NAME] = ArtifactPathReference( + path=f"{BUNDLE_NAME}.h5", + repo_id="fixture/tiny-bundle", + repo_type="dataset", + revision=revision, + sha256=sha256, + ) + + def refuse(*args, **kwargs): + raise AssertionError("a bundle dataset was downloaded") + + return ( + patch.object(manifest_module, "get_release_manifest", lambda _: manifest), + patch.object( + dataset_materialization, "get_release_manifest", lambda _: manifest + ), + patch.object(dataset_materialization.requests, "get", refuse), + ) + + +def test_a_re_pinned_bundle_dataset_cuts_its_year_files_again(tmp_path): + """The ACS-local replacement in miniature: the bundle points a logical + name at a new release, and year files cut from the old one are not + reused, although their names are unchanged.""" + data_folder = tmp_path / "data" + data_folder.mkdir() + # ``materialize_dataset`` keeps a verified copy at the artifact's name in + # the data folder, and downloads only when that copy's sha256 differs. + bundled = data_folder / f"{BUNDLE_NAME}.h5" + _write_entity_tables(bundled) + old_sha256 = sha256_file(bundled) + old = {"uri": f"{BUNDLE_URI}@release-1", "sha256": old_sha256} + + first, second, third = _pinned("release-1", old_sha256) + with first, second, third: + dataset = _ensure(BUNDLE_NAME, data_folder) + assert dataset.filepath == str(data_folder / f"{BUNDLE_NAME}_year_{YEAR}.h5") + assert read_year_file_stamp(dataset.filepath)["source"] == old + assert pd.DataFrame(dataset.data.person)[LIVE].tolist() == DRAW + _assert_reused(BUNDLE_NAME, data_folder) + + # The replacement release stores the infant's draw as True; its verified + # copy replaces the old one, as a download would. + _write_entity_tables(bundled, draw=[False, True, True]) + new_sha256 = sha256_file(bundled) + new = {"uri": f"{BUNDLE_URI}@release-2", "sha256": new_sha256} + assert new_sha256 != old_sha256 + + first, second, third = _pinned("release-2", new_sha256) + with first, second, third: + reason = ( + f"it was cut from {BUNDLE_URI}@release-1 (sha256 {old_sha256}), and " + f"the requested source is {BUNDLE_URI}@release-2 (sha256 {new_sha256})" + ) + assert _problem(dataset.filepath, BUNDLE_NAME) == reason + with pytest.raises(ValueError, match=re.escape(reason)): + load_datasets( + datasets=[BUNDLE_NAME], years=[YEAR], data_folder=str(data_folder) + ) + + regenerated = _ensure(BUNDLE_NAME, data_folder) + + assert regenerated.filepath == dataset.filepath + assert read_year_file_stamp(regenerated.filepath)["source"] == new + assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == [ + False, + True, + True, + ] + loaded = _assert_reused(BUNDLE_NAME, data_folder) + assert _person_outputs(_run(loaded))["wic"].iloc[INFANT] > 0 + + +def test_another_package_version_cuts_year_files_again(tmp_path, monkeypatch): + """Year files cut under one policyengine-us are not reused under another.""" + source, dataset = _create_year_file(tmp_path) + data_folder = tmp_path / "data" + installed = legacy_inputs.installed_package_versions() + assert installed["policyengine-us"] == importlib_metadata.version("policyengine-us") + upgraded = {**installed, "policyengine-us": "99.0.0"} + monkeypatch.setattr( + legacy_inputs, "installed_package_versions", lambda: dict(upgraded) + ) + + reason = ( + f"it was cut with policyengine-us {installed['policyengine-us']}, and " + "policyengine-us 99.0.0 is installed" + ) + assert _problem(dataset.filepath, source) == reason + with pytest.raises(ValueError, match=re.escape(reason)): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + + regenerated = _ensure(source, data_folder) + + assert read_year_file_stamp(regenerated.filepath)["packages"] == upgraded + _assert_reused(source, data_folder) # --- Saved outputs and run records keep the record ---------------------- diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index 7d11eef0..640aa526 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -1,38 +1,47 @@ """Unit and property tests for the stamp on US year files. -``create_datasets`` stamps each year file with the register of renamed stored -inputs it was cut under and a year-file format revision -(``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``). ``ensure_datasets`` cuts a -year file again, and ``load_datasets`` refuses it, unless the stamp matches -this release's. These tests state the check's invariants over generated -stamps, records and registers, without a country engine: +``create_datasets`` stamps each year file with the identity of the source +dataset it was cut from, the policyengine-core and policyengine-us versions +that calculated it, the register of renamed stored inputs it was cut under +and a year-file format revision (``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``). +``ensure_datasets`` cuts a year file again, and ``load_datasets`` refuses it, +unless the stamp matches the one a file cut now from the requested dataset +would carry. These tests state the check's invariants over generated stamps, +records, registers, sources and package versions, without a country engine: - **Exactness.** A year file is current exactly when its stamp has this - release's format revision and register, and its renames record lists only - renames in that register. + release's format revision and register, the requested source and the + installed package versions, and its renames record lists only renames in + that register. - **Totality.** Whatever the stamp and record hold, the check returns a reason or ``None``, and raises only for errors unrelated to the file's contents (no permission, a directory, a lock). - **Round trip.** A stamp reads back as written, and writing again replaces it. -- **Named renames.** The reason names every rename the file's register and - this release's disagree on, from the register alone. +- **Named differences.** The reason names every rename the file's register + and this release's disagree on, both sources when they differ, and every + package whose version differs, and nothing that agrees. - **No partial file is current.** Every truncation of a current year file, every save that stops part-way and every save that does not stamp is out of date. - **Ensure agrees with load.** ``ensure_datasets`` cuts the year files again exactly when one is missing or ``load_datasets`` would refuse one, and - decides from the stamps alone. + decides from the stamps alone: it never reads, downloads or hashes the + source, and never imports policyengine-us. -``test_us_legacy_inputs_integration.py`` runs ``create_datasets`` and +``test_dataset_source_identity.py`` tests the source identities themselves, +and ``test_us_legacy_inputs_integration.py`` runs ``create_datasets`` and ``ensure_datasets`` over a real policyengine-us engine. """ from __future__ import annotations +import builtins import functools +import io import json import os import tempfile +from importlib import metadata as importlib_metadata from pathlib import Path from unittest.mock import MagicMock, patch @@ -44,6 +53,11 @@ from hypothesis import strategies as st from microdf import MicroDataFrame +import policyengine.provenance.dataset_materialization as dataset_materialization +from policyengine.provenance.dataset_materialization import ( + describe_dataset_source_identity, + requested_dataset_source_identity, +) from policyengine.tax_benefit_models.us import datasets as us_datasets from policyengine.tax_benefit_models.us import legacy_inputs from policyengine.tax_benefit_models.us.datasets import ( @@ -56,6 +70,7 @@ from policyengine.tax_benefit_models.us.legacy_inputs import ( LEGACY_INPUT_RENAMES, RENAMES_H5_DATASET, + YEAR_FILE_PACKAGES, YEAR_FILE_STAMP_H5_DATASET, YEAR_FILE_STAMP_REVISION, read_renames_record, @@ -70,15 +85,48 @@ LIVE = "takes_up_wic_if_eligible" REVISION = YEAR_FILE_STAMP_REVISION #: ``hf://`` references are used as they are, and name year files after the -#: file's stem, so no manifest or download is involved. +#: file's stem, so no manifest or download is involved. Its identity is the +#: URI as given. SOURCE = "hf://fixture/tiny.h5" +IDENTITY = {"uri": SOURCE} STEM = "tiny" GROUPS = ("household", "tax_unit", "spm_unit", "family", "marital_unit") +#: The installed versions every test in this module sees (see +#: ``_installed_packages``), so no expectation depends on the environment. +PACKAGES = {"policyengine-core": "3.33.0", "policyengine-us": "2.2.1"} +#: The real reader, kept before ``_installed_packages`` replaces it. +REAL_INSTALLED_PACKAGE_VERSIONS = legacy_inputs.installed_package_versions +MALFORMED = "its stamp of what it was cut from and under is malformed" #: A small alphabet, so generated registers often share names and disagree #: only on some entries. NAMES = st.sampled_from(["alpha", "beta", "gamma", "delta"]) REGISTERS = st.dictionaries(NAMES, NAMES, max_size=3) +#: Small alphabets of source identities in each of their three shapes (see +#: ``dataset_source_identity``), so generated sources often agree. +URIS = st.sampled_from([SOURCE, f"{SOURCE}@r1", f"{SOURCE}@r2"]) +SOURCES = st.one_of( + st.builds(lambda uri: {"uri": uri}, URIS), + st.builds( + lambda uri, sha256: {"uri": uri, "sha256": sha256}, + URIS, + st.sampled_from(["a" * 64, "b" * 64]), + ), + st.builds( + lambda path, size, mtime_ns: { + "path": path, + "size": size, + "mtime_ns": mtime_ns, + }, + st.sampled_from(["/data/tiny.h5", "/elsewhere/tiny.h5"]), + st.sampled_from([10, 11]), + # Before the epoch, at it, and past the year 9999. + st.sampled_from([-1, 0, 1_760_000_000_123_456_789, 10**20]), + ), +) +PACKAGE_NAMES = st.sampled_from([*YEAR_FILE_PACKAGES, "extra-package"]) +VERSIONS = st.one_of(st.none(), st.sampled_from(["1.0", "1.1"])) +PACKAGE_MAPS = st.dictionaries(PACKAGE_NAMES, VERSIONS, max_size=3) #: Raw node contents: ``_ABSENT`` for no node, a JSON value, or a node the #: reader cannot parse as JSON. _ABSENT = object() @@ -92,6 +140,19 @@ ) +@pytest.fixture(autouse=True, scope="module") +def _installed_packages(): + """Pin the installed package versions a stamp records to ``PACKAGES``. + + Module-scoped, so Hypothesis tests can run under it. + """ + with pytest.MonkeyPatch.context() as patcher: + patcher.setattr( + legacy_inputs, "installed_package_versions", lambda: dict(PACKAGES) + ) + yield + + def _frames() -> dict[str, pd.DataFrame]: person = pd.DataFrame( { @@ -160,7 +221,12 @@ def _write_year_file(path, *, stamp=_ABSENT, record=_ABSENT): return Path(path) -def _is_current(stamp, record, register) -> bool: +def _same_json(left, right) -> bool: + """Equal as JSON: ``True`` is not ``1``, and ``1.0`` is not ``1``.""" + return json.dumps(left, sort_keys=True) == json.dumps(right, sort_keys=True) + + +def _is_current(stamp, record, register, source=IDENTITY, packages=PACKAGES) -> bool: """The exactness invariant, stated directly on the stored JSON values.""" def str_map(value): @@ -171,54 +237,109 @@ def str_map(value): return ( isinstance(stamp, dict) - and set(stamp) == {"register", "revision"} + and set(stamp) == {"packages", "register", "revision", "source"} and type(stamp["revision"]) is int and stamp["revision"] == REVISION and str_map(stamp["register"]) and stamp["register"] == register + and _same_json(stamp["source"], source) + and _same_json(stamp["packages"], packages) and str_map(record) and record.items() <= stamp["register"].items() ) -def _stamp(register, revision=REVISION): - return {"register": dict(register), "revision": revision} +def _stamp(register, revision=REVISION, source=IDENTITY, packages=PACKAGES): + return { + "packages": dict(packages), + "register": dict(register), + "revision": revision, + "source": dict(source), + } + + +def _current(): + return year_file_stamp(IDENTITY) # --- The stamp ---------------------------------------------------------- -def test_the_stamp_holds_the_whole_register_and_the_format_revision(): - assert year_file_stamp() == { +def test_the_stamp_holds_the_source_packages_register_and_revision(): + assert _current() == { + "packages": PACKAGES, "register": {LEGACY: LIVE}, - "revision": 1, + "revision": 2, + "source": IDENTITY, } - assert year_file_stamp()["register"] == LEGACY_INPUT_RENAMES + assert _current()["register"] == LEGACY_INPUT_RENAMES -def test_the_stamp_reads_the_register_when_called(monkeypatch): - """datasets.py must not freeze the register at import.""" +def test_the_stamp_reads_the_register_and_packages_when_called(monkeypatch): + """datasets.py must not freeze the register or the versions at import.""" monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", {"b": "c", "a": "d"}) + monkeypatch.setattr( + legacy_inputs, "installed_package_versions", lambda: {"z": None, "y": "1"} + ) - assert year_file_stamp() == {"register": {"a": "d", "b": "c"}, "revision": 1} + assert year_file_stamp(IDENTITY) == { + "packages": {"y": "1", "z": None}, + "register": {"a": "d", "b": "c"}, + "revision": 2, + "source": IDENTITY, + } + + +def test_explicit_packages_override_the_installed_ones(): + assert year_file_stamp(IDENTITY, {"policyengine-us": "9"})["packages"] == { + "policyengine-us": "9" + } def test_the_stamp_is_stored_as_utf8_json(tmp_path): path = _write_year_file(tmp_path / "year.h5") - write_year_file_stamp(path, year_file_stamp()) + write_year_file_stamp(path, _current()) with h5py.File(path, "r") as file: assert file[YEAR_FILE_STAMP_H5_DATASET].asstr()[()] == ( - f'{{"register": {{"{LEGACY}": "{LIVE}"}}, "revision": 1}}' + '{"packages": {"policyengine-core": "3.33.0", "policyengine-us": ' + f'"2.2.1"}}, "register": {{"{LEGACY}": "{LIVE}"}}, "revision": 2, ' + f'"source": {{"uri": "{SOURCE}"}}}}' ) -@settings(max_examples=50, deadline=None) -@given( - first=st.fixed_dictionaries({"register": REGISTERS, "revision": st.integers()}), - second=st.fixed_dictionaries({"register": REGISTERS, "revision": st.integers()}), +def test_installed_package_versions_reads_metadata(monkeypatch): + """The real reader: distribution metadata, ``None`` when absent.""" + monkeypatch.setattr( + legacy_inputs, + "YEAR_FILE_PACKAGES", + ("pytest", "hypothesis", "not-a-distribution-policyengine-test"), + ) + + assert REAL_INSTALLED_PACKAGE_VERSIONS() == { + "pytest": importlib_metadata.version("pytest"), + "hypothesis": importlib_metadata.version("hypothesis"), + "not-a-distribution-policyengine-test": None, + } + + +def test_the_stamp_records_policyengine_core_and_policyengine_us(): + assert YEAR_FILE_PACKAGES == ("policyengine-core", "policyengine-us") + + +STAMPS = st.fixed_dictionaries( + { + "packages": PACKAGE_MAPS, + "register": REGISTERS, + "revision": st.integers(), + "source": SOURCES, + } ) + + +@settings(max_examples=50, deadline=None) +@given(first=STAMPS, second=STAMPS) def test_a_year_file_stamp_reads_back_as_written(first, second): """Round trip: a stamp reads back exactly, a new one replaces it, and the tables and record are left alone.""" @@ -244,7 +365,7 @@ def test_a_missing_or_foreign_file_has_no_stamp(tmp_path): def test_the_stamp_does_not_disturb_pandas_or_the_layout_check(tmp_path): - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) with pd.HDFStore(path, mode="r") as store: assert sorted(store.keys()) == sorted(f"/{key}" for key in _frames()) @@ -261,8 +382,11 @@ def test_the_stamp_does_not_disturb_pandas_or_the_layout_check(tmp_path): @st.composite def stamped_files(draw): - """A current register, and a year file's stamp and record near it.""" + """A current register, source and packages, and a year file's stamp and + record near them.""" register = draw(REGISTERS) + source = draw(SOURCES) + packages = draw(PACKAGE_MAPS) stamped = draw( st.one_of( st.just(dict(register)), @@ -275,7 +399,20 @@ def stamped_files(draw): st.builds(lambda k, v: {**register, k: v}, NAMES, NAMES), ) ) - revision = draw(st.sampled_from([REVISION, REVISION, 0, 2, -1, 99])) + stamped_source = draw(st.one_of(st.just(dict(source)), SOURCES)) + stamped_packages = draw( + st.one_of( + st.just(dict(packages)), + PACKAGE_MAPS, + # One version changed, dropped or added. + st.builds(lambda k, v: {**packages, k: v}, PACKAGE_NAMES, VERSIONS), + st.builds( + lambda name: {k: v for k, v in packages.items() if k != name}, + PACKAGE_NAMES, + ), + ) + ) + revision = draw(st.sampled_from([REVISION, REVISION, 0, 1, 3, -1, 99])) record = draw( st.one_of( st.just(_ABSENT), @@ -287,24 +424,27 @@ def stamped_files(draw): REGISTERS, ) ) - return register, _stamp(stamped, revision), record + stamp = _stamp(stamped, revision, stamped_source, stamped_packages) + return register, source, packages, stamp, record -@settings(max_examples=300, deadline=None) +@settings(max_examples=400, deadline=None) @given(case=stamped_files()) -def test_a_year_file_is_current_only_under_this_releases_exact_stamp(case): +def test_a_year_file_is_current_only_under_the_exact_stamp_requested(case): """Exactness, and refinement of the record check a year file had before. A subset rule would accept a file cut by a newer release under a larger register, which stores live names this release's engine may not define. """ - register, stamp, record = case + register, source, packages, stamp, record = case with tempfile.TemporaryDirectory() as directory: path = _write_year_file(Path(directory) / "year.h5", stamp=stamp, record=record) with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): - problem = year_file_stamp_problem(path) + problem = year_file_stamp_problem(path, source, packages) - assert (problem is None) == _is_current(stamp, record, register) + assert (problem is None) == _is_current( + stamp, record, register, source, packages + ) assert problem is None or (isinstance(problem, str) and problem) if problem is None: # Every file the stamp accepts, the record check accepted too. @@ -323,10 +463,20 @@ def test_a_year_file_is_current_only_under_this_releases_exact_stamp(case): JSON_VALUES, st.fixed_dictionaries( { + "packages": st.one_of(PACKAGE_MAPS, JSON_VALUES), "register": st.one_of(REGISTERS, JSON_VALUES), "revision": st.one_of( st.just(REVISION), st.booleans(), st.floats(), JSON_VALUES ), + "source": st.one_of( + SOURCES, + JSON_VALUES, + # Well-formed keys holding values of the wrong type. + st.fixed_dictionaries( + {"path": JSON_VALUES, "size": JSON_VALUES, "mtime_ns": JSON_VALUES} + ), + st.fixed_dictionaries({"uri": JSON_VALUES, "sha256": JSON_VALUES}), + ), } ), st.tuples(st.just("text"), RAW_TEXT), @@ -337,38 +487,62 @@ def test_a_year_file_is_current_only_under_this_releases_exact_stamp(case): ) -@settings(max_examples=300, deadline=None) -@given(stamp=NODES, record=NODES, register=REGISTERS) +@settings(max_examples=400, deadline=None) +@given( + stamp=NODES, + record=NODES, + register=REGISTERS, + source=SOURCES, + packages=PACKAGE_MAPS, +) def test_the_year_file_problem_is_total_and_exact_for_any_content( - stamp, record, register + stamp, record, register, source, packages ): """Totality: any stamp and record give a reason or ``None``, never an - error, and ``None`` only for this release's exact stamp.""" + error, and ``None`` only for the exact stamp requested.""" with tempfile.TemporaryDirectory() as directory: path = _write_year_file(Path(directory) / "year.h5", stamp=stamp, record=record) with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): - problem = year_file_stamp_problem(path) + problem = year_file_stamp_problem(path, source, packages) assert problem is None or (isinstance(problem, str) and problem) if problem is None: - assert _is_current(stamp, record, register) + assert _is_current(stamp, record, register, source, packages) @pytest.mark.parametrize( "stamp", [ - {"register": {LEGACY: LIVE}, "revision": True}, - {"register": {LEGACY: LIVE}, "revision": 1.0}, - {"register": {LEGACY: LIVE}, "revision": "1"}, - {"register": {LEGACY: LIVE}}, - {"register": {LEGACY: LIVE}, "revision": 1, "extra": None}, - {"register": {LEGACY: 5}, "revision": 1}, - {"register": [[LEGACY, LIVE]], "revision": 1}, - [{"register": {LEGACY: LIVE}, "revision": 1}], + {**_stamp({LEGACY: LIVE}), "revision": True}, + {**_stamp({LEGACY: LIVE}), "revision": 2.0}, + {**_stamp({LEGACY: LIVE}), "revision": "2"}, + {"register": {LEGACY: LIVE}, "revision": 2}, + {k: v for k, v in _stamp({LEGACY: LIVE}).items() if k != "source"}, + {k: v for k, v in _stamp({LEGACY: LIVE}).items() if k != "packages"}, + {**_stamp({LEGACY: LIVE}), "extra": None}, + {**_stamp({LEGACY: LIVE}), "register": {LEGACY: 5}}, + {**_stamp({LEGACY: LIVE}), "register": [[LEGACY, LIVE]]}, + {**_stamp({LEGACY: LIVE}), "packages": {"policyengine-us": 2.2}}, + {**_stamp({LEGACY: LIVE}), "packages": [["policyengine-us", "2.2.1"]]}, + {**_stamp({LEGACY: LIVE}), "source": SOURCE}, + {**_stamp({LEGACY: LIVE}), "source": {}}, + {**_stamp({LEGACY: LIVE}), "source": {"uri": 5}}, + {**_stamp({LEGACY: LIVE}), "source": {"uri": SOURCE, "sha256": None}}, + {**_stamp({LEGACY: LIVE}), "source": {"uri": SOURCE, "revision": "r1"}}, + { + **_stamp({LEGACY: LIVE}), + "source": {"path": "/a", "size": True, "mtime_ns": 1}, + }, + { + **_stamp({LEGACY: LIVE}), + "source": {"path": "/a", "size": 1, "mtime_ns": 1.0}, + }, + {**_stamp({LEGACY: LIVE}), "source": {"path": "/a", "size": 1}}, + [_stamp({LEGACY: LIVE})], ("text", "not json"), ("text", ""), ("bytes", b"\xff\xfe"), - ("number", 1), + ("number", 2), ("array", ["{}"]), ("group", None), ], @@ -376,21 +550,33 @@ def test_the_year_file_problem_is_total_and_exact_for_any_content( def test_a_malformed_stamp_is_out_of_date(tmp_path, stamp): path = _write_year_file(tmp_path / "year.h5", stamp=stamp, record={}) - assert year_file_stamp_problem(path) == ( - "its stamp of the register of renamed inputs it was cut under is malformed" - ) + assert year_file_stamp_problem(path, IDENTITY) == MALFORMED -@pytest.mark.parametrize("revision", [0, 2, -1]) +@pytest.mark.parametrize("revision", [0, 1, 3, -1]) def test_a_year_file_from_another_format_revision_is_out_of_date(tmp_path, revision): - # A later revision may shape the rest of its stamp and record differently, + # Another revision may shape the rest of its stamp and record differently, # and is still reported by revision rather than as malformed. stamp = {"revision": revision, "registers": [[LEGACY, LIVE]]} path = _write_year_file(tmp_path / "year.h5", stamp=stamp) - assert year_file_stamp_problem(path) == ( + assert year_file_stamp_problem(path, IDENTITY) == ( f"it was cut with year-file format revision {revision}, and this " - "version of policyengine.py reads only revision 1" + "version of policyengine.py reads only revision 2" + ) + + +def test_a_year_file_stamped_by_revision_1_is_cut_again_once(tmp_path): + """The stamp #590 wrote holds no source or package versions.""" + path = _write_year_file( + tmp_path / "year.h5", + stamp={"register": dict(LEGACY_INPUT_RENAMES), "revision": 1}, + record={}, + ) + + assert year_file_stamp_problem(path, IDENTITY) == ( + "it was cut with year-file format revision 1, and this version of " + "policyengine.py reads only revision 2" ) @@ -399,12 +585,12 @@ def test_files_without_a_stamp_are_out_of_date_with_distinct_reasons(tmp_path): oldest = _write_year_file(tmp_path / "oldest.h5") record_only = _write_year_file(tmp_path / "record_only.h5", record={LEGACY: LIVE}) - assert year_file_stamp_problem(oldest) == ( + assert year_file_stamp_problem(oldest, IDENTITY) == ( "it has no record of the renamed stored inputs mapped when it was cut, " "so it was cut by an older policyengine.py or not written completely, " f"and may have lost any of: {LEGACY} -> {LIVE}" ) - assert year_file_stamp_problem(record_only) == ( + assert year_file_stamp_problem(record_only, IDENTITY) == ( "it has no stamp of the register of renamed inputs it was cut under, so " f"it may have lost any of: {LEGACY} -> {LIVE}" ) @@ -417,11 +603,11 @@ def test_an_unstamped_file_is_out_of_date_even_with_an_empty_register( oldest = _write_year_file(tmp_path / "oldest.h5") record_only = _write_year_file(tmp_path / "record_only.h5", record={}) - assert year_file_stamp_problem(oldest) == ( + assert year_file_stamp_problem(oldest, IDENTITY) == ( "it has no record of the renamed stored inputs mapped when it was cut, " "so it was cut by an older policyengine.py or not written completely" ) - assert year_file_stamp_problem(record_only) == ( + assert year_file_stamp_problem(record_only, IDENTITY) == ( "it has no stamp of the register of renamed inputs it was cut under" ) @@ -470,36 +656,34 @@ def test_an_unstamped_file_is_out_of_date_even_with_an_empty_register( ], ) def test_a_stamped_file_needs_a_record_within_its_register(tmp_path, record, reason): - path = _write_year_file( - tmp_path / "year.h5", stamp=year_file_stamp(), record=record - ) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record=record) - assert year_file_stamp_problem(path) == reason + assert year_file_stamp_problem(path, IDENTITY) == reason def test_a_stamp_and_record_stored_as_fixed_length_strings_are_read(tmp_path): - stamp = json.dumps(year_file_stamp()).encode() + stamp = json.dumps(_current()).encode() record = json.dumps({LEGACY: LIVE}).encode() path = _write_year_file( tmp_path / "year.h5", stamp=("bytes", stamp), record=("bytes", record) ) - assert year_file_stamp_problem(path) is None - assert read_year_file_stamp(path) == year_file_stamp() + assert year_file_stamp_problem(path, IDENTITY) is None + assert read_year_file_stamp(path) == _current() assert read_renames_record(path) == {LEGACY: LIVE} def test_raising_the_format_revision_makes_older_files_out_of_date( tmp_path, monkeypatch ): - older = _write_year_file(tmp_path / "older.h5", stamp=year_file_stamp(), record={}) + older = _write_year_file(tmp_path / "older.h5", stamp=_current(), record={}) monkeypatch.setattr(legacy_inputs, "YEAR_FILE_STAMP_REVISION", 7) - newer = _write_year_file(tmp_path / "newer.h5", stamp=year_file_stamp(), record={}) + newer = _write_year_file(tmp_path / "newer.h5", stamp=_current(), record={}) - assert year_file_stamp()["revision"] == 7 - assert year_file_stamp_problem(newer) is None - assert year_file_stamp_problem(older) == ( - "it was cut with year-file format revision 1, and this version of " + assert _current()["revision"] == 7 + assert year_file_stamp_problem(newer, IDENTITY) is None + assert year_file_stamp_problem(older, IDENTITY) == ( + "it was cut with year-file format revision 2, and this version of " "policyengine.py reads only revision 7" ) @@ -516,22 +700,18 @@ def test_raising_the_format_revision_makes_older_files_out_of_date( @pytest.mark.parametrize( "stamp, record, reason", [ + (("text", DEEP_JSON), {}, MALFORMED), ( - ("text", DEEP_JSON), - {}, - "its stamp of the register of renamed inputs it was cut under is malformed", - ), - ( - year_file_stamp(), + _current(), ("text", DEEP_JSON), "its record of the renamed stored inputs mapped when it was cut is " "malformed", ), ( - {"revision": 2}, + {"revision": 3}, ("text", DEEP_JSON), - "it was cut with year-file format revision 2, and this version of " - "policyengine.py reads only revision 1", + "it was cut with year-file format revision 3, and this version of " + "policyengine.py reads only revision 2", ), ], ) @@ -540,7 +720,7 @@ def test_json_nested_too_deeply_is_out_of_date_not_an_error( ): path = _write_year_file(tmp_path / "year.h5", stamp=stamp, record=record) - assert year_file_stamp_problem(path) == reason + assert year_file_stamp_problem(path, IDENTITY) == reason if isinstance(stamp, tuple): # Unreadable (None) or parsed into nested lists, but never a stamp. assert not isinstance(read_year_file_stamp(path), dict) @@ -554,7 +734,7 @@ def test_a_recursion_error_while_parsing_is_unreadable_content( tmp_path, monkeypatch, name ): """The RecursionError path, on every platform (see ``DEEP_JSON``).""" - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) loads = json.loads def deep(text, *args, **kwargs): @@ -565,8 +745,8 @@ def deep(text, *args, **kwargs): monkeypatch.setattr(legacy_inputs.json, "loads", deep) - assert year_file_stamp_problem(path) == ( - "its stamp of the register of renamed inputs it was cut under is malformed" + assert year_file_stamp_problem(path, IDENTITY) == ( + MALFORMED if name == YEAR_FILE_STAMP_H5_DATASET else "its record of the renamed stored inputs mapped when it was cut is " "malformed" @@ -581,7 +761,7 @@ def test_dangling_links_in_place_of_the_stamp_and_record_read_as_absent(tmp_path str(tmp_path / "missing.h5"), "/record" ) - assert year_file_stamp_problem(path).startswith( + assert year_file_stamp_problem(path, IDENTITY).startswith( "it has no record of the renamed stored inputs mapped when it was cut" ) assert read_year_file_stamp(path) is None @@ -589,15 +769,13 @@ def test_dangling_links_in_place_of_the_stamp_and_record_read_as_absent(tmp_path @pytest.mark.parametrize("record", [{}, {LEGACY: LIVE}]) -def test_a_file_stamped_with_this_releases_register_is_current(tmp_path, record): - path = _write_year_file( - tmp_path / "year.h5", stamp=year_file_stamp(), record=record - ) +def test_a_file_stamped_for_the_requested_source_is_current(tmp_path, record): + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record=record) - assert year_file_stamp_problem(path) is None + assert year_file_stamp_problem(path, IDENTITY) is None -# --- Named renames ------------------------------------------------------ +# --- Named differences -------------------------------------------------- @settings(max_examples=200, deadline=None) @@ -610,34 +788,217 @@ def test_a_stale_year_file_names_each_rename_it_may_lack(register, stamped): Path(directory) / "year.h5", stamp=_stamp(stamped), record={} ) with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): - problem = year_file_stamp_problem(path) + problem = year_file_stamp_problem(path, IDENTITY) if stamped == register: assert problem is None return missing = {k: v for k, v in register.items() if stamped.get(k) != v} unknown = {k: v for k, v in stamped.items() if register.get(k) != v} + clauses = problem.split("; ") for legacy, live in missing.items(): - assert f"{legacy} -> {live}" in problem.split("; ")[0] + assert f"{legacy} -> {live}" in clauses[0] for legacy, live in unknown.items(): - assert f"{legacy} -> {live}" in problem.split("; ")[-1] + assert f"{legacy} -> {live}" in clauses[-1] assert ("cut before these renames were registered" in problem) == bool(missing) assert ("does not register" in problem) == bool(unknown) for legacy, live in register.items() & stamped.items(): assert f"{legacy} -> {live}" not in problem +@settings(max_examples=300, deadline=None) +@given( + source=SOURCES, + stamped_source=SOURCES, + packages=PACKAGE_MAPS, + stamped_packages=PACKAGE_MAPS, +) +def test_a_stale_year_file_names_each_source_and_version_that_differs( + source, stamped_source, packages, stamped_packages +): + """The reason has a source clause exactly when the sources differ, naming + both, and one clause for each package whose version differs or that only + one side records, naming it; nothing that agrees is mentioned.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file( + Path(directory) / "year.h5", + stamp=_stamp( + LEGACY_INPUT_RENAMES, REVISION, stamped_source, stamped_packages + ), + record={}, + ) + problem = year_file_stamp_problem(path, source, packages) + + differing = { + name + for name in stamped_packages.keys() | packages.keys() + if name not in stamped_packages + or name not in packages + or stamped_packages[name] != packages[name] + } + if stamped_source == source and not differing: + assert problem is None + return + clauses = problem.split("; ") + source_clauses = [c for c in clauses if c.startswith("it was cut from ")] + if stamped_source == source: + assert source_clauses == [] + else: + assert source_clauses == [ + f"it was cut from {describe_dataset_source_identity(stamped_source)}, " + "and the requested source is " + f"{describe_dataset_source_identity(source)}" + ] + package_clauses = [c for c in clauses if c not in source_clauses] + assert len(package_clauses) == len(differing) + for name, clause in zip(sorted(differing), package_clauses): + assert name in clause + for other in (stamped_packages.keys() | packages.keys()) - {name}: + assert other not in clause + + +@pytest.mark.parametrize( + "stamped_source, source, described", + [ + ( + {"uri": "hf://org/repo/acs.h5@old", "sha256": "a" * 64}, + {"uri": "hf://org/repo/acs.h5@new", "sha256": "b" * 64}, + f"hf://org/repo/acs.h5@old (sha256 {'a' * 64}), and the requested " + f"source is hf://org/repo/acs.h5@new (sha256 {'b' * 64})", + ), + ( + # A re-upload under the same revision name. + {"uri": "hf://org/repo/acs.h5@r", "sha256": "a" * 64}, + {"uri": "hf://org/repo/acs.h5@r", "sha256": "b" * 64}, + f"hf://org/repo/acs.h5@r (sha256 {'a' * 64}), and the requested " + f"source is hf://org/repo/acs.h5@r (sha256 {'b' * 64})", + ), + ( + {"uri": "hf://org/repo/acs.h5@r", "sha256": "a" * 64}, + {"uri": "hf://org/repo/acs.h5@r"}, + f"hf://org/repo/acs.h5@r (sha256 {'a' * 64}), and the requested " + "source is hf://org/repo/acs.h5@r", + ), + ( + { + "path": "/data/tiny.h5", + "size": 10, + "mtime_ns": 1_760_000_000_000_000_001, + }, + { + "path": "/data/tiny.h5", + "size": 12, + "mtime_ns": 1_760_000_060_000_000_000, + }, + "/data/tiny.h5 (10 bytes, modified 2025-10-09T08:53:20.000000001Z), " + "and the requested source is /data/tiny.h5 (12 bytes, modified " + "2025-10-09T08:54:20.000000000Z)", + ), + ( + {"path": "/a/tiny.h5", "size": 10, "mtime_ns": 0}, + {"path": "/b/tiny.h5", "size": 10, "mtime_ns": 0}, + "/a/tiny.h5 (10 bytes, modified 1970-01-01T00:00:00.000000000Z), and " + "the requested source is /b/tiny.h5 (10 bytes, modified " + "1970-01-01T00:00:00.000000000Z)", + ), + ], +) +def test_a_year_file_cut_from_another_source_names_both( + tmp_path, stamped_source, source, described +): + path = _write_year_file( + tmp_path / "year.h5", + stamp=_stamp(LEGACY_INPUT_RENAMES, source=stamped_source), + record={}, + ) + + assert year_file_stamp_problem(path, source) == f"it was cut from {described}" + + +@pytest.mark.parametrize( + "stamped_packages, packages, reason", + [ + ( + {**PACKAGES, "policyengine-us": "2.2.0"}, + PACKAGES, + "it was cut with policyengine-us 2.2.0, and policyengine-us 2.2.1 is " + "installed", + ), + ( + PACKAGES, + {**PACKAGES, "policyengine-us": None}, + "it was cut with policyengine-us 2.2.1, and policyengine-us is not " + "installed", + ), + ( + {**PACKAGES, "policyengine-core": None}, + PACKAGES, + "it was cut with no recorded policyengine-core version, and " + "policyengine-core 3.33.0 is installed", + ), + ( + {"policyengine-us": "2.2.1"}, + PACKAGES, + "it was cut with no recorded policyengine-core version, and " + "policyengine-core 3.33.0 is installed", + ), + ( + {**PACKAGES, "extra-package": "1.0"}, + PACKAGES, + "it records the extra-package version it was cut with, which this " + "version of policyengine.py does not check", + ), + ( + {"policyengine-core": "3.32.0", "policyengine-us": "2.2.0"}, + PACKAGES, + "it was cut with policyengine-core 3.32.0, and policyengine-core " + "3.33.0 is installed; it was cut with policyengine-us 2.2.0, and " + "policyengine-us 2.2.1 is installed", + ), + ], +) +def test_a_year_file_cut_with_other_versions_names_them( + tmp_path, stamped_packages, packages, reason +): + path = _write_year_file( + tmp_path / "year.h5", + stamp=_stamp(LEGACY_INPUT_RENAMES, packages=stamped_packages), + record={}, + ) + + assert year_file_stamp_problem(path, IDENTITY, packages) == reason + + +def test_every_difference_is_named_in_order(tmp_path): + path = _write_year_file( + tmp_path / "year.h5", + stamp=_stamp( + {}, + source={"uri": f"{SOURCE}@old"}, + packages={**PACKAGES, "policyengine-us": "2.2.0"}, + ), + record={}, + ) + + assert year_file_stamp_problem(path, IDENTITY) == ( + "it was cut before these renames were registered, so it may have lost " + f"them: {LEGACY} -> {LIVE}; it was cut from {SOURCE}@old, and the " + f"requested source is {SOURCE}; it was cut with policyengine-us 2.2.0, " + "and policyengine-us 2.2.1 is installed" + ) + + def test_a_future_register_entry_makes_a_stamped_file_out_of_date( tmp_path, monkeypatch ): - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) monkeypatch.setattr( legacy_inputs, "LEGACY_INPUT_RENAMES", {**LEGACY_INPUT_RENAMES, "old_input": "new_input"}, ) - assert year_file_stamp_problem(path) == ( + assert year_file_stamp_problem(path, IDENTITY) == ( "it was cut before these renames were registered, so it may have lost " "them: old_input -> new_input" ) @@ -663,13 +1024,13 @@ def test_retiring_or_retargeting_an_entry_makes_a_stamped_file_out_of_date( tmp_path, monkeypatch, register, reason ): path = _write_year_file( - tmp_path / "year.h5", stamp=year_file_stamp(), record={LEGACY: LIVE} + tmp_path / "year.h5", stamp=_current(), record={LEGACY: LIVE} ) monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", register) # The record of a file cut under the old register lies within that # register, so the register is what makes it out of date. - assert year_file_stamp_problem(path) == reason + assert year_file_stamp_problem(path, IDENTITY) == reason # --- No partial file is current ----------------------------------------- @@ -679,9 +1040,9 @@ def test_retiring_or_retargeting_an_entry_makes_a_stamped_file_out_of_date( def _current_year_file_bytes() -> bytes: with tempfile.TemporaryDirectory() as directory: path = _write_year_file( - Path(directory) / "year.h5", stamp=year_file_stamp(), record={} + Path(directory) / "year.h5", stamp=_current(), record={} ) - assert year_file_stamp_problem(path) is None + assert year_file_stamp_problem(path, IDENTITY) is None return path.read_bytes() @@ -694,20 +1055,20 @@ def test_a_truncated_year_file_is_never_current(data): path = Path(directory) / "year.h5" path.write_bytes(content[:length]) - assert year_file_stamp_problem(path) is not None + assert year_file_stamp_problem(path, IDENTITY) is not None def test_a_file_that_is_not_hdf5_is_out_of_date(tmp_path): path = tmp_path / "year.h5" path.write_text("not an H5 file") - assert year_file_stamp_problem(path) == ( + assert year_file_stamp_problem(path, IDENTITY) == ( "it cannot be read as an HDF5 file (it may be half-written or corrupt)" ) def test_saving_a_year_file_again_drops_its_stamp(tmp_path): - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) dataset = PolicyEngineUSDataset( name="year-file", description="reload", filepath=str(path), year=2024 @@ -718,11 +1079,11 @@ def test_saving_a_year_file_again_drops_its_stamp(tmp_path): # stamp, which only ``create_datasets`` writes. assert read_renames_record(path) == {} assert read_year_file_stamp(path) is None - assert year_file_stamp_problem(path).startswith("it has no stamp") + assert year_file_stamp_problem(path, IDENTITY).startswith("it has no stamp") def test_a_save_that_stops_part_way_leaves_no_stamp(tmp_path, monkeypatch): - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) put = pd.HDFStore.put def put_until_family(self, key, value, *args, **kwargs): @@ -735,14 +1096,14 @@ def put_until_family(self, key, value, *args, **kwargs): _year_dataset(path).save() assert read_year_file_stamp(path) is None - assert year_file_stamp_problem(path) is not None + assert year_file_stamp_problem(path, IDENTITY) is not None # --- Errors unrelated to the file's contents ---------------------------- def _unreadable_by_permission(path) -> Path: - path = _write_year_file(path, stamp=year_file_stamp(), record={}) + path = _write_year_file(path, stamp=_current(), record={}) os.chmod(path, 0) return path @@ -768,7 +1129,7 @@ def test_errors_unrelated_to_the_contents_are_raised(tmp_path, make, error): create = MagicMock() try: with pytest.raises(error): - year_file_stamp_problem(path) + year_file_stamp_problem(path, IDENTITY) with pytest.raises(error): load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) with patch.object(us_datasets, "create_datasets", create): @@ -783,14 +1144,14 @@ def test_errors_unrelated_to_the_contents_are_raised(tmp_path, make, error): def test_a_locked_file_is_raised_not_cut_again(tmp_path, monkeypatch): - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) def locked(*args, **kwargs): raise BlockingIOError(35, "Unable to lock file") monkeypatch.setattr(h5py, "File", locked) with pytest.raises(BlockingIOError): - year_file_stamp_problem(path) + year_file_stamp_problem(path, IDENTITY) @pytest.mark.parametrize( @@ -804,14 +1165,14 @@ def locked(*args, **kwargs): def test_errors_unrelated_to_the_contents_are_raised_while_reading( tmp_path, monkeypatch, error ): - path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record={}) def fail(self, *args, **kwargs): raise error monkeypatch.setattr(h5py.Dataset, "asstr", fail) with pytest.raises(type(error)): - year_file_stamp_problem(path) + year_file_stamp_problem(path, IDENTITY) # --- read_renames_record reads only well-formed records ----------------- @@ -838,13 +1199,17 @@ def test_a_well_formed_renames_record_still_reads_back(tmp_path): #: Each year file state, and whether it is current. STATES = { - "current": (dict(stamp=year_file_stamp(), record={}), True), + "current": (dict(stamp=_current(), record={}), True), "current, mapped": ( - dict(stamp=year_file_stamp(), record=dict(LEGACY_INPUT_RENAMES)), + dict(stamp=_current(), record=dict(LEGACY_INPUT_RENAMES)), True, ), "6.1.1 and earlier": (dict(), False), "record only": (dict(record={}), False), + "revision 1": ( + dict(stamp={"register": dict(LEGACY_INPUT_RENAMES), "revision": 1}, record={}), + False, + ), "missing entry": (dict(stamp=_stamp({}), record={}), False), "extra entry": ( dict(stamp=_stamp({**LEGACY_INPUT_RENAMES, "old": "new"}), record={}), @@ -852,10 +1217,46 @@ def test_a_well_formed_renames_record_still_reads_back(tmp_path): ), "retargeted": (dict(stamp=_stamp({LEGACY: "elsewhere"}), record={}), False), "older revision": (dict(stamp=_stamp(LEGACY_INPUT_RENAMES, 0), record={}), False), - "newer revision": (dict(stamp={"revision": 2}, record={}), False), + "newer revision": (dict(stamp={"revision": 3}, record={}), False), + "other source revision": ( + dict( + stamp=_stamp(LEGACY_INPUT_RENAMES, source={"uri": f"{SOURCE}@old"}), + record={}, + ), + False, + ), + "verified source": ( + dict( + stamp=_stamp( + LEGACY_INPUT_RENAMES, source={"uri": SOURCE, "sha256": "a" * 64} + ), + record={}, + ), + False, + ), + "other policyengine-us": ( + dict( + stamp=_stamp( + LEGACY_INPUT_RENAMES, + packages={**PACKAGES, "policyengine-us": "2.2.0"}, + ), + record={}, + ), + False, + ), + "other policyengine-core": ( + dict( + stamp=_stamp( + LEGACY_INPUT_RENAMES, + packages={**PACKAGES, "policyengine-core": "3.34.0"}, + ), + record={}, + ), + False, + ), "malformed stamp": (dict(stamp=("text", "not json"), record={}), False), - "no record": (dict(stamp=year_file_stamp()), False), - "record outside": (dict(stamp=year_file_stamp(), record={"x": "y"}), False), + "no record": (dict(stamp=_current()), False), + "record outside": (dict(stamp=_current(), record={"x": "y"}), False), "not HDF5": (None, False), "missing": (None, None), } @@ -877,13 +1278,13 @@ def _create_current_files(datasets, years, data_folder, allow_unmanaged=False): for year in years: _write_year_file( Path(data_folder) / f"{STEM}_year_{year}.h5", - stamp=year_file_stamp(), + stamp=_current(), record={}, ) return {"created": True} -@settings(max_examples=60, deadline=None) +@settings(max_examples=80, deadline=None) @given(states=st.lists(st.sampled_from(sorted(STATES)), min_size=1, max_size=3)) def test_ensure_cuts_again_exactly_when_load_would_refuse(states): years = [2024 + index for index in range(len(states))] @@ -905,7 +1306,7 @@ def test_ensure_cuts_again_exactly_when_load_would_refuse(states): (str(path), problem) for year in years if (path := data_folder / f"{STEM}_year_{year}.h5").exists() - and (problem := year_file_stamp_problem(path)) is not None + and (problem := year_file_stamp_problem(path, IDENTITY)) is not None ) create = MagicMock(side_effect=_create_current_files) with ( @@ -939,16 +1340,42 @@ def test_ensure_cuts_again_exactly_when_load_would_refuse(states): assert sorted(loaded) == [f"{STEM}_{year}" for year in years] +@pytest.mark.parametrize("local_source", [False, True]) @pytest.mark.parametrize( - "states", [["current", "current, mapped"], ["current", "record only"]] + "states", + [ + ["current", "current, mapped"], + ["current", "record only"], + ["current", "other source revision"], + ["current", "other policyengine-us"], + ], ) -def test_ensure_decides_from_the_stamps_alone(tmp_path, states): - """Deciding never opens a year file's tables or the source dataset.""" +def test_ensure_decides_from_the_stamps_alone(tmp_path, states, local_source): + """Deciding never opens a year file's tables, and never reads, hashes or + downloads the source dataset or imports policyengine-us. A local source + is only ``stat()``-ed.""" data_folder = tmp_path / "data" years = [2024, 2025] + dataset, identity = SOURCE, IDENTITY + if local_source: + source_file = tmp_path / "sources" / f"{STEM}.h5" + source_file.parent.mkdir() + source_file.write_bytes(b"the source is never read") + dataset, identity = ( + str(source_file), + requested_dataset_source_identity("us", str(source_file)), + ) for year, state in zip(years, states): - _make_state(data_folder / f"{STEM}_year_{year}.h5", state) + written, _ = STATES[state] + stamp = written.get("stamp", _ABSENT) + if isinstance(stamp, dict) and stamp.get("source") == IDENTITY: + stamp = {**stamp, "source": identity} + _write_year_file( + data_folder / f"{STEM}_year_{year}.h5", + **{**written, "stamp": stamp}, + ) load, create = MagicMock(return_value="loaded"), MagicMock(return_value="created") + source_path = Path(dataset).resolve() if local_source else None def refuse(*args, **kwargs): raise AssertionError("deciding opened a table or the source dataset") @@ -966,6 +1393,7 @@ class StampAndRecordOnly: def __init__(self, path, mode="r", **kwargs): assert mode == "r" + assert source_path is None or Path(path).resolve() != source_path self._file = real_file(path, mode, **kwargs) def __enter__(self): @@ -978,6 +1406,22 @@ def get(self, name, default=None): assert name in (YEAR_FILE_STAMP_H5_DATASET, RENAMES_H5_DATASET), name return self._file.get(name, default) + real_open, real_import = builtins.open, builtins.__import__ + + def open_anything_but_the_source(file, *args, **kwargs): + if ( + source_path is not None + and isinstance(file, (str, os.PathLike)) + and Path(file).resolve() == source_path + ): + refuse() + return real_open(file, *args, **kwargs) + + def import_anything_but_policyengine_us(name, *args, **kwargs): + if name == "policyengine_us" or name.startswith("policyengine_us."): + raise AssertionError("deciding imported policyengine-us") + return real_import(name, *args, **kwargs) + with ( patch.object(h5py, "File", StampAndRecordOnly), patch.object(pd, "HDFStore", RefusedStore), @@ -985,12 +1429,17 @@ def get(self, name, default=None): # opens files without h5py. patch("pandas.io.pytables.HDFStore", RefusedStore), patch("tables.open_file", refuse), + patch.object(builtins, "open", open_anything_but_the_source), + patch.object(io, "open", open_anything_but_the_source), + patch.object(builtins, "__import__", import_anything_but_policyengine_us), + patch.object(dataset_materialization, "sha256_file", refuse), + patch.object(dataset_materialization.requests, "get", refuse), patch.object(us_datasets, "materialize_dataset", refuse), patch.object(us_datasets, "load_datasets", load), patch.object(us_datasets, "create_datasets", create), ): result = ensure_datasets( - datasets=[SOURCE], years=years, data_folder=str(data_folder) + datasets=[dataset], years=years, data_folder=str(data_folder) ) stale = any(STATES[state][1] is False for state in states) @@ -1010,3 +1459,45 @@ def test_the_refusal_names_the_file_and_what_it_may_lack(tmp_path): f"were registered, so it may have lost them: {LEGACY} -> {LIVE}. " "Regenerate it with ensure_datasets() or create_datasets()." ) + + +def test_the_refusal_names_the_source_it_was_cut_from(tmp_path): + data_folder = tmp_path / "data" + path = data_folder / f"{STEM}_year_2024.h5" + _write_year_file( + path, + stamp=_stamp(LEGACY_INPUT_RENAMES, source={"uri": f"{SOURCE}@old"}), + record={}, + ) + + with pytest.raises(ValueError) as refused: + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + assert str(refused.value) == ( + f"US year file {path} is out of date: it was cut from {SOURCE}@old, and " + f"the requested source is {SOURCE}. Regenerate it with ensure_datasets() " + "or create_datasets()." + ) + + +def test_load_and_ensure_compare_with_the_installed_versions(tmp_path, monkeypatch): + """Both read the versions when called, not at import.""" + data_folder = tmp_path / "data" + _write_year_file(data_folder / f"{STEM}_year_2024.h5", stamp=_current(), record={}) + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + monkeypatch.setattr( + legacy_inputs, + "installed_package_versions", + lambda: {**PACKAGES, "policyengine-us": "2.3.0"}, + ) + with pytest.raises(ValueError, match="policyengine-us 2.3.0 is installed"): + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + create = MagicMock(return_value="created") + with patch.object(us_datasets, "create_datasets", create): + assert ( + ensure_datasets( + datasets=[SOURCE], years=[2024], data_folder=str(data_folder) + ) + == "created" + ) From 55476935f8589ceab32cde5522d2bfaa14a15d68 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:31:27 -0400 Subject: [PATCH 3/9] Check the source again before stamping, refuse shared stems, read link cycles as unreadable From the first review of #593: - create_datasets writes a dataset's year files, then checks the source again (dataset_source_change: re-hash a bundle file against its verified sha256, re-read a local file's size and mtime) and stamps them only if it still matches. A file replaced in a shared data folder after verification leaves the year files unstamped and raises DatasetMaterializationError. - Two requested datasets that are different sources with one stem would overwrite each other's year files and, with the identity check, be cut again on every call. ensure_datasets, load_datasets and create_datasets now refuse them with ValueError before anything is read or cut. - A soft-link cycle or overlong chain in place of the stamp or record made h5py raise RuntimeError out of year_file_stamp_problem; it now reads as unreadable content. - Docs: the certified default's own URI selects the bundle dataset, a local file is identified by file-system metadata, and create_datasets takes no lock. Refs #591 Co-Authored-By: Claude Opus 5.5 --- docs/microsim.md | 37 +++-- .../provenance/dataset_materialization.py | 75 ++++++++-- .../tax_benefit_models/us/datasets.py | 98 ++++++++++--- .../tax_benefit_models/us/legacy_inputs.py | 17 ++- tests/test_dataset_source_identity.py | 86 +++++++++++ tests/test_us_legacy_inputs_integration.py | 70 ++++++++- tests/test_us_year_file_stamp.py | 134 ++++++++++++++++++ 7 files changed, 466 insertions(+), 51 deletions(-) diff --git a/docs/microsim.md b/docs/microsim.md index 1fc2986f..5ed73022 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -436,11 +436,14 @@ everything else, it stamps each file with what the file was cut from and under - `source`: the identity of the dataset it was cut from (`dataset_materialization.dataset_source_identity`), in one of three shapes: - a bundle dataset: `{"uri", "sha256"}`, the `hf://…@revision` URI the - bundle pins and the sha256 the local copy was verified against; - - an explicit Hugging Face URI (`allow_unmanaged=True`): `{"uri"}`, the URI - as given; + bundle pins and the sha256 the local copy was verified against. A + request of the certified default dataset's own URI selects the bundle + dataset too; + - any other explicit Hugging Face URI (`allow_unmanaged=True`): `{"uri"}`, + the URI as given; - a local file (`allow_unmanaged=True`): `{"path", "size", "mtime_ns"}`, - its resolved path, size in bytes and modification time. + its resolved path, and its size in bytes and modification time from one + `stat()`. `ensure_datasets` reuses a year file only when its stamp equals the one a file cut now from the requested dataset would carry, and its renames record lists @@ -455,9 +458,24 @@ reason, or `None` for a current year file; `source` is Deciding stays cheap. Both functions read only each year file's stamp and record, never its tables. They identify the requested source without downloading or reading it: a bundle dataset by its bundle manifest entry, an -explicit URI by the URI, and a local file by one `stat()`. Package versions -come from the installed distributions' metadata, so policyengine-us is not -imported. +explicit URI by the URI, and a local file by its file-system metadata. +Package versions come from the installed distributions' metadata, so +policyengine-us is not imported. + +`create_datasets` identifies the source before loading it, and checks it +again once the year files are written. A bundle dataset's file is hashed +again; a local file's size and modification time are read again +(`dataset_materialization.dataset_source_change`). Another process can +replace a file in a shared data folder, for example by materializing another +release of the same dataset there. If the source changed, the year files are +left unstamped, so they are cut again, and `DatasetMaterializationError` is +raised. + +Requested datasets that are different sources with one name, such as two +revisions of one Hugging Face file or two local files called `pop.h5`, would +write the same year files. `ensure_datasets`, `load_datasets` and +`create_datasets` refuse them with `ValueError`; request them in separate +calls with separate data folders. A year file is out of date when: @@ -496,7 +514,10 @@ The source check has limits that follow from not reading the source: (`hf://owner/repo/file.h5@`): a URI without a revision, or with a branch, names whatever the branch holds when the files are cut; - an editable or source install of policyengine-us keeps its version string - while its code changes, so its year files are not cut again. + while its code changes, so its year files are not cut again; +- `create_datasets` takes no lock, so two processes cutting the same year + files at once write them at the same time. Cut year files in one process + at a time per data folder. `ensure_datasets` raises, rather than cutting the files again, when a year file cannot be opened for reasons unrelated to its contents: no permission, a diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index d9b73e14..1ac0ce89 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -268,17 +268,21 @@ def dataset_source_identity(source: DatasetSource) -> dict[str, Any]: - a bundle dataset: ``{"uri": ..., "sha256": ...}``, the ``hf://repo/path@revision`` URI the bundle manifest pins and the sha256 - that the materialized file was verified against; - - an explicit Hugging Face URI (``allow_unmanaged=True``): ``{"uri": ...}``, - the URI as given. Nothing verified its contents, so a URI that does not - pin a revision, or pins a branch, names whatever the branch holds; + that the materialized file was verified against. This includes a + request of the certified default dataset's own URI, which selects the + bundle dataset; + - any other explicit Hugging Face URI (``allow_unmanaged=True``): + ``{"uri": ...}``, the URI as given. Nothing verified its contents, so a + URI that does not pin a revision, or pins a branch, names whatever the + branch holds; - a local file (``allow_unmanaged=True``): ``{"path": ..., "size": ..., - "mtime_ns": ...}``, its resolved path, size in bytes and modification - time in nanoseconds, read with one ``stat()``. + "mtime_ns": ...}``, its resolved path, and its size in bytes and + modification time in nanoseconds from one ``stat()``. - None of them reads the source's contents. See + Computing it reads no source contents. See :func:`requested_dataset_source_identity` for the identity a request - gives without materializing it. + gives without materializing it, and :func:`dataset_source_change` to check + a materialized source again. """ if source.bundle_dataset is not None: return { @@ -299,10 +303,11 @@ def requested_dataset_source_identity( This is what :func:`dataset_source_identity` gives for ``materialize_dataset(country_id, dataset, allow_unmanaged=True)``, computed without downloading anything or reading the source's contents: - the bundle manifest entry for a bundle dataset, the URI for an explicit - Hugging Face URI, and one ``stat()`` for a local file. A name that is - neither a bundle dataset nor a local file is looked up in the data - release manifest, as ``materialize_dataset`` does. + the bundle manifest entry for a bundle dataset (or the certified default's + own URI), the URI for any other explicit Hugging Face URI, and file-system + metadata for a local file. A name that is neither a bundle dataset nor a + local file is looked up in the data release manifest, as + ``materialize_dataset`` does. It does not enforce bundle selection: an unmanaged request gives its identity whether or not ``allow_unmanaged`` would be set. @@ -321,6 +326,47 @@ def requested_dataset_source_identity( return _local_file_identity(source_uri) +def dataset_source_change( + source: DatasetSource, identity: dict[str, Any] +) -> Optional[str]: + """Return how a materialized source differs from its identity, or ``None``. + + ``identity`` is what :func:`dataset_source_identity` gave for ``source`` + earlier. Another process can replace a file in a shared data folder in + the meantime, for example by materializing another release of the same + bundle dataset there. A bundle dataset's file is hashed again and compared + with its verified sha256, and a local file's size and modification time + are read again. An explicit URI's identity records nothing about the + file, so it is never reported as changed. + + Returns: + ``None`` when the source still matches; otherwise a clause saying + what changed. + """ + if "sha256" in identity: + try: + sha256 = sha256_file(source.path) + except FileNotFoundError: + return f"{source.path} no longer exists" + if sha256 != identity["sha256"]: + return ( + f"{source.path} now has sha256 {sha256}, not the verified " + f"{identity['sha256']}" + ) + return None + if "path" in identity: + try: + now = _local_file_identity(source.path) + except FileNotFoundError: + return f"{identity['path']} no longer exists" + if now != identity: + return ( + f"{describe_dataset_source_identity(identity)} is now " + f"{describe_dataset_source_identity(now)}" + ) + return None + + def _local_file_identity(path: str | Path) -> dict[str, Any]: resolved = Path(path).expanduser().resolve() status = resolved.stat() @@ -353,9 +399,10 @@ def describe_dataset_source_identity(identity: dict[str, Any]) -> str: """Describe a well-formed source identity in words (see :func:`is_dataset_source_identity`).""" if "uri" in identity: + uri: str = identity["uri"] if "sha256" in identity: - return f"{identity['uri']} (sha256 {identity['sha256']})" - return identity["uri"] + return f"{uri} (sha256 {identity['sha256']})" + return uri return ( f"{identity['path']} ({identity['size']} bytes, modified " f"{_describe_mtime_ns(identity['mtime_ns'])})" diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 26a3f36c..f85d2882 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -14,8 +14,11 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, MaterializedDataset, + dataset_source_change, dataset_source_identity, + describe_dataset_source_identity, materialize_dataset, requested_dataset_source_identity, ) @@ -383,6 +386,28 @@ def _year_file_source(dataset: str) -> tuple[str, dict[str, Any]]: return stem, requested_dataset_source_identity("us", dataset) +def _year_file_sources(datasets: list[str]) -> list[tuple[str, dict[str, Any]]]: + """Return each requested dataset's year-file stem and source identity. + + Raises ``ValueError`` when two of them are different sources with one + stem. They would write the same year files, each cut overwriting the + other's, so a request of both would cut them all again every time. + """ + sources = [_year_file_source(dataset) for dataset in datasets] + first_by_stem: dict[str, tuple[str, dict[str, Any]]] = {} + for dataset, (stem, identity) in zip(datasets, sources): + first, first_identity = first_by_stem.setdefault(stem, (dataset, identity)) + if first_identity != identity: + raise ValueError( + f"Datasets {first!r} and {dataset!r} would share the US year " + f"files {stem}_year_{{year}}.h5, but are different sources: " + f"{describe_dataset_source_identity(first_identity)} and " + f"{describe_dataset_source_identity(identity)}. Request them " + "in separate calls with separate data folders." + ) + return sources + + def create_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2024, 2025, 2026, 2027, 2028], @@ -391,14 +416,24 @@ def create_datasets( ) -> dict[str, PolicyEngineUSDataset]: """Create PolicyEngineUSDataset instances from logical dataset names or URLs. - Each year file is written to ``{data_folder}/{dataset_stem}_year_{year}.h5`` - and stamped last with the identity of the source dataset it was cut from, - the installed policyengine-core and policyengine-us versions, the register - of renamed stored inputs it was cut under and the year-file format - revision (see ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which - ``ensure_datasets`` and ``load_datasets`` check. A local source is - identified before it is loaded, so one that changes while the files are - cut leaves them out of date. + Each year file is written to ``{data_folder}/{dataset_stem}_year_{year}.h5``. + Once a dataset's year files are all written, each is stamped with the + identity of the source dataset it was cut from, the installed + policyengine-core and policyengine-us versions, the register of renamed + stored inputs it was cut under and the year-file format revision (see + ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which ``ensure_datasets`` + and ``load_datasets`` check. + + The source is identified before it is loaded, and checked again after the + year files are written: a bundle dataset's file is hashed again and a + local file's metadata read again (see + ``dataset_materialization.dataset_source_change``). Another process can + replace a file in a shared data folder, for example by materializing + another release of the same dataset there. If the source changed, the + year files are left unstamped, so they are cut again, and + ``DatasetMaterializationError`` is raised. Two requested datasets that are + different sources with one stem are refused with ``ValueError`` before + anything is cut, since they would overwrite each other's year files. Args: datasets: List of logical dataset names or HuggingFace dataset URLs @@ -408,6 +443,9 @@ def create_datasets( Returns: Dictionary mapping dataset keys (e.g., "populace_us_2024") to PolicyEngineUSDataset objects """ + if datasets is not None and len(datasets) > 1: + _year_file_sources(datasets) + from policyengine_us import Microsimulation from policyengine.tax_benefit_models.us.spm import resolve_spm_selection @@ -425,12 +463,14 @@ def create_datasets( # What each year file is cut from and under: the source, identified # before it is loaded, the package versions and the register the # mapping below uses. - stamp = year_file_stamp(dataset_source_identity(source)) + identity = dataset_source_identity(source) + stamp = year_file_stamp(identity) sim = Microsimulation(dataset=source.path, spm=resolve_spm_selection()) # Map stored inputs the engine has since renamed before extracting, # so each year's file stores the live input rather than dropping it. legacy_input_renames = apply_legacy_input_renames_to_microsimulation(sim) + cut: list[PolicyEngineUSDataset] = [] for year in years: # Get all input variables from the simulation # We'll calculate each input variable for the specified year @@ -588,13 +628,26 @@ def create_datasets( ), ) us_dataset.save() - # Stamp last. ``save()`` truncates the file first, so a year file - # whose write stopped part-way has no stamp and is cut again. - write_year_file_stamp(us_dataset.filepath, stamp) + cut.append(us_dataset) dataset_key = f"{dataset_stem}_{year}" result[dataset_key] = us_dataset + # Stamp only if the source still matches what was identified, so no + # stamp names a source the files may not have been cut from. + change = dataset_source_change(source, identity) + if change is not None: + raise DatasetMaterializationError( + f"The source of the US year files {data_folder}/" + f"{dataset_stem}_year_{{year}}.h5 changed while they were cut: " + f"{change}. They were left unstamped, so ensure_datasets() " + "cuts them again." + ) + # Stamp last. ``save()`` truncates each file first, so a year file + # whose write stopped part-way has no stamp and is cut again. + for us_dataset in cut: + write_year_file_stamp(us_dataset.filepath, stamp) + return result @@ -615,6 +668,8 @@ def load_datasets( draw. A file cut from an earlier release of a dataset the bundle has since re-pinned holds that release's data. The error says why, naming the renames the file may lack, the two sources or the two package versions. + Two requested datasets that are different sources with one stem are + refused, since they would share year files. Args: datasets: List of HuggingFace dataset paths (used to derive file names) @@ -626,8 +681,7 @@ def load_datasets( """ datasets = datasets or [get_release_manifest("us").default_dataset] result = {} - for dataset in datasets: - dataset_stem, source = _year_file_source(dataset) + for dataset_stem, source in _year_file_sources(datasets): for year in years: filepath = f"{data_folder}/{dataset_stem}_year_{year}.h5" if Path(filepath).exists(): @@ -1290,11 +1344,14 @@ def ensure_datasets( ``legacy_inputs.year_file_stamp_problem``). Only each file's stamp and renames record are read to decide, never its tables. The source is never downloaded or read: a bundle dataset is identified by its manifest entry, - an explicit Hugging Face URI by the URI and a local file by one - ``stat()`` (see ``dataset_materialization.dataset_source_identity``). If - any year file is missing or out of date, every year file is created again - from the source dataset, and each out-of-date file is logged with the - reason. + an explicit Hugging Face URI by the URI and a local file by its + file-system metadata (see + ``dataset_materialization.requested_dataset_source_identity``). If any + year file is missing or out of date, every year file is created again from + the source dataset, and each out-of-date file is logged with the reason. + Two requested datasets that are different sources with one stem are + refused with ``ValueError``, since they would overwrite each other's year + files. Args: datasets: List of HuggingFace dataset paths @@ -1310,8 +1367,7 @@ def ensure_datasets( # from another source, with other package versions or under another # register of renamed stored inputs is created again rather than reused. all_current = True - for dataset in datasets: - dataset_stem, source = _year_file_source(dataset) + for dataset_stem, source in _year_file_sources(datasets): for year in years: filepath = Path(f"{data_folder}/{dataset_stem}_year_{year}.h5") if not filepath.exists(): diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 42696432..0b5221bc 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -83,14 +83,17 @@ #: H5 dataset stamping a ``create_datasets`` year file, as UTF-8 JSON (see #: :func:`year_file_stamp`):: #: -#: {"packages": {"policyengine-core": str, "policyengine-us": str}, +#: {"packages": {"policyengine-core": str | null, +#: "policyengine-us": str | null}, #: "register": {legacy: live, ...}, #: "revision": int, #: "source": {...}} #: -#: It holds the versions of the packages that calculated the file's values, -#: the whole register the file was cut under, its year-file format revision -#: and the identity of the dataset it was cut from. ``create_datasets`` writes +#: It holds the versions of the packages that calculated the file's values +#: (``null`` for one that was not installed, see +#: :func:`installed_package_versions`), the whole register the file was cut +#: under, its year-file format revision and the identity of the dataset it was +#: cut from. ``create_datasets`` writes #: it after the tables and the renames record, so a year file whose write #: stopped part-way has none. ``ensure_datasets`` cuts a year file again, and #: ``load_datasets`` refuses it, unless its stamp equals the one a file cut @@ -393,6 +396,8 @@ def _read_json_dataset(file: Any, name: str) -> Any: import h5py try: + # A cycle of soft links, or a chain too long to follow, raises + # RuntimeError here. node = file.get(name) if node is None: return _ABSENT @@ -400,11 +405,11 @@ def _read_json_dataset(file: Any, name: str) -> Any: return _UNREADABLE # A numeric dataset raises TypeError; bytes that are not UTF-8, or # not JSON, raise ValueError; JSON nested too deeply raises - # RecursionError. + # RecursionError, a RuntimeError. return json.loads(node.asstr()[()]) except (PermissionError, IsADirectoryError, BlockingIOError): raise - except (KeyError, OSError, RecursionError, TypeError, ValueError): + except (KeyError, OSError, RuntimeError, TypeError, ValueError): return _UNREADABLE diff --git a/tests/test_dataset_source_identity.py b/tests/test_dataset_source_identity.py index d887f87c..566259c8 100644 --- a/tests/test_dataset_source_identity.py +++ b/tests/test_dataset_source_identity.py @@ -22,6 +22,9 @@ of a ``stat()`` check, the quick check ``rsync`` makes by default. - **Shapes.** Every identity produced has a shape ``is_dataset_source_identity`` accepts, and describing one never raises. +- **Checking again.** ``dataset_source_change`` reports a bundle dataset's + file whose bytes no longer match its verified sha256, and a local file + whose size or modification time moved, and nothing else. """ from __future__ import annotations @@ -41,6 +44,7 @@ import policyengine.provenance.dataset_materialization as dataset_materialization import policyengine.provenance.manifest as manifest_module from policyengine.provenance.dataset_materialization import ( + dataset_source_change, dataset_source_identity, describe_dataset_source_identity, is_dataset_source_identity, @@ -531,3 +535,85 @@ def test_describing_an_identity_is_total(identity): ) def test_an_identity_is_described_in_words(identity, described): assert describe_dataset_source_identity(identity) == described + + +# --- Checking a materialized source again -------------------------------- + + +def test_an_unchanged_bundle_source_has_not_changed(bundle): + _, data_dir = bundle + source = materialize_dataset("us", "tiny_default", data_dir=data_dir) + + assert dataset_source_change(source, dataset_source_identity(source)) is None + + +def test_a_bundle_file_replaced_after_verification_has_changed(bundle): + """Another process materializing another release into the same folder.""" + _, data_dir = bundle + source = materialize_dataset("us", "tiny_default", data_dir=data_dir) + identity = dataset_source_identity(source) + + Path(source.path).write_bytes(b"another release") + assert dataset_source_change(source, identity) == ( + f"{source.path} now has sha256 {_sha256(b'another release')}, not the " + f"verified {_sha256(DEFAULT_BYTES)}" + ) + Path(source.path).unlink() + assert dataset_source_change(source, identity) == ( + f"{source.path} no longer exists" + ) + + +def test_a_local_file_rewritten_after_identification_has_changed(tmp_path): + path = _local_file(tmp_path) + os.utime(path, ns=(SECOND, SECOND)) + source = materialize_dataset("us", str(path), allow_unmanaged=True) + identity = dataset_source_identity(source) + assert dataset_source_change(source, identity) is None + + os.utime(path, ns=(2 * SECOND, 2 * SECOND)) + now = requested_dataset_source_identity("us", str(path)) + assert dataset_source_change(source, identity) == ( + f"{describe_dataset_source_identity(identity)} is now " + f"{describe_dataset_source_identity(now)}" + ) + path.unlink() + assert dataset_source_change(source, identity) == ( + f"{identity['path']} no longer exists" + ) + + +def test_an_explicit_uri_source_is_never_reported_as_changed(bundle): + """Its identity records nothing about the file to compare.""" + _, data_dir = bundle + get, _ = _serve(EXPLICIT_BYTES) + with patch.object(dataset_materialization.requests, "get", get): + source = materialize_dataset( + "us", + "hf://fixture/explicit/explicit.h5@pinned", + allow_unmanaged=True, + data_dir=data_dir, + ) + identity = dataset_source_identity(source) + Path(source.path).write_bytes(b"changed") + + assert dataset_source_change(source, identity) is None + + +@settings(max_examples=100, deadline=None) +@given(first=CONTENTS, second=CONTENTS, first_mtime=MTIMES, second_mtime=MTIMES) +def test_a_local_source_has_changed_exactly_when_its_identity_has( + tmp_path_factory, first, second, first_mtime, second_mtime +): + path = tmp_path_factory.mktemp("local") / "tiny_local.h5" + path.write_bytes(first) + os.utime(path, ns=(first_mtime, first_mtime)) + source = materialize_dataset("us", str(path), allow_unmanaged=True) + identity = dataset_source_identity(source) + + path.write_bytes(second) + os.utime(path, ns=(second_mtime, second_mtime)) + + assert (dataset_source_change(source, identity) is None) == ( + requested_dataset_source_identity("us", str(path)) == identity + ) diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 589b3369..de06cb77 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -13,6 +13,7 @@ from __future__ import annotations +import contextlib import os import re from importlib import metadata as importlib_metadata @@ -34,6 +35,7 @@ import policyengine.provenance.dataset_materialization as dataset_materialization # noqa: E402 import policyengine.provenance.manifest as manifest_module # noqa: E402 from policyengine.provenance.dataset_materialization import ( # noqa: E402 + DatasetMaterializationError, describe_dataset_source_identity, requested_dataset_source_identity, ) @@ -814,11 +816,15 @@ def test_ensure_datasets_regenerates_a_year_file_with_only_the_renames_record( _assert_reused(source, data_folder) -def test_create_datasets_stamps_each_year_file_after_saving_it(tmp_path): +def test_create_datasets_stamps_the_year_files_after_saving_and_checking( + tmp_path, +): + """Every year file is saved, the source checked again, then each stamped.""" source = _write_entity_tables(tmp_path / "populace_layout.h5") events = [] save = PolicyEngineUSDataset.save stamp = us_datasets.write_year_file_stamp + change = us_datasets.dataset_source_change def recording_save(self): save(self) @@ -828,9 +834,14 @@ def recording_stamp(path, value): stamp(path, value) events.append(("stamp", path)) + def recording_change(*args): + events.append(("check source",)) + return change(*args) + with ( patch.object(PolicyEngineUSDataset, "save", recording_save), patch.object(us_datasets, "write_year_file_stamp", recording_stamp), + patch.object(us_datasets, "dataset_source_change", recording_change), ): created = create_datasets( datasets=[source], @@ -842,8 +853,9 @@ def recording_stamp(path, value): paths = [dataset.filepath for dataset in created.values()] assert events == [ ("save", paths[0]), - ("stamp", paths[0]), ("save", paths[1]), + ("check source",), + ("stamp", paths[0]), ("stamp", paths[1]), ] for path in paths: @@ -996,6 +1008,60 @@ def test_a_re_pinned_bundle_dataset_cuts_its_year_files_again(tmp_path): assert _person_outputs(_run(loaded))["wic"].iloc[INFANT] > 0 +@pytest.mark.parametrize("kind", ["bundle", "local"]) +def test_a_source_replaced_while_it_is_cut_leaves_the_year_files_unstamped( + tmp_path, kind +): + """Another process replaces the source after it is identified (for a + bundle dataset, by materializing another release into the same data + folder). The year files hold data the stamp would not describe, so they + are left unstamped and cut again.""" + import policyengine_us + + data_folder = tmp_path / "data" + data_folder.mkdir() + if kind == "bundle": + path = data_folder / f"{BUNDLE_NAME}.h5" + _write_entity_tables(path) + request, pins = BUNDLE_NAME, _pinned("release-1", sha256_file(path)) + else: + path = tmp_path / "populace_layout.h5" + _write_entity_tables(path) + request, pins = str(path), () + original = path.read_bytes() + real_microsimulation = policyengine_us.Microsimulation + + def replaced_once_loaded(*args, **kwargs): + simulation = real_microsimulation(*args, **kwargs) + _write_entity_tables(path, draw=[False, True, True]) + later = path.stat().st_mtime_ns + 1_000_000_000 + os.utime(path, ns=(later, later)) + return simulation + + year_file = data_folder / f"{Path(path).stem}_year_{YEAR}.h5" + with contextlib.ExitStack() as stack: + for pin in pins: + stack.enter_context(pin) + with patch.object(policyengine_us, "Microsimulation", replaced_once_loaded): + with pytest.raises( + DatasetMaterializationError, match="changed while they were cut" + ): + _ensure(request, data_folder) + + assert year_file.exists() + assert read_year_file_stamp(year_file) is None + with pytest.raises(ValueError, match="it has no stamp"): + load_datasets( + datasets=[request], years=[YEAR], data_folder=str(data_folder) + ) + + # The pinned release is back in place, as a download would put it. + path.write_bytes(original) + regenerated = _ensure(request, data_folder) + assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW + assert _problem(year_file, request) is None + + def test_another_package_version_cuts_year_files_again(tmp_path, monkeypatch): """Year files cut under one policyengine-us are not reused under another.""" source, dataset = _create_year_file(tmp_path) diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index 640aa526..a2a15924 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -768,6 +768,51 @@ def test_dangling_links_in_place_of_the_stamp_and_record_read_as_absent(tmp_path assert read_renames_record(path) is None +def _link_cycle(file, name) -> None: + """Point ``name`` into a cycle of soft links.""" + file[name] = h5py.SoftLink(f"/{name}_loop") + file[f"{name}_loop"] = h5py.SoftLink(f"/{name}") + + +def _long_link_chain(file, name) -> None: + """Point ``name`` down a chain of soft links too long to follow, to a + node that would be well formed if reached.""" + content = _current() if name == YEAR_FILE_STAMP_H5_DATASET else {} + _write_node(file, f"{name}_end", content) + file[name] = h5py.SoftLink(f"/{name}_0") + for index in range(100): + file[f"{name}_{index}"] = h5py.SoftLink(f"/{name}_{index + 1}") + file[f"{name}_100"] = h5py.SoftLink(f"/{name}_end") + + +@pytest.mark.parametrize("links", [_link_cycle, _long_link_chain]) +@pytest.mark.parametrize("name", [YEAR_FILE_STAMP_H5_DATASET, RENAMES_H5_DATASET]) +def test_links_that_cannot_be_followed_are_unreadable_content(tmp_path, links, name): + """Totality: h5py raises RuntimeError for these, which reads as + unreadable content rather than escaping.""" + path = _write_year_file( + tmp_path / "year.h5", + **{"stamp": _current(), "record": {}, **{_NODE_ARGUMENT[name]: _ABSENT}}, + ) + with h5py.File(path, "a") as file: + links(file, name) + + assert year_file_stamp_problem(path, IDENTITY) == ( + MALFORMED + if name == YEAR_FILE_STAMP_H5_DATASET + else "its record of the renamed stored inputs mapped when it was cut is " + "malformed" + ) + if name == YEAR_FILE_STAMP_H5_DATASET: + assert read_year_file_stamp(path) is None + else: + assert read_renames_record(path) is None + + +#: Which ``_write_year_file`` argument writes each node. +_NODE_ARGUMENT = {YEAR_FILE_STAMP_H5_DATASET: "stamp", RENAMES_H5_DATASET: "record"} + + @pytest.mark.parametrize("record", [{}, {LEGACY: LIVE}]) def test_a_file_stamped_for_the_requested_source_is_current(tmp_path, record): path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record=record) @@ -1461,6 +1506,95 @@ def test_the_refusal_names_the_file_and_what_it_may_lack(tmp_path): ) +# --- Two sources never share year files -------------------------------- + +#: Requests whose stems (``pop`` or ``other``) collide in different ways. +COLLIDING = [ + "hf://fixture/pop.h5@r1", + "hf://fixture/pop.h5@r2", + "hf://elsewhere/pop.h5@r1", + "hf://fixture/other.h5@r1", +] + + +def _stem_of(uri: str) -> str: + return uri.rsplit("/", 1)[-1].split(".")[0] + + +@settings(max_examples=60, deadline=None) +@given(requests=st.lists(st.sampled_from(COLLIDING), min_size=1, max_size=4)) +def test_two_sources_with_one_stem_are_refused_before_anything_is_cut(requests): + """Requests that would share year files are refused by all three + functions, before reading a year file or cutting one; requests of one + source repeated, or of sources with different stems, are not.""" + stems = {} + for request in requests: + stems.setdefault(_stem_of(request), set()).add(request) + collide = any(len(sources) > 1 for sources in stems.values()) + create = MagicMock(side_effect=lambda *a, **k: {}) + + def refuse(*args, **kwargs): + raise AssertionError("a year file was read or cut") + + with tempfile.TemporaryDirectory() as directory: + with ( + patch.object(us_datasets, "materialize_dataset", refuse), + patch.object(us_datasets, "year_file_stamp_problem", refuse), + patch.object(us_datasets, "PolicyEngineUSDataset", refuse), + ): + for call in ( + lambda: load_datasets( + datasets=requests, years=[2024], data_folder=directory + ), + lambda: us_datasets.create_datasets( + datasets=requests, years=[2024], data_folder=directory + ), + ): + if collide: + with pytest.raises(ValueError, match="would share the US year"): + call() + with patch.object(us_datasets, "create_datasets", create): + if collide: + with pytest.raises(ValueError, match="would share the US year"): + ensure_datasets( + datasets=requests, years=[2024], data_folder=directory + ) + create.assert_not_called() + else: + ensure_datasets(datasets=requests, years=[2024], data_folder=directory) + create.assert_called_once() + + +def test_the_collision_refusal_names_both_requests_and_sources(tmp_path): + first, second = "hf://fixture/pop.h5@r1", "hf://fixture/pop.h5@r2" + + with pytest.raises(ValueError) as refused: + ensure_datasets( + datasets=[first, second], years=[2024], data_folder=str(tmp_path) + ) + + assert str(refused.value) == ( + f"Datasets {first!r} and {second!r} would share the US year files " + f"pop_year_{{year}}.h5, but are different sources: {first} and {second}. " + "Request them in separate calls with separate data folders." + ) + + +def test_one_source_requested_by_two_names_is_not_a_collision(tmp_path): + """The default's logical name and its own URI are one source.""" + from policyengine.provenance.manifest import get_release_manifest + + manifest = get_release_manifest("us") + requests = [manifest.default_dataset, manifest.default_dataset_uri] + create = MagicMock(return_value="created") + + with patch.object(us_datasets, "create_datasets", create): + assert ( + ensure_datasets(datasets=requests, years=[2024], data_folder=str(tmp_path)) + == "created" + ) + + def test_the_refusal_names_the_source_it_was_cut_from(tmp_path): data_folder = tmp_path / "data" path = data_folder / f"{STEM}_year_2024.h5" From e664afaf0819b1dd05fbf3029b9a8c185b695e3f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:35:16 -0400 Subject: [PATCH 4/9] Decode raw stamp and record nodes in the totality oracle The totality property checked only that a None result was current, and gave the oracle a raw node's descriptor rather than what the checker reads back. A current stamp with an empty record stored as raw text or bytes is current, but the oracle called it out of date. _stored decodes raw nodes independently of the checker (fixed-length bytes read back without trailing NULs), and the property now asserts exactness in both directions. Found by the #594 review in the sibling oracle. Refs #591 Co-Authored-By: Claude Opus 5.5 --- tests/test_us_year_file_stamp.py | 51 +++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index a2a15924..c6ce7817 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -221,6 +221,33 @@ def _write_year_file(path, *, stamp=_ABSENT, record=_ABSENT): return Path(path) +#: What :func:`_stored` gives for a node no reader can parse as JSON. +_UNREADABLE = object() + + +def _stored(content): + """Decode what :func:`_write_node` stores, as the checker reads it back. + + An oracle independent of the checker's reader: raw text is parsed as + JSON; raw bytes are fixed-length strings, which read back without + trailing NULs; any other raw node, or text that is not JSON, is + unreadable. A JSON value reads back as its JSON round trip. + """ + if content is _ABSENT: + return _ABSENT + if isinstance(content, tuple): + kind, value = content + try: + if kind == "text": + return json.loads(value) + if kind == "bytes": + return json.loads(value.rstrip(b"\x00").decode("utf-8")) + except (RecursionError, ValueError): + pass + return _UNREADABLE + return json.loads(json.dumps(content)) + + def _same_json(left, right) -> bool: """Equal as JSON: ``True`` is not ``1``, and ``1.0`` is not ``1``.""" return json.dumps(left, sort_keys=True) == json.dumps(right, sort_keys=True) @@ -443,7 +470,7 @@ def test_a_year_file_is_current_only_under_the_exact_stamp_requested(case): problem = year_file_stamp_problem(path, source, packages) assert (problem is None) == _is_current( - stamp, record, register, source, packages + _stored(stamp), _stored(record), register, source, packages ) assert problem is None or (isinstance(problem, str) and problem) if problem is None: @@ -499,15 +526,31 @@ def test_the_year_file_problem_is_total_and_exact_for_any_content( stamp, record, register, source, packages ): """Totality: any stamp and record give a reason or ``None``, never an - error, and ``None`` only for the exact stamp requested.""" + error, and ``None`` exactly for the stamp requested, however it is + stored (``_stored`` decodes raw nodes independently of the checker).""" with tempfile.TemporaryDirectory() as directory: path = _write_year_file(Path(directory) / "year.h5", stamp=stamp, record=record) with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): problem = year_file_stamp_problem(path, source, packages) assert problem is None or (isinstance(problem, str) and problem) - if problem is None: - assert _is_current(stamp, record, register, source, packages) + assert (problem is None) == _is_current( + _stored(stamp), _stored(record), register, source, packages + ) + + +@pytest.mark.parametrize( + "record", + [("text", "{}"), ("bytes", b"{}"), ("bytes", b"{}\x00\x00")], +) +def test_a_record_stored_raw_is_read_as_its_json(tmp_path, monkeypatch, record): + """The case a one-directional totality check missed: raw text or bytes + that hold an empty record, under an empty register.""" + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", {}) + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), record=record) + + assert year_file_stamp_problem(path, IDENTITY) is None + assert _is_current(_stored(_current()), _stored(record), {}) @pytest.mark.parametrize( From 716fb0eacc2e3f0073220514b3e301323bd9d836 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:45:40 -0400 Subject: [PATCH 5/9] Snapshot the source file after materializing and require it unchanged before stamping From the second review of #593: an explicit URI's download shares the data folder's overwritable path, so another process could replace it before the cut and the year files would be stamped with the pinned URI anyway; a bundle file swapped out and back passed the re-hash. Now create_datasets records the source file's device, inode, size and mtime right after materializing it (dataset_source_snapshot) and stamps only if the file is still that one, for every kind of source, besides the bundle re-hash. A download replaces a file with a new inode. create_datasets' up-front check for shared year files now resolves requests with its own allow_unmanaged, so a request materialize_dataset would refuse is refused with the same error, without a network lookup. The changelog mentions both new errors. Refs #591 Co-Authored-By: Claude Opus 5.5 --- changelog.d/591.fixed.md | 2 +- docs/microsim.md | 20 +- .../provenance/dataset_materialization.py | 77 ++++++-- .../tax_benefit_models/us/datasets.py | 42 ++-- tests/test_dataset_source_identity.py | 182 +++++++++++++----- tests/test_us_legacy_inputs_integration.py | 66 +++++-- tests/test_us_year_file_stamp.py | 32 ++- 7 files changed, 314 insertions(+), 107 deletions(-) diff --git a/changelog.d/591.fixed.md b/changelog.d/591.fixed.md index ee3602e9..64d6c83b 100644 --- a/changelog.d/591.fixed.md +++ b/changelog.d/591.fixed.md @@ -1 +1 @@ -US year files now record the source dataset they were cut from and the policyengine-core and policyengine-us versions that cut them. `ensure_datasets` creates a year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source or another version is installed. A year file cut from a bundle dataset is tied to the revision and sha256 the bundle pins, so re-pinning a dataset under the same name (such as the ACS-local replacement) no longer reuses year files cut from the old release. A local source is identified by its path, size and modification time, and an explicit Hugging Face URI by the URI. Deciding still reads only each year file's stamp and renames record, and never downloads or reads the source. +US year files now record the source dataset they were cut from and the policyengine-core and policyengine-us versions that cut them. `ensure_datasets` creates a year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source or another version is installed. A year file cut from a bundle dataset is tied to the revision and sha256 the bundle pins, so re-pinning a dataset under the same name (such as the ACS-local replacement) no longer reuses year files cut from the old release. A local source is identified by its path, size and modification time, and an explicit Hugging Face URI by the URI. Deciding still reads only each year file's stamp and renames record, and never downloads or reads the source. `create_datasets` stamps the year files only if the source file is still the one it materialized, and otherwise leaves them unstamped and raises `DatasetMaterializationError`. Requesting two different sources that share a name (and so would share year files) now raises `ValueError`. diff --git a/docs/microsim.md b/docs/microsim.md index 5ed73022..865356df 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -462,14 +462,16 @@ explicit URI by the URI, and a local file by its file-system metadata. Package versions come from the installed distributions' metadata, so policyengine-us is not imported. -`create_datasets` identifies the source before loading it, and checks it -again once the year files are written. A bundle dataset's file is hashed -again; a local file's size and modification time are read again -(`dataset_materialization.dataset_source_change`). Another process can +`create_datasets` records the source file's device, inode, size and +modification time right after materializing it +(`dataset_materialization.dataset_source_snapshot`). It checks them again +once the year files are written, and hashes a bundle dataset's file again +against its verified sha256 (`dataset_source_change`). Another process can replace a file in a shared data folder, for example by materializing another -release of the same dataset there. If the source changed, the year files are -left unstamped, so they are cut again, and `DatasetMaterializationError` is -raised. +release of the same dataset, or another file with the same name, there. A +download replaces a file with a new inode. If the source changed, the year +files are left unstamped, so they are cut again, and +`DatasetMaterializationError` is raised. Requested datasets that are different sources with one name, such as two revisions of one Hugging Face file or two local files called `pop.h5`, would @@ -512,7 +514,9 @@ The source check has limits that follow from not reading the source: resolution) is not detected; - an explicit Hugging Face URI is compared as a string. Pin it to a commit (`hf://owner/repo/file.h5@`): a URI without a revision, or with a - branch, names whatever the branch holds when the files are cut; + branch, names whatever the branch holds when the files are cut. Its + contents are verified against nothing, so a file another process puts in + its place between the download and the snapshot goes unnoticed; - an editable or source install of policyengine-us keeps its version string while its code changes, so its year files are not cut again; - `create_datasets` takes no lock, so two processes cutting the same year diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 1ac0ce89..e9844efb 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -281,8 +281,8 @@ def dataset_source_identity(source: DatasetSource) -> dict[str, Any]: Computing it reads no source contents. See :func:`requested_dataset_source_identity` for the identity a request - gives without materializing it, and :func:`dataset_source_change` to check - a materialized source again. + gives without materializing it, and :func:`dataset_source_snapshot` and + :func:`dataset_source_change` to check a materialized source again. """ if source.bundle_dataset is not None: return { @@ -297,11 +297,13 @@ def dataset_source_identity(source: DatasetSource) -> dict[str, Any]: def requested_dataset_source_identity( country_id: str, dataset: Optional[str] = None, + *, + allow_unmanaged: bool = True, ) -> dict[str, Any]: """Return the identity of the source ``materialize_dataset`` would use. This is what :func:`dataset_source_identity` gives for - ``materialize_dataset(country_id, dataset, allow_unmanaged=True)``, + ``materialize_dataset(country_id, dataset, allow_unmanaged=...)``, computed without downloading anything or reading the source's contents: the bundle manifest entry for a bundle dataset (or the certified default's own URI), the URI for any other explicit Hugging Face URI, and file-system @@ -309,8 +311,10 @@ def requested_dataset_source_identity( local file is looked up in the data release manifest, as ``materialize_dataset`` does. - It does not enforce bundle selection: an unmanaged request gives its - identity whether or not ``allow_unmanaged`` would be set. + A request ``materialize_dataset`` would refuse with the same + ``allow_unmanaged`` raises the same error here. The default, + ``allow_unmanaged=True``, gives an unmanaged request's identity without + enforcing bundle selection. """ manifest = get_release_manifest(country_id) if dataset is None or dataset == manifest.default_dataset_uri: @@ -319,30 +323,64 @@ def requested_dataset_source_identity( bundle_dataset = _resolve_bundle_dataset(country_id, dataset, manifest=manifest) return {"uri": bundle_dataset.source_uri, "sha256": bundle_dataset.sha256} source_uri = resolve_managed_dataset_reference( - country_id, dataset, allow_unmanaged=True + country_id, dataset, allow_unmanaged=allow_unmanaged ) if "://" in source_uri: return {"uri": source_uri} return _local_file_identity(source_uri) +def dataset_source_snapshot(source: DatasetSource) -> dict[str, int]: + """Return the materialized file's device, inode, size and modification time. + + Taken right after ``materialize_dataset`` returns, it lets + :func:`dataset_source_change` tell whether the file at ``source.path`` is + still the one materialized. Another process can replace a file in a shared + data folder, for example by materializing another release of the same + dataset, or another file with the same name, there. A download replaces + the file with ``os.replace``, which always gives it another inode. + """ + status = os.stat(source.path) + return { + "device": status.st_dev, + "inode": status.st_ino, + "size": status.st_size, + "mtime_ns": status.st_mtime_ns, + } + + def dataset_source_change( - source: DatasetSource, identity: dict[str, Any] + source: DatasetSource, + identity: dict[str, Any], + snapshot: dict[str, int], ) -> Optional[str]: - """Return how a materialized source differs from its identity, or ``None``. + """Return how a materialized source has changed, or ``None``. + + ``identity`` is what :func:`dataset_source_identity` gave for ``source``, + and ``snapshot`` what :func:`dataset_source_snapshot` gave right after it + was materialized. The file must still be the one snapshotted: the same + device, inode, size and modification time. A bundle dataset's file must + also still hash to its verified sha256, which covers a replacement after + verification and before the snapshot. A local file must still have its + identity. - ``identity`` is what :func:`dataset_source_identity` gave for ``source`` - earlier. Another process can replace a file in a shared data folder in - the meantime, for example by materializing another release of the same - bundle dataset there. A bundle dataset's file is hashed again and compared - with its verified sha256, and a local file's size and modification time - are read again. An explicit URI's identity records nothing about the - file, so it is never reported as changed. + One gap remains, for an explicit URI only. Its contents are verified + against nothing, so a file another process puts in its place between the + download and the snapshot goes unnoticed. Returns: ``None`` when the source still matches; otherwise a clause saying what changed. """ + try: + now = dataset_source_snapshot(source) + except FileNotFoundError: + return f"{source.path} no longer exists" + if now != snapshot: + return ( + f"{source.path} was replaced or modified (device, inode, size or " + "modification time changed)" + ) if "sha256" in identity: try: sha256 = sha256_file(source.path) @@ -350,19 +388,18 @@ def dataset_source_change( return f"{source.path} no longer exists" if sha256 != identity["sha256"]: return ( - f"{source.path} now has sha256 {sha256}, not the verified " + f"{source.path} has sha256 {sha256}, not the verified " f"{identity['sha256']}" ) - return None if "path" in identity: try: - now = _local_file_identity(source.path) + local = _local_file_identity(source.path) except FileNotFoundError: return f"{identity['path']} no longer exists" - if now != identity: + if local != identity: return ( f"{describe_dataset_source_identity(identity)} is now " - f"{describe_dataset_source_identity(now)}" + f"{describe_dataset_source_identity(local)}" ) return None diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index f85d2882..02ded9d4 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -18,6 +18,7 @@ MaterializedDataset, dataset_source_change, dataset_source_identity, + dataset_source_snapshot, describe_dataset_source_identity, materialize_dataset, requested_dataset_source_identity, @@ -374,26 +375,35 @@ def _load_policyengine_core_h5(path: Path, year: int) -> USYearData: ) -def _year_file_source(dataset: str) -> tuple[str, dict[str, Any]]: +def _year_file_source( + dataset: str, allow_unmanaged: bool = True +) -> tuple[str, dict[str, Any]]: """Return a requested dataset's year-file stem and its source identity. The stem names the year files (``{stem}_year_{year}.h5``) and drops any revision, so files cut from every release of a logical dataset share names; the identity tells them apart. Neither downloads nor reads the - source (see ``requested_dataset_source_identity``). + source (see ``requested_dataset_source_identity``). A request + ``materialize_dataset`` would refuse with the same ``allow_unmanaged`` + raises the same error, before the stem is resolved. """ + identity = requested_dataset_source_identity( + "us", dataset, allow_unmanaged=allow_unmanaged + ) stem = dataset_logical_name(resolve_dataset_reference("us", dataset)) - return stem, requested_dataset_source_identity("us", dataset) + return stem, identity -def _year_file_sources(datasets: list[str]) -> list[tuple[str, dict[str, Any]]]: +def _year_file_sources( + datasets: list[str], allow_unmanaged: bool = True +) -> list[tuple[str, dict[str, Any]]]: """Return each requested dataset's year-file stem and source identity. Raises ``ValueError`` when two of them are different sources with one stem. They would write the same year files, each cut overwriting the other's, so a request of both would cut them all again every time. """ - sources = [_year_file_source(dataset) for dataset in datasets] + sources = [_year_file_source(dataset, allow_unmanaged) for dataset in datasets] first_by_stem: dict[str, tuple[str, dict[str, Any]]] = {} for dataset, (stem, identity) in zip(datasets, sources): first, first_identity = first_by_stem.setdefault(stem, (dataset, identity)) @@ -424,13 +434,13 @@ def create_datasets( ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which ``ensure_datasets`` and ``load_datasets`` check. - The source is identified before it is loaded, and checked again after the - year files are written: a bundle dataset's file is hashed again and a - local file's metadata read again (see - ``dataset_materialization.dataset_source_change``). Another process can - replace a file in a shared data folder, for example by materializing - another release of the same dataset there. If the source changed, the - year files are left unstamped, so they are cut again, and + The source file's device, inode, size and modification time are recorded + right after it is materialized, and checked again once the year files are + written; a bundle dataset's file is also hashed again against its + verified sha256 (see ``dataset_materialization.dataset_source_change``). + Another process can replace a file in a shared data folder, for example + by materializing another release of the same dataset there. If the source + changed, the year files are left unstamped, so they are cut again, and ``DatasetMaterializationError`` is raised. Two requested datasets that are different sources with one stem are refused with ``ValueError`` before anything is cut, since they would overwrite each other's year files. @@ -444,7 +454,7 @@ def create_datasets( Dictionary mapping dataset keys (e.g., "populace_us_2024") to PolicyEngineUSDataset objects """ if datasets is not None and len(datasets) > 1: - _year_file_sources(datasets) + _year_file_sources(datasets, allow_unmanaged) from policyengine_us import Microsimulation @@ -464,6 +474,7 @@ def create_datasets( # before it is loaded, the package versions and the register the # mapping below uses. identity = dataset_source_identity(source) + snapshot = dataset_source_snapshot(source) stamp = year_file_stamp(identity) sim = Microsimulation(dataset=source.path, spm=resolve_spm_selection()) # Map stored inputs the engine has since renamed before extracting, @@ -633,9 +644,10 @@ def create_datasets( dataset_key = f"{dataset_stem}_{year}" result[dataset_key] = us_dataset - # Stamp only if the source still matches what was identified, so no + # Stamp only if the source file is still the one materialized (and, + # for a bundle dataset, still hashes to its verified sha256), so no # stamp names a source the files may not have been cut from. - change = dataset_source_change(source, identity) + change = dataset_source_change(source, identity, snapshot) if change is not None: raise DatasetMaterializationError( f"The source of the US year files {data_folder}/" diff --git a/tests/test_dataset_source_identity.py b/tests/test_dataset_source_identity.py index 566259c8..5d600f98 100644 --- a/tests/test_dataset_source_identity.py +++ b/tests/test_dataset_source_identity.py @@ -22,9 +22,10 @@ of a ``stat()`` check, the quick check ``rsync`` makes by default. - **Shapes.** Every identity produced has a shape ``is_dataset_source_identity`` accepts, and describing one never raises. -- **Checking again.** ``dataset_source_change`` reports a bundle dataset's - file whose bytes no longer match its verified sha256, and a local file - whose size or modification time moved, and nothing else. +- **Checking again.** ``dataset_source_change`` reports a source file that + is no longer the one snapshotted after materializing (device, inode, size + or modification time), for every kind of source, and a bundle dataset's + file whose bytes no longer hash to its verified sha256. """ from __future__ import annotations @@ -46,6 +47,7 @@ from policyengine.provenance.dataset_materialization import ( dataset_source_change, dataset_source_identity, + dataset_source_snapshot, describe_dataset_source_identity, is_dataset_source_identity, materialize_dataset, @@ -540,80 +542,162 @@ def test_an_identity_is_described_in_words(identity, described): # --- Checking a materialized source again -------------------------------- -def test_an_unchanged_bundle_source_has_not_changed(bundle): +def _replace(path: Path, payload: bytes) -> None: + """Replace a file as a download does: a new file renamed over it.""" + temporary = path.with_name(f".{path.name}.download") + temporary.write_bytes(payload) + os.replace(temporary, path) + + +def _materialized(dataset, data_dir, served=None): + get, _ = _serve(served or b"") + with patch.object( + dataset_materialization.requests, "get", get if served else _refuse + ): + source = materialize_dataset( + "us", dataset, allow_unmanaged=True, data_dir=data_dir + ) + return source, dataset_source_identity(source), dataset_source_snapshot(source) + + +def test_the_snapshot_is_the_files_device_inode_size_and_mtime(bundle): _, data_dir = bundle - source = materialize_dataset("us", "tiny_default", data_dir=data_dir) + source, _, snapshot = _materialized("tiny_default", data_dir) + status = os.stat(source.path) - assert dataset_source_change(source, dataset_source_identity(source)) is None + assert snapshot == { + "device": status.st_dev, + "inode": status.st_ino, + "size": status.st_size, + "mtime_ns": status.st_mtime_ns, + } -def test_a_bundle_file_replaced_after_verification_has_changed(bundle): - """Another process materializing another release into the same folder.""" +@pytest.mark.parametrize( + "dataset, served", + [ + ("tiny_default", None), + ("hf://fixture/explicit/explicit.h5@pinned", EXPLICIT_BYTES), + ("{tmp}/sources/tiny_local.h5", None), + ], +) +def test_an_untouched_source_has_not_changed(bundle, tmp_path, dataset, served): _, data_dir = bundle - source = materialize_dataset("us", "tiny_default", data_dir=data_dir) - identity = dataset_source_identity(source) + _local_file(tmp_path) + source, identity, snapshot = _materialized( + dataset.format(tmp=tmp_path), data_dir, served + ) - Path(source.path).write_bytes(b"another release") - assert dataset_source_change(source, identity) == ( - f"{source.path} now has sha256 {_sha256(b'another release')}, not the " - f"verified {_sha256(DEFAULT_BYTES)}" + assert dataset_source_change(source, identity, snapshot) is None + + +@pytest.mark.parametrize( + "dataset, served", + [ + ("tiny_default", None), + ("hf://fixture/explicit/explicit.h5@pinned", EXPLICIT_BYTES), + ("{tmp}/sources/tiny_local.h5", None), + ], +) +def test_a_source_replaced_after_its_snapshot_has_changed( + bundle, tmp_path, dataset, served +): + """Another process materializing another release, or another file with + the same name, into the same folder. An explicit URI is caught too, + although nothing verifies its contents.""" + _, data_dir = bundle + _local_file(tmp_path) + source, identity, snapshot = _materialized( + dataset.format(tmp=tmp_path), data_dir, served + ) + + _replace(Path(source.path), b"another release") + + assert dataset_source_change(source, identity, snapshot) == ( + f"{source.path} was replaced or modified (device, inode, size or " + "modification time changed)" ) Path(source.path).unlink() - assert dataset_source_change(source, identity) == ( - f"{source.path} no longer exists" + assert dataset_source_change(source, identity, snapshot).endswith( + "no longer exists" + ) + + +def test_a_bundle_file_swapped_out_and_back_has_changed(bundle): + """Its bytes hash to the verified sha256 again, but it is not the file + that was loaded: another release was materialized over it, then the + pinned release downloaded again.""" + _, data_dir = bundle + source, identity, snapshot = _materialized("tiny_default", data_dir) + path = Path(source.path) + + # The loader's open handle keeps the original inode allocated, so the + # file system cannot hand its number to the file swapped back in. + with path.open("rb"): + _replace(path, b"another release") + _replace(path, DEFAULT_BYTES) + + assert _sha256(path.read_bytes()) == identity["sha256"] + assert dataset_source_change(source, identity, snapshot) is not None + + +def test_a_bundle_file_replaced_before_its_snapshot_has_changed(bundle): + """Replaced after verification but before the snapshot, so the snapshot + matches; the sha256 does not.""" + _, data_dir = bundle + source, identity, _ = _materialized("tiny_default", data_dir) + _replace(Path(source.path), b"another release") + snapshot = dataset_source_snapshot(source) + + assert dataset_source_change(source, identity, snapshot) == ( + f"{source.path} has sha256 {_sha256(b'another release')}, not the " + f"verified {_sha256(DEFAULT_BYTES)}" ) -def test_a_local_file_rewritten_after_identification_has_changed(tmp_path): +def test_a_local_file_rewritten_before_its_snapshot_has_changed(tmp_path): + """Its identity was read before the rewrite, the snapshot after.""" path = _local_file(tmp_path) os.utime(path, ns=(SECOND, SECOND)) source = materialize_dataset("us", str(path), allow_unmanaged=True) identity = dataset_source_identity(source) - assert dataset_source_change(source, identity) is None - os.utime(path, ns=(2 * SECOND, 2 * SECOND)) - now = requested_dataset_source_identity("us", str(path)) - assert dataset_source_change(source, identity) == ( + snapshot = dataset_source_snapshot(source) + + assert dataset_source_change(source, identity, snapshot) == ( f"{describe_dataset_source_identity(identity)} is now " - f"{describe_dataset_source_identity(now)}" - ) - path.unlink() - assert dataset_source_change(source, identity) == ( - f"{identity['path']} no longer exists" + f"{describe_dataset_source_identity(requested_dataset_source_identity('us', str(path)))}" ) -def test_an_explicit_uri_source_is_never_reported_as_changed(bundle): - """Its identity records nothing about the file to compare.""" - _, data_dir = bundle - get, _ = _serve(EXPLICIT_BYTES) - with patch.object(dataset_materialization.requests, "get", get): - source = materialize_dataset( - "us", - "hf://fixture/explicit/explicit.h5@pinned", - allow_unmanaged=True, - data_dir=data_dir, - ) - identity = dataset_source_identity(source) - Path(source.path).write_bytes(b"changed") - - assert dataset_source_change(source, identity) is None - - @settings(max_examples=100, deadline=None) -@given(first=CONTENTS, second=CONTENTS, first_mtime=MTIMES, second_mtime=MTIMES) -def test_a_local_source_has_changed_exactly_when_its_identity_has( - tmp_path_factory, first, second, first_mtime, second_mtime +@given( + first=CONTENTS, + second=CONTENTS, + first_mtime=MTIMES, + second_mtime=MTIMES, + replace=st.booleans(), +) +def test_a_local_source_has_changed_exactly_when_its_file_has( + tmp_path_factory, first, second, first_mtime, second_mtime, replace ): + """Rewritten in place or replaced: changed exactly when the snapshot + moved, which a replacement always does (a new inode).""" path = tmp_path_factory.mktemp("local") / "tiny_local.h5" path.write_bytes(first) os.utime(path, ns=(first_mtime, first_mtime)) source = materialize_dataset("us", str(path), allow_unmanaged=True) identity = dataset_source_identity(source) + snapshot = dataset_source_snapshot(source) - path.write_bytes(second) + if replace: + _replace(path, second) + else: + path.write_bytes(second) os.utime(path, ns=(second_mtime, second_mtime)) - assert (dataset_source_change(source, identity) is None) == ( - requested_dataset_source_identity("us", str(path)) == identity + assert (dataset_source_change(source, identity, snapshot) is None) == ( + dataset_source_snapshot(source) == snapshot ) + if replace: + assert dataset_source_change(source, identity, snapshot) is not None diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index de06cb77..9f5becf5 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -1008,37 +1008,77 @@ def test_a_re_pinned_bundle_dataset_cuts_its_year_files_again(tmp_path): assert _person_outputs(_run(loaded))["wic"].iloc[INFANT] > 0 -@pytest.mark.parametrize("kind", ["bundle", "local"]) +class _Download: + """A Hugging Face response serving ``payload``.""" + + status_code = 200 + + def __init__(self, payload: bytes): + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def raise_for_status(self): + pass + + def iter_content(self, chunk_size): + yield self.payload + + +@pytest.mark.parametrize("kind", ["bundle", "explicit URI", "local"]) def test_a_source_replaced_while_it_is_cut_leaves_the_year_files_unstamped( tmp_path, kind ): - """Another process replaces the source after it is identified (for a - bundle dataset, by materializing another release into the same data - folder). The year files hold data the stamp would not describe, so they - are left unstamped and cut again.""" + """Another process replaces the source once it is materialized: by + materializing another release, or another file with the same name, into + the same data folder, or by rewriting a local file. The year files may + hold data the stamp would not describe, so they are left unstamped and + cut again.""" import policyengine_us data_folder = tmp_path / "data" data_folder.mkdir() + original = Path(_write_entity_tables(tmp_path / "original.h5")).read_bytes() if kind == "bundle": path = data_folder / f"{BUNDLE_NAME}.h5" - _write_entity_tables(path) + path.write_bytes(original) request, pins = BUNDLE_NAME, _pinned("release-1", sha256_file(path)) + elif kind == "explicit URI": + path = data_folder / "tiny_uri.h5" + request = "hf://fixture/tiny-uri/tiny_uri.h5@pinned" + pins = ( + patch.object( + dataset_materialization.requests, + "get", + lambda *args, **kwargs: _Download(original), + ), + ) else: path = tmp_path / "populace_layout.h5" - _write_entity_tables(path) + path.write_bytes(original) request, pins = str(path), () - original = path.read_bytes() real_microsimulation = policyengine_us.Microsimulation def replaced_once_loaded(*args, **kwargs): simulation = real_microsimulation(*args, **kwargs) - _write_entity_tables(path, draw=[False, True, True]) - later = path.stat().st_mtime_ns + 1_000_000_000 - os.utime(path, ns=(later, later)) + if kind == "local": + # Rewritten in place, with a later modification time of its own. + _write_entity_tables(path, draw=[False, True, True]) + later = path.stat().st_mtime_ns + 1_000_000_000 + os.utime(path, ns=(later, later)) + else: + # Replaced as a download replaces it. + other = _write_entity_tables( + path.with_name(".other.h5"), draw=[False, True, True] + ) + os.replace(other, path) return simulation - year_file = data_folder / f"{Path(path).stem}_year_{YEAR}.h5" + year_file = data_folder / f"{path.stem}_year_{YEAR}.h5" with contextlib.ExitStack() as stack: for pin in pins: stack.enter_context(pin) @@ -1055,7 +1095,7 @@ def replaced_once_loaded(*args, **kwargs): datasets=[request], years=[YEAR], data_folder=str(data_folder) ) - # The pinned release is back in place, as a download would put it. + # The pinned source is back in place, as a download would put it. path.write_bytes(original) regenerated = _ensure(request, data_folder) assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index c6ce7817..d6c0c6ae 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -1590,7 +1590,10 @@ def refuse(*args, **kwargs): datasets=requests, years=[2024], data_folder=directory ), lambda: us_datasets.create_datasets( - datasets=requests, years=[2024], data_folder=directory + datasets=requests, + years=[2024], + data_folder=directory, + allow_unmanaged=True, ), ): if collide: @@ -1623,6 +1626,33 @@ def test_the_collision_refusal_names_both_requests_and_sources(tmp_path): ) +@pytest.mark.parametrize( + "requests, error", + [ + (["not_a_bundle_dataset", SOURCE], "Unknown managed dataset"), + ([SOURCE, "hf://fixture/other.h5"], "Explicit dataset URIs bypass"), + ], +) +def test_create_refuses_what_materialize_would_with_its_error( + tmp_path, requests, error +): + """With several datasets, ``create_datasets`` checks them all for shared + year files first. A request ``materialize_dataset`` would refuse is + refused there with the same error, and nothing is downloaded.""" + + def refuse(*args, **kwargs): + raise AssertionError("a dataset was materialized or downloaded") + + with ( + patch.object(us_datasets, "materialize_dataset", refuse), + patch("requests.get", refuse), + ): + with pytest.raises(ValueError, match=error): + us_datasets.create_datasets( + datasets=requests, years=[2024], data_folder=str(tmp_path) + ) + + def test_one_source_requested_by_two_names_is_not_a_collision(tmp_path): """The default's logical name and its own URI are one source.""" from policyengine.provenance.manifest import get_release_manifest From 7af7ba42fcd3de535331806fa76810aaa30d7d6d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:59:09 -0400 Subject: [PATCH 6/9] Compare the source snapshot by size and mtime only; refuse unsupported schemes early From the third review of #593 (APPROVE, low findings): - dataset_source_snapshot records size and mtime, not device and inode. Some FUSE mounts (sshfs without -o use_ino) renumber inodes, which would fail every cut; a replacement policyengine.py makes always writes a new file, so its mtime moves anyway. - requested_dataset_source_identity refuses a non-hf:// URI with materialize_dataset's "Unsupported explicit dataset URI" error, so create_datasets' up-front check refuses it before cutting anything. - The docs and docstring name the other gap: a file another process restores to the materialized size and mtime is not noticed. - Tests date materialized files in the past, so a replacement's mtime moves whatever the file system's timestamp resolution. Refs #591 Co-Authored-By: Claude Opus 5.5 --- docs/microsim.md | 22 +++---- .../provenance/dataset_materialization.py | 51 ++++++++-------- .../tax_benefit_models/us/datasets.py | 8 +-- tests/test_dataset_source_identity.py | 58 ++++++++++++------- tests/test_us_year_file_stamp.py | 16 +++++ 5 files changed, 96 insertions(+), 59 deletions(-) diff --git a/docs/microsim.md b/docs/microsim.md index 865356df..95b4b71d 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -462,16 +462,15 @@ explicit URI by the URI, and a local file by its file-system metadata. Package versions come from the installed distributions' metadata, so policyengine-us is not imported. -`create_datasets` records the source file's device, inode, size and -modification time right after materializing it -(`dataset_materialization.dataset_source_snapshot`). It checks them again -once the year files are written, and hashes a bundle dataset's file again -against its verified sha256 (`dataset_source_change`). Another process can -replace a file in a shared data folder, for example by materializing another -release of the same dataset, or another file with the same name, there. A -download replaces a file with a new inode. If the source changed, the year -files are left unstamped, so they are cut again, and -`DatasetMaterializationError` is raised. +`create_datasets` records the source file's size and modification time right +after materializing it (`dataset_materialization.dataset_source_snapshot`). +It checks them again once the year files are written, and hashes a bundle +dataset's file again against its verified sha256 (`dataset_source_change`). +Another process can replace a file in a shared data folder, for example by +materializing another release of the same dataset, or another file with the +same name, there. A download writes a new file, so its modification time +moves. If the source changed, the year files are left unstamped, so they are +cut again, and `DatasetMaterializationError` is raised. Requested datasets that are different sources with one name, such as two revisions of one Hugging Face file or two local files called `pop.h5`, would @@ -517,6 +516,9 @@ The source check has limits that follow from not reading the source: branch, names whatever the branch holds when the files are cut. Its contents are verified against nothing, so a file another process puts in its place between the download and the snapshot goes unnoticed; +- a source file that another process restores to the size and modification + time it had when it was materialized (renamed aside and back, or copied + with its times kept) while year files are cut from it is not noticed; - an editable or source install of policyengine-us keeps its version string while its code changes, so its year files are not cut again; - `create_datasets` takes no lock, so two processes cutting the same year diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index e9844efb..0bdd9343 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -325,28 +325,31 @@ def requested_dataset_source_identity( source_uri = resolve_managed_dataset_reference( country_id, dataset, allow_unmanaged=allow_unmanaged ) - if "://" in source_uri: + if source_uri.startswith("hf://"): return {"uri": source_uri} + if "://" in source_uri: + raise DatasetMaterializationError( + f"Unsupported explicit dataset URI: {source_uri!r}." + ) return _local_file_identity(source_uri) def dataset_source_snapshot(source: DatasetSource) -> dict[str, int]: - """Return the materialized file's device, inode, size and modification time. + """Return the materialized file's size and modification time. Taken right after ``materialize_dataset`` returns, it lets - :func:`dataset_source_change` tell whether the file at ``source.path`` is - still the one materialized. Another process can replace a file in a shared - data folder, for example by materializing another release of the same - dataset, or another file with the same name, there. A download replaces - the file with ``os.replace``, which always gives it another inode. + :func:`dataset_source_change` tell whether the file at ``source.path`` has + been replaced or modified since. Another process can replace a file in a + shared data folder, for example by materializing another release of the + same dataset, or another file with the same name, there. A download + writes a new file, so its modification time is when it finished. + + Device and inode numbers are left out: some file systems (FUSE mounts + such as sshfs without ``-o use_ino``) do not keep them stable, and a + replacement policyengine.py makes always moves the modification time. """ status = os.stat(source.path) - return { - "device": status.st_dev, - "inode": status.st_ino, - "size": status.st_size, - "mtime_ns": status.st_mtime_ns, - } + return {"size": status.st_size, "mtime_ns": status.st_mtime_ns} def dataset_source_change( @@ -358,14 +361,16 @@ def dataset_source_change( ``identity`` is what :func:`dataset_source_identity` gave for ``source``, and ``snapshot`` what :func:`dataset_source_snapshot` gave right after it - was materialized. The file must still be the one snapshotted: the same - device, inode, size and modification time. A bundle dataset's file must - also still hash to its verified sha256, which covers a replacement after - verification and before the snapshot. A local file must still have its - identity. - - One gap remains, for an explicit URI only. Its contents are verified - against nothing, so a file another process puts in its place between the + was materialized. The file must still have the snapshot's size and + modification time, whatever the kind of source. A bundle dataset's file + must also still hash to its verified sha256, which covers a replacement + after verification and before the snapshot. A local file must still have + its identity. + + Two gaps remain. A file that another process restores to the snapshot's + size and modification time is not noticed: one renamed aside and back, + or copied with its times kept. And an explicit URI's contents are + verified against nothing, so a file put in its place between the download and the snapshot goes unnoticed. Returns: @@ -378,8 +383,8 @@ def dataset_source_change( return f"{source.path} no longer exists" if now != snapshot: return ( - f"{source.path} was replaced or modified (device, inode, size or " - "modification time changed)" + f"{source.path} was replaced or modified (its size or modification " + "time changed)" ) if "sha256" in identity: try: diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 02ded9d4..3049e625 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -434,10 +434,10 @@ def create_datasets( ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which ``ensure_datasets`` and ``load_datasets`` check. - The source file's device, inode, size and modification time are recorded - right after it is materialized, and checked again once the year files are - written; a bundle dataset's file is also hashed again against its - verified sha256 (see ``dataset_materialization.dataset_source_change``). + The source file's size and modification time are recorded right after it + is materialized, and checked again once the year files are written; a + bundle dataset's file is also hashed again against its verified sha256 + (see ``dataset_materialization.dataset_source_change``). Another process can replace a file in a shared data folder, for example by materializing another release of the same dataset there. If the source changed, the year files are left unstamped, so they are cut again, and diff --git a/tests/test_dataset_source_identity.py b/tests/test_dataset_source_identity.py index 5d600f98..ccbb83a8 100644 --- a/tests/test_dataset_source_identity.py +++ b/tests/test_dataset_source_identity.py @@ -292,6 +292,22 @@ def test_a_bundle_entry_without_a_revision_uses_the_release_revision(): } +@pytest.mark.parametrize("uri", ["s3://bucket/x.h5", "gs://bucket/x.h5"]) +def test_an_unsupported_uri_scheme_is_refused_as_materialize_does(bundle, uri): + _, data_dir = bundle + + with pytest.raises( + dataset_materialization.DatasetMaterializationError, + match="Unsupported explicit dataset URI", + ): + requested_dataset_source_identity("us", uri) + with pytest.raises( + dataset_materialization.DatasetMaterializationError, + match="Unsupported explicit dataset URI", + ): + materialize_dataset("us", uri, allow_unmanaged=True, data_dir=data_dir) + + def test_a_bundle_entry_without_a_sha256_is_refused_as_materialize_does(tmp_path): manifest = _manifest() manifest.datasets["tiny_other"].sha256 = None @@ -550,6 +566,7 @@ def _replace(path: Path, payload: bytes) -> None: def _materialized(dataset, data_dir, served=None): + """Materialize ``dataset``, dated well before anything replaces it.""" get, _ = _serve(served or b"") with patch.object( dataset_materialization.requests, "get", get if served else _refuse @@ -557,20 +574,19 @@ def _materialized(dataset, data_dir, served=None): source = materialize_dataset( "us", dataset, allow_unmanaged=True, data_dir=data_dir ) + # Materialized long ago, so a replacement made now has a later + # modification time, whatever the file system's timestamp resolution. + os.utime(source.path, ns=(SECOND, SECOND)) return source, dataset_source_identity(source), dataset_source_snapshot(source) -def test_the_snapshot_is_the_files_device_inode_size_and_mtime(bundle): +def test_the_snapshot_is_the_files_size_and_mtime(bundle): + """Not its device or inode, which some file systems do not keep stable.""" _, data_dir = bundle source, _, snapshot = _materialized("tiny_default", data_dir) status = os.stat(source.path) - assert snapshot == { - "device": status.st_dev, - "inode": status.st_ino, - "size": status.st_size, - "mtime_ns": status.st_mtime_ns, - } + assert snapshot == {"size": status.st_size, "mtime_ns": status.st_mtime_ns} @pytest.mark.parametrize( @@ -614,8 +630,8 @@ def test_a_source_replaced_after_its_snapshot_has_changed( _replace(Path(source.path), b"another release") assert dataset_source_change(source, identity, snapshot) == ( - f"{source.path} was replaced or modified (device, inode, size or " - "modification time changed)" + f"{source.path} was replaced or modified (its size or modification " + "time changed)" ) Path(source.path).unlink() assert dataset_source_change(source, identity, snapshot).endswith( @@ -631,14 +647,13 @@ def test_a_bundle_file_swapped_out_and_back_has_changed(bundle): source, identity, snapshot = _materialized("tiny_default", data_dir) path = Path(source.path) - # The loader's open handle keeps the original inode allocated, so the - # file system cannot hand its number to the file swapped back in. - with path.open("rb"): - _replace(path, b"another release") - _replace(path, DEFAULT_BYTES) + _replace(path, b"another release") + _replace(path, DEFAULT_BYTES) - assert _sha256(path.read_bytes()) == identity["sha256"] - assert dataset_source_change(source, identity, snapshot) is not None + # The pinned release downloaded again has its bytes, but a modification + # time of its own. + assert _sha256(path.read_bytes()) == identity["sha256"] + assert dataset_source_change(source, identity, snapshot) is not None def test_a_bundle_file_replaced_before_its_snapshot_has_changed(bundle): @@ -678,11 +693,12 @@ def test_a_local_file_rewritten_before_its_snapshot_has_changed(tmp_path): second_mtime=MTIMES, replace=st.booleans(), ) -def test_a_local_source_has_changed_exactly_when_its_file_has( +def test_a_local_source_has_changed_exactly_when_its_size_or_mtime_has( tmp_path_factory, first, second, first_mtime, second_mtime, replace ): - """Rewritten in place or replaced: changed exactly when the snapshot - moved, which a replacement always does (a new inode).""" + """Rewritten in place or replaced: changed exactly when the size or the + modification time moved. Other bytes restored to the same size and time + are not noticed: the intended, documented gap.""" path = tmp_path_factory.mktemp("local") / "tiny_local.h5" path.write_bytes(first) os.utime(path, ns=(first_mtime, first_mtime)) @@ -697,7 +713,5 @@ def test_a_local_source_has_changed_exactly_when_its_file_has( os.utime(path, ns=(second_mtime, second_mtime)) assert (dataset_source_change(source, identity, snapshot) is None) == ( - dataset_source_snapshot(source) == snapshot + len(first) == len(second) and first_mtime == second_mtime ) - if replace: - assert dataset_source_change(source, identity, snapshot) is not None diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index d6c0c6ae..11a03b64 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -1653,6 +1653,22 @@ def refuse(*args, **kwargs): ) +def test_create_refuses_an_unsupported_scheme_before_cutting_anything(tmp_path): + """Even with ``allow_unmanaged``, as ``materialize_dataset`` would.""" + + def refuse(*args, **kwargs): + raise AssertionError("a dataset was materialized or downloaded") + + with patch.object(us_datasets, "materialize_dataset", refuse): + with pytest.raises(ValueError, match="Unsupported explicit dataset URI"): + us_datasets.create_datasets( + datasets=[SOURCE, "s3://bucket/tiny.h5"], + years=[2024], + data_folder=str(tmp_path), + allow_unmanaged=True, + ) + + def test_one_source_requested_by_two_names_is_not_a_collision(tmp_path): """The default's logical name and its own URI are one source.""" from policyengine.provenance.manifest import get_release_manifest From 0e6b20b88b9388c62b3d4894d25ab7b02e965a34 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 08:11:11 -0400 Subject: [PATCH 7/9] Hash explicit downloads and re-hash before stamping; validate hf:// URIs up front From the fourth review of #593: with device and inode out of the snapshot, an explicit URI replaced by a same-size download within one timestamp tick passed the check, and the docs claimed a replacement always moves the mtime. _download_hugging_face_dataset now hashes the bytes as it writes them (DatasetSource.downloaded_sha256), and dataset_source_change re-hashes a downloaded file against the sha256 it was verified against (bundle) or downloaded with (explicit URI). The remaining, documented gap is a local file changed within one tick at the same size, or restored to its size and mtime. requested_dataset_source_identity parses an hf:// URI as materialize_dataset does (shared _parse_hugging_face_uri), so a malformed one is refused before anything is cut. The stamp tests' fixture URI gains its repository segment; the refusal tests assert DatasetMaterializationError. Refs #591 Co-Authored-By: Claude Opus 5.5 --- docs/microsim.md | 26 ++--- .../provenance/dataset_materialization.py | 88 +++++++++------ .../tax_benefit_models/us/datasets.py | 12 +- tests/test_dataset_source_identity.py | 104 ++++++++++++++++-- tests/test_us_year_file_stamp.py | 32 ++++-- 5 files changed, 191 insertions(+), 71 deletions(-) diff --git a/docs/microsim.md b/docs/microsim.md index 95b4b71d..a7fe7057 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -463,14 +463,14 @@ Package versions come from the installed distributions' metadata, so policyengine-us is not imported. `create_datasets` records the source file's size and modification time right -after materializing it (`dataset_materialization.dataset_source_snapshot`). -It checks them again once the year files are written, and hashes a bundle -dataset's file again against its verified sha256 (`dataset_source_change`). -Another process can replace a file in a shared data folder, for example by +after materializing it (`dataset_materialization.dataset_source_snapshot`), +and checks them again once the year files are written. A downloaded file is +also hashed again, against the sha256 a bundle dataset was verified against +or an explicit URI was downloaded with (`dataset_source_change`). Another +process can replace a file in a shared data folder, for example by materializing another release of the same dataset, or another file with the -same name, there. A download writes a new file, so its modification time -moves. If the source changed, the year files are left unstamped, so they are -cut again, and `DatasetMaterializationError` is raised. +same name, there. If the source changed, the year files are left unstamped, +so they are cut again, and `DatasetMaterializationError` is raised. Requested datasets that are different sources with one name, such as two revisions of one Hugging Face file or two local files called `pop.h5`, would @@ -513,12 +513,12 @@ The source check has limits that follow from not reading the source: resolution) is not detected; - an explicit Hugging Face URI is compared as a string. Pin it to a commit (`hf://owner/repo/file.h5@`): a URI without a revision, or with a - branch, names whatever the branch holds when the files are cut. Its - contents are verified against nothing, so a file another process puts in - its place between the download and the snapshot goes unnoticed; -- a source file that another process restores to the size and modification - time it had when it was materialized (renamed aside and back, or copied - with its times kept) while year files are cut from it is not noticed; + branch, names whatever the branch holds when the files are cut; +- while year files are cut from a local file, a change that keeps its size + and lands within one tick of the file system's timestamp resolution, or + that restores its size and modification time (renamed aside and back, or + copied with its times kept), is not noticed. A downloaded source is also + checked by its sha256, so this does not apply to it; - an editable or source install of policyengine-us keeps its version string while its code changes, so its year files are not cut again; - `create_datasets` takes no lock, so two processes cutting the same year diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 0bdd9343..310eb8ae 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime +import hashlib import os import tempfile from dataclasses import dataclass @@ -55,6 +56,15 @@ class DatasetSource(BaseModel): source_uri: str path: str bundle_dataset: Optional[MaterializedDataset] = None + downloaded_sha256: Optional[str] = None + """sha256 of the bytes downloaded for an explicit Hugging Face URI. + + Nothing verifies an explicit URI's contents, but this records what was + downloaded, so a later check can tell whether the file at ``path`` is + still those bytes (see :func:`dataset_source_change`). ``None`` for a + bundle dataset, whose verified sha256 is on ``bundle_dataset``, and for a + local file. + """ @property def name(self) -> str: @@ -183,11 +193,11 @@ def _use_bundle_dataset( ) -def _download_hugging_face_dataset( - dataset_uri: str, - *, - data_dir: Path, -) -> DatasetSource: +def _parse_hugging_face_uri(dataset_uri: str) -> tuple[str, str, str]: + """Split ``hf://owner/repo/path[@revision]`` into repo, path and revision. + + The revision defaults to ``main``. + """ path_with_repo, revision = ( dataset_uri[5:].rsplit("@", maxsplit=1) if "@" in dataset_uri[5:] @@ -199,9 +209,15 @@ def _download_hugging_face_dataset( "Invalid Hugging Face dataset URI. Expected format " f"'hf://owner/repo/path/to/file[@revision]', got {dataset_uri!r}." ) + return f"{parts[0]}/{parts[1]}", parts[2], revision + - repo_id = f"{parts[0]}/{parts[1]}" - repository_path = parts[2] +def _download_hugging_face_dataset( + dataset_uri: str, + *, + data_dir: Path, +) -> DatasetSource: + repo_id, repository_path, revision = _parse_hugging_face_uri(dataset_uri) destination = data_dir / Path(repository_path).name destination.parent.mkdir(parents=True, exist_ok=True) @@ -240,12 +256,18 @@ def _download_hugging_face_dataset( f"Could not find explicit dataset {dataset_uri!r}." ) response.raise_for_status() + digest = hashlib.sha256() with temporary_path.open("wb") as output: for chunk in response.iter_content(chunk_size=1024 * 1024): if chunk: output.write(chunk) + digest.update(chunk) os.replace(temporary_path, destination) - return DatasetSource(source_uri=dataset_uri, path=str(destination)) + return DatasetSource( + source_uri=dataset_uri, + path=str(destination), + downloaded_sha256=digest.hexdigest(), + ) finally: temporary_path.unlink(missing_ok=True) @@ -326,6 +348,7 @@ def requested_dataset_source_identity( country_id, dataset, allow_unmanaged=allow_unmanaged ) if source_uri.startswith("hf://"): + _parse_hugging_face_uri(source_uri) return {"uri": source_uri} if "://" in source_uri: raise DatasetMaterializationError( @@ -339,14 +362,10 @@ def dataset_source_snapshot(source: DatasetSource) -> dict[str, int]: Taken right after ``materialize_dataset`` returns, it lets :func:`dataset_source_change` tell whether the file at ``source.path`` has - been replaced or modified since. Another process can replace a file in a - shared data folder, for example by materializing another release of the - same dataset, or another file with the same name, there. A download - writes a new file, so its modification time is when it finished. - - Device and inode numbers are left out: some file systems (FUSE mounts - such as sshfs without ``-o use_ino``) do not keep them stable, and a - replacement policyengine.py makes always moves the modification time. + been modified or replaced since. Device and inode numbers are left out, + since some file systems (FUSE mounts such as sshfs without + ``-o use_ino``) do not keep them stable. A downloaded file is also checked + by its sha256, so the snapshot matters most for a local file. """ status = os.stat(source.path) return {"size": status.st_size, "mtime_ns": status.st_mtime_ns} @@ -361,17 +380,21 @@ def dataset_source_change( ``identity`` is what :func:`dataset_source_identity` gave for ``source``, and ``snapshot`` what :func:`dataset_source_snapshot` gave right after it - was materialized. The file must still have the snapshot's size and - modification time, whatever the kind of source. A bundle dataset's file - must also still hash to its verified sha256, which covers a replacement - after verification and before the snapshot. A local file must still have + was materialized. Another process can replace a file in a shared data + folder, for example by materializing another release of the same + dataset, or another file with the same name, there. + + The file must still have the snapshot's size and modification time, + whatever the kind of source. A downloaded file must also still hash to + the sha256 it was verified against (a bundle dataset) or downloaded with + (an explicit URI), so a replacement is caught even when it lands within + one timestamp tick, or before the snapshot. A local file must still have its identity. - Two gaps remain. A file that another process restores to the snapshot's - size and modification time is not noticed: one renamed aside and back, - or copied with its times kept. And an explicit URI's contents are - verified against nothing, so a file put in its place between the - download and the snapshot goes unnoticed. + One gap remains, for a local file only: one modified, or replaced by a + file of the same size, within one tick of the file system's timestamp + resolution, or restored to the snapshot's size and modification time + (renamed aside and back, or copied with its times kept), is not noticed. Returns: ``None`` when the source still matches; otherwise a clause saying @@ -386,16 +409,19 @@ def dataset_source_change( f"{source.path} was replaced or modified (its size or modification " "time changed)" ) + expected: Optional[tuple[str, str]] = None if "sha256" in identity: + expected = ("verified", identity["sha256"]) + elif source.downloaded_sha256 is not None: + expected = ("downloaded", source.downloaded_sha256) + if expected is not None: + kind, sha256 = expected try: - sha256 = sha256_file(source.path) + actual = sha256_file(source.path) except FileNotFoundError: return f"{source.path} no longer exists" - if sha256 != identity["sha256"]: - return ( - f"{source.path} has sha256 {sha256}, not the verified " - f"{identity['sha256']}" - ) + if actual != sha256: + return f"{source.path} has sha256 {actual}, not the {kind} {sha256}" if "path" in identity: try: local = _local_file_identity(source.path) diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 3049e625..93d59e01 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -435,9 +435,10 @@ def create_datasets( and ``load_datasets`` check. The source file's size and modification time are recorded right after it - is materialized, and checked again once the year files are written; a - bundle dataset's file is also hashed again against its verified sha256 - (see ``dataset_materialization.dataset_source_change``). + is materialized, and checked again once the year files are written. A + downloaded file is also hashed again, against the sha256 a bundle + dataset was verified against or an explicit URI was downloaded with (see + ``dataset_materialization.dataset_source_change``). Another process can replace a file in a shared data folder, for example by materializing another release of the same dataset there. If the source changed, the year files are left unstamped, so they are cut again, and @@ -645,8 +646,9 @@ def create_datasets( result[dataset_key] = us_dataset # Stamp only if the source file is still the one materialized (and, - # for a bundle dataset, still hashes to its verified sha256), so no - # stamp names a source the files may not have been cut from. + # if it was downloaded, still hashes to the sha256 it was verified + # against or downloaded with), so no stamp names a source the files + # may not have been cut from. change = dataset_source_change(source, identity, snapshot) if change is not None: raise DatasetMaterializationError( diff --git a/tests/test_dataset_source_identity.py b/tests/test_dataset_source_identity.py index ccbb83a8..5ab58ed3 100644 --- a/tests/test_dataset_source_identity.py +++ b/tests/test_dataset_source_identity.py @@ -292,20 +292,28 @@ def test_a_bundle_entry_without_a_revision_uses_the_release_revision(): } -@pytest.mark.parametrize("uri", ["s3://bucket/x.h5", "gs://bucket/x.h5"]) -def test_an_unsupported_uri_scheme_is_refused_as_materialize_does(bundle, uri): +@pytest.mark.parametrize( + "uri, error", + [ + ("s3://bucket/x.h5", "Unsupported explicit dataset URI"), + ("gs://bucket/x.h5", "Unsupported explicit dataset URI"), + ("hf://owner/repo", "Invalid Hugging Face dataset URI"), + ("hf://owner//x.h5@r", "Invalid Hugging Face dataset URI"), + ], +) +def test_a_uri_materialize_refuses_is_refused_with_its_error(bundle, uri, error): _, data_dir = bundle - with pytest.raises( - dataset_materialization.DatasetMaterializationError, - match="Unsupported explicit dataset URI", - ): + with pytest.raises(dataset_materialization.DatasetMaterializationError) as early: requested_dataset_source_identity("us", uri) - with pytest.raises( - dataset_materialization.DatasetMaterializationError, - match="Unsupported explicit dataset URI", - ): - materialize_dataset("us", uri, allow_unmanaged=True, data_dir=data_dir) + with patch.object(dataset_materialization.requests, "get", _refuse): + with pytest.raises( + dataset_materialization.DatasetMaterializationError + ) as materializing: + materialize_dataset("us", uri, allow_unmanaged=True, data_dir=data_dir) + + assert error in str(early.value) + assert str(early.value) == str(materializing.value) def test_a_bundle_entry_without_a_sha256_is_refused_as_materialize_does(tmp_path): @@ -670,6 +678,80 @@ def test_a_bundle_file_replaced_before_its_snapshot_has_changed(bundle): ) +def test_an_explicit_download_records_the_sha256_of_its_bytes(bundle): + _, data_dir = bundle + source, _, _ = _materialized( + "hf://fixture/explicit/explicit.h5@pinned", data_dir, EXPLICIT_BYTES + ) + + assert source.downloaded_sha256 == _sha256(EXPLICIT_BYTES) + assert Path(source.path).read_bytes() == EXPLICIT_BYTES + + +@pytest.mark.parametrize( + "dataset, served, payload, kind", + [ + ("tiny_default", None, DEFAULT_BYTES, "verified"), + ( + "hf://fixture/explicit/explicit.h5@pinned", + EXPLICIT_BYTES, + EXPLICIT_BYTES, + "downloaded", + ), + ], +) +def test_a_downloaded_file_replaced_within_one_timestamp_tick_has_changed( + bundle, dataset, served, payload, kind +): + """Same size, same modification time, other bytes: as when another + download of the same size finishes within one tick of a coarse file + system. The sha256 it was verified against, or downloaded with, catches + it.""" + _, data_dir = bundle + source, identity, snapshot = _materialized(dataset, data_dir, served) + other = bytes(reversed(payload)) + assert len(other) == len(payload) and other != payload + + _replace(Path(source.path), other) + os.utime(source.path, ns=(snapshot["mtime_ns"], snapshot["mtime_ns"])) + + assert dataset_source_snapshot(source) == snapshot + assert dataset_source_change(source, identity, snapshot) == ( + f"{source.path} has sha256 {_sha256(other)}, not the {kind} {_sha256(payload)}" + ) + + +@settings(max_examples=100, deadline=None) +@given(other=st.binary(min_size=0, max_size=24), keep_mtime=st.booleans()) +def test_a_replaced_explicit_download_has_changed_whatever_its_metadata( + tmp_path_factory, other, keep_mtime +): + """For a downloaded file, any other bytes are a change, even at the + same size and modification time.""" + data_dir = tmp_path_factory.mktemp("data") + get, _ = _serve(EXPLICIT_BYTES) + with patch.object(dataset_materialization.requests, "get", get): + source = materialize_dataset( + "us", + "hf://fixture/explicit/explicit.h5@pinned", + allow_unmanaged=True, + data_dir=data_dir, + ) + os.utime(source.path, ns=(SECOND, SECOND)) + identity, snapshot = ( + dataset_source_identity(source), + dataset_source_snapshot(source), + ) + + _replace(Path(source.path), other) + if keep_mtime: + os.utime(source.path, ns=(SECOND, SECOND)) + + assert (dataset_source_change(source, identity, snapshot) is None) == ( + other == EXPLICIT_BYTES and keep_mtime + ) + + def test_a_local_file_rewritten_before_its_snapshot_has_changed(tmp_path): """Its identity was read before the rewrite, the snapshot after.""" path = _local_file(tmp_path) diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index 11a03b64..bcd31cf8 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -55,6 +55,7 @@ import policyengine.provenance.dataset_materialization as dataset_materialization from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, describe_dataset_source_identity, requested_dataset_source_identity, ) @@ -87,7 +88,7 @@ #: ``hf://`` references are used as they are, and name year files after the #: file's stem, so no manifest or download is involved. Its identity is the #: URI as given. -SOURCE = "hf://fixture/tiny.h5" +SOURCE = "hf://fixture/repo/tiny.h5" IDENTITY = {"uri": SOURCE} STEM = "tiny" GROUPS = ("household", "tax_unit", "spm_unit", "family", "marital_unit") @@ -1553,10 +1554,10 @@ def test_the_refusal_names_the_file_and_what_it_may_lack(tmp_path): #: Requests whose stems (``pop`` or ``other``) collide in different ways. COLLIDING = [ - "hf://fixture/pop.h5@r1", - "hf://fixture/pop.h5@r2", - "hf://elsewhere/pop.h5@r1", - "hf://fixture/other.h5@r1", + "hf://fixture/repo/pop.h5@r1", + "hf://fixture/repo/pop.h5@r2", + "hf://elsewhere/repo/pop.h5@r1", + "hf://fixture/repo/other.h5@r1", ] @@ -1612,7 +1613,7 @@ def refuse(*args, **kwargs): def test_the_collision_refusal_names_both_requests_and_sources(tmp_path): - first, second = "hf://fixture/pop.h5@r1", "hf://fixture/pop.h5@r2" + first, second = "hf://fixture/repo/pop.h5@r1", "hf://fixture/repo/pop.h5@r2" with pytest.raises(ValueError) as refused: ensure_datasets( @@ -1630,7 +1631,7 @@ def test_the_collision_refusal_names_both_requests_and_sources(tmp_path): "requests, error", [ (["not_a_bundle_dataset", SOURCE], "Unknown managed dataset"), - ([SOURCE, "hf://fixture/other.h5"], "Explicit dataset URIs bypass"), + ([SOURCE, "hf://fixture/repo/other.h5"], "Explicit dataset URIs bypass"), ], ) def test_create_refuses_what_materialize_would_with_its_error( @@ -1653,16 +1654,25 @@ def refuse(*args, **kwargs): ) -def test_create_refuses_an_unsupported_scheme_before_cutting_anything(tmp_path): - """Even with ``allow_unmanaged``, as ``materialize_dataset`` would.""" +@pytest.mark.parametrize( + "request_, error", + [ + ("s3://bucket/tiny.h5", "Unsupported explicit dataset URI"), + ("hf://owner/repo", "Invalid Hugging Face dataset URI"), + ], +) +def test_create_refuses_a_uri_materialize_would_before_cutting_anything( + tmp_path, request_, error +): + """Even with ``allow_unmanaged``, with ``materialize_dataset``'s error.""" def refuse(*args, **kwargs): raise AssertionError("a dataset was materialized or downloaded") with patch.object(us_datasets, "materialize_dataset", refuse): - with pytest.raises(ValueError, match="Unsupported explicit dataset URI"): + with pytest.raises(DatasetMaterializationError, match=error): us_datasets.create_datasets( - datasets=[SOURCE, "s3://bucket/tiny.h5"], + datasets=[SOURCE, request_], years=[2024], data_folder=str(tmp_path), allow_unmanaged=True, From 8578cc84f6f6323f3cd67ccc9db2d8ce933ba2b0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 08:16:04 -0400 Subject: [PATCH 8/9] Say the swapped-and-restored gap applies to every source From the fifth review of #593: the re-hash checks the file when the year files are stamped, not as it is read, so a source swapped out while it loads and then restored exactly is not noticed for downloaded files either. Only the same-size change within one timestamp tick is limited to local files. Refs #591 Co-Authored-By: Claude Opus 5.5 --- docs/microsim.md | 11 ++++++----- .../provenance/dataset_materialization.py | 9 ++++++--- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/docs/microsim.md b/docs/microsim.md index a7fe7057..aef5afd3 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -514,11 +514,12 @@ The source check has limits that follow from not reading the source: - an explicit Hugging Face URI is compared as a string. Pin it to a commit (`hf://owner/repo/file.h5@`): a URI without a revision, or with a branch, names whatever the branch holds when the files are cut; -- while year files are cut from a local file, a change that keeps its size - and lands within one tick of the file system's timestamp resolution, or - that restores its size and modification time (renamed aside and back, or - copied with its times kept), is not noticed. A downloaded source is also - checked by its sha256, so this does not apply to it; +- the source is checked when the year files are stamped, not as it is read. + A source swapped out while it is loaded and then restored, bytes and + times (renamed aside and back, or copied back with `cp -p`), is not + noticed, whatever its kind. A local file modified, or replaced by a file of + the same size, within one tick of the file system's timestamp resolution + is not noticed either; a downloaded file is also checked by its sha256; - an editable or source install of policyengine-us keeps its version string while its code changes, so its year files are not cut again; - `create_datasets` takes no lock, so two processes cutting the same year diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 310eb8ae..98bb76c4 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -391,10 +391,13 @@ def dataset_source_change( one timestamp tick, or before the snapshot. A local file must still have its identity. - One gap remains, for a local file only: one modified, or replaced by a + Two gaps remain. Whatever the source, a file swapped out while it is + loaded and then restored (renamed aside and back, or copied back with its + times kept) is not noticed: it is checked when the year files are + stamped, not as it is read. And a local file modified, or replaced by a file of the same size, within one tick of the file system's timestamp - resolution, or restored to the snapshot's size and modification time - (renamed aside and back, or copied with its times kept), is not noticed. + resolution is not noticed either; a downloaded file's bytes would not + hash to its sha256. Returns: ``None`` when the source still matches; otherwise a clause saying From c7bb48d9ea36ab48d45c4f681b2bf610e3c26ef8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 09:43:16 -0400 Subject: [PATCH 9/9] Check the source again only when year files are cut from it create_datasets with no years cuts and stamps nothing, so it no longer stats or re-hashes the source. tests/test_dataset_runtime.py mocks materialize_dataset with a path that does not exist and calls create_datasets(years=[]); it failed with FileNotFoundError from the snapshot. Reported by the UK year-file session (#600). The full suite passes locally: 1543 passed, 9 skipped. Refs #591 Co-Authored-By: Claude Opus 5.5 --- src/policyengine/tax_benefit_models/us/datasets.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 93d59e01..fd86ecfe 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -475,7 +475,8 @@ def create_datasets( # before it is loaded, the package versions and the register the # mapping below uses. identity = dataset_source_identity(source) - snapshot = dataset_source_snapshot(source) + # The source is checked again only when year files are cut from it. + snapshot = dataset_source_snapshot(source) if years else None stamp = year_file_stamp(identity) sim = Microsimulation(dataset=source.path, spm=resolve_spm_selection()) # Map stored inputs the engine has since renamed before extracting, @@ -649,7 +650,11 @@ def create_datasets( # if it was downloaded, still hashes to the sha256 it was verified # against or downloaded with), so no stamp names a source the files # may not have been cut from. - change = dataset_source_change(source, identity, snapshot) + change = ( + dataset_source_change(source, identity, snapshot) + if snapshot is not None + else None + ) if change is not None: raise DatasetMaterializationError( f"The source of the US year files {data_folder}/"