Fix the two production-path breakages the node graph left in the UK spine driver (#836)

changelog.d/uk-spi-support-channel-mass-policy-836.fixed.md

@@ -0,0 +1 @@
+Declare the UK `spi_support_channel` graph node's mass change instead of `conserve`: the executor's ledger is weighted person mass, which a stage that stacks differently composed synthetic households at conserved household mass cannot hold, so every full licensed UK spine build failed at that node after the node-graph landed. The kernel now states the person-mass ledger and asserts household-mass conservation itself; a regression test reproduces the class on a two-household fixture.

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

@@ -53,8 +53,21 @@
     {"spi_support_channel", "cgt_incidence_clone", "cgt_band_donors"}
 )
 
+# The executor's mass ledger is weighted *person* mass per stratum
+# (``Frame.stratum_mass``: household weights broadcast through membership), so
+# ``conserve`` is satisfiable only by an expansion that keeps household
+# composition fixed.  CGT cloning does (a clone is its source household at
+# half weight).  The SPI support channel does not: it stacks synthetic
+# households whose person counts differ from the FRS households whose mass
+# they take over, so household mass is conserved exactly (the stage's
+# ``allocate_zero_weight_prior_mass`` declares ``conservation: exact_total``)
+# while person mass moves with the composition change.  On the FRS 2024-25
+# spine that is 68.25m -> 65.44m persons, which ``conserve`` rejects at the
+# node.  The node therefore *declares* its mass change: the kernel states the
+# person-mass ledger the executor verifies and asserts the household-mass
+# invariant itself (``UKExpandStageKernel``).
 _STRUCTURAL_MASS = {
-    "spi_support_channel": "conserve",
+    "spi_support_channel": "declared",
     "cgt_incidence_clone": "conserve",
     "cgt_band_donors": "free",
 }

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

@@ -23,6 +23,7 @@
 from types import MappingProxyType
 from typing import TYPE_CHECKING
 
+import numpy as np
 import pandas as pd
 
 from microcosm.frame import Frame, MassChangeRecord, WeightKind
@@ -859,17 +860,67 @@ def run(self, context: KernelContext) -> KernelResult:
                 f"{after_weights.kind.value!r}, not declared "
                 f"{declared_kind.value!r}."
             )
+        receipt: dict[str, object] = {
+            "stage": self.stage,
+            "frame_mass_log_append": _mass_log_payload(before, after),
+        }
+        if context.node.mass == "declared":
+            receipt["mass"] = _declared_mass_receipt(
+                before, after, stage=self.stage, weight_entity=weight_entity
+            )
         return KernelResult(
             columns=MappingProxyType(columns),
             expand=MappingProxyType(expand),
             weights=after_weights,
-            receipt={
-                "stage": self.stage,
-                "frame_mass_log_append": _mass_log_payload(before, after),
-            },
+            receipt=receipt,
         )
 
 
