Tie UK year files to the source dataset and package versions they were cut from

changelog.d/600.fixed.md

@@ -0,0 +1 @@
+UK year files now record the source dataset they were cut from, the policyengine-core and policyengine-uk versions that cut them, and the year they were cut for, as US year files record their source and versions. `ensure_datasets` creates a UK year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source, another version is installed, or the file holds another year's tables, so re-pinning a dataset such as `enhanced_frs_2024_25` or `populace_uk_2023` under the same name, editing a local source, or upgrading either package no longer reuses year files cut before. Deciding reads only each year file's stamp and whether it records its data year, and never downloads or reads the source. `create_datasets` writes each UK year file beside its path and moves it into place, stamped, only once the source is checked again, so a year file's tables and stamp always come from one cut; if the source changed, it writes none of that dataset's year files and raises `DatasetMaterializationError`. Requesting two different sources that share a name (and so would share year files) now raises `ValueError`. Year files written before this release are created again once. The stamp code the US and UK share now lives in `policyengine.provenance.year_file_stamp`.

docs/microsim.md

@@ -263,12 +263,99 @@ poverty median, and `shareholding`, which spreads corporate taxes across
 households (PolicyEngine/policyengine.py#567). Variables of a person or
 benefit unit that do not depend on those match the national run.
 
-Year files written by earlier releases have no recorded data year.
-`ensure_datasets` writes them again and `load_datasets` refuses them. A year
-file opened directly, as `PolicyEngineUKDataset(filepath=...)`, is not
-checked. A saved UK simulation output records the data year its run anchored
-on; `Simulation.load()` refuses an output saved without one, and
-`Simulation.ensure()` runs it again.
+UK year files are named after the dataset's logical name,
+`{data_folder}/{name}_year_{year}.h5`, which drops the revision, so files cut
+from every release of a dataset share names. `create_datasets` stamps each
+with what it was cut from and with, as UTF-8 JSON in
+`uk.datasets.YEAR_FILE_STAMP_H5_DATASET` (`policyengine_uk_year_file_stamp`).
+US year files store theirs under another name, so a year file of one country
+never passes as one of the other.
+
+```json
+{
+  "packages": {"policyengine-core": "3.33.0", "policyengine-uk": "2.126.12"},
+  "revision": 1,
+  "source": {
+    "sha256": "5ba399ea5ab9a6c7d4e9aaf6c623080beb34d9ad5f330f6244dbad72b8d960b6",
+    "uri": "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2024_25.h5@1.58.0"
+  },
+  "year": 2026
+}
+```
+
+- `packages`: the installed policyengine-core and policyengine-uk versions.
+  policyengine-uk extends a single-year source to later years with
+  `extend_single_year_dataset`, which uprates each year from the one before
+  by policyengine-uk parameters read through policyengine-core's parameter
+  tree, so either version can change a projected file's values;
+- `revision`: the UK year-file format revision;
+- `source`: the identity of the dataset it was cut from, in the three shapes
+  US year files use (see [Year files](#year-files)): a bundle dataset's
+  pinned `hf://…@revision` URI and verified sha256, an explicit Hugging Face
+  URI as given, or a local file's resolved path, size and modification time;
+- `year`: the year whose tables the file holds, so a year file copied or
+  renamed to another year's name is not taken for that year.
+
+`ensure_datasets` reuses a UK year file only when it records its data year
+and its stamp equals the one a file cut now from the requested dataset, for
+the requested year, would carry. Otherwise it logs why and creates the year
+files again. `load_datasets` refuses the file with the same reason, which
+names the year, the source or the package versions it was cut for, from or
+with, and the ones requested or installed.
+`uk.datasets.year_file_stamp_problem(path, source, year)` returns the
+reason, or `None` for a current year file; `source` is
+`dataset_materialization.requested_dataset_source_identity("uk", dataset)`.
+
+So a re-pinned dataset (a new release of `enhanced_frs_2024_25` or
+`populace_uk_2023` under the same name), an edited local source, or a new
+policyengine-uk or policyengine-core version makes the affected year files
+out of date. Year files written by earlier releases have no stamp, and those
+written before policyengine.py kept the data year have no data year either:
+`ensure_datasets` writes them again once, and `load_datasets` refuses them.
+
+Deciding reads only each year file's stamp and whether it has a data year,
+never its tables. As for US year files, the requested source is identified
+without downloading or reading it, and package versions come from the
+installed distributions' metadata, so policyengine-uk is not imported.
+
+`create_datasets` writes each year file to a hidden temporary file beside its
+path (`.{name}_year_{year}.…partial`). It records the source file's size and
+modification time right after materializing it, and checks them again
+(re-hashing a downloaded file) once a dataset's year files are written. Only
+then does it stamp each temporary file and move it onto its path, in one
+step. The tables and the stamp at a year file's path therefore always come
+from one cut, even when two processes cut the same year files at once. If the
+source changed, none of that dataset's year files is moved into place, the
+ones already there are left as they were, and `DatasetMaterializationError`
+is raised. A process killed mid-cut can leave a temporary file behind, which
+is never read and can be deleted. Requested datasets that are different
+sources with one name are refused with `ValueError`.
+
+Moving a year file into place replaces the file at its path rather than
+rewriting it. The cut therefore needs write access to the data folder (and,
+in a folder with the sticky bit set, ownership of the year file it replaces
+or of the folder), and the year file gets the mode and group a file the
+cutting user creates there would get, not those of the file it replaces.
+
+The limits listed under [Year files](#year-files) apply to UK year files too,
+with policyengine-uk in place of policyengine-us, except that UK year files
+are moved into place whole. Three more apply:
+
+- the versions are those installed when the files are cut. A package upgraded
+  in a running process after policyengine-uk was imported is recorded at its
+  new version, though the cut ran the code already imported. Restart the
+  process after upgrading;
+- policyengine-uk draws new student loan holders in each projected year from
+  numpy's random generator, seeded by the year. numpy's version is not
+  recorded;
+- a multi-year source is read as stored, without the projection.
+
+A year file opened directly, as `PolicyEngineUKDataset(filepath=...)`, is not
+checked, and saving it again with `save()` drops its stamp. A saved UK
+simulation output records the data year its run anchored on;
+`Simulation.load()` refuses an output saved without one, and
+`Simulation.ensure()` runs it again. A saved UK output does not record the
+stamp of the year file it was calculated from.
 
 ## Simulations
 

src/policyengine/provenance/year_file_stamp.py

@@ -0,0 +1,275 @@
+"""Stamp year files with what they were cut from and under.
+
+A country's ``create_datasets`` cuts one file per requested year from a
+source dataset and names it ``{data_folder}/{stem}_year_{year}.h5``, after
+the source's logical name. The name drops any revision, so files cut from
+every release of a dataset share names. Each year file is therefore stamped
+with the identity of the source it was cut from and the versions of the
+packages that calculated it, and the country's ``ensure_datasets`` and
+``load_datasets`` compare that stamp with the one a file cut now would carry.
+
+Each country stores its stamp, as UTF-8 JSON, in an H5 dataset of its own
+name, so a year file of one country never passes as one of another:
+``tax_benefit_models.us.legacy_inputs.YEAR_FILE_STAMP_H5_DATASET`` (which
+also records the register of renamed stored inputs) and
+``tax_benefit_models.uk.datasets.YEAR_FILE_STAMP_H5_DATASET`` (which also
+records the year the file was cut for).
+
+This module holds what the countries share:
+
+- reading and writing a JSON value in an H5 dataset
+  (:func:`read_json_dataset`, :func:`read_stored_json`,
+  :func:`write_json_dataset`);
+- the installed versions of the packages a stamp records
+  (:func:`installed_package_versions`);
+- the clauses naming how a stamp's source and package versions differ from
+  the current ones (:func:`source_and_package_differences`);
+- each requested dataset's year-file stem and source identity, refusing two
+  sources that would share year files (:func:`year_file_sources`);
+- writing a year file beside its path and moving it into place only once it
+  is stamped (:func:`temporary_year_file`, :func:`publish_year_file`).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import secrets
+from collections.abc import Iterable, Mapping
+from importlib import metadata as importlib_metadata
+from pathlib import Path
+from typing import Any, Optional
+
+from .dataset_materialization import (
+    describe_dataset_source_identity,
+    requested_dataset_source_identity,
+)
+from .manifest import dataset_logical_name, resolve_dataset_reference
+
+#: What :func:`read_json_dataset` returns for a node a file does not have, and
+#: for one that is not a scalar dataset holding UTF-8 JSON.
+ABSENT = object()
+UNREADABLE = object()
+
+#: Errors opening or reading a year file that say nothing about its
+#: contents: no permission, a directory at its path, or a lock another
+#: process holds. A year-file check raises them rather than reporting the
+#: file out of date, since the file could not be written again either.
+UNRELATED_TO_CONTENTS = (PermissionError, IsADirectoryError, BlockingIOError)
+
+#: The clause a year-file check gives for a file h5py cannot open.
+NOT_HDF5 = "it cannot be read as an HDF5 file (it may be half-written or corrupt)"
+
+
+def installed_package_versions(names: Iterable[str]) -> dict[str, Optional[str]]:
+    """Return the installed version of each named distribution.
+
+    ``None`` stands for a distribution that is not installed. The versions
+    come from the installed distributions' metadata, so no package is
+    imported.
+    """
+    versions: dict[str, Optional[str]] = {}
+    for name in names:
+        try:
+            versions[name] = importlib_metadata.version(name)
+        except importlib_metadata.PackageNotFoundError:
+            versions[name] = None
+    return versions
+
+
+def read_json_dataset(file: Any, name: str) -> Any:
+    """Return the JSON value stored in an open H5 file's ``name`` dataset.
+
+    Returns :data:`ABSENT` when the file has no such node, and
+    :data:`UNREADABLE` when the node is not a scalar dataset holding UTF-8
+    JSON, or cannot be followed. Errors in :data:`UNRELATED_TO_CONTENTS` are
+    raised.
+    """
+    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
+        if not isinstance(node, h5py.Dataset) or node.shape != ():
+            return UNREADABLE
+        # A numeric dataset raises TypeError; bytes that are not UTF-8, or
+        # not JSON, raise ValueError; JSON nested too deeply raises
+        # RecursionError, a RuntimeError.
+        return json.loads(node.asstr()[()])
+    except UNRELATED_TO_CONTENTS:
+        raise
+    except (KeyError, OSError, RuntimeError, TypeError, ValueError):
+        return UNREADABLE
+
+
+def write_json_dataset(path: str | Path, name: str, value: Any) -> None:
+    """Store ``value`` as UTF-8 JSON in an H5 file's ``name`` dataset."""
+    import h5py
+
+    with h5py.File(path, "a") as file:
+        if name in file:
+            del file[name]
+        # UTF-8 JSON in a dataset, as for the SPM receipt of a saved output.
+        file.create_dataset(
+            name,
+            data=json.dumps(value, sort_keys=True),
+            dtype=h5py.string_dtype("utf-8"),
+        )
+
+
+def read_stored_json(path: str | Path, name: str) -> Any | None:
+    """Return the JSON value of an H5 file's ``name`` dataset, well formed or
+    not.
+
+    ``None`` means the file stores no readable JSON there (see
+    :func:`read_json_dataset`) or is not an H5 file.
+    """
+    import h5py
+
+    try:
+        with h5py.File(path, "r") as file:
+            value = read_json_dataset(file, name)
+    except OSError:
+        return None
+    return None if value is ABSENT or value is UNREADABLE else value
+
+
+def temporary_year_file(path: str | Path) -> Path:
+    """Create an empty file beside ``path`` to write a year file to.
+
+    ``create_datasets`` writes a year file's tables there, stamps it once the
+    source has been checked again, and only then moves it onto ``path`` (see
+    :func:`publish_year_file`). The tables and the stamp at ``path`` then
+    always come from one cut, even when two processes cut the same year files
+    at once, and a cut that stops part-way leaves ``path`` as it was. The
+    file's name starts with ``.{stem}.`` and ends with ``.partial``; the
+    caller removes it if the cut fails.
+
+    The file is created with the mode a new file gets (``0o666`` less the
+    umask), as pandas would create the year file itself, so a year file moved
+    into place is as readable by others as one written at its path. (A file
+    from ``tempfile.mkstemp`` would be readable by its owner alone.)
+    """
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
+    for _ in range(100):
+        temporary = path.parent / f".{path.stem}.{secrets.token_hex(8)}.partial"
+        try:
+            descriptor = os.open(temporary, flags, 0o666)
+        except FileExistsError:
+            continue
+        os.close(descriptor)
+        return temporary
+    raise FileExistsError(f"Could not create a temporary file beside {path}")
+
+
+def publish_year_file(
+    temporary: str | Path, path: str | Path, name: str, stamp: Mapping[str, Any]
+) -> None:
+    """Stamp a year file written to ``temporary``, then move it onto ``path``.
+
+    ``name`` is the country's stamp dataset. The move replaces any file at
+    ``path`` in one step (``os.replace`` within one folder), so a reader sees
+    either the old file or the new one, each with its own stamp.
+    """
+    write_json_dataset(temporary, name, dict(stamp))
+    os.replace(temporary, path)
+
+
+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 source_and_package_differences(
+    stamp: Mapping[str, Any], current: Mapping[str, Any]
+) -> list[str]:
+    """Return the clauses naming how a stamp's source and packages differ.
+
+    ``stamp`` is a year file's well-formed stamp and ``current`` the stamp a
+    file cut now would carry: each has a ``"source"`` that
+    ``dataset_materialization.is_dataset_source_identity`` accepts and
+    ``"packages"`` that :func:`is_version_map` accepts. One clause, starting
+    "it", names both sources when they differ, and one more names each
+    package whose version differs, in name order. A package the stamp
+    records that ``current`` does not is named too, and a package recorded
+    in both with one version is not.
+
+    Returns no clause exactly when the sources and the package maps are
+    equal.
+    """
+    reasons = []
+    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 reasons
+
+
+def year_file_sources(
+    country_id: str, datasets: list[str], allow_unmanaged: bool = True
+) -> list[tuple[str, dict[str, Any]]]:
+    """Return each requested dataset's year-file stem and 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 ``dataset_materialization.requested_dataset_source_identity``).
+    A request ``materialize_dataset`` would refuse with the same
+    ``allow_unmanaged`` raises the same error, before its stem is resolved.
+
+    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 = []
+    for dataset in datasets:
+        identity = requested_dataset_source_identity(
+            country_id, dataset, allow_unmanaged=allow_unmanaged
+        )
+        stem = dataset_logical_name(resolve_dataset_reference(country_id, dataset))
+        sources.append((stem, identity))
+    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 "
+                f"{country_id.upper()} year 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