diff --git a/PROGRESS.md b/PROGRESS.md new file mode 100644 index 00000000000..2220073150f --- /dev/null +++ b/PROGRESS.md @@ -0,0 +1,134 @@ +# Simulation isolation repairs (Astra review of country 2.0.1) + +Branch `max/spm-runtime-isolation-20260914` off `upstream/main` (2.0.1). +Source: `astra-review-country-201.md`. PR A of two. + +## State +- [x] Read review; read `policyengine_us/spm.py`, core `simulations/simulation.py`, + `taxbenefitsystems/tax_benefit_system.py`, `reforms/reform.py`, `parameters/*`. +- [x] Measured `system.parameters.clone()`: **1.5-2.0 s** against a 0.03 s warm + household build and a 0.16 s warm calculate. Detaching on every simulation + (the structural reform is re-applied to *every* one) was therefore not an + option; copy-on-write had to be genuinely lazy. +- [x] Reproduced findings 1-5 on this runtime and turned each into a regression + test. Every new test fails on the parent commit `e24e655a04`. +- [x] Fix findings 1-4 in `policyengine_us/spm.py`. +- [x] Fix finding 5 (county input type) and the two documentation defects. +- [x] Full core/unit + microsimulation + representative YAML runs. +- [x] Draft PR #9463 opened against `main`. +- [x] Adversarial review (4 lenses, 27 raised, 18 confirmed after per-finding + verification) and every confirmed finding repaired - see below. +- [ ] Final re-run after the review repairs; update the PR body. + +## Design decisions +1. **Copy-on-write parameter tree.** `share_spm_policy` returns an instance of a + per-base `SharedParameter` subclass whose `parameters` is a property. + `SPMSimulationMixin.apply_reform` and the shared system's own + `apply_reform_set` arm a read barrier for the duration of the reform; the + first read inside that window clones the tree through core and gives this + system a cold private at-instant cache, leaving the lender's tree and warm + caches untouched. `modify_parameters`, `load_extension` and + `add_abolition_parameters` - the three core methods that edit the live tree + in place - detach outright. A variable-only reform never reads `parameters`, + so the structural reform re-applied at each simulation's start instant still + pays nothing. Measured: household build stays 0.05 s; a parameter reform pays + 1.24 s once. Branches created with `clone_system=False` are moved onto the + detached tree, because they share their parent's policy deliberately. +2. **Clone alias routing.** Core keeps `self.calc = self.calculate` and + `self.df = self.calculate_dataframe` as bound methods in the instance dict and + copies them verbatim when cloning. `clone()` now rebinds every copied bound + method by method name, so an alias core adds later is repaired too. +3. **Private entities.** `share_spm_policy` shallow-copies the entities and binds + them to the private system, keeping one object per key across `entities`, + `person_entity` and `group_entities` (`clone_spm_system` now does the same). + An entity resolves variable names through the system it is bound to, so this + is what makes `set_input` find a reform-added variable. +4. **Traced parameter receipts.** Only the *root* node carries core's + request-specific `trace`/`tracer`/`branch_name` and caches the resulting + `TracingParameterNodeAtInstant`; children are wrapped on the fly. So a traced + simulation gets a shallow private root with its own at-instant cache + (microseconds, not a 1.5 s tree clone), primed as traced - core marks the tree + in `_run_formula`, but `_calculate` reads `gov.abolitions` first, so on + core 3.30.2 nothing was ever recorded. Two traced simulations now each record + their own four `spm_unit_fpg` parameter accesses, and the shared tree is left + untraced. +5. **County input type.** Tightening `is_county_fips` alone would have been + inert: `value_type = str` maps to the numpy `object` dtype, so the model + stores an integer column as integers, and both readers stringify before + asking the provider. The provider now records which counties arrived as + something other than text - at the one point that still sees the type - and + rejects them when a county measurement needs them, before its own memo. A + caller who never asks for a county is still never asked for one. + +## Hardening carried from the blast-radius audit +- A shared-policy system clones through an ordinary copy: core's clone writes + the cloned tree into the instance dictionary, where the property would shadow + it and hand back a "clone" still sharing the lender's tree. +- The shared-policy class is published under its own name, so systems built on + it still pickle. +- `tools/pinned_tbs.py` keys its cache on the parameter tree as well as the + system, because a shared system keeps its identity while detaching a tree. + +## Notes +- Installed core here is **3.30.2**, not the 3.32.5 Astra reviewed. `calc`/`df` + sit at `simulation.py:223-224` and the clone copy loop at 1381-1416, matching + Astra's citations. +- Astra's finding 1 sub-case "an existing unrelated simulation changes after + cache invalidation" needs a real invalidation: `delete_arrays("income_tax")` + alone leaves the cached standard deduction, so the contaminated tree is not + reached. The regression test uses `_invalidate_all_caches()`, the same purge + `apply_reform` performs. +- Astra's finding 4 count ("four accesses, then zero") does not reproduce + literally on core 3.30.2, where a traced simulation recorded *no* parameter + accesses at all: `_calculate` reads `gov.abolitions` and caches a non-tracing + at-instant node before `_run_formula` marks the tree. Priming the private root + as traced repairs that as well, and the test asserts the four accesses per + simulation that Astra described. + +## Review repairs (adversarial review of the change itself) + +Confirmed findings and what each cost, all now fixed and pinned by tests that +fail when the fix is reverted: + +1. **Ownership was tracked on the borrower, not on the tree.** A system that + detached a private copy and then lent it to a second simulation went on + writing to it in place. Sharing is now a one-way mark on the tree, and every + simulation's own system carries the barrier - not only the ones borrowing + the shipped policy - because a tree a simulation starts out owning acquires + readers later. Core's `ParameterNode.clone` copies the instance dictionary, + so a private clone is explicitly unmarked. +2. **Traced simulations stranded their shared-policy branches.** Branches were + moved onto a detached tree by root identity, which a traced branch can never + match, because tracing hands it a shallow copy of the root. Since the model's + own formulas branch this way (itemizing vs not, state EITC refundability, + marginal tax rates), a traced `calculate -> apply_reform -> calculate` + reported pre-reform numbers - up to a $2,860 error in the verifier's + constructed case. Branches are matched on the children of their root now. +3. **An adopting branch had its copy-on-write switched off**, so its own reform + rewrote its parent's tree. It keeps sharing instead. +4. **A shallow traced root orphaned its children's parent chain**, so an + in-place edit cleared the original root's at-instant cache and left the + traced copy serving pre-reform values. The root is only copied when the tree + is shared; otherwise it is marked traced where it stands. +5. **A tuple of variable-only reforms paid for a full tree clone**, because + core recurses into `apply_reform` per member and the nested call's read of + `parameters` fell inside the outer call's armed window. 1.29 s -> 0.001 s. +6. **The county input-type record was a monotone, unkeyed blacklist.** Keyed on + `(year, county)` and replaced per year now, so correcting an input clears + it; carried only by `copy_receipts=True` snapshots, so a new simulation on a + lender's system is not tainted; recorded on a reform simulation's baseline + arm as well; skipped entirely for non-county selections; and each write + scans only the period it wrote. +7. **Test-efficacy gaps**: the clone override, the entity de-duplication, the + `load_extension` detach and the pinned-system re-key had no test that would + fail if reverted. All four now do (verified by reverting each). + +Residual, documented rather than fixed: + +- A county rejection is scoped to a county code and year, not to a row, because + the county reaching a provider has already been stringified by the variable + that read the column. A simulation holding a mistyped code fails every + request for that code in that year. +- `holder.set_input` bypasses the country's `set_input`, so a county written + straight through core's holder API after construction is not re-recorded. + That channel behaves exactly as it did before this change. diff --git a/changelog.d/spm-runtime-isolation.fixed.md b/changelog.d/spm-runtime-isolation.fixed.md new file mode 100644 index 00000000000..2b07dcf1c74 --- /dev/null +++ b/changelog.d/spm-runtime-isolation.fixed.md @@ -0,0 +1 @@ +Isolate one simulation's reform, clone and trace from every other simulation: a parameter reform applied after construction no longer changes unrelated simulations, a clone's `calc` and `df` aliases calculate on the clone rather than the original, `set_input` finds a variable a reform added, each traced request records its own parameter accesses, and a `county_fips` input that is not a five-digit string is rejected however it is spelled. diff --git a/docs/usage/microsimulation.md b/docs/usage/microsimulation.md index 95db177874a..c472d8e8fdd 100644 --- a/docs/usage/microsimulation.md +++ b/docs/usage/microsimulation.md @@ -110,12 +110,16 @@ formula-owned SPM outputs such as `spm_unit_spm_threshold`. Any observed Census measurement is retained under a separate report-only name. The legacy files under `hf://policyengine/policyengine-us-data/` predate that -contract, so resource and poverty outputs are not available over them: +contract, so SPM measurements are not available over them: - `cps_2023.h5` stores `spm_unit_spm_threshold`, so the loader rejects it. - `enhanced_cps_2024.h5` loads and computes tax variables, but carries no SPM - independence roles, so SPM units consisting of one 15-to-17-year-old classify - no measurement adult and every resource output over the file fails closed. + independence roles, so 18 of its SPM units - each a lone 15-to-17-year-old - + classify no measurement adult, and `spm_unit_spm_threshold` raises + `SPM_COMPOSITION_REQUIRED` over the file. That is what + `test_legacy_enhanced_cps_lacks_source_backed_spm_independence_roles` + checks. Outputs that no longer reach the threshold, household net income + and benefits among them, do compute over the file. A household simulation that has no county input can select an SPM area explicitly instead: @@ -205,12 +209,13 @@ When running microsimulations, verify that weights produce sensible population t ```python sim = Microsimulation() -# Check population -person_weight = sim.calc("person_weight", map_to="person") +# Check population. The weights are the values here, so read them +# unweighted: a weighted sum would square them. +person_weight = sim.calc("person_weight", map_to="person", use_weights=False) print(f"Total population: {person_weight.sum():,.0f}") # Check household count -household_weight = sim.calc("household_weight") +household_weight = sim.calc("household_weight", use_weights=False) print(f"Total households: {household_weight.sum():,.0f}") # Verify key aggregates against published statistics @@ -221,6 +226,14 @@ total_snap = sim.calc("snap").sum() print(f"Total SNAP benefits: ${total_snap / 1e9:.1f}B") ``` +`calc` returns a weighted series whose `sum()` multiplies each value by its +weight, which is what makes the aggregates above population totals. Summing a +weight variable that way squares the weights: with household weights of 2 and +3, `sim.calc("household_weight").sum()` is 13 rather than 5. Either opt out +with `use_weights=False`, as above, or use `count()`, which sums the weights +themselves - `sim.calc("age", map_to="person").count()` is the same population +total. + Compare these totals against official statistics to validate your analysis: - US population: ~330 million - US households: ~130 million diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py index bedc086632e..8a276d7b062 100644 --- a/policyengine_us/spm.py +++ b/policyengine_us/spm.py @@ -7,10 +7,12 @@ import re from collections.abc import Mapping +from contextlib import contextmanager from copy import copy, deepcopy from functools import lru_cache -from inspect import signature +from inspect import ismethod, signature +from policyengine_core import periods as periods_ from policyengine_core.simulations import Simulation as CoreSimulation from policyengine_core.taxbenefitsystems import TaxBenefitSystem from spm_calculator.errors import SPMInputError @@ -41,19 +43,28 @@ def is_county_fips(value): - """Accept only a five-digit county FIPS code. - - Integers, pandas missing values and truncated codes all arrive here as - strings, because ``county_fips`` is a string variable and the model casts - every input to its own dtype. A within-state CPS code such as ``5`` and a - missing value such as ``nan`` are not county FIPS codes, so they must fail - as an absent county rather than as an unrecognised one. + """Accept only a five-digit county FIPS code, supplied as text. + + The advertised contract is a five-digit *string*. Stringifying whatever + arrived enforced the digit count alone, so an integer whose decimal form + happens to be five digits (``36061``) passed, while the same mistake for a + state whose code carries a leading zero (``6037``, meaning Los Angeles + County ``06037``) failed - one silently accepted, the other reported as an + absent county. Reject every non-text value here instead, so an integer + county code, a pandas missing value and a truncated code all fail the same + way and name the same fix. + + ``numpy.str_`` and ``numpy.bytes_`` subclass ``str`` and ``bytes``, so a + county read back out of the model - which stores ``county_fips`` as text - + is still accepted. A within-state CPS code such as ``5`` and a missing + value such as ``nan`` fail as an absent county rather than an unrecognised + one, whichever type they arrive as. """ - if value is None: - return False if isinstance(value, bytes): value = value.decode() - return COUNTY_FIPS_PATTERN.fullmatch(str(value)) is not None + if not isinstance(value, str): + return False + return COUNTY_FIPS_PATTERN.fullmatch(value) is not None class CountyRequiringSPMProvider(PolicyEngineSPMProvider): @@ -65,15 +76,85 @@ class CountyRequiringSPMProvider(PolicyEngineSPMProvider): request sending an integer or a missing value, would otherwise be told its county assignment is unavailable, which points at the artifact instead of at the input. + + This provider also remembers which counties the model was handed as + something other than text. ``county_fips`` declares ``value_type = str``, + but core maps ``str`` to the numpy ``object`` dtype, so the model stores + whatever it is given and an integer column stays integers - and both + readers of that column stringify before calling a provider. An integer + county code for a state without a leading zero therefore arrived here + indistinguishable from the same code sent correctly as a string, and was + silently accepted, while the same mistake for California (``6037`` for Los + Angeles County ``06037``) was reported as an absent county. Recording the + input's type at the one point that still sees it, and rejecting it when a + county measurement actually needs that county, makes the two fail the same + way without demanding SPM geography from a caller who never asks for it. """ - def calculate_unit(self, *, year, adults, children, tenure, county_fips=None): - if self.geography_kind == "county" and not is_county_fips(county_fips): + def __post_init__(self): + super().__post_init__() + # The base is a frozen dataclass, so this receipt of input types is + # attached rather than declared. It maps (year, county text) to the + # repr of what was actually supplied, and ``record_county_input_types`` + # replaces a year's entries outright, so correcting an input clears it. + object.__setattr__(self, "_untyped_counties", {}) + + def snapshot(self, *, copy_receipts=False): + snapshot = super().snapshot(copy_receipts=copy_receipts) + if copy_receipts: + # A new simulation re-reads its own inputs; only a clone, which + # keeps this simulation's holders, inherits what they held. + snapshot._untyped_counties.update(self._untyped_counties) + return snapshot + + def record_county_input_types(self, year, values): + """Replace this year's record of counties not supplied as text. + + Replacing rather than accumulating is what lets a caller correct the + input: re-setting ``county_fips`` for a period re-derives the record + for that period from what the model now holds. + """ + if self.geography_kind != "county": + # Nothing will ask this provider for a county. + return + untyped = self._untyped_counties + for key in [key for key in untyped if key[0] == year]: + del untyped[key] + for value in values: + if not isinstance(value, (str, bytes)): + untyped[(year, str(value))] = repr(value) + + def require_county_input(self, year, county_fips): + """Reject a county this provider cannot honour as a five-digit string.""" + if self.geography_kind != "county": + return + if is_county_fips(county_fips): + supplied = self._untyped_counties.get((year, county_fips)) + if supplied is None: + return + # The county reaching a provider has already been stringified by + # the variable that read the column, so this is as precise as the + # rejection gets: every request for this county in this year + # fails, not only the household whose input was mistyped. raise SPMInputError( "SPM_GEOGRAPHY_REQUIRED", - f"County selection has no county FIPS input ({county_fips!r}): " - f"{COUNTY_INPUT_FIX}", + f"County {county_fips} was supplied for {year} as {supplied}, " + f"not as text: {COUNTY_INPUT_FIX}", ) + raise SPMInputError( + "SPM_GEOGRAPHY_REQUIRED", + f"County selection has no county FIPS input ({county_fips!r}): " + f"{COUNTY_INPUT_FIX}", + ) + + def _amounts(self, year, adults, children, tenure, county): + # Check before the memo, so a cache entry filled by a correctly typed + # row cannot let an identically spelled untyped row through. + self.require_county_input(year, county) + return super()._amounts(year, adults, children, tenure, county) + + def calculate_unit(self, *, year, adults, children, tenure, county_fips=None): + self.require_county_input(year, county_fips) return super().calculate_unit( year=year, adults=adults, @@ -181,20 +262,305 @@ def spm_config(provider): } +SHARED_TREE_MARK = "_spm_tree_is_shared" + + +def mark_tree_shared(tree): + """Record that more than one system can now read this parameter tree. + + Ownership is tracked on the tree rather than on the systems reading it, + because a system's own history says nothing about who else holds the tree + it is pointing at: a system that took a private copy, and then lent it to a + second simulation, must stop writing to that copy in place. The mark is + one-way. Dropping it would need to know when the last other reader went + away, and a tree wrongly believed private is silently shared state, while a + tree wrongly believed shared only costs one clone. + """ + if tree is not None: + setattr(tree, SHARED_TREE_MARK, True) + return tree + + +def unmark_tree_shared(tree): + """Record that this tree has exactly one reader. + + Core's ``ParameterNode.clone`` copies the node's instance dictionary, so a + clone taken from a shared tree arrives carrying the mark even though it is + brand new and private. + """ + if tree is not None: + setattr(tree, SHARED_TREE_MARK, False) + return tree + + +def tree_is_shared(tree): + return tree is not None and getattr(tree, SHARED_TREE_MARK, False) + + +SHARED_POLICY_BASES = {} + + +def shared_policy_class(base): + """Return the shared-parameter subclass of ``base``, creating it once. + + ``share_spm_policy`` hands back a system whose parameter tree belongs to + someone else, so the tree needs a read barrier that a plain attribute + cannot provide. Subclassing keeps that barrier off every other system: an + independently built ``CountryTaxBenefitSystem`` is untouched. + """ + if getattr(base, "shared_policy_base", None) is not None: + # Already a shared-policy class: sharing a shared system is idempotent. + return base + created = SHARED_POLICY_BASES.get(base) + if created is None: + name = f"SharedParameter{base.__name__}" + if name in globals(): + # Two bases sharing a class name would otherwise publish one class + # under the other's name, and pickle would restore the wrong one. + name = f"{name}{len(SHARED_POLICY_BASES) + 1}" + created = type( + name, (SharedParameterPolicy, base), {"shared_policy_base": base} + ) + # Publish it under its own name so a system built from it still pickles: + # a class only reachable from a dictionary has no import path. + created.__module__ = __name__ + created.__qualname__ = name + globals()[name] = created + SHARED_POLICY_BASES[base] = created + return created + + +class SharedParameterPolicy: + """Policy state that borrows another system's parameter tree. + + Reading parameters through this system costs nothing: the tree and its warm + at-instant caches are the lending system's own. Writing to it must not + happen at all, because core applies a reform to whatever tree the system + hands it - ``Simulation.apply_reform`` calls ``reform.apply(system)`` + directly, bypassing ``Reform.__init__``'s defensive clone - and a mutated + shared tree changes unrelated simulations, including ones already built. + + So the tree is copy-on-write. Reform application arms the read barrier; + the first read of ``parameters`` inside that window takes a private clone + through core, leaving the lending system's tree and warm caches untouched. + Cloning rebuilds every node of a 130,000-parameter tree and takes seconds, + which is why it is deferred to the reforms that actually reach the tree: + the structural reform that every simulation re-applies at its own start + instant only rebinds variables, so it never reads ``parameters`` and never + pays for a clone. + """ + + shared_policy_base = None + + @property + def parameters(self): + if self.__dict__.get("detaching_shared_parameters", False): + self.detach_parameters() + return self.__dict__["shared_parameters"] + + @parameters.setter + def parameters(self, value): + self.__dict__["shared_parameters"] = value + + @property + def shares_parameters(self): + """Whether anything but this system can see writes to this tree.""" + return tree_is_shared(self.__dict__["shared_parameters"]) + + @property + def detaching_shared_parameters(self): + """Whether a reform is being applied to this system right now.""" + return self.__dict__.get("detaching_shared_parameters", False) + + def detach_parameters(self): + """Take a private copy of the shared tree; report whether one was made.""" + tree = self.__dict__["shared_parameters"] + if tree is None or not tree_is_shared(tree): + return False + # Clear the barrier first: cloning reads the tree through core, and + # that read must not re-enter this method. + self.__dict__["detaching_shared_parameters"] = False + self.__dict__["shared_parameters"] = unmark_tree_shared(tree.clone()) + # The lender keeps its warm at-instant cache; this system starts cold + # because its parameter values are about to differ. + self._parameters_at_instant_cache = {} + return True + + @contextmanager + def detaching_parameters(self): + """Arm the copy-on-write barrier for the duration of a reform.""" + previous = self.__dict__.get("detaching_shared_parameters", False) + self.__dict__["detaching_shared_parameters"] = True + try: + yield + finally: + self.__dict__["detaching_shared_parameters"] = previous + + def apply_reform_set(self, reform): + # Core applies a reform to this system's own tree, so detach first + # however the reform arrives. + with self.detaching_parameters(): + return super().apply_reform_set(reform) + + def modify_parameters(self, modifier_function): + # The documented way for a reform to edit parameters, and core hands + # the live tree to the modifier rather than a copy. + self.detach_parameters() + return super().modify_parameters(modifier_function) + + def load_extension(self, extension): + # Core merges the extension's parameters into this tree in place. + self.detach_parameters() + return super().load_extension(extension) + + def add_abolition_parameters(self): + # Core adds an abolition child per variable, in place. + self.detach_parameters() + return super().add_abolition_parameters() + + def clone(self): + # Core's clone writes the cloned tree into the new instance's + # dictionary, where this class's property shadows it, so hand it an + # ordinary instance and let the base decide how to clone. + return type(self).shared_policy_base.clone(self.unshared_copy()) + + def unshared_copy(self): + """A plain shallow copy holding the tree as an ordinary attribute. + + Core's ``TaxBenefitSystem.clone`` writes the cloned tree straight into + the new instance's ``__dict__`` under ``parameters``, which this class's + property would shadow, so anything that clones through core starts from + an ordinary instance instead. + """ + policy = copy(self) + tree = policy.__dict__.pop("shared_parameters", None) + policy.__dict__.pop("shares_parameters_with_lender", None) + policy.__dict__.pop("detaching_shared_parameters", None) + policy.__class__ = type(self).shared_policy_base + policy.parameters = tree + return policy + + +def plain_policy_copy(system): + """Shallow-copy any system into one that holds its tree as an attribute.""" + if isinstance(system, SharedParameterPolicy): + return system.unshared_copy() + return copy(system) + + +def with_parameter_barrier(system): + """Give ``system`` a copy-on-write parameter tree, in place. + + Every simulation's own system carries the barrier, not only the ones that + borrow the shipped policy, because the tree a simulation starts out owning + can acquire other readers later - a branch copies the system, and another + simulation can be built on it - and from then on a reform must clone + before it writes. Whether anyone else is reading is recorded on the tree + (see :func:`mark_tree_shared`), so installing the barrier costs nothing + and changes nothing on its own. + """ + if isinstance(system, SharedParameterPolicy): + return system + system.__class__ = shared_policy_class(type(system)) + system.__dict__["shared_parameters"] = system.__dict__.pop("parameters", None) + system.__dict__["detaching_shared_parameters"] = False + return system + + +def bind_private_entities(policy): + """Give ``policy`` its own Entity objects, bound to its own registry. + + An Entity resolves variable names through the system it is bound to + (``Entity.get_variable``), and a population asks its entity before reading + or creating a holder. Sharing the lending system's entities therefore sent + every holder lookup to the lending system's registry, so a variable that a + reform added to this system's private registry was invisible to + ``set_input``, which raised ``VariableNotFoundError`` for a variable the + system plainly had. + + Core's own constructor and ``clone`` copy entities for the same reason. + Unlike ``clone``, keep ``person_entity`` and ``group_entities`` pointing at + the very objects in ``entities``, so there is exactly one entity object per + key to rebind. + """ + entities = [copy(entity) for entity in policy.entities] + by_key = {entity.key: entity for entity in entities} + policy.entities = entities + policy.person_entity = by_key[policy.person_entity.key] + policy.group_entities = [by_key[entity.key] for entity in policy.group_entities] + policy.group_entity_keys = [entity.key for entity in policy.group_entities] + for entity in entities: + entity.set_tax_benefit_system(policy) + return policy + + +def isolate_parameter_tracing(system, tracer, branch_name): + """Give ``system`` its own root parameter node, primed for ``tracer``. + + Core marks a traced request by writing that request's tracer, trace flag + and branch name onto the root node of the parameter tree, and the root then + caches the resulting ``TracingParameterNodeAtInstant`` in its own + at-instant cache. On a shared tree the second request reuses the first + request's cached tracing node, so its parameter accesses are recorded + against the first request's tracer - or, once that request's trace frame + has closed, recorded nowhere at all - and the shared tree is left holding a + finished request's tracer for every simulation that reads it afterwards. + + Only the root carries those fields: children are wrapped on the fly from + the root's tracing node and cache plain at-instant nodes. A shallow copy of + the root - the same children, its own at-instant cache and its own trace + fields - therefore isolates a request's parameter receipts completely, and + costs microseconds rather than the seconds a full tree clone takes. The + children stay shared, so a later reform still detaches the whole tree. + + Prime the copy as traced rather than waiting for core to mark it. Core sets + those fields in ``_run_formula``, but ``_calculate`` reads + ``parameters(period).gov.abolitions`` first, so the at-instant node for + each period is built and cached before the tree is ever marked as traced + and no parameter access is recorded at all. + """ + root = system.parameters + if root is None: + return None + if tree_is_shared(root): + private = copy(root) + private._at_instant_cache = {} + # The copy shares the original's children, and each child's parent + # still points at the original root, so an edit made through a child + # would clear the original's at-instant cache and leave this one + # stale. Mark the copy shared: any write detaches a real clone first. + mark_tree_shared(private) + system.parameters = private + system._parameters_at_instant_cache = {} + root = private + else: + # Nobody else reads this tree, so it can be marked traced where it + # stands - and leaving it in place keeps each parameter's parent + # pointing at the root whose at-instant cache an edit has to clear. + root._at_instant_cache = {} + system._parameters_at_instant_cache = {} + root.trace = True + root.tracer = tracer + root.branch_name = branch_name + return root + + def share_spm_policy(system): """Isolate receipts and variable registration without rebuilding policy. - An ordinary simulation applies no user reform, so it needs private receipts - and a private variable registry, not a private copy of the policy itself. - It is also new rather than cloned, so no previous receipt belongs to it. - Core's TaxBenefitSystem.clone() rebuilds the whole parameter tree node by - node and empties both at-instant caches, and this country then deep-copies - every variable object on top; doing that per household simulation throws - away the shared instance's warm parameter caches and lands on household API - request latency. - - Share the parameter tree and its at-instant caches, and reuse the variable - objects. Every core operation that a reform performs on a variable + An ordinary simulation applies no user reform, so it needs private + receipts, a private variable registry and private entities to resolve that + registry - not a private copy of the policy itself. It is also new rather + than cloned, so no previous receipt belongs to it. + + Core's ``TaxBenefitSystem.clone()`` rebuilds the whole parameter tree node + by node and empties both at-instant caches, and this country then + deep-copies every variable object on top; doing that per household + simulation throws away the shared instance's warm parameter caches and + lands on household API request latency. + + So share the parameter tree and its at-instant caches, and reuse the + variable objects. Every core operation that a reform performs on a variable (add_variable, replace_variable, update_variable, neutralize_variable, annualize_variable) rebinds ``variables[name]`` to a newly constructed object rather than mutating the registered one, so a private dict is enough @@ -202,9 +568,18 @@ def share_spm_policy(system): re-applied at its own start instant - out of the shared instance. ``test_ordinary_simulation_shares_default_policy_state`` enforces that invariant against the shared instance itself. + + A reform that reaches the parameter tree has no such protection, so the + returned system takes a private copy of the tree the first time a reform + reads it: see :class:`SharedParameterPolicy`. """ policy = copy(system) policy.variables = dict(system.variables) + bind_private_entities(policy) + with_parameter_barrier(policy) + # The lender is now one reader among several, so neither side may write to + # this tree in place - including a lender that had taken a private copy. + mark_tree_shared(policy.__dict__["shared_parameters"]) policy.spm_forecast_provider = system.spm_forecast_provider.snapshot( copy_receipts=False ) @@ -218,10 +593,19 @@ def clone_spm_system(system, *, copy_receipts=True): losing instance changes such as neutralization and inherited reform fields. Let core clone parameters/entities, then copy the actual variable state. """ - policy = copy(system) + policy = plain_policy_copy(system) policy.variables = {} cloned = TaxBenefitSystem.clone(policy) + unmark_tree_shared(cloned.parameters) + # Core copies ``entities``, ``person_entity`` and ``group_entities`` + # separately, leaving two objects per group entity key; keep one, so a + # rebind reaches every reader of that entity. by_key = {entity.key: entity for entity in cloned.entities} + cloned.person_entity = by_key[cloned.person_entity.key] + cloned.group_entities = [by_key[entity.key] for entity in cloned.group_entities] + cloned.group_entity_keys = [entity.key for entity in cloned.group_entities] + for entity in cloned.entities: + entity.set_tax_benefit_system(cloned) memo = {id(system): cloned, id(system.parameters): cloned.parameters} for entity in system.entities: memo[id(entity)] = by_key[entity.key] @@ -241,6 +625,40 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Core switches the baseline system after cloning its populations. self._rebind_holders() + # Core sets ``trace`` before this simulation owns its policy state. + self._isolate_parameter_tracing() + # Every input path - situation builder, dataset loader, direct + # set_input - has landed in the holder by now. + self._record_county_input_types() + + @property + def trace(self): + return CoreSimulation.trace.fget(self) + + @trace.setter + def trace(self, trace): + CoreSimulation.trace.fset(self, trace) + self._isolate_parameter_tracing() + + def _isolate_parameter_tracing(self): + """Stop a traced request writing its tracer onto shared parameters.""" + if not self.trace: + return + policy = getattr(self, "tax_benefit_system", None) + # Core sets ``trace`` before a new simulation owns its policy state, + # both while constructing one and while cloning one, so leave the + # original's system alone until this simulation has claimed its own. + if policy is not None and getattr(policy, "simulation", None) is self: + isolate_parameter_tracing(policy, self.tracer, self.branch_name) + for branch in getattr(self, "branches", {}).values(): + branch._isolate_parameter_tracing() + + def get_branch(self, name="branch", clone_system=False): + branch = super().get_branch(name, clone_system) + # Core names the branch and hands it this simulation's tracer after + # cloning, so re-prime the branch's root with what it ended up holding. + branch._isolate_parameter_tracing() + return branch def _rebind_holders(self): """Bind populations and cached holders to their branch's policy state.""" @@ -267,9 +685,61 @@ def _rebind_holders(self): branch._rebind_holders() def apply_reform(self, reform): - super().apply_reform(reform) + policy = self.tax_benefit_system + detaching = getattr(policy, "detaching_parameters", None) + if detaching is None or policy.detaching_shared_parameters: + # Either an ordinary system, or core recursing through a tuple of + # reforms into this override again. Reading ``parameters`` inside + # the armed window would itself trip the barrier and clone the + # tree, which is exactly what a variable-only reform must not pay. + super().apply_reform(reform) + self._rebind_holders() + return + previous_children = getattr(policy.parameters, "children", None) + with detaching(): + super().apply_reform(reform) + self._adopt_detached_parameters(previous_children) + # A detached tree is a fresh root, carrying none of the trace state + # core wrote onto the one it replaced. + self._isolate_parameter_tracing() self._rebind_holders() + def _adopt_detached_parameters(self, previous_children): + """Move branches off the pre-reform tree onto the detached copy. + + A branch created with ``clone_system=False`` shares its parent's policy + state deliberately, so a reform the parent applies is a reform the + branch runs under. Detaching would otherwise strand the branch on the + unreformed tree. + + Branches are recognised by the children of their root node rather than + by the root itself, because a traced branch holds a shallow copy of the + root with those same children (see ``isolate_parameter_tracing``) and + would never match on identity. An adopting branch keeps sharing - the + tree now has more than one reader, so its own reform must detach rather + than rewrite its parent's policy. + """ + detached = self.tax_benefit_system.parameters + if detached is None or detached.children is previous_children: + return + cache = self.tax_benefit_system._parameters_at_instant_cache + for name, branch in self.branches.items(): + if name == "baseline": + # The baseline branch holds unreformed policy by construction. + continue + policy = branch.tax_benefit_system + tree = getattr(policy, "parameters", None) + if tree is not None and tree.children is previous_children: + # The detached tree now has a second reader. + mark_tree_shared(detached) + policy.parameters = detached + policy._parameters_at_instant_cache = cache + # A traced branch needs its own root back, with its own + # at-instant cache, or its parameter receipts rejoin the + # parent's. + branch._isolate_parameter_tracing() + branch._adopt_detached_parameters(previous_children) + @property def spm_config(self): return spm_config(self.tax_benefit_system.spm_forecast_provider) @@ -287,8 +757,10 @@ def _prepare_spm_system(self, args, kwargs, config, start_instant): reform = arguments.arguments.get("reform") if supplied is None: if reform is not None: - chosen = self.default_tax_benefit_system( - reform=reform, spm=config, start_instant=start_instant + chosen = with_parameter_barrier( + self.default_tax_benefit_system( + reform=reform, spm=config, start_instant=start_instant + ) ) else: # No reform and no supplied system: nothing distinguishes this @@ -298,7 +770,9 @@ def _prepare_spm_system(self, args, kwargs, config, start_instant): elif reform is not None: # Core applies the reform set to whatever system it is handed, so # the caller's own system must not be the one it reforms. - chosen = clone_spm_system(supplied, copy_receipts=False) + chosen = with_parameter_barrier( + clone_spm_system(supplied, copy_receipts=False) + ) else: # A caller that builds one system and runs many households through # it - the household API among them - keeps that system's warm @@ -332,14 +806,71 @@ def clone(self, debug=False, trace=False, clone_tax_benefit_system=True): # Core branches may share policy state, but each calculation's # receipt belongs to that simulation's provider. cloned.tax_benefit_system = copy(self.tax_benefit_system) + # Two systems now read this tree, so neither may write to it in + # place: a branch's reform must not rewrite its parent's policy. + mark_tree_shared(getattr(cloned.tax_benefit_system, "parameters", None)) cloned.tax_benefit_system.spm_forecast_provider = ( self.tax_benefit_system.spm_forecast_provider.snapshot( copy_receipts=True ) ) cloned._rebind_holders() + self._rebind_method_aliases(cloned) + cloned._isolate_parameter_tracing() return cloned + def _rebind_method_aliases(self, cloned): + """Point copied bound-method aliases at the clone, not the original. + + Core keeps backwards-compatibility aliases as bound methods on the + instance (``self.calc = self.calculate``, ``self.df = + self.calculate_dataframe``) and its ``clone`` copies the instance + dictionary verbatim, so every alias on the clone still called the + original simulation. ``clone.calc(...)`` therefore returned the + original's values under the original's policy, and recorded the + original's SPM receipts, while ``clone.calculate(...)`` - documented as + the same call - returned the clone's. Rebind by method name so an alias + core adds later is repaired too. + """ + for name, value in list(cloned.__dict__.items()): + if ismethod(value) and value.__self__ is self: + cloned.__dict__[name] = getattr(cloned, value.__func__.__name__) + + def _record_county_input_types(self, period=None): + """Tell this simulation's providers which counties were not text. + + The county a provider is asked for has already been stringified by the + variable that reads the column, so the type has to be captured here, + where the stored input is still whatever the caller supplied. Record it + on this simulation's own provider and on its branches', because a + reform simulation's baseline arm is handed a separate provider that + never sees an input of its own. + """ + holder = self.get_holder("county_fips") + periods = ( + holder.get_known_periods() if period is None else [periods_.period(period)] + ) + for known_period in periods: + array = holder.get_array(known_period) + if array is None or array.dtype.kind in "SU": + # A real text dtype cannot be hiding an integer. + continue + year = int(known_period.start.year) + for simulation in self._simulation_family(): + record = getattr( + simulation.tax_benefit_system.spm_forecast_provider, + "record_county_input_types", + None, + ) + if record is not None: + record(year, array) + + def _simulation_family(self): + """This simulation and every branch that reads the same inputs.""" + yield self + for branch in getattr(self, "branches", {}).values(): + yield from branch._simulation_family() + def set_input(self, variable_name, period, value): # Core's loader calls set_input for every dataset format. Reject saved # measurement/resource outputs here without loading the population twice. @@ -353,4 +884,7 @@ def set_input(self, variable_name, period, value): f"Dataset supplies formula-owned SPM output {variable_name}. " "Use primitive inputs and retain observed outputs under report-only names." ) - return super().set_input(variable_name, period, value) + result = super().set_input(variable_name, period, value) + if variable_name == "county_fips": + self._record_county_input_types(period) + return result diff --git a/policyengine_us/tests/core/test_spm_simulation_isolation.py b/policyengine_us/tests/core/test_spm_simulation_isolation.py new file mode 100644 index 00000000000..1c88396d872 --- /dev/null +++ b/policyengine_us/tests/core/test_spm_simulation_isolation.py @@ -0,0 +1,488 @@ +"""One simulation's reform, clone or trace must not reach another's. + +Every test here reproduces a defect an external review found on the shipped +2.0.1 runtime, where ordinary simulations share the default policy state. +""" + +import hashlib +from copy import copy + +import pytest +from policyengine_core.periods import YEAR +from policyengine_core.reforms import Reform +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.variables import Variable + +from policyengine_us import Simulation +from policyengine_us.entities import Person +from policyengine_us.spm import clone_spm_system, share_spm_policy +from policyengine_us.system import system +from policyengine_us.tools.pinned_tbs import get_pre_arpa_eitc_tbs + +SINGLE_STANDARD_DEDUCTION = "gov.irs.deductions.standard.amount.SINGLE" +BASELINE_INCOME_TAX = 4_016 +REFORMED_INCOME_TAX = 3_718 + + +def earner_situation(): + return { + "people": {"person": {"age": {2024: 40}, "employment_income": {2024: 50_000}}}, + "households": { + "household": {"members": ["person"], "county_fips": {2024: "06037"}} + }, + } + + +def parameter_fingerprint(policy): + """Digest every authored parameter value, so a shared-tree edit is visible.""" + digest = hashlib.sha256() + for parameter in policy.parameters.get_descendants(): + values = getattr(parameter, "values_list", None) + if values is None: + continue + digest.update(parameter.name.encode()) + for value_at_instant in values: + digest.update( + f"|{value_at_instant.instant_str}={value_at_instant.value}".encode() + ) + return digest.hexdigest() + + +def parameter_accesses(tracer): + """Every parameter name this tracer recorded, at any depth.""" + recorded = [] + + def visit(node): + recorded.extend(parameter.name for parameter in node.parameters) + for child in node.children: + visit(child) + + for tree in tracer.trees: + visit(tree) + return recorded + + +class NeutralizeIncomeTax(Reform): + def apply(self): + self.neutralize_variable("income_tax") + + +class clone_only_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Income input added only by the test reform" + + +class AddedInputReform(Reform): + def apply(self): + self.update_variable(clone_only_income) + + +def test_parameter_reform_leaves_every_other_simulation_unreformed(): + """A reform applied to one simulation must not change the shipped policy. + + Core's ``Simulation.apply_reform`` calls ``reform.apply(system)`` directly + rather than constructing the reform, so it bypasses the defensive parameter + clone in ``Reform.__init__`` and edits whatever tree the system hands it. + On a shared tree that silently re-based every other simulation - including + ones already built and calculated - onto the reformed parameter. + """ + fingerprint_before = parameter_fingerprint(system) + reformed = Simulation(situation=earner_situation()) + unrelated = Simulation(situation=earner_situation()) + assert reformed.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + assert unrelated.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + + reformed.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert reformed.calculate("income_tax", 2024)[0] == REFORMED_INCOME_TAX + # An existing simulation, recomputed from its own inputs. Use the same + # purge apply_reform performs: deleting one variable's arrays leaves the + # cached intermediates that a contaminated tree would have to flow through. + unrelated._invalidate_all_caches() + assert unrelated.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + assert unrelated.calculate("standard_deduction", 2024)[0] < 100_000 + # And one built after the reform. + later = Simulation(situation=earner_situation()) + assert later.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + # Nothing reached the shared instance itself. + assert parameter_fingerprint(system) == fingerprint_before + + +def test_parameter_reform_detaches_the_shared_tree_and_its_warm_caches(): + simulation = Simulation(situation=earner_situation()) + policy = simulation.tax_benefit_system + assert policy.parameters is system.parameters + assert policy._parameters_at_instant_cache is system._parameters_at_instant_cache + + simulation.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert policy.parameters is not system.parameters + assert policy._parameters_at_instant_cache is not ( + system._parameters_at_instant_cache + ) + assert not policy.shares_parameters + + +def test_variable_only_reform_keeps_sharing_the_parameter_tree(): + """Detaching rebuilds 130,000 parameter nodes, so it must stay lazy. + + Every simulation re-applies the structural reform at its own start instant, + and that reform only rebinds variables. Detaching for it would put a full + tree clone on every household API request. + """ + simulation = Simulation(situation=earner_situation()) + policy = simulation.tax_benefit_system + + simulation.apply_reform(NeutralizeIncomeTax) + + assert simulation.calculate("income_tax", 2024)[0] == 0 + assert policy.parameters is system.parameters + assert policy.shares_parameters + + +def test_shared_policy_branch_follows_its_parents_parameter_reform(): + """A branch cloned without its own system shares policy deliberately.""" + simulation = Simulation(situation=earner_situation()) + branch = simulation.get_branch("shared_policy") + + simulation.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert branch.tax_benefit_system.parameters is ( + simulation.tax_benefit_system.parameters + ) + assert branch.calculate("income_tax", 2024)[0] == REFORMED_INCOME_TAX + isolated = simulation.get_branch("separate_policy", clone_system=True) + assert isolated.calculate("income_tax", 2024)[0] == REFORMED_INCOME_TAX + + +def test_clone_method_aliases_calculate_on_the_clone(): + """``calc`` is documented as ``calculate``; core copies it as a bound method. + + Core keeps ``self.calc = self.calculate`` and + ``self.df = self.calculate_dataframe`` in the instance dictionary and its + ``clone`` copies that dictionary verbatim, so every alias on the clone + still called the original simulation, under the original's policy. + """ + original = Simulation(situation=earner_situation()) + assert original.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + clone = original.clone() + clone.apply_reform(NeutralizeIncomeTax) + + assert clone.calculate("income_tax", 2024)[0] == 0 + assert clone.calc("income_tax", 2024)[0] == 0 + assert clone.calc.__self__ is clone + assert clone.df.__self__ is clone + # The original keeps its own policy and its own answer. + assert original.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + + +def test_clone_method_alias_records_receipts_on_the_clone(): + original = Simulation(situation=earner_situation()) + clone = original.clone() + + clone.calc("spm_unit_spm_threshold", 2025) + + assert set(clone.spm_provenance()["years"]) == {"2025"} + assert original.spm_provenance()["years"] == {} + + +def test_reform_added_variable_accepts_input_on_an_ordinary_simulation(): + """A population asks its entity for a variable before it builds a holder. + + The entity resolves the name through the system it is bound to, so a shared + system whose entities still belonged to the shared instance raised + ``VariableNotFoundError`` for a variable its own registry plainly held. + """ + simulation = Simulation(situation=earner_situation()) + simulation.apply_reform(AddedInputReform) + + simulation.set_input("clone_only_income", 2024, [123]) + + assert simulation.calculate("clone_only_income", 2024)[0] == 123 + assert "clone_only_income" not in system.variables + + +@pytest.mark.parametrize("period", [2033, 2034]) +def test_traced_simulations_record_their_own_parameter_accesses(period): + """Core marks the parameter tree with the requesting simulation's tracer. + + On a shared tree the second traced request reused the first request's + cached tracing node, so its parameter accesses were recorded against the + first request's tracer - or, with that request's trace frame closed, + recorded nowhere - and the shared tree kept a finished request's tracer for + every simulation that read it afterwards. + """ + first = Simulation(situation=earner_situation()) + first.trace = True + first.calculate("spm_unit_fpg", period) + first_accesses = parameter_accesses(first.tracer) + + second = Simulation(situation=earner_situation()) + second.trace = True + second.calculate("spm_unit_fpg", period) + + expected = { + "gov.abolitions.spm_unit_fpg", + "gov.abolitions.spm_unit_size", + "gov.hhs.fpg.additional_person", + "gov.hhs.fpg.first_person", + } + assert set(first_accesses) == expected + assert set(parameter_accesses(second.tracer)) == expected + # The first simulation's trace did not grow when the second one calculated. + assert parameter_accesses(first.tracer) == first_accesses + # Each traced request owns the root node core writes its tracer onto. + first_root = first.tax_benefit_system.parameters + second_root = second.tax_benefit_system.parameters + assert first_root is not second_root + assert first_root._at_instant_cache is not second_root._at_instant_cache + assert first_root.tracer is first.tracer + assert second_root.tracer is second.tracer + + +def test_tracing_leaves_the_shared_parameter_tree_untraced(): + traced = Simulation(situation=earner_situation()) + traced.trace = True + traced.calculate("spm_unit_fpg", 2035) + + assert traced.tax_benefit_system.parameters is not system.parameters + assert system.parameters.trace is False + assert system.parameters.tracer is None + assert system.parameters.branch_name is None + # Tracing changes no answer. + untraced = Simulation(situation=earner_situation()) + assert ( + untraced.calculate("spm_unit_fpg", 2035)[0] + == (traced.calculate("spm_unit_fpg", 2035)[0]) + ) + + +def test_cloning_a_shared_policy_system_stops_sharing_its_tree(): + """Core's clone writes the cloned tree into the instance dictionary. + + A shared system resolves ``parameters`` through a property, which would + shadow that write and hand back a "clone" still sharing the lender's tree, + with nothing raised. + """ + shared = Simulation(situation=earner_situation()).tax_benefit_system + assert shared.parameters is system.parameters + + cloned = shared.clone() + + assert cloned.parameters is not system.parameters + assert not getattr(cloned, "shares_parameters", False) + assert "shared_parameters" not in cloned.__dict__ + assert parameter_fingerprint(cloned) == parameter_fingerprint(system) + + +def test_shared_policy_class_has_an_import_path(): + """A class only reachable from a dictionary cannot be pickled or named. + + Sharing policy puts every simulation's system on a class this module + creates at run time, so that class has to be published under its own name - + otherwise anything that pickles a system (a worker pool, joblib, a notebook + checkpoint) fails on the class itself rather than on its contents. + """ + import pickle + + import policyengine_us.spm as spm + + shared_class = type(Simulation(situation=earner_situation()).tax_benefit_system) + + assert shared_class.__module__ == spm.__name__ + assert getattr(spm, shared_class.__qualname__) is shared_class + assert pickle.loads(pickle.dumps(shared_class)) is shared_class + + +def test_parameter_reform_on_a_shared_branch_leaves_its_parent_alone(): + """A branch that reforms policy owns the tree it reformed.""" + simulation = Simulation(situation=earner_situation()) + branch = simulation.get_branch("reformed_branch") + + branch.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert branch.calculate("income_tax", 2024)[0] == REFORMED_INCOME_TAX + assert simulation.calculate("income_tax", 2024)[0] == BASELINE_INCOME_TAX + assert simulation.tax_benefit_system.parameters is system.parameters + + +@pytest.mark.parametrize( + "mutate", + [ + lambda policy: policy.modify_parameters( + {SINGLE_STANDARD_DEDUCTION: {"year:2024:1": 100_000}} + ), + lambda policy: policy.add_abolition_parameters(), + # Core's load_extension ends in ``self.parameters.merge(...)``. Stub + # the base out: a real extension is not needed to pin that this + # country override detaches before core reaches the tree. + lambda policy: policy.load_extension("any_extension"), + ], + ids=["modify_parameters", "add_abolition_parameters", "load_extension"], +) +def test_in_place_parameter_edits_detach_before_they_write(mutate, monkeypatch): + """Core edits these into the live tree rather than a copy of it.""" + monkeypatch.setattr( + TaxBenefitSystem, "load_extension", lambda self, extension: None + ) + fingerprint_before = parameter_fingerprint(system) + policy = Simulation(situation=earner_situation()).tax_benefit_system + + mutate(policy) + + assert policy.parameters is not system.parameters + assert parameter_fingerprint(system) == fingerprint_before + + +def test_shared_policy_over_a_core_cloning_base_clones_without_sharing(): + """A base that clones through core must not defeat the property. + + Core's ``TaxBenefitSystem.clone`` writes the cloned tree straight into the + new instance's dictionary, where ``SharedParameterPolicy.parameters`` + shadows it, so a system built on a base that uses core's clone would hand + back a "clone" still reading the lender's tree, with nothing raised. + """ + base = copy(system) + base.__class__ = type( + "CoreCloningCountrySystem", + (type(system),), + {"clone": TaxBenefitSystem.clone}, + ) + shared = share_spm_policy(base) + assert shared.parameters is system.parameters + + cloned = shared.clone() + + assert cloned.parameters is not system.parameters + assert parameter_fingerprint(cloned) == parameter_fingerprint(system) + + +def test_cloned_systems_hold_one_entity_object_per_key(): + """Core's clone leaves two objects per group-entity key, both bound. + + Rebinding through ``entities`` would then miss the copy that + ``instantiate_entities`` and core's simulation builder read. + """ + cloned = clone_spm_system(system) + + by_key = {entity.key: entity for entity in cloned.entities} + assert len(by_key) == len(cloned.entities) + assert cloned.person_entity is by_key[cloned.person_entity.key] + assert all(entity is by_key[entity.key] for entity in cloned.group_entities) + for entity in cloned.entities: + assert entity._tax_benefit_system is cloned + + +def test_pinned_systems_are_rebuilt_when_their_source_tree_detaches(): + """A shared system keeps its identity while swapping its tree. + + The NY EITC and CTC formulas reuse a cached clone pinned to an earlier + federal vintage. Keyed on the system alone, that clone would survive a + reform the system detached for, and NY's decoupled credits would silently + go on reading pre-reform federal parameters. + """ + simulation = Simulation(situation=earner_situation()) + policy = simulation.tax_benefit_system + pinned = get_pre_arpa_eitc_tbs(policy) + assert get_pre_arpa_eitc_tbs(policy) is pinned + + simulation.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert policy.parameters is not system.parameters + assert get_pre_arpa_eitc_tbs(policy) is not pinned + + +@pytest.mark.parametrize("traced", [False, True]) +def test_shared_branch_follows_its_parent_whether_or_not_it_is_traced(traced): + """Tracing gives a simulation its own root node, not its own policy. + + A branch cloned without its own system shares its parent's policy, so the + detached tree has to reach it however the two roots are spelled. Matching + on root identity missed every traced branch, and the model's own formulas + branch this way - itemizing against not itemizing, the state EITC + refundability branches, marginal tax rates - so a traced simulation + reported pre-reform numbers under a reform. + """ + simulation = Simulation(situation=earner_situation()) + simulation.trace = traced + branch = simulation.get_branch("shared_policy") + + simulation.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert simulation.calculate("income_tax", 2024)[0] == REFORMED_INCOME_TAX + assert branch.calculate("income_tax", 2024)[0] == REFORMED_INCOME_TAX + assert branch.tax_benefit_system.parameters.children is ( + simulation.tax_benefit_system.parameters.children + ) + + +def test_traced_simulation_sees_a_later_parameter_reform(): + """A traced root must not outlive the tree its children belong to. + + A shallow root copy keeps the original's children, and each child's parent + still points at the original root, so a parameter edit clears that root's + at-instant cache and leaves the traced copy serving pre-reform values. + """ + simulation = Simulation( + situation=earner_situation(), + reform=Reform.from_dict({SINGLE_STANDARD_DEDUCTION: {"2024": 20_000}}), + ) + simulation.trace = True + assert simulation.calculate("standard_deduction", 2024)[0] == 20_000 + + simulation.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert simulation.calculate("standard_deduction", 2024)[0] == 100_000 + + +def test_lending_a_detached_tree_stops_the_lender_writing_to_it(): + """Owning a tree once is not owning it forever. + + A simulation that detached a private tree, and then lent it to a second + simulation, has to clone again before its next reform: the borrower can + see everything it writes. + """ + lender = Simulation(situation=earner_situation()) + lender.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + borrower = Simulation( + tax_benefit_system=lender.tax_benefit_system, situation=earner_situation() + ) + assert borrower.calculate("standard_deduction", 2024)[0] == 100_000 + + lender.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 7_777}}) + + assert lender.calculate("standard_deduction", 2024)[0] == 7_777 + borrower._invalidate_all_caches() + assert borrower.calculate("standard_deduction", 2024)[0] == 100_000 + + +def test_a_branch_reform_does_not_rewrite_its_parents_policy(): + simulation = Simulation(situation=earner_situation()) + branch = simulation.get_branch("reforming_branch") + simulation.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 50_000}}) + + branch.apply_reform({SINGLE_STANDARD_DEDUCTION: {"2024": 100_000}}) + + assert branch.calculate("standard_deduction", 2024)[0] == 100_000 + assert simulation.calculate("standard_deduction", 2024)[0] == 50_000 + + +def test_a_tuple_of_variable_only_reforms_keeps_sharing_the_tree(): + """Core recurses into this override for each member of a reform tuple. + + Reading ``parameters`` inside the outer call's armed window would itself + trip the barrier, so every tuple reform - including the purely structural + ones - would have paid for a full tree clone. + """ + simulation = Simulation(situation=earner_situation()) + policy = simulation.tax_benefit_system + + simulation.apply_reform((NeutralizeIncomeTax, AddedInputReform)) + + assert simulation.calculate("income_tax", 2024)[0] == 0 + simulation.set_input("clone_only_income", 2024, [123]) + assert simulation.calculate("clone_only_income", 2024)[0] == 123 + assert policy.parameters is system.parameters + assert policy.shares_parameters diff --git a/policyengine_us/tests/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py index 2b60bcd08ed..9177cbe660d 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -15,7 +15,7 @@ from policyengine_us import Microsimulation, Simulation from policyengine_us.data.dataset_schema import USMultiYearDataset, USSingleYearDataset from policyengine_us.entities import Person -from policyengine_us.spm import create_spm_provider +from policyengine_us.spm import create_spm_provider, is_county_fips from policyengine_us.system import DEFAULT_DATASET, _resolve_dataset_path, system @@ -152,8 +152,22 @@ def test_ordinary_simulation_shares_default_policy_state(): assert policy is not system assert policy.parameters is system.parameters assert policy._parameters_at_instant_cache is system._parameters_at_instant_cache - # Entities stay bound to the shared instance, exactly as before this change. - assert policy.entities is system.entities + # Entities are private and bound to this simulation's own registry: an + # entity resolves variable names through the system it is bound to, so + # sharing the shared instance's entities would send every holder lookup to + # the shared registry. + assert policy.entities is not system.entities + assert {entity.key for entity in policy.entities} == { + entity.key for entity in system.entities + } + for entity in policy.entities: + assert entity._tax_benefit_system is policy + for entity in system.entities: + assert entity._tax_benefit_system is system + # One object per key, so rebinding one reaches every reader of it. + by_key = {entity.key: entity for entity in policy.entities} + assert policy.person_entity is by_key[policy.person_entity.key] + assert all(entity is by_key[entity.key] for entity in policy.group_entities) # A private registry, holding the shared instance's own variable objects. assert policy.variables is not system.variables assert ( @@ -381,23 +395,134 @@ def missing(*args, **kwargs): @pytest.mark.parametrize("simulation_type", [Simulation, Microsimulation]) -def test_integer_county_column_requires_county_fips_instead_of_reporting_unavailable( - simulation_type, +@pytest.mark.parametrize( + "counties", + [ + # A legacy population file storing the CPS within-state code. + [5, 1], + # A county code whose decimal form is five digits, so only its type + # distinguishes it from the documented input. + [36_061, 36_061], + # The same mistake for a state whose code carries a leading zero. + [6_037, 6_037], + [float("nan"), float("nan")], + ], + ids=["within_state", "five_digit_integer", "dropped_leading_zero", "missing"], +) +def test_non_text_county_column_requires_county_fips_however_it_is_spelled( + simulation_type, counties ): - """Legacy population files store the CPS within-state code as an integer.""" + """``county_fips`` is documented as a five-digit *string*. + + Core maps a ``str`` variable onto the numpy ``object`` dtype, so the model + stores an integer column as integers, and every reader stringifies before + asking the forecast provider. An integer county code for a state without a + leading zero was therefore accepted silently, while the identical mistake + for California was reported as an absent county. + """ source = small_dataset() - source.household["county_fips"] = [5, 1] + source.household["county_fips"] = counties simulation = simulation_type(dataset=source) with pytest.raises(SPMInputError) as error: simulation.calculate("spm_unit_spm_threshold", 2024) assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" assert "five-digit string" in str(error.value) assert 'geography_kind="national"' in str(error.value) - # The same population computes once an SPM area is selected explicitly. + # The same population computes once an SPM area is selected explicitly: + # a caller who never asks for a county is never asked for one. national = simulation_type(dataset=source, spm={"geography_kind": "national"}) assert np.all(national.calculate("spm_unit_spm_threshold", 2024) > 0) +def test_correcting_a_county_input_clears_its_rejection(): + """The record has to follow the input, not outlive it.""" + situation = single_person_situation() + situation["households"]["household"]["county_fips"] = {2024: 36_061} + simulation = Simulation(situation=situation) + with pytest.raises(SPMInputError): + simulation.calculate("spm_unit_spm_threshold", 2024) + + simulation.set_input("county_fips", 2024, ["36061"]) + + assert simulation.calculate("spm_unit_spm_threshold", 2024)[0] > 0 + + +def test_a_county_rejection_does_not_follow_the_system_to_a_new_simulation(): + """A new simulation reads its own inputs; only a clone keeps these.""" + situation = single_person_situation() + situation["households"]["household"]["county_fips"] = {2024: 36_061} + mistyped = Simulation(situation=situation) + with pytest.raises(SPMInputError): + mistyped.calculate("spm_unit_spm_threshold", 2024) + + # The same county, sent correctly: the lender's record of its own mistake + # must not reject it. + situation = single_person_situation() + situation["households"]["household"]["county_fips"] = {2024: "36061"} + correct = Simulation( + tax_benefit_system=mistyped.tax_benefit_system, situation=situation + ) + assert correct.calculate("spm_unit_spm_threshold", 2024)[0] > 0 + with pytest.raises(SPMInputError): + mistyped.clone().calculate("spm_unit_spm_threshold", 2024) + + +def test_a_reform_simulations_baseline_arm_rejects_the_same_county_input(): + """Core hands the baseline arm a provider that never sees an input.""" + situation = single_person_situation() + situation["households"]["household"]["county_fips"] = {2024: 36_061} + simulation = Simulation( + situation=situation, + reform=Reform.from_dict( + {"gov.irs.credits.ctc.amount.base[0].amount": {"2024": 0}} + ), + ) + for arm in (simulation, simulation.baseline): + with pytest.raises(SPMInputError) as error: + arm.calculate("spm_unit_spm_threshold", 2024) + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" + + +def test_a_national_selection_records_no_county_input_types(): + """Nothing will ask that provider for a county, so nothing is scanned.""" + source = small_dataset() + source.household["county_fips"] = [36_061, 36_061] + national = Microsimulation(dataset=source, spm={"geography_kind": "national"}) + provider = national.tax_benefit_system.spm_forecast_provider + assert provider._untyped_counties == {} + assert np.all(national.calculate("spm_unit_spm_threshold", 2024) > 0) + + +def test_text_county_column_still_resolves_its_county(): + source = small_dataset() + source.household["county_fips"] = ["06037", "36061"] + simulation = Microsimulation(dataset=source) + assert np.all(simulation.calculate("spm_unit_spm_threshold", 2024) > 0) + assert len(simulation.spm_provenance()["geographies"]) == 2 + + +@pytest.mark.parametrize( + ("value", "accepted"), + [ + ("06037", True), + (b"06037", True), + (np.str_("36061"), True), + (np.bytes_(b"36061"), True), + ("6037", False), + ("", False), + ("360610", False), + (36_061, False), + (np.int64(36_061), False), + (36_061.0, False), + (float("nan"), False), + (None, False), + (True, False), + ], +) +def test_county_fips_accepts_five_digit_text_only(value, accepted): + assert is_county_fips(value) is accepted + + def small_dataset(): groups = ("household", "tax_unit", "spm_unit", "family", "marital_unit") return USSingleYearDataset( diff --git a/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py b/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py index 0439c6ab2be..29d0e786128 100644 --- a/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py +++ b/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py @@ -267,6 +267,8 @@ def mock_super_init(self, *args, **kwargs): for population in self.populations.values(): population.simulation = self self._user_input_keys = set() + # Core's __init__ sets these before the country's __init__ resumes. + self.trace = False self.branches = {} self.baseline = None self.is_over_dataset = True diff --git a/policyengine_us/tests/microsimulation/test_microsim.py b/policyengine_us/tests/microsimulation/test_microsim.py index 7fa111a3754..25f764eee7e 100644 --- a/policyengine_us/tests/microsimulation/test_microsim.py +++ b/policyengine_us/tests/microsimulation/test_microsim.py @@ -24,10 +24,16 @@ def test_legacy_enhanced_cps_lacks_source_backed_spm_independence_roles(): Its county column does hold five-digit FIPS codes, so geography is not what fails. 18 of its 43,134 SPM units are a lone 15-to-17-year-old carrying no source-backed SPM independence role, so those units classify no measurement - adult, and every output that reaches the threshold - household net income, - benefits, poverty and marginal tax rates included - fails closed over the - file. Assert on the threshold itself: it raises straight off the + adult and every output that reaches the threshold fails closed over the + file. Only the threshold is asserted here: it raises straight off the composition, without building the whole resource chain over 43,134 units. + + Resource outputs are no longer among the outputs that reach it. The + capped housing subsidy asks for the SPM housing portion only for units + receiving housing assistance, and none of these 18 units does - which is + asserted below - so household net income and benefits compute over the + file. Computing them here would add two minutes to this suite, so they + are left to the resource tests. """ import numpy as np from spm_calculator.errors import SPMInputError @@ -41,6 +47,11 @@ def test_legacy_enhanced_cps_lacks_source_backed_spm_independence_roles(): adults = np.asarray(simulation.calculate("spm_measurement_adults", 2024)) assert (adults < 1).sum() == 18 + assisted = np.asarray( + simulation.calculate("receives_housing_assistance", 2024, map_to="spm_unit") + ) + assert assisted.sum() == 870 + assert not assisted[adults < 1].any() with pytest.raises(SPMInputError) as error: simulation.calculate("spm_unit_spm_threshold", 2024) assert error.value.code == "SPM_COMPOSITION_REQUIRED" diff --git a/policyengine_us/tools/pinned_tbs.py b/policyengine_us/tools/pinned_tbs.py index 2e9c1318836..8f2ce50f787 100644 --- a/policyengine_us/tools/pinned_tbs.py +++ b/policyengine_us/tools/pinned_tbs.py @@ -14,22 +14,32 @@ from policyengine_core.parameters import Parameter from policyengine_core.periods import instant -# pin name -> (weakref to base TBS, pinned clone). +# pin name -> (weakref to base TBS, weakref to its parameter tree, pinned clone). # The clone is held strongly so it survives across simulations; the # weakref to the base is used to detect that the base system changed # (e.g. a reformed system in another test), in which case the entry is # rebuilt. Using a weakref (rather than id() alone) also guards against # id reuse after the base system is garbage-collected. +# +# The parameter tree is tracked separately because a shared-policy system +# keeps its identity while swapping its tree: a reform detaches a private +# copy in place (see policyengine_us.spm.SharedParameterPolicy), so a pin +# built before that reform would otherwise be reused after it. _PINNED_TBS_CACHE = {} def _get_pinned_tbs(base_tbs, pin_name, pin_fn): entry = _PINNED_TBS_CACHE.get(pin_name) - if entry is not None and entry[0]() is base_tbs: - return entry[1] + parameters = base_tbs.parameters + if entry is not None and entry[0]() is base_tbs and entry[1]() is parameters: + return entry[2] pinned = base_tbs.clone() pin_fn(pinned) - _PINNED_TBS_CACHE[pin_name] = (weakref.ref(base_tbs), pinned) + _PINNED_TBS_CACHE[pin_name] = ( + weakref.ref(base_tbs), + weakref.ref(parameters), + pinned, + ) return pinned