diff --git a/changelog.d/591.fixed.md b/changelog.d/591.fixed.md new file mode 100644 index 00000000..64d6c83b --- /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 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 1e63432a..aef5afd3 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -409,17 +409,74 @@ 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. 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, 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 +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 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 size and modification time right +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. 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: @@ -434,16 +491,46 @@ 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; +- 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 + 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 +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/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 35824357..98bb76c4 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -2,11 +2,13 @@ from __future__ import annotations +import datetime +import hashlib 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 @@ -54,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: @@ -182,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:] @@ -198,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) @@ -239,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) @@ -257,6 +280,216 @@ 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. 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, and its size in bytes and + modification time in nanoseconds from one ``stat()``. + + 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_snapshot` and + :func:`dataset_source_change` to check a materialized source again. + """ + 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, + *, + 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=...)``, + 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 + 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. + + 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: + 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=allow_unmanaged + ) + if source_uri.startswith("hf://"): + _parse_hugging_face_uri(source_uri) + 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 size and modification time. + + Taken right after ``materialize_dataset`` returns, it lets + :func:`dataset_source_change` tell whether the file at ``source.path`` has + 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} + + +def dataset_source_change( + source: DatasetSource, + identity: dict[str, Any], + snapshot: dict[str, int], +) -> Optional[str]: + """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. 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. 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 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 + 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 (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: + actual = sha256_file(source.path) + except FileNotFoundError: + return f"{source.path} no longer exists" + 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) + except FileNotFoundError: + return f"{identity['path']} no longer exists" + if local != identity: + return ( + f"{describe_dataset_source_identity(identity)} is now " + f"{describe_dataset_source_identity(local)}" + ) + return None + + +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: + uri: str = identity["uri"] + if "sha256" in identity: + return f"{uri} (sha256 {identity['sha256']})" + return 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..fd86ecfe 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -14,8 +14,14 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, MaterializedDataset, + dataset_source_change, + dataset_source_identity, + dataset_source_snapshot, + describe_dataset_source_identity, materialize_dataset, + requested_dataset_source_identity, ) from policyengine.provenance.manifest import ( dataset_logical_name, @@ -369,6 +375,49 @@ def _load_policyengine_core_h5(path: Path, year: int) -> USYearData: ) +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``). 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, identity + + +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, 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)) + 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], @@ -377,12 +426,26 @@ 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 register of renamed stored inputs it was cut - under and the year-file format revision (see + 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 file's size and modification time are recorded right after it + 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 + ``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 years: List of years to extract data for @@ -391,6 +454,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, allow_unmanaged) + from policyengine_us import Microsimulation from policyengine.tax_benefit_models.us.spm import resolve_spm_selection @@ -405,13 +471,19 @@ 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. + identity = dataset_source_identity(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()) - # 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) + cut: list[PolicyEngineUSDataset] = [] for year in years: # Get all input variables from the simulation # We'll calculate each input variable for the specified year @@ -569,13 +641,32 @@ 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 file is still the one materialized (and, + # 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 snapshot is not None + else None + ) + 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 @@ -586,13 +677,18 @@ 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. + 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) @@ -604,13 +700,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) + 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(): - 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 +1357,20 @@ 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 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 @@ -1281,18 +1383,16 @@ 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) + 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(): 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..0b5221bc 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,53 @@ #: 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 | null, +#: "policyengine-us": str | null}, +#: "register": {legacy: live, ...}, +#: "revision": int, +#: "source": {...}} +#: +#: 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 +#: 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 +176,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 +236,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 +261,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 +274,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 +320,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 +360,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) @@ -295,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 @@ -302,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 @@ -331,6 +434,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())) diff --git a/tests/test_dataset_source_identity.py b/tests/test_dataset_source_identity.py new file mode 100644 index 00000000..5ab58ed3 --- /dev/null +++ b/tests/test_dataset_source_identity.py @@ -0,0 +1,799 @@ +"""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. +- **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 + +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_change, + dataset_source_identity, + dataset_source_snapshot, + 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), + } + + +@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) as early: + requested_dataset_source_identity("us", uri) + 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): + 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 + + +# --- Checking a materialized source again -------------------------------- + + +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): + """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 + ): + 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_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 == {"size": status.st_size, "mtime_ns": status.st_mtime_ns} + + +@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 + _local_file(tmp_path) + source, identity, snapshot = _materialized( + dataset.format(tmp=tmp_path), data_dir, served + ) + + 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 (its size or modification " + "time changed)" + ) + Path(source.path).unlink() + 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) + + _replace(path, b"another release") + _replace(path, DEFAULT_BYTES) + + # 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): + """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_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) + os.utime(path, ns=(SECOND, SECOND)) + source = materialize_dataset("us", str(path), allow_unmanaged=True) + identity = dataset_source_identity(source) + os.utime(path, ns=(2 * SECOND, 2 * SECOND)) + 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(requested_dataset_source_identity('us', str(path)))}" + ) + + +@settings(max_examples=100, deadline=None) +@given( + first=CONTENTS, + second=CONTENTS, + first_mtime=MTIMES, + second_mtime=MTIMES, + replace=st.booleans(), +) +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 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)) + source = materialize_dataset("us", str(path), allow_unmanaged=True) + identity = dataset_source_identity(source) + snapshot = dataset_source_snapshot(source) + + if replace: + _replace(path, second) + else: + path.write_bytes(second) + os.utime(path, ns=(second_mtime, second_mtime)) + + assert (dataset_source_change(source, identity, snapshot) is None) == ( + len(first) == len(second) and first_mtime == second_mtime + ) diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 182d39af..9f5becf5 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -13,6 +13,11 @@ from __future__ import annotations +import contextlib +import os +import re +from importlib import metadata as importlib_metadata +from pathlib import Path from unittest.mock import patch import h5py @@ -27,6 +32,17 @@ 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 + DatasetMaterializationError, + 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 +65,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 +121,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 +579,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 +600,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 +721,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 +736,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 +745,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 +762,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 +782,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 +792,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,16 +811,20 @@ 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) -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) @@ -784,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], @@ -798,12 +853,13 @@ 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: - 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 +881,250 @@ 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 + + +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 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" + 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" + path.write_bytes(original) + request, pins = str(path), () + real_microsimulation = policyengine_us.Microsimulation + + def replaced_once_loaded(*args, **kwargs): + simulation = real_microsimulation(*args, **kwargs) + 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.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 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 + 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) + 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..bcd31cf8 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,12 @@ from hypothesis import strategies as st from microdf import MicroDataFrame +import policyengine.provenance.dataset_materialization as dataset_materialization +from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, + 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 +71,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 +86,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. -SOURCE = "hf://fixture/tiny.h5" +#: file's stem, so no manifest or download is involved. Its identity is the +#: URI as given. +SOURCE = "hf://fixture/repo/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 +141,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 +222,39 @@ def _write_year_file(path, *, stamp=_ABSENT, record=_ABSENT): return Path(path) -def _is_current(stamp, record, register) -> bool: +#: 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) + + +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 +265,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 +393,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 +410,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 +427,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 +452,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( + _stored(stamp), _stored(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 +491,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 +515,78 @@ 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`` 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) + 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 (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( "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 +594,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 +629,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 +647,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 +700,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 +744,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 +764,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 +778,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 +789,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,23 +805,66 @@ 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 assert read_renames_record(path) is None -@pytest.mark.parametrize("record", [{}, {LEGACY: LIVE}]) -def test_a_file_stamped_with_this_releases_register_is_current(tmp_path, record): +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=year_file_stamp(), record=record + 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) is None + 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 -# --- Named renames ------------------------------------------------------ +#: 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) + + assert year_file_stamp_problem(path, IDENTITY) is None + + +# --- Named differences -------------------------------------------------- @settings(max_examples=200, deadline=None) @@ -610,34 +877,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 +1113,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 +1129,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 +1144,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 +1168,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 +1185,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 +1218,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 +1233,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 +1254,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 +1288,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 +1306,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 +1367,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 +1395,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 +1429,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 +1482,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 +1495,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 +1518,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 +1548,189 @@ 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()." ) + + +# --- Two sources never share year files -------------------------------- + +#: Requests whose stems (``pop`` or ``other``) collide in different ways. +COLLIDING = [ + "hf://fixture/repo/pop.h5@r1", + "hf://fixture/repo/pop.h5@r2", + "hf://elsewhere/repo/pop.h5@r1", + "hf://fixture/repo/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, + allow_unmanaged=True, + ), + ): + 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/repo/pop.h5@r1", "hf://fixture/repo/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." + ) + + +@pytest.mark.parametrize( + "requests, error", + [ + (["not_a_bundle_dataset", SOURCE], "Unknown managed dataset"), + ([SOURCE, "hf://fixture/repo/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) + ) + + +@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(DatasetMaterializationError, match=error): + us_datasets.create_datasets( + datasets=[SOURCE, request_], + 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 + + 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" + _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" + )