From 67fbece28ef84a02b8ef550ac2de277d72f0ee05 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 10:11:48 -0400 Subject: [PATCH 1/9] Evaluate anchored recursions in full and drop values a cut reaches max_spiral_loops capped how many uncached frames of one variable the stack could hold, returning the default where the cap fell. Since it counted uncached frames only, a recursion's result depended on what had been calculated before: with an input of 1 in 2010 and recursive(t) = recursive(t - 1) + 1, 2021 was 10 in a new simulation and 12 after calculating 2020. policyengine-us weeks_worked (formula_2025 reading the year before, input in 2024) is cut this way from 2035 on. A recursion that reaches the limit while heading towards an input, or a period with no formula, now unwinds to the outermost calculate, which calculates the deepest period first and starts again; the chain is evaluated all the way down, max_spiral_loops frames at a time. One with no anchor is still cut, but no value calculated while a cut is on the stack outlives that calculation (values above a cut used to be kept), spiral purges delete exactly the period under the calculating branch's name (they deleted contained periods under the default branch), and each clone has its own invalidated set. Co-Authored-By: Claude Opus 5.5 --- .../fix-spiral-anchored-recursion.fixed.md | 1 + .../data_storage/in_memory_storage.py | 11 + .../data_storage/on_disk_storage.py | 6 + policyengine_core/holders/holder.py | 9 + policyengine_core/simulations/simulation.py | 315 +++++++++++++++++- tests/core/test_cycles.py | 11 +- tests/core/test_fast_cache_guards.py | 2 +- tests/core/test_spiral_order.py | 275 +++++++++++++++ tests/core/test_spiral_order_property.py | 66 ++++ tests/fixtures/spirals.py | 85 +++++ 10 files changed, 765 insertions(+), 16 deletions(-) create mode 100644 changelog.d/fix-spiral-anchored-recursion.fixed.md create mode 100644 tests/core/test_spiral_order.py create mode 100644 tests/core/test_spiral_order_property.py create mode 100644 tests/fixtures/spirals.py diff --git a/changelog.d/fix-spiral-anchored-recursion.fixed.md b/changelog.d/fix-spiral-anchored-recursion.fixed.md new file mode 100644 index 000000000..4fe33153a --- /dev/null +++ b/changelog.d/fix-spiral-anchored-recursion.fixed.md @@ -0,0 +1 @@ +Evaluate a recursion over periods that reaches `max_spiral_loops` in full when it heads towards an input or a period with no formula, instead of cutting it where the limit falls, and keep no value calculated from a cut recursion once its calculation ends, so recursive variables no longer depend on calculation order (for example, policyengine-us `weeks_worked` from 2035 on in a new simulation). diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index a1bddf76b..a3de6799f 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -157,6 +157,17 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: } self._shared.intersection_update(self._arrays) + def delete_exact(self, period: Period, branch_name: str = "default") -> None: + """Drop the array stored for exactly ``period``. + + Unlike ``delete``, this leaves the periods ``period`` contains. + """ + if self.is_eternal: + period = periods.period(periods.ETERNITY) + key = f"{branch_name}:{periods.period(period)}" + self._arrays.pop(key, None) + self._shared.discard(key) + def get_known_periods(self) -> list: # Split on the first colon only: an anchored period's string form # itself contains colons (e.g. "default:year:2027-11"). diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3563c805d..b375599fd 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -102,6 +102,12 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: if not period_item == f"{branch_name}_{period}" } + def delete_exact(self, period: Period, branch_name: str = "default") -> None: + """Drop the file recorded for exactly ``period``.""" + if self.is_eternal: + period = periods.period(periods.ETERNITY) + self._files.pop(f"{branch_name}_{periods.period(period)}", None) + def get_known_periods(self) -> list: return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50b..954238b52 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -106,6 +106,14 @@ def delete_arrays( if self._disk_storage: self._disk_storage.delete(period, branch_name) + def delete_array(self, period: Period, branch_name: str = "default") -> None: + """ + Remove the value stored for exactly ``period``, leaving the periods it contains (unlike ``delete_arrays``). + """ + self._memory_storage.delete_exact(period, branch_name) + if self._disk_storage: + self._disk_storage.delete_exact(period, branch_name) + def _get_array_from_storage( self, period: Period, branch_name: str = "default" ) -> ArrayLike: @@ -389,6 +397,7 @@ def put_in_cache( return self._set(period, value, branch_name) + self.simulation._note_cached_value(self.variable.name, period) def default_array(self) -> ArrayLike: """ diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d9723..03000c923 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,7 +1,7 @@ import hashlib import tempfile from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union import numpy as np @@ -62,6 +62,70 @@ def __init__(self, simulation: "Simulation"): ) +# How many deeper periods one outermost calculation may finish first (see +# ``_DeeperPeriodFirst``) before its recursion is cut where it stands, as a +# chain with no anchor is. +MAX_SPIRAL_DEFERRALS = 10_000 + + +class _DeeperPeriodFirst(BaseException): + """A recursion over periods reached ``max_spiral_loops`` and can end. + + ``_check_for_cycle`` raises this, instead of cutting the recursion with a + ``SpiralError``, when the chain of uncached periods is anchored: a + stored input, or a period with no formula, lies further along it (see + ``Simulation._can_finish_deeper_period_first``). The outermost + ``calculate`` catches it once every frame above has unwound, calculates + ``variable`` for ``period`` on ``simulation`` first, as an outermost + calculation with an empty stack, and then starts its own request again, + which now reaches that cached value. + + Where the recursion is cut therefore no longer depends on which periods + happen to be cached: an anchored chain is evaluated all the way down, + however long it is, in at most ``max_spiral_loops`` frames at a time. + + It derives from ``BaseException`` so that a formula that catches + ``Exception`` around a calculation does not swallow it. + """ + + def __init__(self, simulation: "Simulation", variable: str, period: Period): + super().__init__(variable, period) + self.simulation = simulation + self.variable = variable + self.period = period + + @property + def key(self) -> tuple: + return (id(self.simulation), self.variable, self.period) + + def calculate(self) -> None: + self.simulation._calculate_traced(self.variable, self.period) + + def is_cached(self) -> bool: + holder = self.simulation.get_holder(self.variable) + return holder.get_array(self.period, self.simulation.branch_name) is not None + + +@dataclass +class _SpiralDeferrals: + """Deeper periods finished first during one outermost calculation. + + ``futile`` holds the ``_DeeperPeriodFirst.key`` of each one that was not + cached when finished, which ``_check_for_cycle`` then cuts instead. + """ + + count: int = 0 + futile: set = field(default_factory=set) + # Tracers whose stack was active when a recursion was cut and has not + # emptied since: until it does, what is cached may depend on the cut. + cut_tracers: list = field(default_factory=list) + + +_spiral_deferrals: ContextVar[Optional[_SpiralDeferrals]] = ContextVar( + "_spiral_deferrals", default=None +) + + def _stable_hash_to_seed(value: str) -> int: """Deterministically hash a string to an int suitable for numpy.random.seed. @@ -663,6 +727,41 @@ def calculate( if _cached is not None: return _cached + if getattr(self.tracer, "stack", None) or _spiral_deferrals.get() is not None: + # Inside another calculation: a recursion that reaches + # ``max_spiral_loops`` unwinds to the outermost one. + return self._calculate_traced(variable_name, period, map_to, decode_enums) + + deferrals = _SpiralDeferrals() + token = _spiral_deferrals.set(deferrals) + try: + # Deeper periods to calculate before the request, deepest last. + pending = [] + while True: + try: + if not pending: + return self._calculate_traced( + variable_name, period, map_to, decode_enums + ) + pending[-1].calculate() + deferral = pending.pop() + if not deferral.is_cached(): + # Cut further down, or never stored: starting again + # would reach the same frame, so cut the chain there. + deferrals.futile.add(deferral.key) + except _DeeperPeriodFirst as deferral: + deferrals.count += 1 + pending.append(deferral) + finally: + _spiral_deferrals.reset(token) + + def _calculate_traced( + self, + variable_name: str, + period: Period, + map_to: str = None, + decode_enums: bool = False, + ) -> ArrayLike: self.tracer.record_calculation_start(variable_name, period, self.branch_name) # No per-variable RNG seeding: formulas may not use randomness at all @@ -683,6 +782,7 @@ def calculate( return result finally: self.tracer.record_calculation_end() + self._leave_frame() self.purge_cache_of_invalid_values() def map_result( @@ -999,13 +1099,33 @@ def purge_cache_of_invalid_values(self) -> None: # We wait for the end of calculate(), signalled by an empty stack, before purging the cache if self.tracer.stack: return + # A traced branch shares this simulation's stack, so its frames can + # only be purged now that the stack is empty. + branches = list(getattr(self, "branches", {}).values()) + for branch in branches: + branch.purge_cache_of_invalid_values() _fast_cache = getattr(self, "_fast_cache", None) invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: return + descendants = [] + while branches: + branch = branches.pop() + descendants.append(branch) + branches.extend(getattr(branch, "branches", {}).values()) + branch_name = getattr(self, "branch_name", "default") for _name, _period in invalidated_caches: holder = self.get_holder(_name) - holder.delete_arrays(_period) + # Delete the value this simulation calculated: the one stored + # under its own branch name, not another branch's, and for that + # period only, not the inputs or values of the periods within it. + holder.delete_array(_period, branch_name) + # Branches made since hold it too, under this branch's name. + for branch in descendants: + population = branch.get_variable_population(_name) + branch_holder = population._holders.get(_name) + if branch_holder is not None: + branch_holder.delete_array(_period, branch_name) if _fast_cache is not None: _fast_cache.pop((_name, _period), None) self.invalidated_caches = set() @@ -1289,13 +1409,124 @@ def _check_for_cycle(self, variable: str, period: Period) -> None: ) ) spiral = len(previous_periods) >= self.max_spiral_loops - if spiral: + if spiral and self._has_calculation(variable, period): + if self._can_finish_deeper_period_first(variable, period): + raise _DeeperPeriodFirst(self, variable, period) self.invalidate_spiral_variables(variable) message = "Quasicircular definition detected on formula {}@{} involving {}".format( variable, period, self.tracer.stack ) raise SpiralError(message, variable) + def _has_calculation(self, variable_name: str, period: Period) -> bool: + """Whether calculating the variable for ``period`` can calculate others. + + With no formula for ``period`` and no ``adds`` or ``subtracts``, the + value is an uprated or carried-over value or the default, which read + stored values only. Such a frame ends a recursion instead of + continuing it, so it is never a spiral. + """ + tax_benefit_system = getattr(self, "tax_benefit_system", None) + if tax_benefit_system is None: + return True + variable = tax_benefit_system.get_variable(variable_name) + if variable is None: + return True + return bool(variable.adds or variable.subtracts or variable.get_formula(period)) + + def _can_finish_deeper_period_first( + self, variable_name: str, period: Period + ) -> bool: + """Whether a recursion that reached ``max_spiral_loops`` is anchored. + + ``period`` is the newest frame of ``variable_name``. The frames from + its previous frame down to this one are one turn of the recursion. + The chain is anchored, and is then evaluated in full (see + ``_DeeperPeriodFirst``), when a variable it passes through is heading + towards something that stops it whatever is cached: + + - going back in time, a period before the variable's first formula, + when its formulas all start after the earliest date + (``formula_2010``, not ``formula``) and it has no ``adds`` or + ``subtracts``: its value there is an input, an uprated or + carried-over value or its default, none of which recurse; + - going forward, likewise a period after its ``end``; + - either way, an input stored for the variable at a period of its own + unit further along. + + Any other chain, such as one whose only base case is a condition + inside a formula, is cut at ``max_spiral_loops`` frames as before. + Being anchored depends on the variables and the inputs, not on what + has been calculated, so whether a chain is cut does not depend on + calculation order either. + """ + deferrals = _spiral_deferrals.get() + if ( + # Only an outermost ``calculate`` can finish the deeper period + # first; called any other way, the recursion is cut as before. + deferrals is None + or deferrals.count >= MAX_SPIRAL_DEFERRALS + or (id(self), variable_name, period) in deferrals.futile + ): + return False + tax_benefit_system = getattr(self, "tax_benefit_system", None) + if tax_benefit_system is None: + return False + stack = self.tracer.stack + previous = max( + ( + index + for index, frame in enumerate(stack[:-1]) + if frame["name"] == variable_name + and frame["branch_name"] == self.branch_name + ), + default=None, + ) + if previous is None: # max_spiral_loops is 0 + return False + previous_period = stack[previous]["period"] + if period.start < previous_period.start: + backward = True + elif period.start > previous_period.start: + backward = False + else: + return False + turn = stack[previous + 1 :] + input_keys = getattr(self, "_user_input_keys", ()) + visible_branches = self._get_visible_branch_names() + for frame in turn: + variable = tax_benefit_system.get_variable(frame["name"]) + frame_period = frame["period"] + if variable is None or not isinstance(frame_period, Period): + continue + if self._has_formula_boundary_ahead(variable, frame_period, backward): + return True + for input_name, branch_name, input_period in input_keys: + if ( + input_name == variable.name + and branch_name in visible_branches + and input_period.unit == variable.definition_period + and ( + input_period.start < frame_period.start + if backward + else input_period.start > frame_period.start + ) + ): + return True + return False + + @staticmethod + def _has_formula_boundary_ahead( + variable: Variable, period: Period, backward: bool + ) -> bool: + """Whether ``variable`` stops having a formula beyond ``period``.""" + if variable.adds or variable.subtracts or not variable.formulas: + return False + if backward: + first_start = next(iter(variable.formulas)) + return "0001-01-01" < first_start <= str(period.start) + return variable.end is not None and period.start.date <= variable.end + def invalidate_cache_entry(self, variable: str, period: Period) -> None: invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: @@ -1304,17 +1535,72 @@ def invalidate_cache_entry(self, variable: str, period: Period) -> None: invalidated_caches.add((variable, period)) def invalidate_spiral_variables(self, variable: str) -> None: - # Visit the stack, from the bottom (most recent) up; we know that we'll find - # the variable implicated in the spiral (max_spiral_loops+1) times; we keep the - # intermediate values computed (to avoid impacting performance) but we mark them - # for deletion from the cache once the calculation ends. - count = 0 + """Purge, once their calculation ends, the values a cut recursion reaches. + + The value a recursion cut at ``max_spiral_loops`` frames gets depends + on where the cut fell, and so on which frames were on the stack. A + later calculation reading it from the cache, from another depth, + would get a different result from calculating it afresh. So no value + calculated while the cut is on the stack outlives that calculation: + the frames on the stack are marked now, and everything cached until + each active stack empties is marked as it is cached (see + ``_note_cached_value``). They are still reused until then. + """ for frame in reversed(self.tracer.stack): - self.invalidate_cache_entry(frame["name"], frame["period"]) - if frame["name"] == variable: - count += 1 - if count > self.max_spiral_loops: - break + # A traced branch shares its parent's stack, so a frame can be + # the parent's (or another ancestor's). + owner = self._simulation_for_branch(frame.get("branch_name")) + if owner is not None: + owner.invalidate_cache_entry(frame["name"], frame["period"]) + state = _spiral_deferrals.get() + if state is None: + return + # The cut value also reaches every calculation in progress in a + # related simulation (a branch's caller, say). + for simulation in self._related_simulations(): + tracer = getattr(simulation, "tracer", None) + if getattr(tracer, "stack", None) and not any( + tracer is cut_tracer for cut_tracer in state.cut_tracers + ): + state.cut_tracers.append(tracer) + + def _related_simulations(self) -> List["Simulation"]: + """This simulation, its ancestor branches and all their branches.""" + root = self + while getattr(root, "parent_branch", None) is not None: + root = root.parent_branch + found, pending = [], [root] + while pending: + simulation = pending.pop() + found.append(simulation) + pending.extend(getattr(simulation, "branches", {}).values()) + return found + + def _note_cached_value(self, variable_name: str, period: Period) -> None: + """Mark a value cached while a cut recursion's calculation runs.""" + state = _spiral_deferrals.get() + if state is not None and state.cut_tracers: + self.invalidate_cache_entry(variable_name, period) + + def _leave_frame(self) -> None: + """Stop marking cached values once every stack a cut reached is empty.""" + state = _spiral_deferrals.get() + if state is not None and state.cut_tracers: + state.cut_tracers = [ + tracer for tracer in state.cut_tracers if getattr(tracer, "stack", None) + ] + + def _simulation_for_branch(self, branch_name: Optional[str]) -> "Simulation": + """This simulation or the ancestor branch named ``branch_name``.""" + simulation = self + while simulation is not None: + if ( + branch_name is None + or getattr(simulation, "branch_name", "default") == branch_name + ): + return simulation + simulation = getattr(simulation, "parent_branch", None) + return None # ----- Methods to access stored values ----- # @@ -1502,6 +1788,9 @@ def clone( ): new_dict[key] = value new._fast_cache = {} + # Its own set: values a spiral marks in one simulation are purged + # from that simulation's holders only. + new.invalidated_caches = set() # Only pass ``share_arrays`` when sharing, so a population or holder # ``clone`` override with the earlier signature still deep-copies. diff --git a/tests/core/test_cycles.py b/tests/core/test_cycles.py index a488765b1..793caeac3 100644 --- a/tests/core/test_cycles.py +++ b/tests/core/test_cycles.py @@ -145,9 +145,16 @@ def test_spiral_heuristic(simulation: Simulation, reference_period): def test_spiral_cache(simulation, reference_period): - simulation.calculate("variable7", period=reference_period) + # variable7 reads variable5, whose recursion is cut, so its value depends + # on where the cut fell. It is not kept once the calculation ends, so + # that a later request gets the same value whatever ran before it. + variable7 = simulation.calculate("variable7", period=reference_period) cached_variable7 = simulation.get_holder("variable7").get_array(reference_period) - assert cached_variable7 is not None + assert cached_variable7 is None + simulation.calculate("variable5", period=reference_period.last_month) + tools.assert_near( + simulation.calculate("variable7", period=reference_period), variable7 + ) def test_cotisation_1_level(simulation, reference_period): diff --git a/tests/core/test_fast_cache_guards.py b/tests/core/test_fast_cache_guards.py index 77dfcaba0..45715653d 100644 --- a/tests/core/test_fast_cache_guards.py +++ b/tests/core/test_fast_cache_guards.py @@ -72,7 +72,7 @@ def test_purge_cache_of_invalid_values_without_fast_cache_attribute(): sim.tracer = types.SimpleNamespace(stack=[]) sim.invalidated_caches = {("variable_name", "2024")} sim.get_holder = lambda name: types.SimpleNamespace( - delete_arrays=lambda period: None + delete_array=lambda period, branch_name="default": None ) sim.purge_cache_of_invalid_values() diff --git a/tests/core/test_spiral_order.py b/tests/core/test_spiral_order.py new file mode 100644 index 000000000..3e672aeca --- /dev/null +++ b/tests/core/test_spiral_order.py @@ -0,0 +1,275 @@ +"""A recursion over periods gives the same result whatever is cached. + +``max_spiral_loops`` caps how many frames of one variable the stack holds. +A recursion longer than that used to be cut where the cap fell, returning +the default there. Since the cap counts uncached frames only, the result +depended on what had been calculated before: with an input of 1 in 2010 and +``recursive(t) = recursive(t - 1) + 1``, 2021 was 10 on a new simulation and +12 after calculating 2020. + +Now a recursion heading towards an input or a period with no formula (an +anchor) finishes its deeper periods first and is evaluated all the way down. +One with no anchor is still cut, and nothing calculated from the cut value +is kept once the calculation ends. +""" + +import pytest + +from policyengine_core.errors import CycleError +from policyengine_core.simulations import simulation as simulation_module +from tests.fixtures.spirals import build, make_recurrence, run, stored_years + + +def fresh_and_after(variables, inputs, target, prior, loops=None): + fresh = build(variables, inputs, loops) + used = build(variables, inputs, loops) + for request in prior: + run(used, *request) + return run(fresh, *target), run(used, *target) + + +def counting_up(start=2010): + return make_recurrence("recursive", [("recursive", 1)], start=start) + + +@pytest.mark.parametrize( + "loops, target, prior, expected", + [(1, 2012, 2011, 3.0), (None, 2021, 2020, 12.0)], +) +def test_anchored_recursion_ignores_what_is_cached(loops, target, prior, expected): + # The review's examples: an input of 1 in 2010 anchors the recursion. + fresh, after = fresh_and_after( + [counting_up()], + {("recursive", 2010): 1}, + ("recursive", target), + [("recursive", prior)], + loops, + ) + assert fresh == after == [expected] + + +def test_recursion_is_evaluated_down_to_the_first_formula(): + # No input: 2009 has no formula, so it is 0 and 2010 is 1. + fresh, after = fresh_and_after( + [counting_up()], {}, ("recursive", 2060), [("recursive", 2040)] + ) + assert fresh == after == [51.0] + + +def test_recursion_is_evaluated_down_to_an_input(): + # An undated formula, anchored by the input in 2010 only. + fresh, after = fresh_and_after( + [counting_up(start=None)], + {("recursive", 2010): 1}, + ("recursive", 2045), + [("recursive", 2030)], + ) + assert fresh == after == [36.0] + + +def test_forward_recursion_is_evaluated_up_to_the_end(): + ahead = make_recurrence("ahead", [("ahead", -1)], end=2030) + fresh, after = fresh_and_after([ahead], {}, ("ahead", 2000), [("ahead", 2020)]) + assert fresh == after == [31.0] + + +def test_long_recursion_through_another_variable_is_evaluated_in_full(): + # first(t) = 1 + second(t), second(t) = 1 + first(t - 1); second's dated + # formula anchors the chain. + first = make_recurrence("first", [("second", 0)]) + second = make_recurrence("second", [("first", 1)], start=2010) + fresh, after = fresh_and_after( + [first, second], {}, ("first", 2040), [("first", 2030)] + ) + # second(2009) = 0, first(2009) = 1; each later year adds 2. + assert fresh == after == [1.0 + 2 * 31] + + +def test_unanchored_recursion_is_still_cut_at_the_limit(): + # No input and an undated formula: nothing ends the recursion, so it is + # cut after max_spiral_loops frames, as before. + unanchored = counting_up(start=None) + fresh, after = fresh_and_after( + [unanchored], {}, ("recursive", 2021), [("recursive", 2020)] + ) + assert fresh == after == [10.0] + + +def test_values_reached_by_a_cut_are_not_kept(): + simulation = build([counting_up(start=None)]) + assert run(simulation, "recursive", 2021) == [10.0] + assert stored_years(simulation, "recursive") == [] + + +def test_values_above_a_cut_are_not_kept(): + # anchored(t) = 1 + unanchored(t) + anchored(t - 1): anchored's own chain + # ends in 2009, but every value reads a cut recursion. + unanchored = make_recurrence("unanchored", [("unanchored", 1)]) + anchored = make_recurrence( + "anchored", [("unanchored", 0), ("anchored", 1)], start=2010 + ) + fresh, after = fresh_and_after( + [unanchored, anchored], + {}, + ("anchored", 2030), + [("anchored", 2025), ("unanchored", 2028)], + ) + assert fresh == after + simulation = build([unanchored, anchored]) + run(simulation, "anchored", 2030) + assert stored_years(simulation, "anchored") == [] + + +def test_skipped_input_does_not_loop(): + # skip(t) = 1 + skip(t - 2) passes the 2011 input by when it starts from + # an even year, so it can't end there; it is cut instead of retried. + skip = make_recurrence("skip", [("skip", 2)]) + fresh, after = fresh_and_after( + [skip], {("skip", 2011): 5}, ("skip", 2040), [("skip", 2030)] + ) + assert fresh == after == [10.0] + + +def test_cycle_is_still_an_error(): + first = make_recurrence("first", [("second", 0)]) + second = make_recurrence("second", [("first", 0)]) + assert run(build([first, second]), "first", 2020) == CycleError.__name__ + + +def test_formula_catching_exception_does_not_swallow_the_deeper_period(): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class guarded(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Recursion inside try/except Exception" + + def formula_2010(person, period): + try: + return person("guarded", period.last_year) + 1 + except Exception: + return person.filled_array(-1000.0) + + fresh, after = fresh_and_after( + [guarded], {}, ("guarded", 2040), [("guarded", 2035)] + ) + assert fresh == after == [31.0] + + +def test_cap_on_deeper_periods_falls_back_to_the_cut(monkeypatch): + monkeypatch.setattr(simulation_module, "MAX_SPIRAL_DEFERRALS", 0) + simulation = build([counting_up()], {("recursive", 2010): 1}) + assert run(simulation, "recursive", 2021) == [10.0] + + +def test_spiral_in_a_branch_is_purged_from_the_branch(): + simulation = build([counting_up(start=None)]) + branch = simulation.get_branch("other") + assert run(branch, "recursive", 2021) == [10.0] + assert stored_years(branch, "recursive", "other") == [] + assert run(branch, "recursive", 2022) == run( + build([counting_up(start=None)]).get_branch("other"), "recursive", 2022 + ) + + +def test_anchored_recursion_in_a_traced_branch_called_from_a_formula(): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class from_branch(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Reads recursive on a branch" + + def formula(person, period): + branch = person.simulation.get_branch("other") + return branch.calculate("recursive", period) + + traced = build([counting_up(), from_branch], {("recursive", 2010): 1}) + traced.trace = True + plain = build([counting_up(), from_branch], {("recursive", 2010): 1}) + assert run(traced, "from_branch", 2040) == run(plain, "from_branch", 2040) + assert run(plain, "from_branch", 2040) == [31.0] + + +def test_traced_branch_cut_from_a_formula_is_purged(): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class from_branch(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Reads recursive on a branch" + + def formula(person, period): + branch = person.simulation.get_branch("other") + return branch.calculate("recursive", period) + + simulation = build([counting_up(start=None), from_branch]) + simulation.trace = True + run(simulation, "from_branch", 2021) + branch = simulation.branches["other"] + assert stored_years(branch, "recursive", "other") == [] + assert stored_years(simulation, "from_branch") == [] + + +def test_branch_made_after_a_cut_does_not_keep_the_cut_values(): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class then_branch(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Reads a cut recursion, then makes a branch" + + def formula(person, period): + value = person("recursive", period) + person.simulation.get_branch("late") + return value + + simulation = build([counting_up(start=None), then_branch]) + run(simulation, "then_branch", 2021) + branch = simulation.branches["late"] + assert stored_years(branch, "recursive", "default") == [] + + +def test_untraced_branch_cut_from_a_formula_is_purged_from_the_caller(): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class from_branch(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Reads recursive on a branch" + + def formula(person, period): + branch = person.simulation.get_branch("other") + return branch.calculate("recursive", period) + + simulation = build([counting_up(start=None), from_branch]) + assert run(simulation, "from_branch", 2021) == [10.0] + assert stored_years(simulation, "from_branch") == [] + assert stored_years(simulation.branches["other"], "recursive", "other") == [] + + +def test_branch_spiral_does_not_purge_the_parent(): + # A cut in a branch marks values for purging in that branch only. A + # parent sharing the branch's set would later delete its own values at + # those periods, an input set after branching included. + simulation = build([counting_up(start=None)]) + branch = simulation.get_branch("other") + simulation.set_input("recursive", "2015", [100]) + assert run(branch, "recursive", 2021) == [10.0] + assert run(simulation, "recursive", 2016) == [101.0] + assert simulation.get_array("recursive", "2015").tolist() == [100.0] diff --git a/tests/core/test_spiral_order_property.py b/tests/core/test_spiral_order_property.py new file mode 100644 index 000000000..41b5c1e8b --- /dev/null +++ b/tests/core/test_spiral_order_property.py @@ -0,0 +1,66 @@ +"""Property: recursions over periods give results independent of order. + +For any system of one or two yearly variables, each equal to 1 plus other +variables (itself included) at lags of 0 to 2 years, with dated or undated +formulas, inputs at random years and any ``max_spiral_loops``, every request +returns the same value, or raises the same error, after any earlier requests +as it does in a new simulation. +""" + +import pytest + +from tests.fixtures.spirals import build, make_recurrence, run + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +NAMES = ["r0", "r1"] +YEARS = st.integers(2000, 2035) + + +@st.composite +def systems(draw): + names = NAMES[: draw(st.integers(1, 2))] + variables = [] + for name in names: + terms = draw( + st.lists( + st.tuples(st.sampled_from(names), st.integers(0, 2)), + min_size=1, + max_size=2, + ) + ) + # Lag 0 on itself is a plain cycle, tested elsewhere. + terms = [(other, lag or (1 if other == name else 0)) for other, lag in terms] + start = draw(st.sampled_from([None, 2004, 2008])) + variables.append((name, tuple(terms), start)) + inputs = draw( + st.dictionaries( + st.tuples(st.sampled_from(names), st.integers(2000, 2012)), + st.integers(1, 9), + max_size=2, + ) + ) + loops = draw(st.sampled_from([1, 2, 3, 10])) + requests = st.tuples(st.sampled_from(names), YEARS) + prior = draw(st.lists(requests, max_size=4)) + target = draw(requests) + return variables, inputs, loops, prior, target + + +@hypothesis.settings(max_examples=200, deadline=None) +@hypothesis.given(systems()) +def test_earlier_requests_do_not_change_a_result(system): + variables, inputs, loops, prior, target = system + + def new(): + return build( + [make_recurrence(name, terms, start) for name, terms, start in variables], + inputs, + loops, + ) + + used = new() + for request in prior: + run(used, *request) + assert run(used, *target) == run(new(), *target) diff --git a/tests/fixtures/spirals.py b/tests/fixtures/spirals.py new file mode 100644 index 000000000..b69ee2e57 --- /dev/null +++ b/tests/fixtures/spirals.py @@ -0,0 +1,85 @@ +"""Variables that read themselves, or each other, at other periods. + +Shared by ``tests/core/test_spiral_order.py`` and its Hypothesis property +module, which pytest must be able to skip on its own when Hypothesis is not +installed. +""" + +from typing import Dict, Iterable, Optional, Tuple + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + + +def make_recurrence( + name: str, + terms: Iterable[Tuple[str, int]], + start: Optional[int] = None, + end: Optional[int] = None, +) -> type: + """A yearly variable equal to 1 plus other variables at lagged years. + + Each of ``terms`` is ``(variable, lag)``: the variable's value ``lag`` + years earlier (later when negative). ``start`` dates the formula + (``formula_``), so that earlier years have none, and ``end`` is the + last year the variable has a formula for. + """ + terms = tuple(terms) + + def formula(person, period): + value = person.filled_array(1.0) + for variable, lag in terms: + value = value + person(variable, period.offset(-lag, "year")) + return value + + attributes = dict( + value_type=float, + entity=Person, + definition_period=periods.YEAR, + label=f"Recurrence {name}", + ) + attributes["formula" if start is None else f"formula_{start}"] = formula + if end is not None: + attributes["end"] = f"{end}-12-31" + return type(name, (Variable,), attributes) + + +def build( + variables: Iterable[type], + inputs: Optional[Dict[Tuple[str, int], float]] = None, + max_spiral_loops: Optional[int] = None, +): + """One-person simulation of ``variables`` with ``inputs``. + + ``inputs`` maps ``(variable, year)`` to a value. + """ + system = CountryTaxBenefitSystem() + system.add_variables(*variables) + simulation = SimulationBuilder().build_default_simulation(system, count=1) + for (variable, year), value in (inputs or {}).items(): + simulation.set_input(variable, str(year), np.array([value])) + if max_spiral_loops is not None: + simulation.max_spiral_loops = max_spiral_loops + return simulation + + +def run(simulation, variable: str, year: int): + """The value as a list, or the name of the error the calculation raises.""" + try: + return simulation.calculate(variable, str(year)).tolist() + except Exception as error: + return type(error).__name__ + + +def stored_years(simulation, variable: str, branch_name: str = "default"): + """The years ``variable`` has a value stored for under ``branch_name``.""" + return sorted( + period.start.year + for branch, period in simulation.get_holder(variable).get_known_branch_periods() + if branch == branch_name + ) From 416852cc8c33767a5c996b43925d4becf1af4b4c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 10:03:00 -0400 Subject: [PATCH 2/9] Fix the round-2 review's eight regressions in deferred spiral recursion Branch and copy cleanup (findings 1 and 4). A cut in a simulation's calculation no longer deletes its branch-name key from every descendant, which deleted a descendant's own input stored under the same name. A copy (clone or branch) instead drops, when it is made, the values the simulation it copies is waiting to purge: every such value is marked by the time it is cached, and a copy never sees what is cached after. Purges delete only values the simulation calculated (derived), never an input. Swallowed requests (2). Each request to finish a deeper period first is recorded as outstanding when raised. A formula that returns while one is outstanding (a bare except caught it) has its value dropped and the request raised again; an error raised while one is outstanding is treated as the request. requires_computation_after (3). A deferred period keeps the names of the variables that needed it, so its caller still counts while it is calculated with an empty stack. Tracing and storage (6, 7). The flat trace prefers a completed node to an abandoned one. Deferred values are kept for the outermost calculation in its own cache, which _calculate reads, so deferral works the same traced or not and when the holder stores nothing; values calculated under a cut are marked for purging when written to the fast cache too. Termination (8). MAX_SPIRAL_DEFERRALS is gone: it cut long chains depending on what was cached. Each deferred period starts at most once per outermost calculation, and only on the requested simulation or its registered branches (a copy made in a formula is made anew on every retry), so the calculation still ends. Finding 5 is settled by master's #562 (only inputs carry over); its regression test passes on the merge commit. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/in_memory_storage.py | 12 +- .../data_storage/on_disk_storage.py | 12 +- policyengine_core/holders/holder.py | 15 +- policyengine_core/simulations/simulation.py | 280 +++++++++--- policyengine_core/tracers/flat_trace.py | 17 +- tests/core/test_fast_cache_guards.py | 2 +- tests/core/test_spiral_order.py | 404 +++++++++++++++++- 7 files changed, 666 insertions(+), 76 deletions(-) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 6335f97c8..7add3dabd 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -320,14 +320,22 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: self._stop_sharing_dropped_keys() self._unmark_dropped_keys() - def delete_exact(self, period: Period, branch_name: str = "default") -> None: + def delete_exact( + self, + period: Period, + branch_name: str = "default", + derived_only: bool = False, + ) -> None: """Drop the array stored for exactly ``period``. - Unlike ``delete``, this leaves the periods ``period`` contains. + Unlike ``delete``, this leaves the periods ``period`` contains. With + ``derived_only``, an input stored there stays. """ if self.is_eternal: period = periods.period(periods.ETERNITY) key = f"{branch_name}:{periods.period(period)}" + if derived_only and key in self._inputs: + return self._arrays.pop(key, None) self._stop_sharing(key) self._unmark_input(key) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index d87c2d96b..3e1f30fa3 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -350,11 +350,19 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: } self._derived.intersection_update(self._files) - def delete_exact(self, period: Period, branch_name: str = "default") -> None: - """Drop the file recorded for exactly ``period``.""" + def delete_exact( + self, + period: Period, + branch_name: str = "default", + derived_only: bool = False, + ) -> None: + """Drop the file recorded for exactly ``period``. With + ``derived_only``, an input stored there stays.""" if self.is_eternal: period = periods.period(periods.ETERNITY) key = f"{branch_name}_{periods.period(period)}" + if derived_only and key in self._files and key not in self._derived: + return self._files.pop(key, None) self._derived.discard(key) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 7a212dc63..d50f9087d 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -110,13 +110,20 @@ def delete_arrays( if self._disk_storage: self._disk_storage.delete(period, branch_name) - def delete_array(self, period: Period, branch_name: str = "default") -> None: + def delete_array( + self, + period: Period, + branch_name: str = "default", + derived_only: bool = False, + ) -> None: """ - Remove the value stored for exactly ``period``, leaving the periods it contains (unlike ``delete_arrays``). + Remove the value stored for exactly ``period`` under ``branch_name``, leaving the periods it contains (unlike ``delete_arrays``). + + With ``derived_only``, remove it only if the simulation calculated it (see ``put_in_cache``): an input stored there stays. """ - self._memory_storage.delete_exact(period, branch_name) + self._memory_storage.delete_exact(period, branch_name, derived_only) if self._disk_storage: - self._disk_storage.delete_exact(period, branch_name) + self._disk_storage.delete_exact(period, branch_name, derived_only) def _get_array_from_storage( self, period: Period, branch_name: str = "default" diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 7a532e8a7..d1058969d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -66,12 +66,6 @@ def __init__(self, simulation: "Simulation"): ) -# How many deeper periods one outermost calculation may finish first (see -# ``_DeeperPeriodFirst``) before its recursion is cut where it stands, as a -# chain with no anchor is. -MAX_SPIRAL_DEFERRALS = 10_000 - - class _DeeperPeriodFirst(BaseException): """A recursion over periods reached ``max_spiral_loops`` and can end. @@ -82,48 +76,102 @@ class _DeeperPeriodFirst(BaseException): ``calculate`` catches it once every frame above has unwound, calculates ``variable`` for ``period`` on ``simulation`` first, as an outermost calculation with an empty stack, and then starts its own request again, - which now reaches that cached value. + which now reaches that value. Where the recursion is cut therefore no longer depends on which periods happen to be cached: an anchored chain is evaluated all the way down, however long it is, in at most ``max_spiral_loops`` frames at a time. It derives from ``BaseException`` so that a formula that catches - ``Exception`` around a calculation does not swallow it. + ``Exception`` around a calculation does not swallow it. One that catches + everything (a bare ``except``) can, so each request is recorded as + outstanding when raised (``_SpiralDeferrals.outstanding``), and a formula + that returns while one is outstanding has its value dropped and the + request raised again (see ``Simulation._calculate``). """ - def __init__(self, simulation: "Simulation", variable: str, period: Period): + def __init__( + self, + simulation: "Simulation", + variable: str, + period: Period, + callers: frozenset = frozenset(), + ): super().__init__(variable, period) self.simulation = simulation self.variable = variable self.period = period + # The names of the variables whose calculation needed this one: the + # callers of the deeper period being calculated first when it was + # raised, if any, and the frames on the stack. Calculated first with + # an empty stack, the period still has them as callers, as + # ``requires_computation_after`` requires. + self.callers = callers @property def key(self) -> tuple: return (id(self.simulation), self.variable, self.period) - def calculate(self) -> None: - self.simulation._calculate_traced(self.variable, self.period) - - def is_cached(self) -> bool: - holder = self.simulation.get_holder(self.variable) - return holder.get_array(self.period, self.simulation.branch_name) is not None + def calculate(self) -> ArrayLike: + return self.simulation._calculate_traced(self.variable, self.period) @dataclass class _SpiralDeferrals: - """Deeper periods finished first during one outermost calculation. - - ``futile`` holds the ``_DeeperPeriodFirst.key`` of each one that was not - cached when finished, which ``_check_for_cycle`` then cuts instead. + """The deeper periods one outermost calculation finishes first. + + Each deeper period's calculation starts at most once (``started``), and + only on the simulation the calculation was requested on or one of its + branches (``root``), so the outermost calculation ends: every round + starts a period not started before, or ends one, or ends the request, + and the periods a deterministic calculation can reach are finite. No + count of rounds cuts a long chain, since where such a count fell would + depend on what was cached before. """ - count: int = 0 - futile: set = field(default_factory=set) + # The simulation the outermost calculation was requested on, or the + # simulation it is a branch of (see ``Simulation._branch_root``). + root: Any = None + # ``_DeeperPeriodFirst.key`` of each period whose calculation has + # started. A recursion reaching one again is cut there: it is under way + # (a cycle through the deferrals), or it was calculated with a recursion + # cut on the way, or it is finished, and then never reached. + started: set = field(default_factory=set) + # Of those, the ones whose calculation has ended. + done: set = field(default_factory=set) + # The value of each period calculated first without a recursion cut on + # the way, by simulation: ``(simulation, {(variable, period): value})`` + # under ``id(simulation)``. ``_calculate`` reads it as a stored value, + # so starting again reaches it however the holder stores values + # (``_do_not_store``, ``opt_out_cache``) and whether or not the + # calculation is traced. It goes when the calculation ends. + finished: dict = field(default_factory=dict) + # How many recursions have been cut so far. + cuts: int = 0 + # Requests raised and not yet back in the outermost ``calculate``: a + # formula caught them, or they are on their way up. + outstanding: list = field(default_factory=list) + # The callers of the deeper period being calculated first, if any (see + # ``_DeeperPeriodFirst.callers``). + callers: frozenset = frozenset() # Tracers whose stack was active when a recursion was cut and has not # emptied since: until it does, what is cached may depend on the cut. cut_tracers: list = field(default_factory=list) + def finish(self, deferral: _DeeperPeriodFirst, value: ArrayLike) -> None: + """Record ``value``, calculated for ``deferral`` with no cut.""" + simulation = deferral.simulation + entry = self.finished.get(id(simulation)) + if entry is None: + entry = self.finished[id(simulation)] = (simulation, {}) + entry[1][(deferral.variable, deferral.period)] = value + + def finished_value( + self, simulation: "Simulation", variable: str, period: Period + ) -> Optional[ArrayLike]: + entry = self.finished.get(id(simulation)) + return None if entry is None else entry[1].get((variable, period)) + _spiral_deferrals: ContextVar[Optional[_SpiralDeferrals]] = ContextVar( "_spiral_deferrals", default=None @@ -864,29 +912,74 @@ def calculate( # ``max_spiral_loops`` unwinds to the outermost one. return self._calculate_traced(variable_name, period, map_to, decode_enums) - deferrals = _SpiralDeferrals() - token = _spiral_deferrals.set(deferrals) + state = _SpiralDeferrals(root=self._branch_root()) + token = _spiral_deferrals.set(state) try: # Deeper periods to calculate before the request, deepest last. pending = [] while True: + deferral = pending[-1] if pending else None + if deferral is not None: + if deferral.key in state.done: + # Requested twice before it started (see ``_defer``). + pending.pop() + continue + state.started.add(deferral.key) + state.callers = frozenset() if deferral is None else deferral.callers + cuts = state.cuts try: - if not pending: - return self._calculate_traced( + if deferral is None: + result = self._calculate_traced( variable_name, period, map_to, decode_enums ) - pending[-1].calculate() - deferral = pending.pop() - if not deferral.is_cached(): - # Cut further down, or never stored: starting again - # would reach the same frame, so cut the chain there. - deferrals.futile.add(deferral.key) - except _DeeperPeriodFirst as deferral: - deferrals.count += 1 - pending.append(deferral) + else: + result = deferral.calculate() + except _DeeperPeriodFirst as deeper: + pending.extend(self._take_outstanding(state, deeper)) + continue + except Exception: + # A formula caught a request to finish a deeper period + # first and raised an error of its own instead: the error + # comes from the request, so carry it out. If the error + # does not, it comes back once that period is done. + if not state.outstanding: + raise + pending.extend(self._take_outstanding(state)) + continue + if state.outstanding: + # Caught on the way and never raised again, which + # ``_calculate`` prevents for formulas: still carry it out. + pending.extend(self._take_outstanding(state)) + continue + if deferral is None: + return result + pending.pop() + state.done.add(deferral.key) + if state.cuts == cuts: + state.finish(deferral, result) + # Otherwise a recursion was cut on the way, so the value + # depends on where. Starting again reaches this period again, + # and, started once already, it is cut there. finally: _spiral_deferrals.reset(token) + @staticmethod + def _take_outstanding( + state: _SpiralDeferrals, raised: Optional[_DeeperPeriodFirst] = None + ) -> List[_DeeperPeriodFirst]: + """The requests to finish a deeper period first that the outermost + ``calculate`` is to carry out: those raised since it last took them, + in the order raised, whether they came back up or a formula caught + them, and ``raised``, the one that came back.""" + taken = list(state.outstanding) + state.outstanding.clear() + if raised is not None and not any(raised is other for other in taken): + taken.append(raised) + for deferral in taken: + # Keep no frame of the abandoned attempt alive while it waits. + deferral.__traceback__ = deferral.__context__ = None + return taken + def _calculate_traced( self, variable_name: str, @@ -1049,6 +1142,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: cached_array = holder.get_array(period, self.branch_name) if cached_array is not None: return cached_array + # Then for one this calculation finished first (see + # ``_SpiralDeferrals.finished``), which the holder may not store. + state = _spiral_deferrals.get() + if state is not None and state.finished: + finished = state.finished_value(self, variable_name, period) + if finished is not None: + return finished # Check if cache can be used, if available, check if path exists is_cache_available = self.check_macro_cache(variable_name, str(period)) @@ -1073,6 +1173,10 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if variable.requires_computation_after is not None: variables_in_stack = [node.get("name") for node in self.tracer.stack] + if state is not None: + # A deeper period calculated first, with an empty stack, + # keeps the callers that needed it. + variables_in_stack = [*state.callers, *variables_in_stack] variable_in_stack = ( variable.requires_computation_after in variables_in_stack ) @@ -1122,6 +1226,12 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: try: self._check_for_cycle(variable.name, period) array = self._run_formula(variable, population, period) + if state is not None and state.outstanding: + # The formula caught a request to finish a deeper period + # first (with a bare ``except``) and returned anyway: its + # value may be its fallback. Raise the request again instead + # of keeping it. + raise state.outstanding[-1] # If no result, use the default value and cache it if array is None: @@ -1283,6 +1393,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: smc.set_cache_value(cache_path, array) if hasattr(self, "_fast_cache"): + # Purged with what the holder stores if a cut recursion reached + # it, even when the holder stores nothing (``_do_not_store``). + self._note_cached_value(variable_name, period) self._fast_cache[(variable_name, period)] = array return array @@ -1293,31 +1406,22 @@ def purge_cache_of_invalid_values(self) -> None: return # A traced branch shares this simulation's stack, so its frames can # only be purged now that the stack is empty. - branches = list(getattr(self, "branches", {}).values()) - for branch in branches: + for branch in list(getattr(self, "branches", {}).values()): branch.purge_cache_of_invalid_values() _fast_cache = getattr(self, "_fast_cache", None) invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: return - descendants = [] - while branches: - branch = branches.pop() - descendants.append(branch) - branches.extend(getattr(branch, "branches", {}).values()) branch_name = getattr(self, "branch_name", "default") for _name, _period in invalidated_caches: holder = self.get_holder(_name) # Delete the value this simulation calculated: the one stored # under its own branch name, not another branch's, and for that # period only, not the inputs or values of the periods within it. - holder.delete_array(_period, branch_name) - # Branches made since hold it too, under this branch's name. - for branch in descendants: - population = branch.get_variable_population(_name) - branch_holder = population._holders.get(_name) - if branch_holder is not None: - branch_holder.delete_array(_period, branch_name) + # An input stored there stays. A branch or copy made before now + # dropped what it took of these values when it was made (see + # ``clone``), and never sees what this simulation stores after. + holder.delete_array(_period, branch_name, derived_only=True) if _fast_cache is not None: _fast_cache.pop((_name, _period), None) self.invalidated_caches = set() @@ -1609,7 +1713,7 @@ def _check_for_cycle(self, variable: str, period: Period) -> None: spiral = len(previous_periods) >= self.max_spiral_loops if spiral and self._has_calculation(variable, period): if self._can_finish_deeper_period_first(variable, period): - raise _DeeperPeriodFirst(self, variable, period) + raise self._defer(variable, period) self.invalidate_spiral_variables(variable) message = "Quasicircular definition detected on formula {}@{} involving {}".format( variable, period, self.tracer.stack @@ -1657,14 +1761,26 @@ def _can_finish_deeper_period_first( Being anchored depends on the variables and the inputs, not on what has been calculated, so whether a chain is cut does not depend on calculation order either. + + Each period's calculation starts once per outermost calculation: + one reached again after that (a cycle through the deferrals, or a + period calculated with a recursion cut on the way) is cut. And only + on the simulation the calculation was requested on and its branches: + a copy made while it runs is made anew when it starts again, and + would never reach what was finished on the first one. """ deferrals = _spiral_deferrals.get() - if ( + if deferrals is None: # Only an outermost ``calculate`` can finish the deeper period # first; called any other way, the recursion is cut as before. - deferrals is None - or deferrals.count >= MAX_SPIRAL_DEFERRALS - or (id(self), variable_name, period) in deferrals.futile + return False + key = (id(self), variable_name, period) + if ( + key in deferrals.started + # Requested already by this attempt at the calculation, and the + # request caught: the formula tries again. + or any(key == other.key for other in deferrals.outstanding) + or self._branch_root() is not deferrals.root ): return False tax_benefit_system = getattr(self, "tax_benefit_system", None) @@ -1725,6 +1841,37 @@ def _has_formula_boundary_ahead( return "0001-01-01" < first_start <= str(period.start) return variable.end is not None and period.start.date <= variable.end + def _defer(self, variable_name: str, period: Period) -> _DeeperPeriodFirst: + """A request to calculate ``variable_name`` for ``period`` first, + outstanding until the outermost ``calculate`` takes it. + + The period may be requested already and not started: a formula + caught that request and the calculation went on to another. This + request puts it first. + """ + state = _spiral_deferrals.get() + callers = state.callers.union(frame["name"] for frame in self.tracer.stack[:-1]) + deferral = _DeeperPeriodFirst(self, variable_name, period, callers) + state.outstanding.append(deferral) + return deferral + + def _branch_root(self) -> "Simulation": + """The simulation this one is a branch of, through ``parent_branch`` + and as ``get_branch`` returns it, or this one if it is not a branch + (a copy of a branch is not).""" + simulation = self + while True: + parent = getattr(simulation, "parent_branch", None) + if ( + parent is None + or getattr(parent, "branches", {}).get( + getattr(simulation, "branch_name", None) + ) + is not simulation + ): + return simulation + simulation = parent + def invalidate_cache_entry(self, variable: str, period: Period) -> None: invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: @@ -1753,6 +1900,7 @@ def invalidate_spiral_variables(self, variable: str) -> None: state = _spiral_deferrals.get() if state is None: return + state.cuts += 1 # The cut value also reaches every calculation in progress in a # related simulation (a branch's caller, say). for simulation in self._related_simulations(): @@ -2028,7 +2176,8 @@ def clone( if hasattr(self, "_user_input_keys"): new._user_input_keys = set(self._user_input_keys) # Its own set: values a spiral marks in one simulation are purged - # from that simulation's holders only. + # from that simulation's holders only. It starts with nothing to + # purge: it drops what it takes of this simulation's below. new.invalidated_caches = set() # Only pass ``share_arrays`` when sharing, so a population or holder @@ -2051,9 +2200,32 @@ def clone( new.tax_benefit_system = self.tax_benefit_system new.debug = debug new.trace = trace + new._drop_invalid_values(self) return new + def _drop_invalid_values(self, source: "Simulation") -> None: + """Drop the values this copy of ``source`` took that ``source`` is to + purge once its calculation ends (see ``invalidate_spiral_variables``). + + A value a cut recursion reaches is marked when it is cached, or + before, so every one ``source`` holds now is marked now, and a copy + never sees what ``source`` caches later. Only the copy's own index + changes; ``source`` purges its own values when its calculation ends. + Inputs stay, as they do in ``source``. + """ + invalidated_caches = getattr(source, "invalidated_caches", None) + if not invalidated_caches: + return + branch_name = getattr(source, "branch_name", "default") + for name, invalid_period in invalidated_caches: + variable = self.tax_benefit_system.get_variable(name) + if variable is None: + continue + holder = self.populations[variable.entity.key]._holders.get(name) + if holder is not None: + holder.delete_array(invalid_period, branch_name, derived_only=True) + def get_branch( self, name: str = "branch", clone_system: bool = False ) -> "Simulation": diff --git a/policyengine_core/tracers/flat_trace.py b/policyengine_core/tracers/flat_trace.py index 7cebeb829..f5c90851f 100644 --- a/policyengine_core/tracers/flat_trace.py +++ b/policyengine_core/tracers/flat_trace.py @@ -33,14 +33,15 @@ def get_trace(self) -> dict: # We don't want cache read to overwrite data about the initial # calculation. # - # We therefore use a non-overwriting update. - trace.update( - { - key: node_trace - for key, node_trace in self._get_flat_trace(node).items() - if key not in trace - } - ) + # We therefore use a non-overwriting update, except over a + # calculation that never recorded a result: one abandoned to + # calculate a deeper period first (see ``Simulation.calculate``), + # or ended by an error. The first that completes replaces it. + for key, node_trace in self._get_flat_trace(node).items(): + if key not in trace or ( + trace[key]["value"] is None and node_trace["value"] is not None + ): + trace[key] = node_trace return trace diff --git a/tests/core/test_fast_cache_guards.py b/tests/core/test_fast_cache_guards.py index 45715653d..7efe9104b 100644 --- a/tests/core/test_fast_cache_guards.py +++ b/tests/core/test_fast_cache_guards.py @@ -72,7 +72,7 @@ def test_purge_cache_of_invalid_values_without_fast_cache_attribute(): sim.tracer = types.SimpleNamespace(stack=[]) sim.invalidated_caches = {("variable_name", "2024")} sim.get_holder = lambda name: types.SimpleNamespace( - delete_array=lambda period, branch_name="default": None + delete_array=lambda period, branch_name="default", derived_only=False: None ) sim.purge_cache_of_invalid_values() diff --git a/tests/core/test_spiral_order.py b/tests/core/test_spiral_order.py index 3e672aeca..e229c95f2 100644 --- a/tests/core/test_spiral_order.py +++ b/tests/core/test_spiral_order.py @@ -16,7 +16,6 @@ import pytest from policyengine_core.errors import CycleError -from policyengine_core.simulations import simulation as simulation_module from tests.fixtures.spirals import build, make_recurrence, run, stored_years @@ -159,10 +158,60 @@ def formula_2010(person, period): assert fresh == after == [31.0] -def test_cap_on_deeper_periods_falls_back_to_the_cut(monkeypatch): - monkeypatch.setattr(simulation_module, "MAX_SPIRAL_DEFERRALS", 0) - simulation = build([counting_up()], {("recursive", 2010): 1}) - assert run(simulation, "recursive", 2021) == [10.0] +def test_recursion_changing_direction_ends(): + # zigzag(t) = 1 + zigzag(t + 2) in odd years and 1 + zigzag(t - 1) in + # even ones, until 2020: it heads back towards its first formula, then + # forward towards its end, in turn. Each deeper period is started once, + # so the calculation ends, and with the same value however it starts. + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class zigzag(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Recursion heading back and forward in turn" + end = "2020-12-31" + + def formula_2010(person, period): + lag = 1 if period.start.year % 2 == 0 else -2 + return person("zigzag", period.offset(-lag, "year")) + 1 + + # Evaluated in full: 2021 has no formula and is 0, 2009 likewise. + expected = {2010: 1, 2011: 5, 2012: 6, 2013: 4, 2014: 5, 2015: 3} + expected.update({2016: 4, 2017: 2, 2018: 3, 2019: 1, 2020: 2}) + for loops in (1, 2, 10): + for year, value in expected.items(): + fresh, after = fresh_and_after( + [zigzag], {}, ("zigzag", year), [("zigzag", 2017)], loops + ) + assert fresh == after == [float(value)] + + +def test_copy_made_in_a_formula_ends_its_recursion(): + # A formula that makes a copy of the simulation and calculates on it + # makes a new copy each time the calculation starts again, so the copy's + # deeper periods are never finished first: its recursion is cut, as + # before, instead of starting again for ever. + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class on_copy(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Calculates recursive on a new copy of the simulation" + + def formula(person, period): + return person.simulation.clone().calculate("recursive", period) + + simulation = build([counting_up(), on_copy], {}, 1) + assert run(simulation, "on_copy", 2015) == [1.0] + # On the copy itself, outside a calculation, it is evaluated in full. + copy = simulation.clone() + assert run(copy, "recursive", 2015) == [6.0] def test_spiral_in_a_branch_is_purged_from_the_branch(): @@ -273,3 +322,348 @@ def test_branch_spiral_does_not_purge_the_parent(): assert run(branch, "recursive", 2021) == [10.0] assert run(simulation, "recursive", 2016) == [101.0] assert simulation.get_array("recursive", "2015").tolist() == [100.0] + + +# Round-2 review regressions. Each case compares a calculation with the same +# calculation in a new simulation, or with the full evaluation of an anchored +# chain, which is what a new simulation returns. + + +def reads(name, target, lag=0): + """A yearly variable equal to ``target`` ``lag`` years earlier.""" + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + def formula(person, period): + return person(target, period.offset(-lag, "year")) + + return type( + name, + (Variable,), + dict( + value_type=float, + entity=Person, + definition_period=YEAR, + label=f"Reads {target}", + formula=formula, + ), + ) + + +@pytest.mark.parametrize("traced", [False, True]) +def test_cut_in_the_parent_keeps_a_branch_input_under_the_same_name(traced): + # A branch of a branch can be named "default", like the simulation it + # descends from. A cut in the simulation's calculation used to delete the + # simulation's branch-name key from every descendant, and with it this + # branch's own input. + simulation = build( + [counting_up(start=None), reads("wrapper", "recursive")], max_spiral_loops=1 + ) + simulation.trace = traced + branch = simulation.get_branch("one").get_branch("default") + branch.set_input("wrapper", "2021", [100]) + assert run(simulation, "wrapper", 2021) == [1.0] + assert branch.get_array("wrapper", "2021").tolist() == [100.0] + assert run(branch, "wrapper", 2021) == [100.0] + + +@pytest.mark.parametrize("traced", [False, True]) +def test_cut_keeps_an_input_set_on_a_branch_during_the_calculation(traced): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class via(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Sets an input on a branch, then reads a cut recursion there" + + def formula(person, period): + branch = person.simulation.get_branch("one").get_branch("default") + branch.set_input("via", period, [100]) + return branch.calculate("recursive", period) + + simulation = build([counting_up(start=None), via], max_spiral_loops=1) + simulation.trace = traced + assert run(simulation, "via", 2021) == [1.0] + branch = simulation.branches["one"].branches["default"] + assert branch.get_array("via", "2021").tolist() == [100.0] + assert run(branch, "via", 2021) == [100.0] + + +def catching(name, handler): + """A dated recursion whose formula reads last year inside a bare except.""" + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + def formula_2010(person, period): + try: + return person(name, period.last_year) + 1 + except: # noqa: E722 - the case under test + return handler(person, period) + + return type( + name, + (Variable,), + dict( + value_type=float, + entity=Person, + definition_period=YEAR, + label="Recursion inside a bare except", + formula_2010=formula_2010, + ), + ) + + +def _fallback(person, period): + return person.filled_array(-1000.0) + + +def _reraise(person, period): + raise ValueError("the formula turned the exception into its own") + + +def _fallback_calculation(person, period): + return person("plain", period) - 1000 + + +@pytest.mark.parametrize("handler", [_fallback, _reraise, _fallback_calculation]) +@pytest.mark.parametrize("prior", [[], [2011], [2010, 2011]]) +def test_bare_except_in_a_formula_does_not_swallow_the_deeper_period(handler, prior): + # The chain is anchored (2009 has no formula), so 2012 is 3 whatever was + # calculated before. A bare except used to catch the request to finish + # the deeper period first and cache its fallback. + variables = [catching("catches_all", handler), make_recurrence("plain", [])] + fresh, after = fresh_and_after( + variables, + {}, + ("catches_all", 2012), + [("catches_all", year) for year in prior], + loops=1, + ) + assert fresh == after == [3.0] + + +def requires_caller(): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class caller(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Calls dependent" + + def formula(person, period): + return person("dependent", period) + + class dependent(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Requires caller to be requested first" + requires_computation_after = "caller" + + def formula_2010(person, period): + return person("dependent", period.last_year) + 1 + + return [caller, dependent] + + +@pytest.mark.parametrize( + "loops, year, prior, expected", + [ + (1, 2012, [], 3.0), + (1, 2012, [2011], 3.0), + (None, 2021, [], 12.0), + (None, 2021, [2015], 12.0), + ], +) +def test_deeper_periods_keep_the_caller_requires_computation_after_needs( + loops, year, prior, expected +): + # dependent may only be calculated while caller is: the deeper periods + # finished first used to run with an empty stack, and raised ValueError. + fresh, after = fresh_and_after( + requires_caller(), + {}, + ("caller", year), + [("caller", earlier) for earlier in prior], + loops, + ) + assert fresh == after == [expected] + + +def test_requires_computation_after_still_refuses_a_direct_request(): + assert run(build(requires_caller()), "dependent", 2012) == ValueError.__name__ + + +def copying(copies, branch=False): + """A variable that reads recursive, then copies the simulation.""" + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + def formula(person, period): + value = person("recursive", period) + simulation = person.simulation + if branch: + copies.append(simulation.get_branch(f"late{len(copies)}")) + else: + copies.append(simulation.clone()) + return value + + return type( + "copy_after_cut", + (Variable,), + dict( + value_type=float, + entity=Person, + definition_period=YEAR, + label="Copies the simulation after a cut", + formula=formula, + ), + ) + + +@pytest.mark.parametrize("branch", [False, True]) +@pytest.mark.parametrize("loops", [1, None]) +@pytest.mark.parametrize("year", [2015, 2021, 2022]) +def test_copy_made_during_a_cut_does_not_keep_the_cut_values(branch, loops, year): + # The copy is made while the cut recursion's values are cached and + # waiting to be purged; it used to keep them, so its later results + # depended on the cut. + copies = [] + simulation = build([counting_up(start=None), copying(copies, branch)], {}, loops) + run(simulation, "copy_after_cut", 2021) + copy = copies[0] + name = copy.branch_name + assert stored_years(copy, "recursive", "default") == [] + fresh = build([counting_up(start=None)], {}, loops) + if branch: + fresh = fresh.get_branch(name) + assert run(copy, "recursive", year) == run(fresh, "recursive", year) + + +@pytest.mark.parametrize("loops", [1, 2, 10]) +@pytest.mark.parametrize("end", [2012, 2020]) +@pytest.mark.parametrize("carry", [False, True]) +def test_forward_anchor_does_not_read_carried_calculations(carry, end, loops): + # forward(t) = forward(t + 1) + 1 until its end; 2008 and 2009 are 7. + # Only inputs carry over, so the year after the end is 0 whatever was + # calculated before. + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR + from policyengine_core.variables import Variable + + class forward(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Forward recursion with an earlier constant formula" + + def formula_2008(person, period): + return person.filled_array(7) + + def formula_2010(person, period): + return person("forward", period.offset(1, "year")) + 1 + + forward.end = f"{end}-12-31" + + def new(): + simulation = build([forward], {}, loops) + simulation.tax_benefit_system.auto_carry_over_input_variables = carry + return simulation + + fresh, used = new(), new() + run(used, "forward", 2008) + assert run(fresh, "forward", 2010) == run(used, "forward", 2010) + assert run(fresh, "forward", 2010) == [end - 2010 + 1.0] + + +@pytest.mark.parametrize("loops, year", [(1, 2011), (None, 2021)]) +def test_traced_deeper_periods_serialize_the_completed_calculation(loops, year): + # The first attempt at the request is abandoned when it reaches the + # limit; its trace nodes have no value. The flat trace used to keep them. + simulation = build([counting_up()], {}, loops) + simulation.trace = True + expected = [year - 2009.0] + assert run(simulation, "recursive", year) == expected + trace = simulation.tracer.get_serialized_flat_trace() + for earlier in range(2010, year + 1): + node = trace[f"recursive<{earlier}, (default)>"] + assert node["value"] == [earlier - 2009.0] + assert trace[f"recursive<{year}, (default)>"]["dependencies"] == [ + f"recursive<{year - 1}, (default)>" + ] + assert simulation.tracer.stack == [] + + +def without_storage(simulation, method): + if method == "drop": + simulation.get_holder("recursive")._do_not_store = True + elif method == "blacklist": + simulation.opt_out_cache = True + simulation.tax_benefit_system.cache_blacklist = {"recursive"} + + +@pytest.mark.parametrize("traced", [False, True]) +@pytest.mark.parametrize("method", ["drop", "blacklist"]) +@pytest.mark.parametrize( + "start, loops, year, prior, expected", + [ + # Anchored: evaluated in full, whatever the storage or tracing. + (2010, 1, 2012, [], 3.0), + (2010, 1, 2012, [2011], 3.0), + (2010, None, 2021, [], 12.0), + (2010, None, 2021, [2015], 12.0), + # Not anchored: cut at the limit; nothing calculated from the cut is + # reused afterwards, even untraced from the fast cache. + (None, 1, 2012, [], 1.0), + (None, 1, 2012, [2011], 1.0), + (None, None, 2021, [], 10.0), + (None, None, 2021, [2020], 10.0), + ], +) +def test_deeper_periods_do_not_depend_on_storage_or_tracing( + traced, method, start, loops, year, prior, expected +): + def new(): + simulation = build([counting_up(start)], {}, loops) + simulation.trace = traced + without_storage(simulation, method) + return simulation + + fresh, used = new(), new() + for earlier in prior: + run(used, "recursive", earlier) + assert run(fresh, "recursive", year) == [expected] + assert run(used, "recursive", year) == [expected] + + +def test_long_anchored_recursion_is_evaluated_in_full_however_long(): + # 10,002 monthly formula periods at max_spiral_loops 1: past what a + # budget of 10,000 deeper periods allowed, after which the chain was + # cut, or not, depending on what was cached. + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import MONTH + from policyengine_core.variables import Variable + + class monthly(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Dated monthly backward recurrence" + + def formula_2010(person, period): + return person("monthly", period.offset(-1, "month")) + 1 + + fresh = build([monthly], {}, 1) + warm = build([monthly], {}, 1) + assert warm.calculate("monthly", "2010-01").tolist() == [1.0] + assert fresh.calculate("monthly", "2843-06").tolist() == [10002.0] + assert warm.calculate("monthly", "2843-06").tolist() == [10002.0] From e0a11a21646eed0d2e7d3cdeed77bde647a7a9b6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:19:16 -0400 Subject: [PATCH 3/9] Add bounded scheduler and exact-deletion differential coverage Retain the original recurrence order generator. Compare finite anchored systems against an independent iterative NumPy evaluator across caller authorization, bare exception handlers, tracing, cache policies, copies and descendant inputs. Compare exact-deletion snapshots and provenance across copied memory, shared memory and disk storage. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_spiral_scheduler_property.py | 287 +++++++++++++++++++ tests/core/test_spiral_storage.py | 184 ++++++++++++ 2 files changed, 471 insertions(+) create mode 100644 tests/core/test_spiral_scheduler_property.py create mode 100644 tests/core/test_spiral_storage.py diff --git a/tests/core/test_spiral_scheduler_property.py b/tests/core/test_spiral_scheduler_property.py new file mode 100644 index 000000000..2efac4c0e --- /dev/null +++ b/tests/core/test_spiral_scheduler_property.py @@ -0,0 +1,287 @@ +"""Finite anchored recurrences agree with an independent iterative evaluator. + +Every dependency moves one or two years in the same direction. Backward +formulas start in 2010; forward formulas end in 2024. This gives the oracle +a finite temporal order without using the simulation's recursion scheduler, +cache, tracer, or exception handling. The existing broader order property +continues to cover mixed graphs and unanchored recurrences. +""" + +from typing import NamedTuple + +import numpy as np +import pytest + +from policyengine_core.country_template.entities import Person +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable +from tests.fixtures.spirals import build, make_recurrence + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +FIRST_YEAR = 2010 +LAST_YEAR = 2024 + + +class Recurrence(NamedTuple): + name: str + terms: tuple + guarded: bool + + +class Scenario(NamedTuple): + recurrences: tuple + direction: int + inputs: dict + loops: int + prior: tuple + target: tuple + + +@st.composite +def anchored_scenarios(draw): + names = ("r0", "r1")[: draw(st.integers(1, 2))] + direction = draw(st.sampled_from((-1, 1))) + recurrences = [] + for name in names: + terms = draw( + st.lists( + st.tuples(st.sampled_from(names), st.integers(1, 2)), + min_size=1, + max_size=2, + ) + ) + recurrences.append( + Recurrence( + name, + tuple((other, direction * distance) for other, distance in terms), + draw(st.booleans()), + ) + ) + requests = st.tuples(st.sampled_from(names), st.integers(FIRST_YEAR, LAST_YEAR)) + inputs = draw( + st.dictionaries( + requests, + st.floats( + min_value=-9, + max_value=9, + allow_nan=False, + allow_infinity=False, + width=32, + ), + max_size=3, + ) + ) + return Scenario( + tuple(recurrences), + direction, + inputs, + draw(st.sampled_from((1, 2, 3, 10))), + tuple(draw(st.lists(requests, max_size=3))), + draw(requests), + ) + + +def iterative_values(scenario): + """Fill the recurrence table in dependency order, including its boundary. + + ``filled_array(1.0)`` uses NumPy's inferred float64 dtype. Each formula + adds its already-cast float32 dependencies in order and then casts its + completed result to the float variable's float32 dtype. Casting only + the final mathematical answer would miss rounding at earlier years. + """ + values = {} + years = range(FIRST_YEAR - 2, LAST_YEAR + 3) + if scenario.direction > 0: + years = reversed(years) + for year in years: + outside_formula = ( + year < FIRST_YEAR if scenario.direction < 0 else year > LAST_YEAR + ) + for recurrence in scenario.recurrences: + key = (recurrence.name, year) + if key in scenario.inputs: + value = np.array([scenario.inputs[key]], dtype=np.float32) + elif outside_formula: + value = np.zeros(1, dtype=np.float32) + else: + value = np.full(1, 1.0) + for other, offset in recurrence.terms: + value = value + values[(other, year + offset)] + value = value.astype(np.float32) + values[key] = value + return values + + +def recurrence_variable(recurrence, direction, requires_caller): + def formula(person, period): + value = person.filled_array(1.0) + for other, offset in recurrence.terms: + value = value + person(other, period.offset(offset, "year")) + return value + + def guarded_formula(person, period): + try: + return formula(person, period) + except: # noqa: E722 - exercises a swallowed scheduling request + return person.filled_array(-1000.0) + + attributes = dict( + value_type=float, + entity=Person, + definition_period=YEAR, + label="Finite anchored recurrence", + ) + formula_name = f"formula_{FIRST_YEAR}" if direction < 0 else "formula" + attributes[formula_name] = guarded_formula if recurrence.guarded else formula + if direction > 0: + attributes["end"] = f"{LAST_YEAR}-12-31" + if requires_caller: + attributes["requires_computation_after"] = "caller" + return type(recurrence.name, (Variable,), attributes) + + +def new_simulation(scenario, traced, storage, requires_caller=False): + variables = [ + recurrence_variable(recurrence, scenario.direction, requires_caller) + for recurrence in scenario.recurrences + ] + if requires_caller: + + class caller(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Authorizes the recursive calculation" + + def formula(person, period): + return person("r0", period) + + variables.append(caller) + simulation = build(variables, scenario.inputs, scenario.loops) + simulation.tax_benefit_system.auto_carry_over_input_variables = False + simulation.trace = traced + names = [variable.__name__ for variable in variables] + for name in names: + assert simulation.tax_benefit_system.get_variable(name).dtype == np.float32 + if storage == "drop": + for name in names: + simulation.get_holder(name)._do_not_store = True + elif storage == "blacklist": + simulation.opt_out_cache = True + simulation.tax_benefit_system.cache_blacklist = set(names) + return simulation + + +def assert_request(simulation, name, year, expected, traced): + actual = simulation.calculate(name, str(year)) + assert actual.dtype == expected.dtype + np.testing.assert_array_equal(actual, expected) + if traced: + trace = simulation.tracer.get_serialized_flat_trace() + assert trace[f"{name}<{year}, (default)>"]["value"] == expected.tolist() + assert simulation.tracer.stack == [] + + +GUARDED_BACKWARD = Scenario( + (Recurrence("r0", (("r0", -1),), True),), + -1, + {}, + 1, + (("r0", 2023),), + ("r0", 2024), +) +GUARDED_FORWARD = Scenario( + (Recurrence("r0", (("r0", 1),), True),), + 1, + {("r0", 2024): 0.25}, + 1, + (("r0", 2011),), + ("r0", 2010), +) + + +@hypothesis.settings(max_examples=60, deadline=None, derandomize=True) +@hypothesis.given( + scenario=anchored_scenarios(), + traced=st.booleans(), + storage=st.sampled_from(("store", "drop", "blacklist")), +) +@hypothesis.example(scenario=GUARDED_BACKWARD, traced=True, storage="drop") +@hypothesis.example(scenario=GUARDED_FORWARD, traced=False, storage="blacklist") +def test_anchored_scheduler_matches_iterative_values(scenario, traced, storage): + expected = iterative_values(scenario) + used = new_simulation(scenario, traced, storage) + for name, year in scenario.prior: + assert_request(used, name, year, expected[(name, year)], traced) + name, year = scenario.target + assert_request(used, name, year, expected[(name, year)], traced) + # Compare with both tracing modes on a fresh simulation. Neither the + # stored work nor tracing may change the independently evaluated value. + fresh = new_simulation(scenario, not traced, storage) + assert_request(fresh, name, year, expected[(name, year)], not traced) + + +@hypothesis.settings(max_examples=40, deadline=None, derandomize=True) +@hypothesis.given( + scenario=anchored_scenarios(), + traced=st.booleans(), + storage=st.sampled_from(("store", "drop", "blacklist")), +) +@hypothesis.example(scenario=GUARDED_BACKWARD, traced=True, storage="drop") +def test_deferred_periods_keep_caller_authorization(scenario, traced, storage): + expected = iterative_values(scenario) + used = new_simulation(scenario, traced, storage, requires_caller=True) + for _, year in scenario.prior: + assert_request(used, "caller", year, expected[("r0", year)], traced) + _, year = scenario.target + assert_request(used, "caller", year, expected[("r0", year)], traced) + fresh = new_simulation(scenario, not traced, storage, requires_caller=True) + assert_request(fresh, "caller", year, expected[("r0", year)], not traced) + + +@hypothesis.settings(max_examples=24, deadline=None, derandomize=True) +@hypothesis.given( + loops=st.sampled_from((1, 2, 3, 10)), + year=st.integers(FIRST_YEAR, LAST_YEAR), + later=st.integers(1, 3), + traced=st.booleans(), + branch=st.booleans(), + input_value=st.integers(50, 150), +) +def test_cut_copies_and_descendant_inputs_are_independent( + loops, year, later, traced, branch, input_value +): + copies = [] + + class wrapper(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Copies the simulation while a cut is pending" + + def formula(person, period): + value = person("unanchored", period) + simulation = person.simulation + copies.append( + simulation.get_branch("copy") if branch else simulation.clone() + ) + return value + + recursive = make_recurrence("unanchored", (("unanchored", 1),)) + simulation = build([recursive, wrapper], max_spiral_loops=loops) + simulation.trace = traced + descendant = simulation.get_branch("one").get_branch("default") + descendant.set_input("wrapper", str(year), [input_value]) + expected = np.array([loops], dtype=np.float32) + assert_request(simulation, "wrapper", year, expected, traced) + assert descendant.get_array("wrapper", str(year)).tolist() == [input_value] + assert descendant.calculate("wrapper", str(year)).tolist() == [input_value] + copy = copies[0] + actual = copy.calculate("unanchored", str(year + later)) + np.testing.assert_array_equal(actual, expected) + fresh = build([recursive], max_spiral_loops=loops) + np.testing.assert_array_equal( + actual, fresh.calculate("unanchored", str(year + later)) + ) diff --git a/tests/core/test_spiral_storage.py b/tests/core/test_spiral_storage.py new file mode 100644 index 000000000..e66d2e9ea --- /dev/null +++ b/tests/core/test_spiral_storage.py @@ -0,0 +1,184 @@ +"""Exact spiral cleanup preserves inputs and independent storage snapshots. + +Run the same deletion scenarios against memory and disk storage and compare +their observable values and provenance with explicit expected snapshots. +The shared-memory cases delete before reading the clone, exercising the +metadata cleanup without first materializing its shared arrays. +""" + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage + + +@pytest.fixture(params=["memory", "shared-memory", "disk"]) +def storage_family(request, tmp_path): + def make(is_eternal=False): + if request.param == "disk": + return OnDiskStorage( + str(tmp_path), + is_eternal=is_eternal, + preserve_storage_dir=True, + ) + return InMemoryStorage(is_eternal=is_eternal) + + def clone(storage): + if isinstance(storage, InMemoryStorage): + return storage.clone(share_arrays=request.param == "shared-memory") + return storage.clone() + + return make, clone + + +def _put(storage, period, branch, value, derived=False): + storage.put(np.array([value]), period, branch, derived=derived) + + +def _assert_snapshot(storage, expected, probes): + """Check values, input provenance and keys against the expected snapshot. + + ``expected`` maps (branch, period) to (value, derived), using storage's + public period representation rather than its encoded internal keys. + """ + expected = { + (branch, periods.period(period)): (value, derived) + for (branch, period), (value, derived) in expected.items() + } + assert set(storage.get_known_branch_periods()) == set(expected) + for branch, period in probes: + canonical = periods.period(periods.ETERNITY if storage.is_eternal else period) + entry = expected.get((branch, canonical)) + assert storage.has(period, branch) == (entry is not None) + assert storage.is_derived(period, branch) == (entry is not None and entry[1]) + if entry is None: + assert storage.get(period, branch) is None + else: + np.testing.assert_array_equal(storage.get(period, branch), [entry[0]]) + + +def _assert_metadata(storage): + """Deleted entries retain neither sharing nor provenance marks.""" + if isinstance(storage, InMemoryStorage): + assert storage._shared <= storage._arrays.keys() + assert storage._inputs <= storage._arrays.keys() + else: + assert storage._derived <= storage._files.keys() + + +def test_exact_cleanup_keeps_descendant_inputs_and_contained_periods(storage_family): + make, clone = storage_family + source = make() + original = { + ("default", "2021"): (10, True), + ("default", "2021-01"): (20, False), + ("other", "2021"): (30, True), + } + for (branch, period), (value, derived) in original.items(): + _put(source, period, branch, value, derived) + copy = clone(source) + descendant = clone(copy) + # A descendant replaces an inherited calculation under the same name + # with its own input. Cleanup must protect that input independently. + _put(descendant, "2021", "default", 100) + + source.delete_exact("2021", "default", derived_only=True) + _assert_metadata(source) + remaining = dict(original) + del remaining[("default", "2021")] + probes = list(original) + _assert_snapshot(source, remaining, probes) + _assert_snapshot(copy, original, probes) + + copy.delete_exact("2021", "default", derived_only=True) + descendant.delete_exact("2021", "default", derived_only=True) + _assert_metadata(copy) + _assert_metadata(descendant) + _assert_snapshot(copy, remaining, probes) + own_input = dict(original) + own_input[("default", "2021")] = (100, False) + _assert_snapshot(descendant, own_input, probes) + + +def test_exact_deletion_releases_clone_marks_without_changing_its_source( + storage_family, +): + make, clone = storage_family + source = make() + _put(source, "2021", "default", 10) + _put(source, "2022", "default", 20, derived=True) + copy = clone(source) + descendant = clone(copy) + probes = [("default", "2021"), ("default", "2022")] + original = {probes[0]: (10, False), probes[1]: (20, True)} + + # Unconditional deletion still removes inputs, and dropping every + # shared key must work with the empty-set sentinels on a later retry. + copy.delete_exact("2021") + copy.delete_exact("2022", derived_only=True) + copy.delete_exact("2021") + copy.delete_exact("2022", derived_only=True) + _assert_metadata(copy) + if isinstance(copy, InMemoryStorage): + assert "_inputs" not in vars(copy) + assert "_shared" not in vars(copy) + _assert_snapshot(copy, {}, probes) + _assert_snapshot(source, original, probes) + _assert_snapshot(descendant, original, probes) + + # Reusing deleted keys changes only this copy, including provenance; + # disk storage must not overwrite files the other snapshots still read. + _put(copy, "2021", "default", 101, derived=True) + _put(copy, "2022", "default", 202) + _assert_snapshot(copy, {probes[0]: (101, True), probes[1]: (202, False)}, probes) + _assert_snapshot(source, original, probes) + _assert_snapshot(descendant, original, probes) + + +def test_derived_only_cleanup_keeps_an_input_and_its_clone_metadata(storage_family): + make, clone = storage_family + source = make() + _put(source, "2021", "default", 10) + copy = clone(source) + copy.delete_exact("2021", derived_only=True) + _assert_metadata(copy) + + expected = {("default", "2021"): (10, False)} + probes = list(expected) + _assert_snapshot(copy, expected, probes) + _assert_snapshot(source, expected, probes) + # A protected input can still be removed explicitly from the copy. + copy.delete_exact("2021") + _assert_metadata(copy) + _assert_snapshot(copy, {}, probes) + _assert_snapshot(source, expected, probes) + + +def test_eternal_exact_cleanup_normalizes_the_period_and_isolates_branches( + storage_family, +): + make, clone = storage_family + source = make(is_eternal=True) + _put(source, "2021", "default", 10) + _put(source, "2021-01", "other", 20, derived=True) + copy = clone(source) + original = { + ("default", periods.ETERNITY): (10, False), + ("other", periods.ETERNITY): (20, True), + } + probes = [ + (branch, period) + for branch in ("default", "other") + for period in ("2021", "2022-02", periods.ETERNITY) + ] + + copy.delete_exact("2022-02", "default", derived_only=True) + copy.delete_exact("2022", "other", derived_only=True) + _assert_metadata(copy) + _assert_snapshot(copy, {("default", periods.ETERNITY): (10, False)}, probes) + _assert_snapshot(source, original, probes) + copy.delete_exact("2023-03", "default") + _assert_metadata(copy) + _assert_snapshot(copy, {}, probes) + _assert_snapshot(source, original, probes) From 661926cc1c9d2c4361ea4ec1fbba9e7acec4308f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 10:23:15 -0400 Subject: [PATCH 4/9] Clarify exact storage deletion documentation Describe the shared behavior without claiming disk delete uses contained-period deletion. Co-Authored-By: Claude Opus 5.5 --- docs/python_api/data_storage.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 77f72adf9..d7d631448 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -3,8 +3,8 @@ The `policyengine_core.data_storage` module contains the classes that are used to handle the storage of data in simulations. Both storage classes provide `delete_exact(period, branch_name, derived_only)` -to remove a single stored period. Unlike `delete`, it leaves contained periods -in place. With `derived_only=True`, an input at that key remains stored; this is +to remove a single stored period, leaving contained periods in place. +With `derived_only=True`, an input at that key remains stored; this is used to discard provisional results from cut period recursions. ## InMemoryStorage From 8d7ea9cafd47f0d5cd08be31250e0d934d813391 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 10:45:25 -0400 Subject: [PATCH 5/9] Ignore deleted and derived input anchors during recursion scheduling Use readable storage provenance instead of stale input registry keys. Preserve replay bookkeeping and inherited inputs. Add traced and untraced deletion regressions and a bounded provenance property; clarify deferred-key retry invariants. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 37 +++++++------- tests/core/test_spiral_order.py | 47 ++++++++++++++++++ tests/core/test_spiral_scheduler_property.py | 51 +++++++++++++++++++- 3 files changed, 115 insertions(+), 20 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index a0150a473..e1a40a8fe 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -121,12 +121,14 @@ def calculate(self) -> ArrayLike: class _SpiralDeferrals: """The deeper periods one outermost calculation finishes first. - Each deeper period's calculation starts at most once (``started``), and - only on the simulation the calculation was requested on or one of its - registered branches (``root``). For a fixed branch graph and a finite - set of reachable calculation keys, every round starts a new period, - finishes one, or returns the request. No count of rounds cuts a long - chain, since where such a count fell would depend on earlier caching. + Once a deferred key starts evaluation (``started``), recursion cannot + defer it again. Its pending calculation can still be retried after its + deeper dependencies finish. Deferrals run only on the simulation the + calculation was requested on or its registered branches (``root``). + For a fixed branch graph and finite reachable calculation keys, every + round schedules a new key, finishes pending work, or returns the request. + No count of rounds cuts a long chain, since where such a count fell + would depend on earlier caching. """ # The simulation the outermost calculation was requested on, or the @@ -1865,7 +1867,7 @@ def _can_finish_deeper_period_first( has been calculated, so whether a chain is cut does not depend on calculation order either. - Each period's calculation starts once per outermost calculation: + A period cannot be deferred again after its evaluation has started: one reached again after that (a cycle through the deferrals, or a period calculated with a recursion cut on the way) is cut. And only on the simulation the calculation was requested on and its branches: @@ -1909,8 +1911,6 @@ def _can_finish_deeper_period_first( else: return False turn = stack[previous + 1 :] - input_keys = getattr(self, "_user_input_keys", ()) - visible_branches = self._get_visible_branch_names() for frame in turn: variable = tax_benefit_system.get_variable(frame["name"]) frame_period = frame["period"] @@ -1918,16 +1918,15 @@ def _can_finish_deeper_period_first( continue if self._has_formula_boundary_ahead(variable, frame_period, backward): return True - for input_name, branch_name, input_period in input_keys: - if ( - input_name == variable.name - and branch_name in visible_branches - and input_period.unit == variable.definition_period - and ( - input_period.start < frame_period.start - if backward - else input_period.start > frame_period.start - ) + # The set_input registry can outlive deleted inputs. Use the + # provenance of the value this branch actually reads, including + # ancestor fallback and replacements by calculated values. + holder = self.get_holder(variable.name) + for input_period in holder.get_input_periods(self.branch_name): + if input_period.unit == variable.definition_period and ( + input_period.start < frame_period.start + if backward + else input_period.start > frame_period.start ): return True return False diff --git a/tests/core/test_spiral_order.py b/tests/core/test_spiral_order.py index b921e9fd7..1e34b801c 100644 --- a/tests/core/test_spiral_order.py +++ b/tests/core/test_spiral_order.py @@ -66,6 +66,53 @@ def test_recursion_is_evaluated_down_to_an_input(): assert fresh == after == [36.0] +@pytest.mark.parametrize( + "removal", ["period", "all", "invalidate", "derived_replacement"] +) +@pytest.mark.parametrize("traced", [False, True]) +def test_deleted_input_does_not_keep_authorizing_deeper_periods(removal, traced): + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import YEAR, period + from policyengine_core.variables import Variable + + class recursive(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Undated recursion with an error before a former input" + + def formula(person, period): + if period.start.year == 2011: + raise ValueError("This period is beyond the unanchored depth limit") + return person("recursive", period.last_year) + 1 + + simulation = build([recursive], {("recursive", 2010): 40}, 1) + simulation.trace = traced + input_period = period("2010") + if removal == "all": + simulation.delete_arrays("recursive") + elif removal == "invalidate": + simulation.invalidate_cache_entry("recursive", input_period) + simulation.purge_cache_of_invalid_values() + else: + simulation.delete_arrays("recursive", input_period) + if removal == "derived_replacement": + # A calculated value at the old input's key has no input + # provenance and cannot authorize the scheduler to go deeper. + simulation.get_holder("recursive").put_in_cache( + [7], input_period, derived=True + ) + assert simulation.get_holder("recursive").get_input_periods() == [] + + fresh = build([recursive], max_spiral_loops=1) + fresh.trace = traced + assert fresh.calculate("recursive", "2020").tolist() == [1.0] + # With no actual input left, the recursion is cut before 2011. A stale + # set_input registry entry must not expose that unreachable error. + assert simulation.calculate("recursive", "2020").tolist() == [1.0] + assert simulation.tracer.stack == [] + + def test_forward_recursion_is_evaluated_up_to_the_end(): ahead = make_recurrence("ahead", [("ahead", -1)], end=2030) fresh, after = fresh_and_after([ahead], {}, ("ahead", 2000), [("ahead", 2020)]) diff --git a/tests/core/test_spiral_scheduler_property.py b/tests/core/test_spiral_scheduler_property.py index 2efac4c0e..d83e3f4ef 100644 --- a/tests/core/test_spiral_scheduler_property.py +++ b/tests/core/test_spiral_scheduler_property.py @@ -13,7 +13,7 @@ import pytest from policyengine_core.country_template.entities import Person -from policyengine_core.periods import YEAR +from policyengine_core.periods import YEAR, period from policyengine_core.variables import Variable from tests.fixtures.spirals import build, make_recurrence @@ -241,6 +241,55 @@ def test_deferred_periods_keep_caller_authorization(scenario, traced, storage): assert_request(fresh, "caller", year, expected[("r0", year)], not traced) +@hypothesis.settings(max_examples=24, deadline=None, derandomize=True) +@hypothesis.given( + input_year=st.integers(2010, 2014), + loops=st.sampled_from((1, 2, 3, 10)), + distance=st.integers(1, 5), + traced=st.booleans(), + removal=st.sampled_from(("period", "all", "invalidate", "derived_replacement")), +) +def test_scheduler_uses_current_input_provenance( + input_year, loops, distance, traced, removal +): + error_year = input_year + 1 + + class recursive(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Undated recurrence with an unreachable error" + + def formula(person, period): + if period.start.year == error_year: + raise ValueError("This period is beyond the unanchored depth limit") + return person("recursive", period.last_year) + 1 + + used = build([recursive], {("recursive", input_year): 40}, loops) + used.trace = traced + input_period = period(str(input_year)) + if removal == "all": + used.delete_arrays("recursive") + elif removal == "invalidate": + used.invalidate_cache_entry("recursive", input_period) + used.purge_cache_of_invalid_values() + else: + used.delete_arrays("recursive", input_period) + if removal == "derived_replacement": + used.get_holder("recursive").put_in_cache([7], input_period, derived=True) + assert used.get_holder("recursive").get_input_periods() == [] + + # Every generated request stops above the error in an unanchored + # simulation. Neither deleted input history nor a derived replacement + # may authorize the scheduler to reach it. + target = error_year + loops + distance + fresh = build([recursive], max_spiral_loops=loops) + fresh.trace = not traced + expected = fresh.calculate("recursive", str(target)) + np.testing.assert_array_equal(expected, np.array([loops], dtype=np.float32)) + assert_request(used, "recursive", target, expected, traced) + + @hypothesis.settings(max_examples=24, deadline=None, derandomize=True) @hypothesis.given( loops=st.sampled_from((1, 2, 3, 10)), From 0b241a5fbbc80638690cafc4443ba4fbc34f4ce3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 11:06:24 -0400 Subject: [PATCH 6/9] Keep bounded recurrence properties independent of generation timing Retain all example counts and semantic assertions. Suppress only the slow-generation health check for the bounded two-variable generator after a 47-second draw interrupted otherwise passing properties on the shared host. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_spiral_scheduler_property.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/core/test_spiral_scheduler_property.py b/tests/core/test_spiral_scheduler_property.py index d83e3f4ef..4b5ee4dad 100644 --- a/tests/core/test_spiral_scheduler_property.py +++ b/tests/core/test_spiral_scheduler_property.py @@ -202,7 +202,15 @@ def assert_request(simulation, name, year, expected, traced): ) -@hypothesis.settings(max_examples=60, deadline=None, derandomize=True) +# The generator is bounded to two variables, two terms and three inputs. +# Host pauses can trip generation timing even when the same draw takes +# milliseconds, so validate every example without a wall-clock health gate. +@hypothesis.settings( + max_examples=60, + deadline=None, + derandomize=True, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) @hypothesis.given( scenario=anchored_scenarios(), traced=st.booleans(), @@ -223,7 +231,12 @@ def test_anchored_scheduler_matches_iterative_values(scenario, traced, storage): assert_request(fresh, name, year, expected[(name, year)], not traced) -@hypothesis.settings(max_examples=40, deadline=None, derandomize=True) +@hypothesis.settings( + max_examples=40, + deadline=None, + derandomize=True, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) @hypothesis.given( scenario=anchored_scenarios(), traced=st.booleans(), From eac7ec566431f7b5ea9bdffefda10b2048e7ac58 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 04:09:30 -0400 Subject: [PATCH 7/9] Keep descendant explicit invalidations owned by their branch --- policyengine_core/simulations/simulation.py | 15 ++++++--- tests/core/test_clone_invalidated_caches.py | 34 +++++++++++++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index d1cff7235..48c27b091 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2279,20 +2279,25 @@ def _cache_result( return array def purge_cache_of_invalid_values(self) -> None: + self._purge_invalid_values(spiral_only=False) + + def _purge_invalid_values(self, *, spiral_only: bool) -> None: # We wait for the end of calculate(), signalled by an empty stack, before purging the cache if self.tracer.stack: return - # A traced branch shares this simulation's stack, so its frames can - # only be purged now that the stack is empty. + # A traced branch shares this simulation's stack, so its cut values + # can only be purged now that the stack is empty. Explicit invalidations + # belong to that branch and wait for its own calculation or purge. for branch in list(getattr(self, "branches", {}).values()): - branch.purge_cache_of_invalid_values() + branch._purge_invalid_values(spiral_only=True) _fast_cache = getattr(self, "_fast_cache", None) invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: return branch_name = getattr(self, "branch_name", "default") spiral_invalidations = getattr(self, "_spiral_invalidated_caches", ()) - for _name, _period in invalidated_caches: + purged = set(spiral_invalidations if spiral_only else invalidated_caches) + for _name, _period in purged: holder = self.get_holder(_name) if (_name, _period) in spiral_invalidations: # A cut only invalidates its calculated value, at exactly @@ -2323,7 +2328,7 @@ def purge_cache_of_invalid_values(self) -> None: ) ): _fast_cache.pop(key, None) - self.invalidated_caches = set() + self.invalidated_caches = invalidated_caches - purged self._spiral_invalidated_caches = set() def calculate_add( diff --git a/tests/core/test_clone_invalidated_caches.py b/tests/core/test_clone_invalidated_caches.py index 5a4cbd015..25d40fde3 100644 --- a/tests/core/test_clone_invalidated_caches.py +++ b/tests/core/test_clone_invalidated_caches.py @@ -139,3 +139,37 @@ def test_a_pending_invalidation_purges_each_copy_once(system, make): assert source.invalidated_caches == {("salary", MONTH)} source.purge_cache_of_invalid_values() assert source.get_holder("salary").get_array(MONTH) is None + + +@pytest.mark.parametrize("depth", [1, 2]) +@pytest.mark.parametrize("trace", [False, True]) +@pytest.mark.parametrize("spiral_pending", [False, True]) +def test_parent_purge_preserves_a_descendants_explicit_invalidation( + system, depth, trace, spiral_pending +): + source = _simulation(system) + source.trace = trace + source.set_input("salary", MONTH, [INPUT]) + branch = source + for index in range(depth): + branch = branch.get_branch(f"branch_{index}") + holder = branch.get_holder("income_tax") + value = np.array([42.0], dtype=holder.variable.dtype) + holder.put_in_cache(value, MONTH, derived=True) + pending = ("income_tax", MONTH) + branch.invalidate_cache_entry(*pending) + if spiral_pending: + cut_holder = branch.get_holder("spiral_a") + cut_holder.put_in_cache(value, MONTH, branch.branch_name, derived=True) + branch._invalidate_spiral_cache_entry("spiral_a", MONTH) + + source.calculate("income_tax", MONTH) + + assert branch.invalidated_caches == {pending} + np.testing.assert_array_equal(holder.get_array(MONTH), value) + if spiral_pending: + assert cut_holder._memory_storage.get(MONTH, branch.branch_name) is None + # The branch still consumes its explicit invalidation when it calculates. + branch.calculate("salary", MONTH) + assert branch.invalidated_caches == set() + assert holder.get_array(MONTH) is None From 1d8e8e255a21395c93ce8539524a8269e51ecfbb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 04:16:49 -0400 Subject: [PATCH 8/9] Preserve input and cut cache invariants across the master merge --- .../fix-spiral-anchored-recursion.fixed.md | 2 +- docs/python_api/data_storage.md | 13 +- policyengine_core/simulations/simulation.py | 151 ++++++++- tests/core/test_clone_invalidated_caches.py | 40 +++ tests/core/test_spiral_merge_integration.py | 288 ++++++++++++++++++ tests/core/test_spiral_storage.py | 59 ++++ 6 files changed, 538 insertions(+), 15 deletions(-) create mode 100644 tests/core/test_spiral_merge_integration.py diff --git a/changelog.d/fix-spiral-anchored-recursion.fixed.md b/changelog.d/fix-spiral-anchored-recursion.fixed.md index 13b757d6f..62a2e8e21 100644 --- a/changelog.d/fix-spiral-anchored-recursion.fixed.md +++ b/changelog.d/fix-spiral-anchored-recursion.fixed.md @@ -1 +1 @@ -Evaluate finite anchored period recursions in full regardless of earlier calculations, tracing or cache policy, and discard calculated values from cut recursions without deleting inputs or retaining invalid values in clones. +Evaluate finite anchored period recursions in full regardless of earlier calculations, tracing or cache policy, and discard calculated values from cut recursions without deleting inputs or retaining invalid values in clones, preserving branch-owned explicit invalidations and input and option-cache safety after merging master. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 11856463f..fe41d3fcc 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -5,7 +5,18 @@ The `policyengine_core.data_storage` module contains the classes that are used t Both storage classes provide `delete_exact(period, branch_name, derived_only)` to remove a single stored period, leaving contained periods in place. With `derived_only=True`, an input at that key remains stored; this is -used to discard provisional results from cut period recursions. +used to discard provisional results from cut period recursions. Exact holder +deletion also evicts the corresponding fast-cache entry and storage sequence +metadata. An explicit invalidation remains owned by its simulation: a parent's +cleanup discards only a descendant's cut-derived values, leaving its explicit +invalidation pending even when both refer to the same key. + +The period-recursion scheduler's temporary results are discarded after input +writes or computed-value invalidation. Retried calculations retain foreign +input dependencies, and repeated input changes use the existing calculation +retry limit; stable anchored chains have no scheduling count limit. ADD and +DIVIDE aggregates derived from a recursion cut are returned without being +cached, while inputs supplied during the calculation still take precedence. ## InMemoryStorage diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 48c27b091..b5654d676 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -154,6 +154,14 @@ class _SpiralDeferrals: # (``_do_not_store``, ``opt_out_cache``) and whether or not the # calculation is traced. It goes when the calculation ends. finished: dict = field(default_factory=dict) + # Mutations invalidate temporary results just as they invalidate holder + # caches. Reopen completed keys only; pending keys still guard cycles. + input_generation: int = field(default_factory=lambda: _spiral_input_generation) + input_restarts: int = 0 + current_key: Optional[tuple] = None + completed_frame: Any = None + # Every cut-touched key, including those with a pending explicit invalidation. + invalidated: dict = field(default_factory=dict) # How many recursions have been cut so far. cuts: int = 0 # Requests raised and not yet back in the outermost ``calculate``: a @@ -172,13 +180,53 @@ def finish(self, deferral: _DeeperPeriodFirst, value: ArrayLike) -> None: entry = self.finished.get(id(simulation)) if entry is None: entry = self.finished[id(simulation)] = (simulation, {}) - entry[1][(deferral.variable, deferral.period)] = value + frame = self.completed_frame + entry[1][(deferral.variable, deferral.period)] = ( + value, + tuple(frame.reads.items()) if frame is not None else (), + tuple(frame.registrations) if frame is not None else (), + ) + self.completed_frame = None + + def refresh_inputs(self) -> bool: + """Discard temporary values after writes, without reopening active cycles. + + The retry limit applies only to input mutations, as in ``_Frame``; + stable anchored chains retain their unlimited iterative scheduling. + """ + generation = _spiral_input_generation + if generation == self.input_generation: + return False + self.input_generation = generation + keys = { + (identity, variable, period) + for identity, (_, values) in self.finished.items() + for variable, period in values + } + if keys: + if self.input_restarts < _RERUNS_AFTER_INPUT_CHANGE: + self.started.difference_update(keys) + self.done.difference_update(keys) + self.input_restarts += 1 + self.finished.clear() + return True def finished_value( self, simulation: "Simulation", variable: str, period: Period ) -> Optional[ArrayLike]: entry = self.finished.get(id(simulation)) - return None if entry is None else entry[1].get((variable, period)) + completed = None if entry is None else entry[1].get((variable, period)) + if completed is None: + return None + value, reads, registrations = completed + frames = _calculation_frames.get() + if frames: + frame = frames[-1] + for other, epoch in reads: + if other is not frame.simulation: + frame.reads.setdefault(other, epoch) + frame.registrations.update(registrations) + return value _spiral_deferrals: ContextVar[Optional[_SpiralDeferrals]] = ContextVar( @@ -318,6 +366,18 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: # changing ordinary frames' scoped read tracking or convergence decisions. _restricted_frames = set() _restricted_frames_lock = Lock() +# Dedicated generations exclude ordinary reads and completed calculations. +# Process-wide counters also observe writes/cuts in context-free worker calls. +_spiral_input_generation = 0 +_spiral_cut_generation = 0 + + +def _note_spiral_input_change() -> None: + global _spiral_input_generation + with _restricted_frames_lock: + _spiral_input_generation += 1 + + # Activity between restricted calculations must also invalidate their shared # attempt cache, even when no restricted child is currently running. _unobserved_activity_generation = 0 @@ -738,6 +798,7 @@ def _note_input_change( value: Optional[ArrayLike] = None, ) -> None: """Record actual stores and deletes in each frame observing this context.""" + _note_spiral_input_change() frames = _calculation_frames.get() if not frames: _taint_waiting_frames(simulation) @@ -1587,7 +1648,10 @@ def calculate( # Deeper periods to calculate before the request, deepest last. pending = [] while True: + state.refresh_inputs() deferral = pending[-1] if pending else None + state.current_key = None if deferral is None else deferral.key + state.completed_frame = None if deferral is not None: if deferral.key in state.done: # Requested twice before it started (see ``_defer``). @@ -1596,6 +1660,8 @@ def calculate( state.started.add(deferral.key) state.callers = frozenset() if deferral is None else deferral.callers cuts = state.cuts + input_generation = _spiral_input_generation + input_restarts = state.input_restarts attempt_tracer = ( self.tracer if deferral is None else deferral.simulation.tracer ) @@ -1637,9 +1703,24 @@ def calculate( ): attempt_tracer.trees[attempt_index]._spiral_provisional = False return result + changed = input_generation != _spiral_input_generation + state.refresh_inputs() + if changed: + # Refresh may already have counted this mutation wave. + if state.input_restarts == input_restarts: + state.input_restarts += 1 + if state.input_restarts <= _RERUNS_AFTER_INPUT_CHANGE: + # Finish the pending request at settled inputs. A result + # from an exhausted input retry is never reusable. + continue pending.pop() state.done.add(deferral.key) - if state.cuts == cuts: + completed = state.completed_frame + if ( + state.cuts == cuts + and not changed + and (completed is None or completed.may_keep()) + ): state.finish(deferral, result) # Otherwise a recursion was cut on the way, so the value # depends on where. Starting again reaches this period again, @@ -1751,6 +1832,13 @@ def _calculate_traced( variable_name ).entity.key result = self.map_result(result, source_entity, map_to) + if state is not None and state.current_key == ( + id(self), + variable_name, + period, + ): + frames = _calculation_frames.get() + state.completed_frame = frames[-1] if frames else None return result finally: try: @@ -1875,6 +1963,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: """ if variable_name not in self.tax_benefit_system.variables: raise ValueError(f"Variable {variable_name} does not exist.") + state = _spiral_deferrals.get() + if state is not None: + state.refresh_inputs() input_state = self._calculation_start() # Lets a custom ``set_input`` handler tell whether it calculated. self._calculations_started = self._calculations_started + 1 @@ -2247,6 +2338,8 @@ def _cache_result( array: ArrayLike, period: Period, input_state: Tuple[int, int, int], + *, + allow_store: bool = True, ) -> ArrayLike: """Cache a calculated value, and return the value to use for it. @@ -2262,6 +2355,9 @@ def _cache_result( stored_input = self._input_set_meanwhile(holder, period, started_at) if stored_input is not None: return stored_input + if not allow_store: + holder._record_store(period, next_sequence_number()) + return array frames = _calculation_frames.get() may_keep = not frames or frames[-1].may_keep() if epoch == self._input_epoch and may_keep: @@ -2295,11 +2391,15 @@ def _purge_invalid_values(self, *, spiral_only: bool) -> None: if invalidated_caches is None: return branch_name = getattr(self, "branch_name", "default") - spiral_invalidations = getattr(self, "_spiral_invalidated_caches", ()) - purged = set(spiral_invalidations if spiral_only else invalidated_caches) + spiral_invalidations = set(getattr(self, "_spiral_invalidated_caches", ())) + state = _spiral_deferrals.get() + touched = set() if state is None else state.invalidated.pop(self, set()) + purged = ( + spiral_invalidations | touched if spiral_only else set(invalidated_caches) + ) for _name, _period in purged: holder = self.get_holder(_name) - if (_name, _period) in spiral_invalidations: + if spiral_only or (_name, _period) in spiral_invalidations: # A cut only invalidates its calculated value, at exactly # this period and under the calculating branch's name. holder.delete_array(_period, branch_name, derived_only=True) @@ -2328,7 +2428,8 @@ def _purge_invalid_values(self, *, spiral_only: bool) -> None: ) ): _fast_cache.pop(key, None) - self.invalidated_caches = invalidated_caches - purged + consumed = spiral_invalidations if spiral_only else purged + self.invalidated_caches = invalidated_caches - consumed self._spiral_invalidated_caches = set() def calculate_add( @@ -2381,6 +2482,7 @@ def total(): try: with frame if frame is not None else nullcontext(): input_state = self._calculation_start() + cut_generation = _spiral_cut_generation result = total() previous_transition = None for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): @@ -2395,8 +2497,11 @@ def total(): # obsolete (see ``calculate``). Sum again. frame.restart() input_state = self._calculation_start() + cut_generation = _spiral_cut_generation result = total() - return self._cache_option_result(variable, period, result, input_state) + return self._cache_option_result( + variable, period, result, input_state, cut_generation + ) finally: self._calculations_in_flight -= 1 @@ -2431,8 +2536,11 @@ def calculate_divide( with frame if frame is not None else nullcontext(): input_state = self._calculation_start() computation_period = period.this_year + cut_generation = _spiral_cut_generation result = self.calculate(variable_name, period=computation_period) / 12.0 - return self._cache_option_result(variable, period, result, input_state) + return self._cache_option_result( + variable, period, result, input_state, cut_generation + ) elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -2448,6 +2556,7 @@ def _cache_option_result( period: Period, result: ArrayLike, input_state: Tuple[int, int, int], + cut_generation: int, ) -> ArrayLike: """Cache an ADD or DIVIDE result at ``period`` if a plain read would return it. @@ -2491,8 +2600,14 @@ def _cache_option_result( ) if not routed or np.asarray(result).dtype != variable.dtype: return result + # A direct option outlives each term's scheduler. Skip storing a cut + # aggregate, retaining the input-set-meanwhile rule in ``_cache_result``. return self._cache_result( - self.get_holder(variable.name), result, period, input_state + self.get_holder(variable.name), + result, + period, + input_state, + allow_store=cut_generation == _spiral_cut_generation, ) def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: @@ -2865,6 +2980,9 @@ def invalidate_cache_entry(self, variable: str, period: Period) -> None: def _invalidate_spiral_cache_entry(self, variable: str, period: Period) -> None: """Mark an exact calculated value, preserving explicit invalidations.""" key = (variable, period) + state = _spiral_deferrals.get() + if state is not None: + state.invalidated.setdefault(self, set()).add(key) invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: invalidated_caches = self.invalidated_caches = set() @@ -2887,6 +3005,9 @@ def invalidate_spiral_variables(self, variable: str) -> None: each active stack empties is marked as it is cached (see ``_note_cached_value``). They are still reused until then. """ + global _spiral_cut_generation + with _restricted_frames_lock: + _spiral_cut_generation += 1 for frame in reversed(self.tracer.stack): # A traced branch shares its parent's stack, so a frame can be # the parent's (or another ancestor's). @@ -3198,6 +3319,7 @@ def _drop_values_that_may_depend_on( def _drop_computed(self, since: Optional[int] = None) -> int: """Drop every non-input value numbered ``since`` or later (all, without it).""" + _note_spiral_input_change() dropped = 0 for population in self.populations.values(): for holder in population._holders.values(): @@ -3450,7 +3572,10 @@ def _drop_invalid_values(self, source: "Simulation") -> None: changes; ``source`` purges its own values when its calculation ends. Inputs stay, as they do in ``source``. """ - invalidated_caches = getattr(source, "_spiral_invalidated_caches", None) + spiral_invalidations = set(getattr(source, "_spiral_invalidated_caches", ())) + state = _spiral_deferrals.get() + touched = set() if state is None else state.invalidated.get(source, set()) + invalidated_caches = spiral_invalidations | touched if not invalidated_caches: return branch_name = getattr(source, "branch_name", "default") @@ -3461,8 +3586,8 @@ def _drop_invalid_values(self, source: "Simulation") -> None: holder = self.populations[variable.entity.key]._holders.get(name) if holder is not None: holder.delete_array(invalid_period, branch_name, derived_only=True) - self.invalidated_caches.difference_update(invalidated_caches) - self._spiral_invalidated_caches.difference_update(invalidated_caches) + self.invalidated_caches.difference_update(spiral_invalidations) + self._spiral_invalidated_caches.difference_update(spiral_invalidations) def get_branch( self, name: str = "branch", clone_system: bool = False diff --git a/tests/core/test_clone_invalidated_caches.py b/tests/core/test_clone_invalidated_caches.py index 25d40fde3..ae04207bc 100644 --- a/tests/core/test_clone_invalidated_caches.py +++ b/tests/core/test_clone_invalidated_caches.py @@ -173,3 +173,43 @@ def test_parent_purge_preserves_a_descendants_explicit_invalidation( branch.calculate("salary", MONTH) assert branch.invalidated_caches == set() assert holder.get_array(MONTH) is None + + +@pytest.mark.parametrize("clone_after_cut", [False, True]) +def test_branch_cut_cleanup_preserves_an_overlapping_explicit_invalidation( + clone_after_cut, +): + copies = [] + + class from_branch(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "Reads a cut recursion on a traced branch" + + def formula(person, period): + branch = person.simulation.get_branch("other") + value = branch.calculate("spiral_a", period) + if clone_after_cut: + # The shared tracer still holds the parent's formula, so + # the branch has not yet purged its cut-derived values. + copies.append(branch.clone()) + return value + + system = CountryTaxBenefitSystem() + system.add_variables(spiral_a, spiral_b, from_branch) + source = _simulation(system) + source.trace = True + branch = source.get_branch("other") + pending = ("spiral_a", MONTH) + branch.invalidate_cache_entry(*pending) + + np.testing.assert_array_equal(source.calculate("from_branch", MONTH), [22.0]) + + # Keeping the explicit request pending must not keep the calculated + # value reached by the cut, even when both concern exactly the same key. + for simulation in [branch, *copies]: + assert simulation.invalidated_caches == {pending} + holder = simulation.get_holder("spiral_a") + assert holder._memory_storage.get(MONTH, simulation.branch_name) is None + assert pending not in simulation._fast_cache diff --git a/tests/core/test_spiral_merge_integration.py b/tests/core/test_spiral_merge_integration.py new file mode 100644 index 000000000..08b74953a --- /dev/null +++ b/tests/core/test_spiral_merge_integration.py @@ -0,0 +1,288 @@ +"""Recursion scheduling preserves input invalidation and option-cache rules.""" + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template.entities import Person +from policyengine_core.variables import QuantityType, Variable +from tests.fixtures.spirals import build + + +@pytest.mark.parametrize("traced", [False, True]) +@pytest.mark.parametrize("store", [True, False]) +def test_finished_deeper_periods_follow_an_input_changed_by_their_caller(traced, store): + class source(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Input used by the recursive calculation" + + class recursive(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Anchored recurrence over a mutable input" + + def formula_2010(person, period): + return person("recursive", period.last_year) + person("source", "2010") + + class caller(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Caller that replaces the recursive calculation's input" + + def formula(person, period): + result = person("recursive", period) + # The outer calculation must run again with this input. Deeper + # periods completed by the scheduler must run again as well. + person.simulation.set_input("source", "2010", np.array([2.0])) + return result + + simulation = build([source, recursive, caller], {("source", 2010): 1}, 1) + simulation.trace = traced + branch = simulation.get_branch("changed") + if not store: + for name in ("recursive", "caller"): + branch.get_holder(name)._do_not_store = True + + fresh = build([source, recursive, caller], {("source", 2010): 2}, 1) + fresh.trace = traced + expected = fresh.calculate("caller", "2024") + + np.testing.assert_array_equal(expected, [30.0]) + np.testing.assert_array_equal(branch.calculate("caller", "2024"), expected) + assert branch.tracer.stack == [] + assert branch._calculations_in_flight == 0 + + +@pytest.mark.parametrize("traced", [False, True]) +def test_finished_deeper_periods_follow_a_foreign_input_change(traced): + class source(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Input read in another simulation" + + class recursive(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Anchored recurrence with a foreign input" + + def formula_2010(person, period): + inputs = person.simulation.get_branch("inputs") + return person("recursive", period.last_year) + inputs.calculate( + "source", "2010" + ) + + class caller(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Caller that replaces a foreign input" + + def formula(person, period): + result = person("recursive", period) + person.simulation.get_branch("inputs").set_input( + "source", "2010", np.array([2.0]) + ) + return result + + simulation = build([source, recursive, caller], {("source", 2010): 1}, 1) + simulation.trace = traced + simulation.get_branch("inputs") + # Isolate the scheduler's temporary values. Persistent values from + # already-ended calculations intentionally survive foreign mutations. + del simulation._fast_cache + for name in ("recursive", "caller"): + simulation.get_holder(name)._do_not_store = True + + np.testing.assert_array_equal(simulation.calculate("caller", "2024"), [30.0]) + assert simulation.tracer.stack == [] + assert simulation._calculations_in_flight == 0 + + +@pytest.mark.parametrize("traced", [False, True]) +@pytest.mark.parametrize("settles", [True, False]) +def test_repeated_input_changes_settle_or_stop_without_reusing_stale_values( + traced, settles +): + calls = [] + + class source(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Input changed on successive attempts" + + class recursive(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Anchored recurrence whose input changes repeatedly" + + def formula_2010(person, period): + return person("recursive", period.last_year) + person("source", "2010") + + class caller(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Caller that repeatedly changes a recursive input" + + def formula(person, period): + calls.append(str(period)) + # Fail promptly if reopening completed keys bypasses the + # existing bound on retries after input changes. + assert len(calls) <= 200, "Input changes kept reopening deferred periods" + result = person("recursive", period) + old = person("source", "2010") + new = np.minimum(old + 1, 7) if settles else 3 - old + person.simulation.set_input("source", "2010", new) + return result + + simulation = build([source, recursive, caller], {("source", 2010): 1}, 1) + simulation.trace = traced + branch = simulation.get_branch("changed") + + result = branch.calculate("caller", "2024") + + if settles: + np.testing.assert_array_equal(result, [105.0]) + np.testing.assert_array_equal(branch.calculate("recursive", "2024"), [105.0]) + else: + # A continuously changing input has no stable numerical answer. + # The retry limit still returns a finite result and keeps no caller + # result calculated before its last input change. + assert np.all(np.isfinite(result)) + assert branch.get_array("caller", "2024") is None + assert branch.tracer.stack == [] + assert branch._calculations_in_flight == 0 + + +@pytest.mark.parametrize("traced", [False, True]) +def test_an_unsettled_deferred_period_is_never_published_as_finished(traced): + calls = [] + + class source(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Input continuously changed inside a deferred calculation" + + class recursive(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Recurrence with an unsettled deeper period" + + def formula_2010(person, period): + if period.start.year == 2023: + calls.append(str(period)) + assert len(calls) <= 200, "An unsettled deferred period kept retrying" + result = person("source", "2010") + person.simulation.set_input("source", "2010", 3 - result) + # A nested calculation observes the mutation generation + # before the scheduler receives this still-stale result. + person("source", "2010") + return result + return person("recursive", period.last_year) + 1 + + simulation = build([source, recursive], {("source", 2010): 1}, 1) + simulation.trace = traced + branch = simulation.get_branch("changed") + + result = branch.calculate("recursive", "2024") + + assert calls + assert np.all(np.isfinite(result)) + # The input never settles. Neither the exhausted deferred result nor + # a caller using it may be kept as though calculated at settled inputs. + assert branch.get_array("recursive", "2023") is None + assert branch.get_array("recursive", "2024") is None + assert branch.tracer.stack == [] + assert branch._calculations_in_flight == 0 + + +@pytest.mark.parametrize("traced", [False, True]) +@pytest.mark.parametrize("operation", ["add", "divide"]) +def test_direct_option_does_not_keep_a_value_derived_from_a_recursion_cut( + traced, operation +): + definition_period = periods.MONTH if operation == "add" else periods.YEAR + + class recursive(Variable): + value_type = float + entity = Person + quantity_type = QuantityType.FLOW + label = "Unanchored recurrence used through an option" + + def formula(person, period): + previous = period.offset(-1, definition_period) + return person("recursive", previous) + 1 + + recursive.definition_period = definition_period + used = build([recursive], max_spiral_loops=1) + used.trace = traced + fresh = build([recursive], max_spiral_loops=1) + fresh.trace = traced + target = "2020" if operation == "add" else "2020-01" + method = used.calculate_add if operation == "add" else used.calculate_divide + + result = method("recursive", target) + + # The direct option has a different recursion depth from a plain read + # of the converted period. Its incomplete result must not answer that + # subsequent plain read, nor remain available for cloning or export. + expected_option = [12.0] if operation == "add" else [1.0 / 12] + np.testing.assert_allclose(result, expected_option) + assert used.get_array("recursive", target) is None + np.testing.assert_array_equal( + used.calculate("recursive", target), fresh.calculate("recursive", target) + ) + assert used.get_array("recursive", target) is None + assert used.tracer.stack == [] + assert used._calculations_in_flight == 0 + + +@pytest.mark.parametrize("traced", [False, True]) +@pytest.mark.parametrize("operation", ["add", "divide"]) +def test_a_cut_option_keeps_an_input_set_meanwhile_at_its_requested_period( + traced, operation +): + definition_period = periods.MONTH if operation == "add" else periods.YEAR + target = "2020" if operation == "add" else "2020-01" + + class recursive(Variable): + value_type = float + entity = Person + quantity_type = QuantityType.FLOW + set_input = None + label = "Cut recurrence that supplies the requested option value" + + def formula(person, period): + previous = period.offset(-1, definition_period) + result = person("recursive", previous) + 1 + if period.start.year == 2020 and period.start.month == 1: + person.simulation.set_input("recursive", target, np.array([120.0])) + return result + + recursive.definition_period = definition_period + simulation = build([recursive], max_spiral_loops=1) + simulation.trace = traced + method = ( + simulation.calculate_add if operation == "add" else simulation.calculate_divide + ) + + result = method("recursive", target) + + # Suppressing an incomplete aggregate's cache write must retain the + # input-first rule: an input supplied for that aggregate while it ran + # supplies its answer, even though the computed terms were cut. + np.testing.assert_array_equal(result, [120.0]) + np.testing.assert_array_equal(simulation.get_array("recursive", target), [120.0]) + assert not simulation.get_holder("recursive").is_derived(periods.period(target)) + assert simulation.tracer.stack == [] + assert simulation._calculations_in_flight == 0 diff --git a/tests/core/test_spiral_storage.py b/tests/core/test_spiral_storage.py index e66d2e9ea..df5edde65 100644 --- a/tests/core/test_spiral_storage.py +++ b/tests/core/test_spiral_storage.py @@ -11,6 +11,7 @@ from policyengine_core import periods from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from tests.fixtures.disk_storage import build_simulation @pytest.fixture(params=["memory", "shared-memory", "disk"]) @@ -63,8 +64,10 @@ def _assert_metadata(storage): if isinstance(storage, InMemoryStorage): assert storage._shared <= storage._arrays.keys() assert storage._inputs <= storage._arrays.keys() + assert storage._sequence_numbers.keys() <= storage._arrays.keys() else: assert storage._derived <= storage._files.keys() + assert storage._sequence_numbers.keys() <= storage._files.keys() def test_exact_cleanup_keeps_descendant_inputs_and_contained_periods(storage_family): @@ -182,3 +185,59 @@ def test_eternal_exact_cleanup_normalizes_the_period_and_isolates_branches( _assert_metadata(copy) _assert_snapshot(copy, {}, probes) _assert_snapshot(source, original, probes) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_holder_exact_cleanup_drops_only_its_cached_period(on_disk): + simulation = build_simulation(on_disk=on_disk) + january = periods.period("2021-01") + february = periods.period("2021-02") + simulation.set_input("salary", january, [1000]) + simulation.set_input("salary", february, [2000]) + first = simulation.calculate("income_tax", january) + other = simulation.calculate("income_tax", february) + holder = simulation.get_holder("income_tax") + + holder.delete_array(january, derived_only=True) + + assert holder.get_array(january) is None + assert holder._stored_sequence_number(january) is None + assert ("income_tax", january) not in simulation._fast_cache + assert simulation._fast_cache[("income_tax", february)] is other + np.testing.assert_array_equal(holder.get_array(february), other) + np.testing.assert_array_equal(simulation.calculate("income_tax", january), first) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_holder_exact_deletion_forgets_only_an_input_it_removes(on_disk): + simulation = build_simulation(on_disk=on_disk) + january = periods.period("2021-01") + february = periods.period("2021-02") + simulation.set_input("income_tax", january, [7]) + simulation.set_input("income_tax", february, [11]) + simulation.calculate("income_tax", january) + holder = simulation.get_holder("income_tax") + january_key = ("income_tax", "default", january) + february_key = ("income_tax", "default", february) + + holder.delete_array(january, derived_only=True) + + np.testing.assert_array_equal(holder.get_array(january), [7]) + assert january_key in simulation._user_input_keys + assert holder._stores_user_input(january, "default") + + holder.delete_array(january) + + assert holder.get_array(january) is None + assert january_key not in simulation._user_input_keys + assert ("income_tax", january) not in simulation._fast_cache + assert february_key in simulation._user_input_keys + assert holder._stores_user_input(february, "default") + simulation.calculate("income_tax", january) + assert holder.is_derived(january) + assert january_key not in simulation._user_input_keys + + simulation._invalidate_all_caches() + + assert holder.get_array(january) is None + np.testing.assert_array_equal(holder.get_array(february), [11]) From 7bc9a284efd96915ca1242e5e84d2dcd7e573722 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 04:20:11 -0400 Subject: [PATCH 9/9] Scope deferred invalidation to its input dependencies --- policyengine_core/simulations/simulation.py | 59 +++++++++++++++----- tests/core/test_branch_input_invalidation.py | 10 ++-- tests/core/test_spiral_merge_integration.py | 37 ++++++++++++ 3 files changed, 87 insertions(+), 19 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index b5654d676..2acc24d1c 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -185,6 +185,10 @@ def finish(self, deferral: _DeeperPeriodFirst, value: ArrayLike) -> None: value, tuple(frame.reads.items()) if frame is not None else (), tuple(frame.registrations) if frame is not None else (), + tuple( + (other, other._spiral_input_generation) + for other in {simulation, *(frame.reads if frame is not None else ())} + ), ) self.completed_frame = None @@ -198,18 +202,29 @@ def refresh_inputs(self) -> bool: if generation == self.input_generation: return False self.input_generation = generation - keys = { - (identity, variable, period) - for identity, (_, values) in self.finished.items() - for variable, period in values - } + keys = set() + for identity, (_, values) in self.finished.items(): + stale = [ + key + for key, (_, _, registrations, inputs) in values.items() + if any( + other._spiral_input_generation != version + for other, version in inputs + ) + or any( + parent.branches.get(name) is not branch + for parent, name, branch in registrations + ) + ] + for variable, period in stale: + del values[(variable, period)] + keys.add((identity, variable, period)) if keys: if self.input_restarts < _RERUNS_AFTER_INPUT_CHANGE: self.started.difference_update(keys) self.done.difference_update(keys) self.input_restarts += 1 - self.finished.clear() - return True + return bool(keys) def finished_value( self, simulation: "Simulation", variable: str, period: Period @@ -218,7 +233,19 @@ def finished_value( completed = None if entry is None else entry[1].get((variable, period)) if completed is None: return None - value, reads, registrations = completed + value, reads, registrations, _ = completed + if any( + parent.branches.get(name) is not branch + for parent, name, branch in registrations + ): + # Registrations can change without writing an input. + del entry[1][(variable, period)] + key = (id(simulation), variable, period) + if self.input_restarts < _RERUNS_AFTER_INPUT_CHANGE: + self.started.discard(key) + self.done.discard(key) + self.input_restarts += 1 + return None frames = _calculation_frames.get() if frames: frame = frames[-1] @@ -372,9 +399,10 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: _spiral_cut_generation = 0 -def _note_spiral_input_change() -> None: +def _note_spiral_input_change(simulation: "Simulation") -> None: global _spiral_input_generation with _restricted_frames_lock: + simulation._spiral_input_generation += 1 _spiral_input_generation += 1 @@ -798,7 +826,7 @@ def _note_input_change( value: Optional[ArrayLike] = None, ) -> None: """Record actual stores and deletes in each frame observing this context.""" - _note_spiral_input_change() + _note_spiral_input_change(simulation) frames = _calculation_frames.get() if not frames: _taint_waiting_frames(simulation) @@ -867,6 +895,7 @@ class Simulation: # Counters read on every calculation (see ``set_input``). Class defaults, # so a subclass that skips ``__init__`` reads them too; incrementing one # gives the simulation its own. + _spiral_input_generation: int = 0 # writes/deletes affecting temporary results _input_epoch: int = 0 # drops on this simulation _inputs_set: int = 0 # inputs stored through ``Holder.set_input`` _calculations_in_flight: int = 0 # ``calculate``/``calculate_add`` running @@ -1660,7 +1689,8 @@ def calculate( state.started.add(deferral.key) state.callers = frozenset() if deferral is None else deferral.callers cuts = state.cuts - input_generation = _spiral_input_generation + calculating = self if deferral is None else deferral.simulation + input_generation = calculating._spiral_input_generation input_restarts = state.input_restarts attempt_tracer = ( self.tracer if deferral is None else deferral.simulation.tracer @@ -1703,7 +1733,7 @@ def calculate( ): attempt_tracer.trees[attempt_index]._spiral_provisional = False return result - changed = input_generation != _spiral_input_generation + changed = input_generation != calculating._spiral_input_generation state.refresh_inputs() if changed: # Refresh may already have counted this mutation wave. @@ -2272,7 +2302,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: except SpiralError: array = holder.default_array() - # Not stored, but what reads it is (see ``set_input``). + # Record this source even though cut-derived caches are purged + # when the calculation ends (see ``set_input``). holder._record_store(period, next_sequence_number()) except RecursionError as e: if isinstance(self.tracer, FullTracer): @@ -3319,7 +3350,7 @@ def _drop_values_that_may_depend_on( def _drop_computed(self, since: Optional[int] = None) -> int: """Drop every non-input value numbered ``since`` or later (all, without it).""" - _note_spiral_input_change() + _note_spiral_input_change(self) dropped = 0 for population in self.populations.values(): for holder in population._holders.values(): diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index ec86d82ba..84300c22d 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -533,14 +533,14 @@ def test_value_a_holder_does_not_keep_is_tracked(): def test_value_a_spiral_defaults_is_tracked(): - """The default a spiral returns is not stored, but what reads it is.""" + """A spiral's default stays in dependency history after cut caches are purged.""" simulation = synthetic_simulation(ROOT_INPUTS) branch = simulation.get_branch("branch") branch.calculate("p_spiral", "2015") - deepest = min( - periods.period(key.split(":", 1)[1]) - for key in branch.get_holder("p_spiral")._memory_storage._arrays - ).last_year + assert not branch.get_holder("p_spiral")._memory_storage._arrays + history = branch._get_store_history() + deepest = min(history._first_stored["p_spiral"]) + assert history.earliest_dependency("p_spiral", deepest) is not None branch.set_input("p_spiral", deepest, np.array([100.0, 100.0, 100.0])) diff --git a/tests/core/test_spiral_merge_integration.py b/tests/core/test_spiral_merge_integration.py index 08b74953a..c3dfa4496 100644 --- a/tests/core/test_spiral_merge_integration.py +++ b/tests/core/test_spiral_merge_integration.py @@ -9,6 +9,43 @@ from tests.fixtures.spirals import build +@pytest.mark.parametrize("traced", [False, True]) +@pytest.mark.parametrize("storage", ["drop", "blacklist"]) +def test_unrelated_input_writes_do_not_spend_the_recursion_input_retry_budget( + traced, storage +): + class scratch_input(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Input in an independent temporary simulation" + + class recursive(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Anchored recurrence independent of scratch inputs" + + def formula_2010(person, period): + # This temporary simulation is never read by the recurrence. + # Its writes must not discard completed periods or exhaust the + # retry budget of an unchanged simulation's anchored chain. + build([scratch_input], {("scratch_input", 2010): 1}) + return person("recursive", period.last_year) + 1 + + simulation = build([recursive], max_spiral_loops=1) + simulation.trace = traced + if storage == "drop": + simulation.get_holder("recursive")._do_not_store = True + else: + simulation.opt_out_cache = True + simulation.tax_benefit_system.cache_blacklist = {"recursive"} + + np.testing.assert_array_equal(simulation.calculate("recursive", "2024"), [15.0]) + assert simulation.tracer.stack == [] + assert simulation._calculations_in_flight == 0 + + @pytest.mark.parametrize("traced", [False, True]) @pytest.mark.parametrize("store", [True, False]) def test_finished_deeper_periods_follow_an_input_changed_by_their_caller(traced, store):