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 00000000..13b757d6 --- /dev/null +++ b/changelog.d/fix-spiral-anchored-recursion.fixed.md @@ -0,0 +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. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 4307934f..d7d63144 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -2,6 +2,11 @@ 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, 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 ```{eval-rst} diff --git a/docs/python_api/tracers.md b/docs/python_api/tracers.md index 6c72e330..ab894cea 100644 --- a/docs/python_api/tracers.md +++ b/docs/python_api/tracers.md @@ -13,6 +13,12 @@ The `policyengine_core.tracers` module contains classes used to represent tracer ## FlatTrace +Flat traces retain the first completed calculation's dependencies and parameters +when a later request reads its cache. When the period-recursion scheduler retries +a calculation, a completed accepted result replaces an abandoned or provisional +cut result for the same variable, period, and branch. The full computation and +performance logs still contain the retry attempts. + ```{eval-rst} .. autoclass:: policyengine_core.tracers.flat_trace.FlatTrace :members: diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index f1c3304f..7add3dab 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -320,6 +320,26 @@ 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", + derived_only: bool = False, + ) -> None: + """Drop the array stored for exactly ``period``. + + 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) + 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 3c5dbbb3..3e1f30fa 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -350,6 +350,22 @@ 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", + 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) + 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 324e3d7e..d50f9087 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -110,6 +110,21 @@ def delete_arrays( if self._disk_storage: self._disk_storage.delete(period, branch_name) + def delete_array( + self, + period: Period, + branch_name: str = "default", + derived_only: bool = False, + ) -> None: + """ + 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, derived_only) + if self._disk_storage: + self._disk_storage.delete_exact(period, branch_name, derived_only) + def _get_array_from_storage( self, period: Period, branch_name: str = "default" ) -> ArrayLike: @@ -423,6 +438,7 @@ def put_in_cache( return self._set(period, value, branch_name, derived=derived) + 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 123a5f1a..e1a40a8f 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2,7 +2,7 @@ import os import types 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 @@ -67,6 +67,119 @@ def __init__(self, simulation: "Simulation"): ) +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 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. 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, + 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) -> ArrayLike: + return self.simulation._calculate_traced(self.variable, self.period) + + +@dataclass +class _SpiralDeferrals: + """The deeper periods one outermost calculation finishes first. + + 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 + # 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 +) + + def _stable_hash_to_seed(value: str) -> int: """Deterministically hash a string to an int suitable for numpy.random.seed. @@ -296,6 +409,7 @@ def __init__( self.is_over_dataset = dataset is not None self.invalidated_caches = set() + self._spiral_invalidated_caches = set() self._fast_cache: dict = {} # ``set_input`` records each (variable_name, branch_name, period) it # populates so ``_invalidate_all_caches`` can tell user-provided @@ -443,6 +557,7 @@ def _invalidate_all_caches(self) -> None: """ self._fast_cache = {} self.invalidated_caches = set() + self._spiral_invalidated_caches = set() # Snapshot user-provided inputs before wiping so they can be # replayed into the fresh storage. Use the storage API instead of # hand-building keys, since ETERNITY variables canonicalize every @@ -852,7 +967,104 @@ 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) + + 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 + attempt_tracer = ( + self.tracer if deferral is None else deferral.simulation.tracer + ) + attempt_index = ( + len(attempt_tracer.trees) + if isinstance(attempt_tracer, FullTracer) + else None + ) + try: + if deferral is None: + result = self._calculate_traced( + variable_name, period, map_to, decode_enums + ) + 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: + # This root is the value actually returned to the + # caller, ahead of cut results from abandoned retries. + if ( + attempt_index is not None + and len(attempt_tracer.trees) > attempt_index + ): + attempt_tracer.trees[attempt_index]._spiral_provisional = False + 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, + period: Period, + map_to: str = None, + decode_enums: bool = False, + ) -> ArrayLike: self.tracer.record_calculation_start(variable_name, period, self.branch_name) + state = _spiral_deferrals.get() + cuts = 0 if state is None else state.cuts # No per-variable RNG seeding: formulas may not use randomness at all # (enforced statically at variable registration by @@ -863,6 +1075,12 @@ def calculate( result = self._calculate(variable_name, period) if isinstance(result, EnumArray) and decode_enums: result = result.decode_to_str() + if isinstance(self.tracer, FullTracer) and state is not None: + node = self.tracer._current_node + if node is not None: + node._spiral_provisional = state.cuts != cuts or bool( + state.cut_tracers + ) self.tracer.record_calculation_result(result) if map_to is not None: source_entity = self.tax_benefit_system.get_variable( @@ -872,6 +1090,7 @@ def calculate( return result finally: self.tracer.record_calculation_end() + self._leave_frame() self.purge_cache_of_invalid_values() def map_result( @@ -1006,6 +1225,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)) @@ -1030,6 +1256,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 ) @@ -1079,6 +1309,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: @@ -1240,6 +1476,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 @@ -1248,16 +1487,49 @@ 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. + 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 + branch_name = getattr(self, "branch_name", "default") + spiral_invalidations = getattr(self, "_spiral_invalidated_caches", ()) for _name, _period in invalidated_caches: holder = self.get_holder(_name) - holder.delete_arrays(_period) - if _fast_cache is not None: - _fast_cache.pop((_name, _period), None) + if (_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) + if _fast_cache is not None: + _fast_cache.pop((_name, _period), None) + else: + # Explicit invalidation retains its public contract: inputs + # and contained periods in default storage are deleted too. + holder.delete_arrays(_period) + if _fast_cache is not None: + invalid_period = periods.period(_period) + eternal = ( + getattr( + getattr(holder, "variable", None), "definition_period", None + ) + == periods.ETERNITY + ) + for key in list(_fast_cache): + if key[0] == _name and ( + eternal + or key[1] == _period + or ( + isinstance(key[1], Period) + and invalid_period.start <= key[1].start + and _end_order(key[1]) <= _end_order(invalid_period) + ) + ): + _fast_cache.pop(key, None) self.invalidated_caches = set() + self._spiral_invalidated_caches = set() def calculate_add( self, @@ -1544,32 +1816,256 @@ 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 self._defer(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. + + 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: + 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 deferrals is None: + # Only an outermost ``calculate`` can finish the deeper period + # first; called any other way, the recursion is cut as before. + 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) + 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 :] + 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 + # 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 + + @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 _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: + """Mark an explicit invalidation, including inputs and contained periods.""" invalidated_caches = getattr(self, "invalidated_caches", None) if invalidated_caches is None: self.invalidated_caches = {(variable, period)} - return - invalidated_caches.add((variable, period)) + else: + invalidated_caches.add((variable, period)) + spiral_invalidations = getattr(self, "_spiral_invalidated_caches", None) + if spiral_invalidations is not None: + spiral_invalidations.discard((variable, period)) + + def _invalidate_spiral_cache_entry(self, variable: str, period: Period) -> None: + """Mark an exact calculated value, preserving explicit invalidations.""" + key = (variable, period) + invalidated_caches = getattr(self, "invalidated_caches", None) + if invalidated_caches is None: + invalidated_caches = self.invalidated_caches = set() + if key not in invalidated_caches: + spiral_invalidations = getattr(self, "_spiral_invalidated_caches", None) + if spiral_invalidations is None: + spiral_invalidations = self._spiral_invalidated_caches = set() + spiral_invalidations.add(key) + invalidated_caches.add(key) 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_spiral_cache_entry(frame["name"], frame["period"]) + 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(): + 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_spiral_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 ----- # @@ -1746,7 +2242,10 @@ def clone( Copy the simulation just enough to be able to run the copy without modifying the original simulation. The copy records its own cache invalidations (``invalidated_caches``, - starting from this simulation's pending ones). If ``baseline`` is a + starting from this simulation's pending ones). It removes copied + values from cut recursions before exposing the copy, preserving + inputs; explicit invalidations remain pending until its next purge. + If ``baseline`` is a branch of this simulation, as a reform simulation's is, the copy gets a copy of that branch as its own: a branch of the copy, under the same name, with this baseline's tax-benefit system, traced in the copy if @@ -1830,6 +2329,9 @@ def clone( # over: the clone's cached arrays start as copies of its arrays. if getattr(self, "invalidated_caches", None) is not None: new.invalidated_caches = set(self.invalidated_caches) + new._spiral_invalidated_caches = set( + getattr(self, "_spiral_invalidated_caches", ()) + ) # Only pass ``share_arrays`` when sharing, so a population or holder # ``clone`` override with the earlier signature still deep-copies. @@ -1865,6 +2367,7 @@ def clone( new.tax_benefit_system = self.tax_benefit_system new.debug = debug new.trace = trace + new._drop_invalid_values(self) # A branch shares its parent's baseline: formulas that run in a # branch read the parent's baseline values through it. That holds for @@ -1903,6 +2406,30 @@ def clone( 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 calculated 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, "_spiral_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) + self.invalidated_caches.difference_update(invalidated_caches) + self._spiral_invalidated_caches.difference_update(invalidated_caches) + 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 7cebeb82..9e6c764d 100644 --- a/policyengine_core/tracers/flat_trace.py +++ b/policyengine_core/tracers/flat_trace.py @@ -28,19 +28,19 @@ def key(self, node: tracers.TraceNode) -> str: def get_trace(self) -> dict: trace = {} + qualities = {} for node in self._full_tracer.browse_trace(): - # 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 - } - ) + # Preserve the first completed calculation over later cache + # reads. A scheduler retry can leave either no result or a + # provisional result from a cut recursion; the first accepted + # calculation replaces either. Keep provisional results when + # no accepted calculation exists, so dependencies still resolve. + quality = 0 if node.value is None else 1 if node._spiral_provisional else 2 + for key, node_trace in self._get_flat_trace(node).items(): + if key not in trace or quality > qualities[key]: + trace[key] = node_trace + qualities[key] = quality return trace diff --git a/policyengine_core/tracers/trace_node.py b/policyengine_core/tracers/trace_node.py index 50164877..8009ca80 100644 --- a/policyengine_core/tracers/trace_node.py +++ b/policyengine_core/tracers/trace_node.py @@ -24,6 +24,7 @@ class TraceNode: value: typing.Optional[Array] = None start: float = 0 end: float = 0 + _spiral_provisional: bool = False def calculation_time(self, round_: bool = True) -> Time: result = self.end - self.start diff --git a/tests/core/test_cycles.py b/tests/core/test_cycles.py index a488765b..793caeac 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 77dfcaba..61640479 100644 --- a/tests/core/test_fast_cache_guards.py +++ b/tests/core/test_fast_cache_guards.py @@ -72,7 +72,8 @@ 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_arrays=lambda period: 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 new file mode 100644 index 00000000..1e34b801 --- /dev/null +++ b/tests/core/test_spiral_order.py @@ -0,0 +1,759 @@ +"""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 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] + + +@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)]) + 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_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(): + 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] + + +# 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] + + +@pytest.mark.parametrize("spiral_first", [False, True]) +def test_explicit_invalidation_keeps_precedence_over_spiral_cleanup(spiral_first): + from policyengine_core.periods import period + + simulation = build([counting_up()], {("recursive", 2021): 40}) + key = ("recursive", period("2021")) + markers = [ + simulation._invalidate_spiral_cache_entry, + simulation.invalidate_cache_entry, + ] + if not spiral_first: + markers.reverse() + for mark in markers: + mark(*key) + simulation.purge_cache_of_invalid_values() + assert simulation.get_array(*key) is None + + +def test_spiral_cleanup_and_copy_preserve_an_input_at_the_marked_period(): + from policyengine_core.periods import period + + simulation = build([counting_up()], {("recursive", 2021): 40}) + key = ("recursive", period("2021")) + simulation._invalidate_spiral_cache_entry(*key) + copy = simulation.clone() + assert copy.invalidated_caches == set() + assert copy.get_array(*key).tolist() == [40] + simulation.purge_cache_of_invalid_values() + assert simulation.get_array(*key).tolist() == [40] + + +def test_explicit_invalidation_removes_contained_inputs_and_fast_values(): + from policyengine_core.periods import period + + simulation = build([]) + simulation.set_input("salary", "2020-06", [100]) + assert simulation.calculate("salary", "2020-06").tolist() == [100] + simulation.invalidate_cache_entry("salary", period("2020")) + simulation.purge_cache_of_invalid_values() + assert simulation.get_array("salary", "2020-06") is None + assert simulation.calculate("salary", "2020-06").tolist() == [0] diff --git a/tests/core/test_spiral_order_property.py b/tests/core/test_spiral_order_property.py new file mode 100644 index 00000000..41b5c1e8 --- /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/core/test_spiral_scheduler_property.py b/tests/core/test_spiral_scheduler_property.py new file mode 100644 index 00000000..4b5ee4da --- /dev/null +++ b/tests/core/test_spiral_scheduler_property.py @@ -0,0 +1,349 @@ +"""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, period +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), +) + + +# 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(), + 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, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@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( + 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)), + 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 00000000..e66d2e9e --- /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) diff --git a/tests/core/test_spiral_trace_retries.py b/tests/core/test_spiral_trace_retries.py new file mode 100644 index 00000000..dd85825b --- /dev/null +++ b/tests/core/test_spiral_trace_retries.py @@ -0,0 +1,94 @@ +"""Flat traces retain accepted calculations across spiral retries.""" + +import numpy as np +import pytest + +from policyengine_core.tracers import FullTracer +from tests.fixtures.spirals import build, make_recurrence + + +@pytest.mark.parametrize("storage", ["store", "drop", "blacklist"]) +def test_flat_trace_matches_returned_result_after_a_cut_retry(storage): + # Reduced from a counterexample executed by the prior review: a cut + # child recorded 0 before the accepted request returned 1. + variable = make_recurrence("r", [("r", 1), ("r", -1)], end=2020) + plain = build([variable], max_spiral_loops=1) + traced = build([variable], max_spiral_loops=1) + traced.trace = True + for simulation in (plain, traced): + if storage == "drop": + simulation.get_holder("r")._do_not_store = True + elif storage == "blacklist": + simulation.opt_out_cache = True + simulation.tax_benefit_system.cache_blacklist = {"r"} + + expected = plain.calculate("r", "2000").tolist() + result = traced.calculate("r", "2000").tolist() + assert result == expected == [1.0] + trace = traced.tracer.get_serialized_flat_trace() + assert trace["r<2000, (default)>"]["value"] == result + assert traced.tracer.stack == [] + + +def test_flat_trace_keeps_first_accepted_calculation_over_cache_reads(): + tracer = FullTracer() + tracer.record_calculation_start("r", "2000") + tracer.record_calculation_end() # An abandoned retry has no result. + + tracer.record_calculation_start("r", "2000") + tracer.record_calculation_result(np.array([0.0])) + tracer._current_node._spiral_provisional = True + tracer.record_calculation_end() + + tracer.record_calculation_start("r", "2000") + tracer.record_calculation_start("dependency", "2000") + tracer.record_calculation_result(np.array([1.0])) + tracer.record_calculation_end() + tracer.record_calculation_result(np.array([1.0])) + tracer.record_calculation_end() + + tracer.record_calculation_start("r", "2000") + tracer.record_calculation_result(np.array([1.0])) + tracer.record_calculation_end() # A cache read has no dependencies. + + entry = tracer.get_serialized_flat_trace()["r<2000, (default)>"] + assert entry["value"] == [1.0] + assert entry["dependencies"] == ["dependency<2000, (default)>"] + + +def test_stable_child_metadata_survives_an_abandoned_parent_retry(): + tracer = FullTracer() + tracer.record_calculation_start("retry", "2000") + tracer._current_node._spiral_provisional = True + tracer.record_calculation_start("stable", "2000") + tracer.record_parameter_access("rate", "2000-01-01", "default", 2.0) + tracer.record_calculation_start("dependency", "2000") + tracer.record_calculation_result(np.array([3.0])) + tracer.record_calculation_end() + tracer.record_calculation_result(np.array([6.0])) + tracer.record_calculation_end() + tracer.record_calculation_end() # The parent's deeper request unwinds. + + tracer.record_calculation_start("stable", "2000") + tracer.record_calculation_result(np.array([6.0])) + tracer.record_calculation_end() + + entry = tracer.get_serialized_flat_trace()["stable<2000, (default)>"] + assert entry["value"] == [6.0] + assert entry["dependencies"] == ["dependency<2000, (default)>"] + assert entry["parameters"] == {"rate<2000-01-01, (default)>": 2.0} + + +def test_provisional_dependencies_remain_in_the_flat_trace(): + tracer = FullTracer() + tracer.record_calculation_start("r", "2000") + tracer.record_calculation_start("cut", "1999") + tracer.record_calculation_result(np.array([0.0])) + tracer._current_node._spiral_provisional = True + tracer.record_calculation_end() + tracer.record_calculation_result(np.array([1.0])) + tracer.record_calculation_end() + + trace = tracer.get_serialized_flat_trace() + dependency = trace["r<2000, (default)>"]["dependencies"][0] + assert trace[dependency]["value"] == [0.0] diff --git a/tests/fixtures/spirals.py b/tests/fixtures/spirals.py new file mode 100644 index 00000000..b69ee2e5 --- /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 + )