+#: Relative tolerance for the weight-entity mass invariant a ``declared`` UK
+#: expansion must still hold; equals the executor's own ledger tolerance.
+_WEIGHT_ENTITY_MASS_RTOL = 1e-9
+
+
+def _declared_mass_receipt(
+    before: Frame,
+    after: Frame,
+    *,
+    stage: str,
+    weight_entity: str,
+) -> dict[str, object]:
+    """State the person-mass ledger of a ``declared`` expansion.
+
+    The executor's mass ledger is weighted person mass per stratum, which an
+    expansion that changes household composition cannot conserve even when it
+    conserves the mass of the entity it reweights.  A ``declared`` UK expansion
+    therefore states the person-mass ledger for the executor to verify and
+    asserts here the invariant that is actually its contract: the weight
+    entity's total mass is unchanged.
+    """
+
+    before_entity = float(before.weights_for(weight_entity).total)
+    after_entity = float(after.weights_for(weight_entity).total)
+    if not np.isclose(
+        after_entity, before_entity, rtol=_WEIGHT_ENTITY_MASS_RTOL, atol=0.0
+    ):
+        raise ValueError(
+            f"UK EXPAND stage {stage!r} declares its person-mass change but must "
+            f"conserve {weight_entity!r} mass: {before_entity!r} -> {after_entity!r}."
+        )
+    before_mass = before.stratum_mass()
+    after_mass = after.stratum_mass()
+    return {
+        "policy": "declared",
+        "before": float(before_mass.sum()),
+        "after": float(after_mass.sum()),
+        "stratum_before": {key: float(value) for key, value in before_mass.items()},
+        "stratum_after": {key: float(value) for key, value in after_mass.items()},
+        "weight_entity": weight_entity,
+        "weight_entity_mass_before": before_entity,
+        "weight_entity_mass_after": after_entity,
+    }
+
+
 def build_uk_registry(
     graph: Graph,
     implementations: Mapping[str, object],

packages/microcosm-build/tests/test_uk_graph.py

@@ -260,3 +260,213 @@ def test_uk_graph_json_round_trip_is_canonical() -> None:
 
     assert graph_from_json(serialized) == graph
     assert graph_to_json(graph_from_json(serialized)) == serialized
+
+
+def _mixed_size_population() -> Population:
+    """Two households of different size: one person in 10, two in 20."""
+
+    frame = Frame(
+        {
+            "person": pd.DataFrame(
+                {
+                    "person_id": pd.Series([1, 2, 3], dtype="int64"),
+                    "person_benunit_id": pd.Series([100, 200, 200], dtype="int64"),
+                    "person_household_id": pd.Series([10, 20, 20], dtype="int64"),
+                }
+            ),
+            "benunit": pd.DataFrame(
+                {"benunit_id": pd.Series([100, 200], dtype="int64")}
+            ),
+            "household": pd.DataFrame(
+                {
+                    "household_id": pd.Series([10, 20], dtype="int64"),
+                    "region": pd.Series(["LONDON", "WALES"], dtype="string"),
+                }
+            ),
+        },
+        EntitySchema(group_entities=("benunit", "household")),
+        {
+            "household": Weights(
+                np.array([1.0, 2.0], dtype=np.float64), WeightKind.DESIGN
+            )
+        },
+        pd.Series(["base", "base", "base"], dtype="string", name="stratum"),
+        metadata={"time_period": "2024"},
+    )
+    return Population.from_frame(frame, "root")
+
+
+def _mass_shifting_expand_result(*, declared: bool) -> KernelResult:
+    """Clone the one-person household and move mass onto it from the larger one.
+
+    Household mass is conserved (1 + 2 == 0.5 + 1.5 + 1.0) while person mass is
+    not (1 + 2*2 = 5 against 0.5 + 1.5*2 + 1.0 = 4.5): the shape of the SPI
+    support channel, whose prior-mass allocation moves half the household mass
+    onto stacked households whose composition differs from the FRS households
+    it is taken from.
+    """
+
+    receipt: dict[str, object] = {
+        "frame_mass_log_append": [
+            {
+                "entity": "household",
+                "old_total": 3.0,
+                "new_total": 3.0,
+                "declared_factor": None,
+                "reason": "test stack conserves household mass, not person mass",
+            }
+        ]
+    }
+    if declared:
+        receipt["mass"] = {
+            "policy": "declared",
+            "before": 5.0,
+            "after": 4.5,
+            "stratum_before": {"base": 5.0},
+            "stratum_after": {"base": 4.5},
+        }
+    return KernelResult(
+        expand={
+            "person": pd.Series(
+                [1], index=pd.Index([4], name="person_id"), dtype="int64"
+            ),
+            "benunit": pd.Series(
+                [100], index=pd.Index([300], name="benunit_id"), dtype="int64"
+            ),
+            "household": pd.Series(
+                [10], index=pd.Index([30], name="household_id"), dtype="int64"
+            ),
+        },
+        columns={
+            ("household", "is_clone"): pd.Series(
+                [False, False, True],
+                index=pd.Index([10, 20, 30], name="household_id"),
+                dtype="bool",
+            ),
+        },
+        weights=Weights(
+            np.array([0.5, 1.5, 1.0], dtype=np.float64),
+            WeightKind.IMPORTANCE,
+        ),
+        receipt=receipt,
+    )
+
+
+def test_conserve_rejects_a_mass_shift_across_household_sizes() -> None:
+    # The executor's ledger is person mass: an expansion that conserves the
+    # weight entity's mass but shifts it between households of different size
+    # cannot pass ``conserve``.  This is the class that refused the SPI support
+    # channel on the licensed FRS 2024-25 spine (68.25m -> 65.44m persons).
+    with pytest.raises(PopulationError, match="changed stratum"):
+        patch(
+            _mixed_size_population(),
+            _expand_node(),
+            _mass_shifting_expand_result(declared=False),
+        )
+
+
+def test_declared_accepts_the_same_expansion_with_the_kernel_ledger() -> None:
+    node = Node(
+        id="stack",
+        kernel="uk.stage.expand.test@1",
+        structural=StructuralDelta.EXPAND,
+        base="root",
+        params=_expand_node().params,
+        mass="declared",
+    )
+    expanded = patch(
+        _mixed_size_population(),
+        node,
+        _mass_shifting_expand_result(declared=True),
+    )
+
+    assert expanded.frame.table("person")["person_household_id"].tolist() == [
+        10,
+        20,
+        20,
+        30,
+    ]
+    assert expanded.frame.weights_for("household").total == pytest.approx(3.0)
+    record = expanded.mass_ledger[-1]
+    assert record.policy == "declared"
+    assert record.before_total == pytest.approx(5.0)
+    assert record.after_total == pytest.approx(4.5)
+
+
+def test_spi_support_channel_declares_its_mass_change_and_cgt_clones_conserve() -> None:
+    graph = uk_spine_graph(load_country_spec("uk"))
+
+    assert graph.node("spi_support_channel").mass == "declared"
+    assert graph.node("cgt_incidence_clone").mass == "conserve"
+    assert graph.node("cgt_band_donors").mass == "free"
+
+
+@pytest.mark.requires_uk
+def test_driver_projects_a_stage_record_for_every_graph_stage_on_the_fixture(
+    tmp_path,
+) -> None:
+    """The driver's record projection must cover every declared output.
+
+    ``frs_spine`` declares the entity ids and memberships among its outputs,
+    but the executor carries those outside owned cells, so the root node
+    exposes no artifact for them.  The first full licensed run through the
+    graph completed every stage and then died here, on ``person_id``; this
+    test runs the projection on the hermetic H2 fixture so the class fails
+    in CI's engine lane instead.
+    """
+
+    import importlib.util
+    from pathlib import Path
+
+    from microcosm.build.uk_runtime.graph_kernels import fixture_stage_plan_inputs
+    from microcosm.graph import ContentStore, run_graph
+
+    root = Path(__file__).resolve().parents[3]
+    fixture = root / "packages/microcosm-graph/tests/fixtures/parity/uk_spine"
+    if not fixture.exists():
+        pytest.skip("UK spine parity fixture is not present")
+    spec = importlib.util.spec_from_file_location(
+        "build_uk_frs_spine", root / "tools" / "build_uk_frs_spine.py"
+    )
+    driver = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(driver)
+
+    country = load_country_spec("uk")
+    stages = [
+        stage
+        for stage in country.sources.stages
+        if stage.stage not in UK_SPINE_EXCLUSIONS
+    ]
+    _, implementations = fixture_stage_plan_inputs(fixture / "sources")
+    graph = uk_spine_graph()
+    compiled = compile_graph(graph)
+    store = ContentStore(tmp_path / "store")
+    manifest = run_graph(
+        compiled,
+        sources={"frs": fixture / "sources"},
+        store=store,
+        kernels=uk_registry(dict(implementations)),
+        resume="forbid",
+        decisions=(),
+    )
+    final = manifest.population(compiled.versions[compiled.order[-1]])
+
+    records = driver._graph_stage_records(
+        manifest=manifest, store=store, stages=stages, frame=final
+    )
+
+    assert [record.stage for record in records] == [stage.stage for stage in stages]
+    by_stage = {record.stage: record for record in records}
+    for stage in stages:
+        assert set(by_stage[stage.stage].nonzero_share) == set(stage.outputs), (
+            stage.stage
+        )
+    root_shares = by_stage["frs_spine"].nonzero_share
+    for column in (
+        "person_id",
+        "person_benunit_id",
+        "person_household_id",
+        "benunit_id",
+        "household_id",
+    ):
+        assert root_shares[column] == 1.0

tools/build_uk_frs_spine.py

@@ -633,9 +633,18 @@ def _graph_stage_records(
     manifest,
     store: ContentStore,
     stages,
+    frame,
 ) -> tuple[StageRecord, ...]:
-    """Project immediate node artifacts onto the legacy record schema."""
+    """Project immediate node artifacts onto the legacy record schema.
+
+    Entity ids and memberships are executor-carried context, not owned cells,
+    so the root node exposes no artifact for them although ``frs_spine``
+    declares them as outputs.  Their share is read from the final population
+    instead, which is what the legacy plan recorded (identity columns are
+    never zero, so the value is 1.0 on every vintage).
+    """
 
+    structural = _structural_columns(frame)
     records: list[StageRecord] = []
     for stage in stages:
         output_node = (
@@ -651,6 +660,9 @@ def _graph_stage_records(
                 for coordinate, key in output_receipt.artifacts.items()
                 if coordinate[1] == column
             ]
+            if not matches and column in structural:
+                shares[column] = _nonzero_shares(frame, [column])[column]
+                continue
             if len(matches) != 1:
                 raise RuntimeError(
                     f"graph stage {stage.stage!r} exposes {len(matches)} artifacts "
@@ -670,6 +682,15 @@ def _graph_stage_records(
     return tuple(records)
 
 
+def _structural_columns(frame) -> frozenset[str]:
+    """Entity id and membership columns the executor carries outside owned cells."""
+
+    schema = frame.schema
+    columns = {schema.entity_id_column(entity) for entity in frame.entities}
+    columns.update(schema.membership_column(group) for group in schema.group_entities)
+    return frozenset(columns)
+
+
 def _new_build_id(timestamp: datetime) -> str:
     return f"uk-frs-spine-{timestamp.strftime('%Y%m%dT%H%M%SZ')}"
 
@@ -1253,6 +1274,7 @@ def main(argv: list[str] | None = None) -> int:
             manifest=graph_manifest,
             store=graph_store,
             stages=stages,
+            frame=frame,
         )
         sampling = sampled_root.sampling
         if spine_battery is not None: