UK national staging onto the outer stage runtime, increment 3 of #612

changelog.d/612-uk-stage-runtime.changed.md

@@ -0,0 +1,23 @@
+UK national staging onto the outer stage runtime (#612 increment 3). The
+descent fences between the retained-leaves and SPI stages are now
+content-addressed (`uk_frame_content_identity`) instead of Python object
+identity, so the certified-candidate descent guarantee survives a process
+boundary; `StageRuntime.load` gains an additive `frame_metadata_key` that
+restores caller-bound frame metadata from the validated run-context record,
+and UK stage checkpoints round-trip `time_period` through it.
+`build_uk_national_dataset` gains a checkpointed mode: with
+`--checkpoint-dir` each stage boundary persists a lossless Frame checkpoint
+through the outer stage runtime, completed stages resume from their
+checkpoints (transforms rehydrate their downstream evidence — retained-leaves
+descent identities, SPI fit-weight audit records — from the run-context
+record), and the run is pinned by a content-addressed run config (certified
+candidate digest, raw-source digests, seeds); a changed configuration is
+refused. The monolith path is untouched, and the staged build's output is
+content-identical to it. `_UKSourceFileFingerprint` is scope-reduced to the
+mid-read race guard and same-process candidate re-binding. The US PUF
+support tool's private `_builder_code_identity` is promoted to
+`microcosm.build.code_identity.builder_code_identity`; the promoted
+function raises on a repo root that is not a real checkout where the old
+helper silently produced a hollow identity, and
+`build_us_puf_support_base.py` inherits that refusal (its run-config
+contents on a real checkout are unchanged).

packages/microcosm-build/src/microcosm/build/code_identity.py

@@ -0,0 +1,75 @@
+"""Content-addressed builder code identity for resumable runs.
+
+Promoted from ``tools/build_us_puf_support_base.py``'s private helper (the
+``_stage_run_config`` idiom): a run config that pins only input digests and
+seeds still resumes across a code or dependency change, silently blending
+old-code checkpoints with new-code stages into one attested artifact. The
+fingerprint here covers every packaged source file plus the versions of the
+numeric dependencies whose behavior the checkpoints depend on, so a resume
+after a ``git pull`` or an environment upgrade is refused by the runtime's
+run-config equality check instead of blended.
+
+Only meaningful from a repository checkout (where ``packages/*/src``
+exists); resumable builders are checkout-run tools, and the function raises
+on a root without packaged sources rather than returning a hollow identity.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import sys
+from importlib import metadata as importlib_metadata
+from pathlib import Path
+
+__all__ = ["builder_code_identity"]
+
+
+def builder_code_identity(
+    repo_root: str | Path,
+    *,
+    tool_path: str | Path,
+    distributions: tuple[str, ...],
+) -> dict[str, object]:
+    """Fingerprint executable sources and dependency versions for safe resume.
+
+    Args:
+        repo_root: The repository checkout root (carries ``packages/``).
+        tool_path: The invoking tool's own file, included in the digest.
+        distributions: Installed distributions whose versions to record —
+            the numeric stack the run's outputs depend on.
+    """
+
+    root = Path(repo_root).resolve()
+    packages = root / "packages"
+    if not packages.is_dir():
+        raise ValueError(
+            f"builder code identity requires a repository checkout; {root} "
+            "carries no packages/ directory."
+        )
+    candidates = [Path(tool_path).resolve(), root / "pyproject.toml", root / "uv.lock"]
+    for source_root in sorted(packages.glob("*/src")):
+        candidates.extend(
+            path
+            for path in source_root.rglob("*")
+            if path.is_file()
+            and path.suffix in {".json", ".py", ".toml", ".yaml", ".yml"}
+        )
+    digest = hashlib.sha256()
+    for source_path in sorted(set(candidates)):
+        relative = source_path.relative_to(root).as_posix().encode("utf-8")
+        content = source_path.read_bytes()
+        digest.update(len(relative).to_bytes(8, "little"))
+        digest.update(relative)
+        digest.update(len(content).to_bytes(8, "little"))
+        digest.update(content)
+    dependency_versions: dict[str, str | None] = {}
+    for distribution in distributions:
+        try:
+            dependency_versions[distribution] = importlib_metadata.version(distribution)
+        except importlib_metadata.PackageNotFoundError:
+            dependency_versions[distribution] = None
+    return {
+        "dependency_versions": dependency_versions,
+        "python": sys.version,
+        "source_sha256": digest.hexdigest(),
+    }

packages/microcosm-build/src/microcosm/build/outer_stage_runtime.py

@@ -371,8 +371,22 @@ def complete_without_frame(
         self._append_record(record)
         return self._root / record.checkpoint_filename
 
-    def load(self, stage_name: str) -> LoadedStageCheckpoint:
-        """Load and validate the Frame checkpoint bound to a completed stage."""
+    def load(
+        self,
+        stage_name: str,
+        *,
+        frame_metadata_key: str | None = None,
+    ) -> LoadedStageCheckpoint:
+        """Load and validate the Frame checkpoint bound to a completed stage.
+
+        ``frame_metadata_key`` names a stage-metadata entry whose mapping
+        value is restored as :class:`~microcosm.frame.Frame` metadata during
+        the loader's one Frame construction (checkpoints do not serialize
+        frame metadata). The entry travels in the run context and was
+        normalized at :meth:`complete`, so restoring from it keeps the
+        metadata bound to the recorded stage rather than to a side channel.
+        Naming a key that is absent or not a mapping fails closed.
+        """
 
         context = self._load_context()
         stage_index = self._pipeline.index(stage_name)
@@ -384,6 +398,15 @@ def load(self, stage_name: str) -> LoadedStageCheckpoint:
                 f"run context record {stage_index} is {record.stage!r}, not "
                 f"{stage_name!r}."
             )
+        frame_metadata: Mapping[str, object] | None = None
+        if frame_metadata_key is not None:
+            value = record.metadata.get(frame_metadata_key)
+            if not isinstance(value, Mapping):
+                raise ValueError(
+                    f"stage {record.stage!r} metadata carries no mapping under "
+                    f"{frame_metadata_key!r} to restore as frame metadata."
+                )
+            frame_metadata = value
         checkpoint_path = self._root / record.checkpoint_filename
         actual_checkpoint_sha256 = _file_sha256(checkpoint_path)
         if actual_checkpoint_sha256 != record.checkpoint_sha256:
@@ -392,7 +415,7 @@ def load(self, stage_name: str) -> LoadedStageCheckpoint:
                 f"expected {record.checkpoint_sha256}, got "
                 f"{actual_checkpoint_sha256}."
             )
-        loaded = load_frame_checkpoint(checkpoint_path)
+        loaded = load_frame_checkpoint(checkpoint_path, frame_metadata=frame_metadata)
         expected_checkpoint_index = self._pipeline.index(record.checkpoint_stage)
         expected_metadata = {
             "artifact_kind": _ARTIFACT_KIND,
@@ -422,14 +445,22 @@ def load(self, stage_name: str) -> LoadedStageCheckpoint:
             metadata=_normalize_json_mapping(record.metadata, label="metadata"),
         )
 
-    def load_predecessor(self, stage_name: str) -> LoadedStageCheckpoint | None:
+    def load_predecessor(
+        self,
+        stage_name: str,
+        *,
+        frame_metadata_key: str | None = None,
+    ) -> LoadedStageCheckpoint | None:
         """Validate readiness and load the immediate predecessor, if one exists."""
 
         self.require_ready(stage_name)
         stage_index = self._pipeline.index(stage_name)
         if stage_index == 0:
             return None
-        return self.load(self._pipeline.names[stage_index - 1])
+        return self.load(
+            self._pipeline.names[stage_index - 1],
+            frame_metadata_key=frame_metadata_key,
+        )
 
     def _checkpoint_path(self, stage_index: int, stage_name: str) -> Path:
         return self._root / f"{stage_index:03d}_{stage_name}.frame.h5"

packages/microcosm-build/src/microcosm/build/uk_runtime/__init__.py

@@ -9,6 +9,9 @@
     materialize_uk_cgt_calibration_frame,
     uk_cgt_annual_exempt_amount,
 )
+from microcosm.build.uk_runtime.content_identity import (
+    uk_frame_content_identity,
+)
 from microcosm.build.uk_runtime.diagnostics import (
     UK_DIAGNOSTICS_SCHEMA_VERSION,
     UK_TARGET_GEOGRAPHY_LEVELS,
@@ -361,6 +364,12 @@
     support_clone_index_column,
     support_source_id_column,
 )
+from microcosm.build.uk_runtime.stage_checkpoints import (
+    UK_FRAME_METADATA_KEY,
+    load_uk_stage_checkpoint,
+    load_uk_stage_predecessor,
+    uk_stage_metadata,
+)
 from microcosm.build.uk_runtime.terminal_gates import (
     UK_DEFAULT_ZERO_WEIGHT_STRATA,
     UK_MAX_TARGET_ABS_RELATIVE_ERROR,
@@ -395,6 +404,7 @@
     "ARTIFACT_CLONE_INDEX_COLUMN",
     "UKRowwiseDoctrineSolve",
     "UK_CGT_ANNUAL_EXEMPT_AMOUNTS",
+    "UK_FRAME_METADATA_KEY",
     "UK_LOCAL_MAX_WEIGHT_RATIO",
     "UK_LOCAL_SOLVE_DOCTRINE",
     "UK_LOCAL_TARGET_LOSS_CAP",
@@ -403,6 +413,8 @@
     "UK_CGT_TAXPAYER_COUNT_COLUMN",
     "UKCGTTargetMaterialization",
     "ladder_clone_index_column",
+    "load_uk_stage_checkpoint",
+    "load_uk_stage_predecessor",
     "materialize_uk_cgt_calibration_frame",
     "rowwise_calibration_mass_reason",
     "uk_cgt_annual_exempt_amount",
@@ -670,11 +682,13 @@
     "uk_firm_source_data_from_frames",
     "uk_firm_source_data_from_ledger_facts",
     "uk_calibration_diagnostics_payload",
+    "uk_frame_content_identity",
     "uk_geography_ladder_assignment_summary",
     "uk_geography_ladder_gate",
     "uk_release_input_coverage_gate",
     "uk_release_input_coverage_required_columns",
     "uk_release_input_coverage_reviewed_exclusions",
+    "uk_stage_metadata",
     "uk_weight_summary",
     "uk_zero_weight_strata",
     "update_england_wales_lad_codes",

packages/microcosm-build/src/microcosm/build/uk_runtime/content_identity.py

@@ -0,0 +1,152 @@
+"""Content identity for UK national frames (#612 increment 3).
+
+The pre-increment-3 descent fences compared Python object identity
+(``retained.frame is frame``), which only holds while every stage runs in
+one process on the very same objects. Content identity makes the same
+guarantee survive a process boundary: two frames carry the same identity
+exactly when their schema, tables (column order, dtypes, index, values),
+typed weights, mass log, and metadata agree. A checkpoint-rehydrated frame
+is content-identical to the frame that was checkpointed, so resumable
+staged builds can keep the certified-candidate descent fence.
+
+This is deliberately not
+:func:`microcosm.build.outer_stage_runtime.frame_identity`: that identity is
+structural only (ids, memberships, provenance columns) and ignores payload
+values and weights — sufficient for row-order guarantees, too weak for a
+substitution fence.
+
+Two boundaries the digest declares rather than hides. Identities are
+comparable only within a pinned environment: the table bytes ride on
+``pd.util.hash_pandas_object``, which is not guaranteed stable across
+pandas versions — the UK run config pins the environment via
+``builder_code_identity``, and that pin is what makes cross-process
+comparison sound; do not reuse this as a cross-version artifact identity.
+And the digest header carries a version (``:v1``): bump it whenever the
+covered surface changes, so digests from different definitions can never
+compare equal by accident.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+
+import numpy as np
+import pandas as pd
+
+from microcosm.frame import Frame
+
+__all__ = ["uk_frame_content_identity"]
+
+_IDENTITY_HEADER = "microcosm-uk-frame-content-identity:v1"
+
+
+def uk_frame_content_identity(frame: Frame) -> str:
+    """Return a sha256 hex digest over the frame's full content.
+
+    Covers, in deterministic order: the entity set, each entity table's
+    column order, dtypes, index, and cell values, each weighted entity's
+    typed weight kind and vector, the strata labels, the mass log, and the
+    frame metadata. Structural-only changes (a renamed column, a reordered
+    column) move the identity just as value changes do. Link tables are
+    outside the digest: the UK national schema declares none
+    (``validate_uk_national_frame`` enforces linklessness) — extend this
+    before reusing it on a linked schema.
+    """
+
+    if not isinstance(frame, Frame):
+        raise TypeError("content identity requires a microcosm Frame.")
+    digest = hashlib.sha256()
+    digest.update(_IDENTITY_HEADER.encode("utf-8"))
+    for entity in frame.entities:
+        table = frame.table(entity)
+        digest.update(f"\x00entity\x1f{entity}".encode())
+        digest.update(
+            json.dumps(
+                {
+                    "columns": [str(column) for column in table.columns],
+                    "dtypes": [str(dtype) for dtype in table.dtypes],
+                    "index_dtype": str(table.index.dtype),
+                    "rows": int(len(table)),
+                },
+                sort_keys=True,
+            ).encode("utf-8")
+        )
+        row_hashes = pd.util.hash_pandas_object(table, index=True)
+        digest.update(np.ascontiguousarray(row_hashes.to_numpy()).tobytes())
+    for entity in frame.weighted_entities:
+        weights = frame.weights_for(entity)
+        digest.update(f"\x00weights\x1f{entity}\x1f{weights.kind.name}".encode())
+        digest.update(np.ascontiguousarray(weights.values, dtype=np.float64).tobytes())
+    strata = frame.strata
+    digest.update(b"\x00strata")
+    digest.update(
+        np.ascontiguousarray(
+            pd.util.hash_pandas_object(strata, index=True).to_numpy()
+        ).tobytes()
+    )
+    mass_log_payload = [
+        {
+            "entity": record.entity,
+            "old_total": record.old_total,
+            "new_total": record.new_total,
+            "declared_factor": record.declared_factor,
+            "reason": record.reason,
+        }
+        for record in frame.mass_log
+    ]
+    digest.update(b"\x00mass_log")
+    try:
+        digest.update(
+            json.dumps(mass_log_payload, sort_keys=True, allow_nan=False).encode(
+                "utf-8"
+            )
+        )
+    except ValueError as exc:
+        raise ValueError(
+            "the frame cannot be content-identified: its mass log carries a "
+            "non-finite value."
+        ) from exc
+    digest.update(b"\x00metadata")
+    try:
+        digest.update(
+            json.dumps(
+                _jsonable_metadata(frame.metadata),
+                sort_keys=True,
+                allow_nan=False,
+            ).encode("utf-8")
+        )
+    except ValueError as exc:
+        raise ValueError(
+            "the frame cannot be content-identified: its metadata carries a "
+            "non-finite value."
+        ) from exc
+    return digest.hexdigest()
+
+
+def _jsonable_metadata(value: object) -> object:
+    """Coerce frozen frame metadata into a canonically serializable shape.
+
+    Set members sort by ``repr`` — the same canonical order
+    ``stage_checkpoints._thawed`` uses, so a set that rides a checkpoint
+    round trip (JSON has no set type; it returns as a sequence) keeps its
+    content identity. Anything outside the JSON-shaped vocabulary is
+    refused: digesting an arbitrary object's ``repr`` could fold a memory
+    address into the identity, making it differ across processes in exactly
+    the dimension the fence exists to make robust.
+    """
+
+    if isinstance(value, dict) or hasattr(value, "items"):
+        return {str(key): _jsonable_metadata(item) for key, item in value.items()}
+    if isinstance(value, str | bool | int | float) or value is None:
+        return value
+    if isinstance(value, tuple | list | set | frozenset):
+        items = [_jsonable_metadata(item) for item in value]
+        if isinstance(value, set | frozenset):
+            return sorted(items, key=repr)
+        return items
+    raise TypeError(
+        f"frame metadata value of type {type(value).__name__} cannot be "
+        "content-identified; metadata must be composed of mappings, "
+        "sequences, sets, and scalar values."
+    )