diff --git a/changelog.d/fix-test-runner-memory.fixed.md b/changelog.d/fix-test-runner-memory.fixed.md new file mode 100644 index 000000000..76101c289 --- /dev/null +++ b/changelog.d/fix-test-runner-memory.fixed.md @@ -0,0 +1 @@ +Stop the YAML test runner keeping per-case state alive for the whole run: finished cases drop runner-owned simulation references (also when the simulation failed to build) and any tracer, the reform-system cache keeps the 2 most recently used (`--reform-cache-size`), and neither a traced parameter tree nor the tax-benefit system's cache of it at each instant (`get_parameters_at_instant`) keeps tracer-bound wrappers. diff --git a/changelog.d/lazy-parameter-node-at-instant.changed.md b/changelog.d/lazy-parameter-node-at-instant.changed.md new file mode 100644 index 000000000..00e96589b --- /dev/null +++ b/changelog.d/lazy-parameter-node-at-instant.changed.md @@ -0,0 +1 @@ +Resolve `ParameterNodeAtInstant` children on first read instead of building the whole parameter tree for every instant, which cost policyengine-us 16 to 33 MB and several seconds per instant on every tax-benefit system. diff --git a/changelog.d/modify-parameters-cached-instants.fixed.md b/changelog.d/modify-parameters-cached-instants.fixed.md new file mode 100644 index 000000000..5be9d168e --- /dev/null +++ b/changelog.d/modify-parameters-cached-instants.fixed.md @@ -0,0 +1 @@ +`modify_parameters` drops the values a parameter tree has cached at each instant once its modifier has run, so a modifier's in-place changes to values, brackets or children are read even where the system had already read those parameters. diff --git a/changelog.d/parameter-cache-invalidation.fixed.md b/changelog.d/parameter-cache-invalidation.fixed.md new file mode 100644 index 000000000..937dcea13 --- /dev/null +++ b/changelog.d/parameter-cache-invalidation.fixed.md @@ -0,0 +1 @@ +`Parameter.update` on a bracket of a freshly loaded scale and `ParameterNode.add_child` clear the cached parameters of every node above them, `ParameterNode.clear_instant_caches` clears them after direct edits to values, brackets or children, and a tax-benefit system clears them once its parameter operations have run. diff --git a/changelog.d/parameters-at-instant-cache-follows-updates.fixed.md b/changelog.d/parameters-at-instant-cache-follows-updates.fixed.md new file mode 100644 index 000000000..78559e977 --- /dev/null +++ b/changelog.d/parameters-at-instant-cache-follows-updates.fixed.md @@ -0,0 +1 @@ +`TaxBenefitSystem.get_parameters_at_instant` returns the parameters as they are after a `Parameter.update`, instead of the node it had cached before the change. diff --git a/changelog.d/reform-modify-parameters-own-tree.changed.md b/changelog.d/reform-modify-parameters-own-tree.changed.md new file mode 100644 index 000000000..332266b20 --- /dev/null +++ b/changelog.d/reform-modify-parameters-own-tree.changed.md @@ -0,0 +1 @@ +`Reform.modify_parameters` no longer deep-copies the baseline's parameter tree. The modifier changes the reform's own copy, as it already did for a reform applied with `apply_reform_set`, so changes made earlier in `apply` are kept and building such a reform skips a full copy of the tree. diff --git a/changelog.d/test-runner-perf.added.md b/changelog.d/test-runner-perf.added.md new file mode 100644 index 000000000..900355107 --- /dev/null +++ b/changelog.d/test-runner-perf.added.md @@ -0,0 +1 @@ +The YAML test runner sets a case's dotted parameter inputs on a reform that copies only those parameters' paths instead of the whole tree, builds reform systems on the shared reform-free system so each holds one copy of the parameter tree, adds a `--read-only-arrays` option that fails a formula writing into an array it read, and `policyengine_core.tools.forked_test_runner.run_batches_in_forks` runs test batches in children forked from one process that built the tax-benefit system. diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 86e214485..b20f02e35 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -36,3 +36,50 @@ broader behavioral change, run the relevant focused tests and formatting check. - Prefer small synthetic fixtures for regression tests. - When fixing a bug, add a regression test that fails without the fix and passes with it unless the change is documentation-only. + +## YAML suites and memory + +`policyengine-core test -c ` runs every case in one +process. Its memory is bounded by design, and tests must keep it that way: + +- A finished case releases its simulation (`YamlItem.teardown`). pytest keeps + every collected item until the session ends, so anything stored on the item + lives for the whole run. +- A case that names `reforms` or `extensions` runs on a system built for + that combination on the reform-free system: a reform copies the parameter + tree before it changes it, so each holds one copy of the tree. The runner + caches the `--reform-cache-size` most recently used of them (default 2) plus + the reform-free system, and drops one before building another when full. + Do not add a cache to the runner that grows with the number of cases or + combinations. +- A case that sets a parameter with a dotted input key runs on a system of its + own (`_ParameterInputReform`) that copies only the nodes on the paths to the + parameters it sets (`policyengine_core.parameters.path_copy`) and shares the + rest of the tree, with its cached at-instant values, with the system its + `reforms` select. Nothing may change a shared node in place. +- Each system still caches its parameter tree at every distinct instant its + cases read (`ParameterNode.get_at_instant`, + `TaxBenefitSystem.get_parameters_at_instant`), so memory grows with the + number of distinct dates a suite reads, not with the number of cases. + Neither cache may hold a tracer: a read of a traced tree is wrapped for that + read only. +- `tests/core/tools/test_runner/test_runner_memory.py` holds the regression + tests: a few hundred cases in one process must not grow traced memory or + the number of live simulations and systems, and traced cases at distinct + dates must leave no tracer or recorded result alive. + +Run one large suite at a time. A country package's whole YAML tree in one +process still holds the baseline system, the reform-free copy and the cached +reform systems at once; use the country package's own batch runner where it +has one (policyengine-us: `policyengine_us/tests/test_batched.py` and the +`make test-yaml-*` targets). + +`policyengine_core.tools.forked_test_runner.run_batches_in_forks` runs batches +in children forked from a process that has already built the system, so a +batch skips the import and build, and the children share the parent's system +in memory until they write to it. `tests/core/tools/test_runner/test_forked_test_runner.py` +covers its exit codes, worker limit and timeout. + +`--read-only-arrays` (option `read_only_arrays`) stores every array read-only +for the run, so a formula that writes into an array it read from another +variable fails at that line (`test_read_only_arrays.py`). diff --git a/docs/python_api/tools.md b/docs/python_api/tools.md index 4c8473ab3..47fe40021 100644 --- a/docs/python_api/tools.md +++ b/docs/python_api/tools.md @@ -8,6 +8,9 @@ The `policyengine_core.tools` module contains miscellaneous utility functions, i .. autofunction:: policyengine_core.tools.test_runner.run_tests ``` +The same runner backs `policyengine-core test -c `, +where `--reform-cache-size N` sets `reform_cache_size`. + ## dump_simulation ```{eval-rst} diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 423d19b9a..3aa8ee015 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -85,6 +85,14 @@ class InMemoryStorage: # holder per branch for inputs most branches never change. _inputs: Union[Set[str], FrozenSet[str]] = _NO_INPUTS is_eternal: bool + # Whether every storage stores and hands out read-only views of numpy + # arrays, so that code writing into an array it read (``x += y`` on the + # result of ``tax_unit("x", period)``) raises at that line instead of + # changing ``x`` for every later reader. Off by default; the YAML test + # runner turns it on for a run with the ``read_only_arrays`` option. The + # caller's own array stays writeable: only the stored view is protected. + # Masked arrays are left as they are, since a view shares their mask. + read_only_arrays: bool = False def __init__(self, is_eternal: bool): self._arrays = {} @@ -182,6 +190,8 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: # First read of an array shared by ``clone``: from here on the # caller may write into it, so it needs to be this storage's own. values = values.copy() + if self.read_only_arrays: + values = _read_only_view(values) self._arrays[key] = values self._stop_sharing(key) return values @@ -304,6 +314,8 @@ def put( "string form is lossy (see policyengine-core#526)." ) key = f"{branch_name}:{period}" + if self.read_only_arrays and _can_share(value): + value = _read_only_view(value) self._arrays[key] = value self._stop_sharing(key) if derived: diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 0df0e387b..997a809a7 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -182,6 +182,8 @@ def add_child(self, name: str, child: Union["ParameterNode", Parameter]): self.children[name] = child setattr(self, name, child) child.parent = self + # Nodes at instant cached here and above were built without it. + self.clear_parent_cache() def __repr__(self) -> str: result = os.linesep.join( @@ -214,31 +216,60 @@ def clone(self) -> "ParameterNode": return clone def _get_at_instant(self, instant: Instant) -> ParameterNodeAtInstant: - if instant in self._at_instant_cache: - return self._at_instant_cache[instant] - node_at_instant = ParameterNodeAtInstant(self.name, self, instant) + node_at_instant = self._at_instant_cache.get(instant) + if node_at_instant is None: + node_at_instant = ParameterNodeAtInstant(self.name, self, instant) + self._at_instant_cache[instant] = node_at_instant if self.trace: - at_instant = TracingParameterNodeAtInstant( + # Wrapped on each read and never cached: a cached wrapper kept the + # tracer of whichever simulation first read this instant alive for + # the life of the parameters, and recorded later reads into it, + # while a node cached before tracing was turned on went untraced. + return TracingParameterNodeAtInstant( node_at_instant, self.tracer, self.branch_name ) - else: - at_instant = node_at_instant - self._at_instant_cache[instant] = at_instant - return at_instant + return node_at_instant def attach_to_parent(self, parent: "ParameterNode"): self.parent = parent def clear_parent_cache(self): - self._at_instant_cache.clear() - if self.parent is not None: - self.parent.clear_parent_cache() + """Clear the nodes at instant cached by this node and every node above + it (scales cache nothing and are passed through). + + Walks up without recursing and stops at a node it has already + cleared: while a node is loaded, a child named ``parent`` (as + policyengine-us's ``gov.hhs.medicaid.eligibility.categories.parent``) + replaces its ``parent`` attribute and points back at it until the node + is attached to its own parent. + """ + node, cleared = self, set() + while node is not None and id(node) not in cleared: + cleared.add(id(node)) + cache = getattr(node, "_at_instant_cache", None) + if cache is not None: + cache.clear() + node = getattr(node, "parent", None) def mark_as_modified(self): self.modified = True if self.parent is not None: self.parent.mark_as_modified() + def clear_instant_caches(self) -> None: + """Drop the nodes at instant cached by this node, every node below it + and every node above it. + + ``Parameter.update`` and ``add_child`` clear what they change. Call + this after changing ``values_list``, a scale's ``brackets`` or + ``children`` directly, so that nodes at instant requested afterwards + read the change. + """ + self.clear_parent_cache() + for descendant in self.get_descendants(): + if isinstance(descendant, ParameterNode): + descendant._at_instant_cache.clear() + def get_child(self, path: str) -> "ParameterNode": node = self for name in path.split("."): diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index b55889eb6..5699a92fe 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -1,10 +1,11 @@ import os import sys +import threading from typing import TYPE_CHECKING, Iterable, Union import numpy -from policyengine_core import parameters, tools +from policyengine_core import parameters, periods, tools from policyengine_core.errors import ParameterNotFoundError from policyengine_core.parameters import helpers @@ -16,9 +17,90 @@ ) +# The one instance attribute a node keeps for itself; a parameter of this +# name is rejected. +_STATE = "_ParameterNodeAtInstant__state" +# Attributes a node sets on itself, which a child of the same name shadows. +_OWN_NAMES = frozenset({"_name", "_instant_str", "_children", "_vectorial_node"}) +_MISSING = object() + +# Held while a child is resolved for the first time, so that concurrent +# readers never see a half-resolved node or resolve one child twice. Reads of +# a resolved child never take it. Reentrant because resolving a child can +# build that child's own node. +_RESOLUTION_LOCK = threading.RLock() + + +def _reset_resolution_lock() -> None: + """Give a forked child a lock no thread holds. A thread of the parent that + held it at the fork does not exist in the child, so nothing would release + it, and the child's first resolution would wait for ever.""" + global _RESOLUTION_LOCK + _RESOLUTION_LOCK = threading.RLock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_resolution_lock) + + +class _LazyChildren: + """What a node at instant knows about its children. A shallow copy of the + node shares it, as copies of a node that built every child up front shared + their ``_children`` mapping.""" + + __slots__ = ("instant", "source", "resolved", "absent", "added", "children") + + def __init__(self, instant: str, source: dict): + self.instant = instant + # Each child's parameter node as at construction, in order; None once + # every child is resolved. + self.source = source + # Children resolved from ``source`` so far. + self.resolved = {} + # Children of ``source`` with no value at the instant. + self.absent = set() + # Children given with ``add_child``, in call order. + self.added = {} + # Every child in order, once all of ``source`` is resolved. + self.children = None + + # Pickle protocols 0 and 1 refuse a class with ``__slots__`` that does not + # define these, and a parameter tree pickles the nodes it has cached. + def __getstate__(self) -> dict: + # A pickle or a deep copy is a snapshot, as when every child was built + # up front. Resolving what is left drops ``source``, whose parameter + # nodes reach the whole tree through their parents: the copy would + # otherwise carry the tree with it. + if self.source is not None: + _resolve_remaining(self) + return {name: getattr(self, name) for name in self.__slots__} + + def __setstate__(self, state: dict) -> None: + for name, value in state.items(): + setattr(self, name, value) + + class ParameterNodeAtInstant: """ Parameter node of the legislation, at a given instant. + + Children resolve on first read and are kept from then on. Building every + descendant up front made each instant a model is asked about cost a copy of + the whole parameter tree (16 to 33 MB and several seconds per instant for + policyengine-us), which every tax-benefit system then kept for its + lifetime. Which children a node has, and in what order, is fixed when it is + built. A child's value is taken when the child is first read, so a node + held across an in-place change to its parameters reads the changed value + for a child it has not read yet; once read (including found to have no + value), a child keeps what that read found. + + A node requested after a change reads the change: ``Parameter.update`` and + ``ParameterNode.add_child`` clear the cached nodes at instant of every + node above the one they change. After changing ``values_list``, a scale's + ``brackets`` or ``children`` directly, call + ``ParameterNode.clear_instant_caches``. A pickle or deep copy of a node + resolves every child first, so it is a snapshot that does not carry the + parameter tree. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -28,26 +110,85 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): :param instant_str: A date in the format `YYYY-MM-DD`. """ + # Built up front, a node at an instant that is not a date (such as + # ``parameters(None)``, which reads "None") failed on its first scale + # and was never cached. Built lazily, it would be cached and fail, or + # read the latest values, only when a child is read or copied. + periods.instant(instant_str) + # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - self._children = {} - - for child_name, child in node.children.items(): - child_at_instant = child._get_at_instant(instant_str) - if child_at_instant is not None: - self.add_child(child_name, child_at_instant) + source = dict(node.children) + if _STATE in source: + raise ValueError(f"{_STATE!r} cannot name a parameter.") + state = _LazyChildren(instant_str, source) + self.__dict__[_STATE] = state + cls = type(self) + if cls.add_child is not ParameterNodeAtInstant.add_child: + # A subclass's add_child sees every child, as before. + for child_name, child in source.items(): + child_at_instant = child._get_at_instant(instant_str) + if child_at_instant is not None: + self.add_child(child_name, child_at_instant) + with _RESOLUTION_LOCK: + state.source = None + state.children = dict(state.added) + return + for child_name in source: + # A child whose name the node itself answers to shadows it, as + # when every child was set here. + if child_name in _OWN_NAMES or hasattr(cls, child_name): + child_at_instant = _resolve(state, child_name) + if child_at_instant is not _MISSING: + setattr(self, child_name, child_at_instant) def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): - self._children[child_name] = child_at_instant + if child_name == _STATE: + raise ValueError(f"{_STATE!r} cannot name a parameter.") + state = self.__dict__[_STATE] + with _RESOLUTION_LOCK: + state.added[child_name] = child_at_instant + if state.children is not None: + state.children[child_name] = child_at_instant setattr(self, child_name, child_at_instant) def __getattr__(self, key: str): - if key == "_name": + # Reached only for a name that is not an instance attribute: a child + # not read yet by this instance, ``_children`` before it is set, or a + # name this node does not have. + state = self.__dict__.get(_STATE) + if state is None or key == "_name": + # An instance still being rebuilt by copy or pickle. raise AttributeError(key) + if key == "_children": + return _materialise(self, state) + child_at_instant = _resolve(state, key) + if child_at_instant is not _MISSING: + # Attributes hold what the parameters give; ``add_child`` on + # another copy changes only that copy's attribute, as before. + self.__dict__.setdefault(key, child_at_instant) + return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) + def __dir__(self): + children = self._children + names = set(super().__dir__()) | set(map(str, children)) + names.discard(_STATE) + return list(names) + + def __setstate__(self, state: dict): + self.__dict__.update(state) + if _STATE not in state: + # Pickled when every child was built up front, all in ``_children``. + children = state.get("_children") + if children is None: + children = self.__dict__["_children"] = {} + lazy = _LazyChildren(state.get("_instant_str"), None) + lazy.resolved = lazy.children = children + self.__dict__[_STATE] = lazy + def __getitem__( self, key: str ) -> Union["ParameterNodeAtInstant", VectorialParameterNodeAtInstant]: @@ -68,7 +209,15 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - return self._children[key] + state = self.__dict__[_STATE] + child_at_instant = state.added.get(key, _MISSING) + if child_at_instant is _MISSING: + child_at_instant = state.resolved.get(key, _MISSING) + if child_at_instant is _MISSING: + child_at_instant = _resolve(state, key) + if child_at_instant is _MISSING: + raise KeyError(key) + return child_at_instant def __iter__(self) -> Iterable: return iter(self._children) @@ -81,3 +230,77 @@ def __repr__(self) -> str: ] ) return result + + +def _resolve(state: _LazyChildren, child_name): + """A child as the parameters give it, resolving it if no read has yet; + ``_MISSING`` if the node has no such child or it has no value at the + instant.""" + with _RESOLUTION_LOCK: + child_at_instant = state.resolved.get(child_name, _MISSING) + if child_at_instant is not _MISSING or child_name in state.absent: + return child_at_instant + source = state.source + child = None if source is None else source.get(child_name) + if child is None: + return _MISSING + child_at_instant = child._get_at_instant(state.instant) + # A read made while resolving (the same thread re-entering) may have + # recorded this child already; the first recorded result stands. + if child_name in state.absent: + return _MISSING + recorded = state.resolved.get(child_name, _MISSING) + if recorded is not _MISSING: + child_at_instant = recorded + elif child_at_instant is None: + state.absent.add(child_name) + else: + state.resolved[child_name] = child_at_instant + if state.source is not None and len(state.resolved) + len(state.absent) == len( + state.source + ): + _complete(state) + return _MISSING if child_at_instant is None else child_at_instant + + +def _complete(state: _LazyChildren): + """Order every child once all of ``source`` is resolved; let go of the + parameter nodes. The caller holds ``_RESOLUTION_LOCK``.""" + resolved, added = state.resolved, state.added + children = { + name: added.get(name, resolved[name]) + for name in state.source + if name in resolved + } + for name, child_at_instant in added.items(): + # Added children the parameters do not have (or have no value for) + # follow, in the order they were added. + children.setdefault(name, child_at_instant) + if not added: + state.resolved = children + state.children = children + state.source = None + state.absent = set() + + +def _resolve_remaining(state: _LazyChildren): + """Resolve every child not read yet, and complete the state.""" + with _RESOLUTION_LOCK: + if state.children is None: + for child_name in list(state.source): + if state.children is not None: + break + _resolve(state, child_name) + if state.children is None: + _complete(state) + + +def _materialise(node_at_instant: ParameterNodeAtInstant, state: _LazyChildren): + """Resolve every child and set ``_children``.""" + with _RESOLUTION_LOCK: + _resolve_remaining(state) + children = state.children + if "_children" not in node_at_instant.__dict__: + # Unless a child named ``_children`` took the name. + node_at_instant.__dict__["_children"] = children + return children diff --git a/policyengine_core/parameters/parameter_scale.py b/policyengine_core/parameters/parameter_scale.py index 9b92a75b8..50d39d0d8 100644 --- a/policyengine_core/parameters/parameter_scale.py +++ b/policyengine_core/parameters/parameter_scale.py @@ -53,6 +53,9 @@ def __init__(self, name: str, data: dict, file_path: str): bracket = parameters.ParameterScaleBracket( name=bracket_name, data=bracket_data, file_path=file_path ) + # As ``clone`` does, so that ``Parameter.update`` on a bracket's + # rate or threshold clears the caches above the scale. + bracket.parent = self brackets.append(bracket) self.brackets: typing.List[parameters.ParameterScaleBracket] = brackets self.propagate_uprating() diff --git a/policyengine_core/parameters/path_copy.py b/policyengine_core/parameters/path_copy.py new file mode 100644 index 000000000..3e1987264 --- /dev/null +++ b/policyengine_core/parameters/path_copy.py @@ -0,0 +1,92 @@ +"""Copy only the parts of a parameter tree that a change will touch.""" + +from typing import Iterable, Set, Union + +from policyengine_core import commons + +from .parameter import Parameter +from .parameter_node import ParameterNode +from .parameter_scale import ParameterScale + +TreeNode = Union[ParameterNode, ParameterScale, Parameter] + + +def copy_paths(root: ParameterNode, paths: Iterable[str]) -> ParameterNode: + """Return a tree that shares every node of ``root`` except those on ``paths``. + + Each path names a parameter as ``set_parameter`` and + ``ParameterNode.get_child`` do: dotted child names, with ``name[i]`` for + bracket ``i`` of a scale. The returned root, every node between it and each + named parameter, and each named parameter are new; every other node, with + the at-instant values it has cached, is ``root``'s own. + + A full ``clone`` rebuilds all of a large tree (policyengine-us has over + 130,000 nodes) and drops every cached at-instant value, which costs + seconds and a copy of the tree. Copying the paths instead costs a few + nodes per path. + + Change the returned tree only at the named parameters, for example with + ``Parameter.update``: an update clears the at-instant caches of the new + nodes above it and marks them modified, and stops there. A node shared + with ``root`` still names its parent in ``root``'s tree, so changing it + would change ``root``'s tree too. + + Raises the ``ValueError`` of ``ParameterNode.get_child`` if a path names + no parameter, before copying anything. + """ + paths = list(paths) + for path in paths: + root.get_child(path) + + copied_root = _copy_node(root) + copies = {id(copied_root)} + for path in paths: + node = copied_root + for name in path.split("."): + if "[" in name: + scale_name, index = name.split("[") + index = int(index[:-1]) + scale = _copied_child(node, scale_name, copies) + bracket = scale.brackets[index] + if id(bracket) not in copies: + bracket = _copy_node(bracket) + copies.add(id(bracket)) + scale.brackets[index] = bracket + bracket.parent = scale + node = bracket + else: + node = _copied_child(node, name, copies) + return copied_root + + +def _copied_child(node: ParameterNode, name: str, copies: Set[int]) -> TreeNode: + """``node``'s child ``name``, replaced by a copy unless it already is one.""" + child = node.children[name] + if id(child) in copies: + return child + child = _copy_node(child) + copies.add(id(child)) + node.children[name] = child + setattr(node, name, child) + child.parent = node + return child + + +def _copy_node(node: TreeNode) -> TreeNode: + """A copy of ``node`` that shares its children and their caches. + + A parameter is cloned outright, as it has no children to share. A node or + scale gets its own children mapping or bracket list, holding the same + children, and a node an empty at-instant cache, since a value below it is + about to change. + """ + if isinstance(node, Parameter): + return node.clone() + copied = commons.empty_clone(node) + copied.__dict__ = node.__dict__.copy() + if isinstance(node, ParameterScale): + copied.brackets = list(node.brackets) + else: + copied.children = dict(node.children) + copied._at_instant_cache = {} + return copied diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index 9358ee5c1..2465013cc 100644 --- a/policyengine_core/reforms/reform.py +++ b/policyengine_core/reforms/reform.py @@ -1,7 +1,6 @@ from __future__ import annotations -import copy -from typing import Callable, Union, TYPE_CHECKING +from typing import Union, TYPE_CHECKING from policyengine_core.parameters import ParameterNode, Parameter from policyengine_core.taxbenefitsystems import TaxBenefitSystem @@ -152,12 +151,11 @@ def __init__(self, baseline: TaxBenefitSystem): super().__init__(baseline.entities) self.baseline = baseline # Clone the parameter tree so in-place mutations inside ``apply()`` - # (``set_parameter(...)`` or ``self.parameters.path.update(...)``) - # don't leak into the baseline. ``modify_parameters()`` was already - # deep-copying, but plenty of reforms mutate ``self.parameters`` - # directly (bug C4). + # (``set_parameter(...)``, ``self.parameters.path.update(...)`` or a + # ``modify_parameters()`` modifier) don't leak into the baseline + # (bug C4). if baseline.parameters is not None: - self.parameters = baseline.parameters.clone() + self.parameters = self._copy_parameters(baseline.parameters) else: self.parameters = None # Use a fresh at-instant cache so reform mutations aren't shadowed @@ -172,6 +170,16 @@ def __init__(self, baseline: TaxBenefitSystem): ) self.apply() + def _copy_parameters(self, baseline_parameters: ParameterNode) -> ParameterNode: + """The tree ``apply`` starts from and may change in place, directly or + through ``modify_parameters``. + + A full clone of the baseline's tree, since a reform can change any + node. A reform that changes only parameters it can name in advance can + copy less (see ``policyengine_core.parameters.path_copy``). + """ + return baseline_parameters.clone() + def __getattr__(self, attribute): return getattr(self.baseline, attribute) @@ -186,29 +194,6 @@ def full_key(self) -> str: key = ".".join([baseline_full_key, key]) return key - def modify_parameters( - self, modifier_function: Callable[[ParameterNode], ParameterNode] - ) -> None: - """Make modifications on the parameters of the legislation. - - Call this function in `apply()` if the reform asks for legislation parameter modifications. - - Args: - modifier_function: A function that takes a :obj:`.ParameterNode` and should return an object of the same type. - """ - baseline_parameters = self.baseline.parameters - baseline_parameters_copy = copy.deepcopy(baseline_parameters) - reform_parameters = modifier_function(baseline_parameters_copy) - if not isinstance(reform_parameters, ParameterNode): - return ValueError( - "modifier_function {} in module {} must return a ParameterNode".format( - modifier_function.__name__, - modifier_function.__module__, - ) - ) - self.parameters = reform_parameters - self._parameters_at_instant_cache = {} - @staticmethod def from_dict( parameter_values: dict, diff --git a/policyengine_core/scripts/policyengine_command.py b/policyengine_core/scripts/policyengine_command.py index 7aa85c888..94835d323 100644 --- a/policyengine_core/scripts/policyengine_command.py +++ b/policyengine_core/scripts/policyengine_command.py @@ -9,6 +9,13 @@ """ +def non_negative_int(value: str) -> int: + number = int(value) + if number < 0: + raise argparse.ArgumentTypeError(f"must be 0 or more, got {number}") + return number + + def get_parser(): parser = argparse.ArgumentParser() @@ -91,6 +98,18 @@ def build_test_parser(parser): default=None, help="variables to ignore. If specified, do not test the given variables.", ) + parser.add_argument( + "--reform-cache-size", + type=non_negative_int, + default=None, + help="how many reform systems (built for tests with reforms or extensions) to keep cached; each is a full copy of the tax-benefit system. Defaults to 2; 0 caches none.", + ) + parser.add_argument( + "--read-only-arrays", + action="store_true", + default=False, + help="store every calculated and input array read-only, so a formula that writes into an array it read from another variable fails at that line instead of changing that variable for later readers", + ) return parser diff --git a/policyengine_core/scripts/run_test.py b/policyengine_core/scripts/run_test.py index 5eb4c7c7e..2fcb40c74 100644 --- a/policyengine_core/scripts/run_test.py +++ b/policyengine_core/scripts/run_test.py @@ -30,6 +30,8 @@ def main(parser): "only_variables": args.only_variables, "ignore_variables": args.ignore_variables, "visualize": args.visualize, + "reform_cache_size": args.reform_cache_size, + "read_only_arrays": args.read_only_arrays, } paths = [os.path.abspath(path) for path in args.path] diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 212de058d..4ee9efd6e 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -50,12 +50,19 @@ ) from policyengine_core.periods import Instant, Period from policyengine_core.populations import GroupPopulation, Population +from policyengine_core.tracers import TracingParameterNodeAtInstant from policyengine_core.variables import Variable from policyengine_core.variables.variable import NUMERIC_VALUE_TYPES log = logging.getLogger(__name__) +def _clear_instant_caches(parameters) -> None: + """Drop the nodes at instant every node of a parameter tree has cached.""" + if isinstance(parameters, ParameterNode): + parameters.clear_instant_caches() + + class TaxBenefitSystem: """ Represents the legislation. @@ -154,6 +161,10 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: self.parameters = uprate_parameters(self.parameters) self.parameters = propagate_parameter_metadata(self.parameters) self.add_abolition_parameters() + # The operations above change values in place, after the reforms + # have run and possibly read the parameters. + self.parameters.clear_instant_caches() + self._parameters_at_instant_cache = {} self.add_modelled_policy_metadata() @@ -609,10 +620,29 @@ def get_parameters_at_instant(self, instant: Instant) -> ParameterNodeAtInstant: ) ) + parameters = self.parameters parameters_at_instant = self._parameters_at_instant_cache.get(instant) - if parameters_at_instant is None and self.parameters is not None: - parameters_at_instant = self.parameters.get_at_instant(str(instant)) + if parameters_at_instant is not None and not _tree_still_caches( + parameters, str(instant), parameters_at_instant + ): + # A change to the tree (``Parameter.update``, ``add_child``) + # clears the tree's own caches up to its root, not this one. + parameters_at_instant = None + if parameters_at_instant is None and parameters is not None: + parameters_at_instant = parameters.get_at_instant(str(instant)) + if isinstance(parameters_at_instant, TracingParameterNodeAtInstant): + parameters_at_instant = parameters_at_instant.parameter_node_at_instant self._parameters_at_instant_cache[instant] = parameters_at_instant + if parameters_at_instant is not None and getattr(parameters, "trace", False): + # Wrapped on each read and never cached, as the tree's own reads + # are (``ParameterNode._get_at_instant``): a cached wrapper kept + # the tracer of whichever simulation first read this instant + # alive for the life of the system, and recorded later reads into + # it, while a node cached before tracing was turned on went + # untraced. + return TracingParameterNodeAtInstant( + parameters_at_instant, parameters.tracer, parameters.branch_name + ) return parameters_at_instant def get_package_metadata(self) -> dict: @@ -751,6 +781,11 @@ def modify_parameters( Call this function in `apply()` if the reform asks for legislation parameter modifications. + The modifier is given this system's own parameter tree, with any + change ``apply()`` has made so far, and may change it in place. A + reform built with ``Reform(baseline)`` has its own copy of the + baseline's tree, so the baseline is left unchanged. + Args: modifier_function: A function that takes a :obj:`.ParameterNode` and should return an object of the same type. """ @@ -764,6 +799,11 @@ def modify_parameters( ) else: reform_parameters = modifier_function(self.parameters) + # A modifier may change values, brackets or children in place + # rather than through ``update``, which is what drops the values a + # node has cached at each instant. + self._parameters_at_instant_cache = {} + _clear_instant_caches(self.parameters) if not isinstance(reform_parameters, ParameterNode): return ValueError( "modifier_function {} in module {} must return a ParameterNode".format( @@ -771,6 +811,8 @@ def modify_parameters( modifier_function.__module__, ) ) + if reform_parameters is not self.parameters: + _clear_instant_caches(reform_parameters) self.parameters = reform_parameters self._parameters_at_instant_cache = {} return self @@ -784,3 +826,11 @@ def add_modelled_policy_metadata(self): with open(self.modelled_policies, "r") as f: modelled_policies = yaml.load(f, Loader=yaml.FullLoader) self.modelled_policies = modelled_policies + + +def _tree_still_caches(parameters, instant_str: str, parameters_at_instant) -> bool: + """Whether the root of ``parameters`` still caches this node at instant; + ``Parameter.update`` clears the caches from the parameter it changes up to + the root.""" + tree_cache = getattr(parameters, "_at_instant_cache", None) + return tree_cache is None or tree_cache.get(instant_str) is parameters_at_instant diff --git a/policyengine_core/tools/forked_test_runner.py b/policyengine_core/tools/forked_test_runner.py new file mode 100644 index 000000000..7e97de2d8 --- /dev/null +++ b/policyengine_core/tools/forked_test_runner.py @@ -0,0 +1,228 @@ +"""Run batches of YAML tests in processes forked from one that built the system. + +Starting a process for each batch of a large country suite repeats the same +work every time: policyengine-us takes about 40 seconds to import and build +its tax-benefit system before the first case runs, and over half a minute to +tear its interpreter down afterwards. Forking from a process that has already +built the system skips both, and the children share the parent's copy of the +system in memory until they write to it, so several batches can run at once +in the memory one used to take. +""" + +import gc +import os +import signal +import sys +import tempfile +import time +import traceback +from collections import deque +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Sequence + +from policyengine_core.tools.test_runner import _get_tax_benefit_system, run_tests + +# pytest's exit code for an internal error, used when a child fails outside +# pytest (for example, opening its log). +PYTEST_INTERNAL_ERROR = 3 +# How long a child gets to exit after SIGTERM before it is killed. +TERMINATE_GRACE_SECONDS = 5 +# How often the parent checks on its children. +POLL_SECONDS = 0.1 + + +@dataclass +class BatchResult: + """How one batch ended. + + ``exit_code`` is pytest's exit code, negative for a child ended by a + signal (as ``subprocess`` reports it). ``peak_rss_mb`` is the child's + peak resident set, which counts the pages it still shares with the parent. + """ + + paths: List[str] + exit_code: int + elapsed: float + peak_rss_mb: Optional[float] + output: str + timed_out: bool = False + + +def run_batches_in_forks( + tax_benefit_system, + batches: Sequence[Sequence[str]], + options: Optional[dict] = None, + workers: int = 1, + timeout: Optional[float] = None, + on_finish: Optional[Callable[[int, BatchResult], None]] = None, + on_heartbeat: Optional[Callable[[Dict[int, float]], None]] = None, + heartbeat_seconds: float = 60, +) -> List[BatchResult]: + """Run each batch of test paths in its own child process, ``workers`` at once. + + Each child is forked from this process, runs ``run_tests`` on its batch + with ``tax_benefit_system`` and ``options``, and exits as soon as pytest + returns. Its output goes to a file and is returned whole in its result, + so concurrent batches never interleave. Batches start in the order given. + + ``timeout`` bounds each batch in seconds: a batch still running then is + sent SIGTERM, then SIGKILL, and reported with ``timed_out``. + ``on_finish(index, result)`` is called as each batch ends, and + ``on_heartbeat({index: seconds running})`` every ``heartbeat_seconds`` + while any runs. + + Call this from a process with no other threads: a thread holding a lock + when the process forks leaves that lock held forever in the child. That + includes a numerical library's thread pool, so set + ``OPENBLAS_NUM_THREADS=1`` (and ``OMP_NUM_THREADS``, ``MKL_NUM_THREADS``) + before the process imports numpy. + + Before forking, this builds the reform-free system the runner gives every + case without reforms (see ``_get_tax_benefit_system``), so the children + share it rather than each building their own, and moves every object the + parent holds to the garbage collector's permanent generation + (``gc.freeze``), so a child's collections do not write to, and copy, the + pages it shares with the parent. + """ + if not hasattr(os, "fork"): + raise NotImplementedError("Running batches in forks needs os.fork.") + if workers < 1: + raise ValueError(f"workers must be 1 or more, got {workers}") + if options is None: + options = {} + + _get_tax_benefit_system(tax_benefit_system, [], []) + + pending = deque(enumerate(batches)) + running: Dict[int, tuple] = {} + results: List[Optional[BatchResult]] = [None] * len(batches) + last_heartbeat = time.monotonic() + log_directory = tempfile.mkdtemp(prefix="policyengine-batches-") + gc.collect() + gc.freeze() + try: + while pending or running: + while pending and len(running) < workers: + index, paths = pending.popleft() + log_path = os.path.join(log_directory, f"batch-{index}.log") + sys.stdout.flush() + sys.stderr.flush() + pid = os.fork() + if pid == 0: + _run_child(tax_benefit_system, list(paths), options, log_path) + running[pid] = (index, list(paths), log_path, time.monotonic(), None) + + any_finished = False + for pid, (index, paths, log_path, started, terminated_at) in list( + running.items() + ): + waited_pid, status, usage = os.wait4(pid, os.WNOHANG) + now = time.monotonic() + if waited_pid == 0: + if terminated_at is None: + if timeout is not None and now - started >= timeout: + os.kill(pid, signal.SIGTERM) + running[pid] = (index, paths, log_path, started, now) + elif now - terminated_at >= TERMINATE_GRACE_SECONDS: + os.kill(pid, signal.SIGKILL) + continue + any_finished = True + del running[pid] + result = BatchResult( + paths=paths, + exit_code=os.waitstatus_to_exitcode(status), + elapsed=now - started, + peak_rss_mb=_peak_rss_mb(usage), + output=_read_and_remove(log_path), + timed_out=terminated_at is not None, + ) + results[index] = result + if on_finish is not None: + on_finish(index, result) + + now = time.monotonic() + if ( + on_heartbeat is not None + and running + and (now - last_heartbeat >= heartbeat_seconds) + ): + on_heartbeat( + { + index: now - started + for index, _, _, started, _ in running.values() + } + ) + last_heartbeat = now + if not any_finished: + time.sleep(POLL_SECONDS) + finally: + # Reached with children still running only if the parent was + # interrupted: do not leave them behind. + for pid in running: + try: + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + except (ChildProcessError, ProcessLookupError): + pass + gc.unfreeze() + for name in os.listdir(log_directory): + os.remove(os.path.join(log_directory, name)) + os.rmdir(log_directory) + return results + + +def _run_child(tax_benefit_system, paths, options, log_path) -> None: + """Run one batch in a forked child, then exit without returning.""" + exit_code = PYTEST_INTERNAL_ERROR + try: + log = os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644) + # The standard descriptors, and whatever ``sys.stdout`` and + # ``sys.stderr`` write to if something (pytest's capture) replaced them. + for descriptor in {1, 2, *_stream_descriptors()}: + os.dup2(log, descriptor) + os.close(log) + exit_code = int(run_tests(tax_benefit_system, paths, options)) + except BaseException: + traceback.print_exc() + finally: + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + # Skip interpreter teardown: freeing a large country's systems + # one object at a time took over half a minute, and nothing the + # child holds needs releasing before the process goes. + os._exit(exit_code) + + +def _stream_descriptors() -> List[int]: + descriptors = [] + for stream in (sys.stdout, sys.stderr): + try: + descriptors.append(stream.fileno()) + except (AttributeError, OSError, ValueError): + pass + return descriptors + + +def _peak_rss_mb(usage) -> Optional[float]: + """A finished child's peak resident set in MB, from ``os.wait4``.""" + if usage is None or not usage.ru_maxrss: + return None + # Linux reports kilobytes and macOS bytes. + if sys.platform == "darwin": + return usage.ru_maxrss / 1e6 + return usage.ru_maxrss / 1e3 + + +def _read_and_remove(log_path: str) -> str: + try: + with open(log_path, encoding="utf-8", errors="replace") as log: + return log.read() + except FileNotFoundError: + return "" + finally: + try: + os.remove(log_path) + except FileNotFoundError: + pass diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index 6b8d4aefc..9470d6078 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -16,12 +16,14 @@ eval_expression, ) from policyengine_core.simulations import SimulationBuilder +from policyengine_core.data_storage import InMemoryStorage from policyengine_core.errors import ( SituationParsingError, VariableNotFoundError, ) from policyengine_core.scripts import build_tax_benefit_system from policyengine_core.reforms import Reform, set_parameter +from policyengine_core.parameters.path_copy import copy_paths from policyengine_core.populations import ADD, DIVIDE log = logging.getLogger(__name__) @@ -67,9 +69,20 @@ def import_yaml(): # reuses after GC — a collected baseline could produce the same id as a # completely unrelated new baseline and hit a stale cache entry (bug H9). import weakref +from collections import OrderedDict _tax_benefit_system_cache: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +# How many reform systems (cases with ``reforms`` or ``extensions``) each +# baseline keeps, least recently used evicted first. Every one is a full clone +# of the baseline's parameters and variables, which is over a gigabyte for +# policyengine-us, and an unbounded cache made a long run grow by that much for +# every distinct combination it met. The reform-free system that most cases +# share is always kept and does not count. Dotted parameter inputs do not use +# the cache: each such case gets a ``_ParameterInputReform`` of its own, which +# shares nearly all of its tree with the cached system it is built on. +DEFAULT_REFORM_CACHE_SIZE = 2 + def run_tests(tax_benefit_system, paths, options=None): """ @@ -94,6 +107,22 @@ def run_tests(tax_benefit_system, paths, options=None): +-------------------------------+-----------+ See :any:`openfisca_test` options doc + | name_filter | ``str`` | | +-------------------------------+-----------+-------------------------------------------+ + | reform_cache_size | ``int`` | Reform systems kept cached (default 2) | + +-------------------------------+-----------+-------------------------------------------+ + | read_only_arrays | ``bool`` | Store cached arrays read-only, so a | + | | | formula writing into one it read fails | + +-------------------------------+-----------+-------------------------------------------+ + + **Memory**: a case that names ``reforms`` or ``extensions`` runs on a + system built for that combination, which is a full copy of + ``tax_benefit_system``. The runner keeps the ``reform_cache_size`` most + recently used of them and the reform-free system every other case shares. + A case that sets a parameter with a dotted input key runs on a system of + its own that copies only the paths to the parameters it sets and shares + everything else with the system its reforms select. Each case releases its + simulation and that system when it finishes, so memory does not grow with + the number of cases or of distinct combinations. Each system still caches + its parameters at every distinct instant the cases read. """ @@ -111,10 +140,16 @@ def run_tests(tax_benefit_system, paths, options=None): if options is None: options = {} - return pytest.main( - [*argv, *paths], - plugins=[OpenFiscaPlugin(tax_benefit_system, options)], - ) + previously_read_only = InMemoryStorage.read_only_arrays + if options.get("read_only_arrays"): + InMemoryStorage.read_only_arrays = True + try: + return pytest.main( + [*argv, *paths], + plugins=[OpenFiscaPlugin(tax_benefit_system, options)], + ) + finally: + InMemoryStorage.read_only_arrays = previously_read_only class YamlFile(pytest.File): @@ -183,6 +218,7 @@ def __init__(self, *, baseline_tax_benefit_system, test, options, **kwargs): self.test = test self.simulation = None self.tax_benefit_system = None + self._system_simulation = None def runtest(self): self.name = self.test.get("name", "") @@ -207,53 +243,38 @@ def runtest(self): unsafe_input = {} period = self.test.get("period") input = {} - inline_reforms = [] - parametric_reform_items = [] + parameter_inputs = [] for key, value in unsafe_input.items(): if "." in key: - inline_reforms += [ - set_parameter( - key, - value, - return_modifier=True, - period=f"year:2000:40", - ) - ] - parametric_reform_items.append((key, value)) + parameter_inputs.append((key, value)) else: input[key] = value - if len(inline_reforms) == 0: - inline_reform = [] - else: - - class inline_reform_class(Reform): - def apply(self): - for modifier in inline_reforms: - self.parameters = modifier(self.parameters) - - inline_reform = [inline_reform_class] - reforms = self.test.get("reforms", []) if isinstance(reforms, str): reforms = [reforms] - if not inline_reforms: - inline_reform = [] - self.tax_benefit_system = _get_tax_benefit_system( self.baseline_tax_benefit_system, - reforms + inline_reform, + reforms, self.test.get("extensions", []), - reform_key="=".join( - [f"{key}:{value}" for key, value in parametric_reform_items] - ), + cache_size=self.options.get("reform_cache_size"), ) + if parameter_inputs: + # Built for this case alone, and cheap: it shares every parameter + # it does not set with the cached system it is built on. + self.tax_benefit_system = _ParameterInputReform( + self.tax_benefit_system, parameter_inputs + ) verbose = self.options.get("verbose") performance_graph = self.options.get("performance_graph") performance_tables = self.options.get("performance_tables") visualize = self.options.get("visualize") + # ``Simulation.__init__`` registers itself on the system before the + # builder returns, so note what was there to restore it at teardown + # even if building fails. + self._system_simulation = getattr(self.tax_benefit_system, "simulation", None) try: builder.set_default_period(period) self.simulation = builder.build_from_dict(self.tax_benefit_system, input) @@ -288,6 +309,33 @@ def apply(self): if visualize: self.generate_variable_graph(tracer) + def teardown(self): + # pytest keeps every collected item until the session ends, so a + # finished case must not keep its simulation, or the system it ran on, + # alive: over a run of thousands of cases that retained every + # simulation and every reform system any case used. + system, simulation = self.tax_benefit_system, self.simulation + if system is not None: + if getattr(system, "simulation", None) is not self._system_simulation: + # A simulation built for this case, including one whose build + # failed, registered itself on the cached system. + system.simulation = self._system_simulation + tracer = getattr(simulation, "tracer", None) + parameters = getattr(system, "parameters", None) + if tracer is not None and getattr(parameters, "tracer", None) is tracer: + # Core's simulations trace through a per-call view and leave + # the tree unmarked, but code outside core can mark the cached + # system's tree with this case's tracer (policyengine-us's + # ``isolate_parameter_tracing`` does). Left marked, the tree + # would keep this case's tracer alive, and later untraced + # cases would record their parameter reads in it. + parameters.trace = False + parameters.tracer = None + parameters.branch_name = None + self.simulation = None + self.tax_benefit_system = None + self._system_simulation = None + def print_computation_log(self, tracer): print("Computation log:") # noqa T001 tracer.print_computation_log() @@ -419,32 +467,117 @@ def pytest_collect_file(self, parent, file_path): ) +class _CachedSystems: + """One baseline's cached systems: the reform-free one, kept for the + baseline's lifetime, and an LRU of reform systems.""" + + def __init__(self): + self.reform_free = None + self.reforms = OrderedDict() + + +class _ParameterInputReform(Reform): + """Sets a case's dotted parameter inputs on top of ``baseline``. + + Each input sets its parameter from 2000 for 40 years, as + ``set_parameter`` does. The reform copies only the nodes on the paths of + the parameters it sets (see ``policyengine_core.parameters.path_copy``), + so it shares the rest of ``baseline``'s tree, including the at-instant + values ``baseline`` has cached. A full copy of policyengine-us's tree took + seconds and over half a gigabyte for every distinct set of values, and + most cases that set parameters use values no other case does. + """ + + def __init__(self, baseline, parameter_inputs): + self.parameter_inputs = list(parameter_inputs) + super().__init__(baseline) + + def _copy_parameters(self, baseline_parameters): + return copy_paths( + baseline_parameters, [key for key, _ in self.parameter_inputs] + ) + + def apply(self): + for key, value in self.parameter_inputs: + set_parameter(key, value, period="year:2000:40", return_modifier=True)( + self.parameters + ) + + def _get_tax_benefit_system( baseline, reforms, extensions, reform_key=None, + cache_size=None, ): + """Return the system a YAML case runs on. + + ``cache_size`` bounds how many reform systems stay cached for ``baseline`` + (default ``DEFAULT_REFORM_CACHE_SIZE``; 0 caches none). The reform-free + system is always cached. + """ + if cache_size is None: + cache_size = DEFAULT_REFORM_CACHE_SIZE + if cache_size < 0: + raise ValueError(f"cache_size must be 0 or more, got {cache_size}") if not isinstance(reforms, list): reforms = [reforms] if not isinstance(extensions, list): extensions = [extensions] - # Inner key is (reforms in order, reform_key, extensions as a frozenset). - # The outer cache is keyed on the baseline *object* so its entries are - # tied to the baseline's lifetime (see bug H9 above). + # The cache is keyed on the baseline *object* so its entries are tied to + # the baseline's lifetime (see bug H9 above). + cached_systems = _tax_benefit_system_cache.get(baseline) + if cached_systems is None: + cached_systems = _CachedSystems() + _tax_benefit_system_cache[baseline] = cached_systems + + # This call's bound applies to what earlier calls cached under a larger + # one, least recently used first. + reform_systems = cached_systems.reforms + while len(reform_systems) > cache_size: + reform_systems.popitem(last=False) + + if cached_systems.reform_free is None: + cached_systems.reform_free = baseline.clone() + if not reforms and not extensions and not reform_key: + return cached_systems.reform_free + + # Key: (reforms in order, reform_key, extensions as a frozenset). The + # runner rebuilds its inline reform class for every case, so with a + # ``reform_key`` (the parameter values it sets) a class reform keys by + # that; without one, by the class itself. inner_key = ( - ":".join([reform if isinstance(reform, str) else "" for reform in reforms]), + tuple( + reform if isinstance(reform, str) or not reform_key else "" + for reform in reforms + ), reform_key, frozenset(extensions), ) - baseline_entry = _tax_benefit_system_cache.get(baseline) - if baseline_entry is not None: - cached = baseline_entry.get(inner_key) - if cached is not None: - return cached - - current_tax_benefit_system = baseline.clone() + cached = reform_systems.get(inner_key) + if cached is not None: + reform_systems.move_to_end(inner_key) + return cached + + # Make room before building, not after, so the cache and the system being + # built never add up to more than ``cache_size`` reform systems. + while reform_systems and len(reform_systems) >= cache_size: + reform_systems.popitem(last=False) + + if (reforms or extensions) and all( + _copies_its_baseline(reform) for reform in reforms + ): + # Every reform copies the parameters it changes and leaves the system + # it is applied to as it was, and each extension below is loaded into + # a clone, so build on the shared reform-free system: what comes back + # is a new system either way. A clone of the baseline here doubled + # every reform system's size, the clone kept alive as the reform's + # ``baseline``. + current_tax_benefit_system = cached_systems.reform_free + else: + current_tax_benefit_system = baseline.clone() for reform_path in reforms: if isinstance(reform_path, str): @@ -459,14 +592,22 @@ def _get_tax_benefit_system( current_tax_benefit_system = current_tax_benefit_system.clone() current_tax_benefit_system.load_extension(extension) - if baseline_entry is None: - baseline_entry = {} - _tax_benefit_system_cache[baseline] = baseline_entry - baseline_entry[inner_key] = current_tax_benefit_system + if cache_size > 0: + reform_systems[inner_key] = current_tax_benefit_system return current_tax_benefit_system +def _copies_its_baseline(reform) -> bool: + """Whether applying ``reform`` leaves the system it is applied to as it + was: a reform named by its path (``apply_reform`` requires a ``Reform``) + or a ``Reform`` class, whose constructor copies the baseline's parameter + tree, entities and variable registry before ``apply`` changes them.""" + return isinstance(reform, str) or ( + isinstance(reform, type) and issubclass(reform, Reform) + ) + + def assert_near( value, target_value, diff --git a/tests/core/parameters/test_parameter_cache_invalidation.py b/tests/core/parameters/test_parameter_cache_invalidation.py new file mode 100644 index 000000000..ae25473f5 --- /dev/null +++ b/tests/core/parameters/test_parameter_cache_invalidation.py @@ -0,0 +1,141 @@ +"""A change to a parameter tree reaches every node at instant requested after it. + +``ParameterNode`` caches its node at each instant, and a node at instant +resolves its children on first read. A change therefore has to clear the +cache of every node above the one it changes, or a node requested afterwards +would read old values for children an earlier read had resolved and new ones +for the rest. These tests pin each way of changing a tree, with and without +an earlier read. +""" + +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import Parameter, ParameterNode +from policyengine_core.reforms import Reform +from tests.fixtures.parameter_nodes_at_instant import build_tree, materialise, snapshot + +INSTANT = "2017-03-01" + + +@pytest.mark.parametrize("read_before", [False, True]) +def test_a_bracket_update_reaches_a_node_requested_afterwards(read_before): + """A freshly loaded scale's brackets know their scale, so updating a + bracket's rate clears the caches above the scale too.""" + tree = build_tree() + if read_before: + assert tree.get_at_instant(INSTANT).scale.rates == [0.1, 0.2] + + rate = tree.children["scale"].brackets[1].children["rate"] + rate.update(period="year:2017:1", value=0.5) + + assert tree.get_at_instant(INSTANT).scale.rates == [0.1, 0.5] + + +@pytest.mark.parametrize("read_before", [False, True]) +def test_add_child_reaches_a_node_requested_afterwards(read_before): + tree = build_tree() + if read_before: + assert tree.get_at_instant(INSTANT).group.x == 0.1 + + tree.children["group"].add_child( + "added", Parameter("group.added", data={"values": {"2010-01-01": 9}}) + ) + + group = tree.get_at_instant(INSTANT).group + assert group.added == 9 + assert "added" in list(group) + + +@pytest.mark.parametrize("read_before", [False, True]) +def test_clear_instant_caches_after_a_direct_edit(read_before): + """``values_list`` edits bypass ``Parameter.update``; clearing the caches + afterwards makes them visible.""" + tree = build_tree() + if read_before: + assert tree.get_at_instant(INSTANT).group.x == 0.1 + + for value_at_instant in tree.children["group"].children["x"].values_list: + value_at_instant.value = 0.7 + tree.children["group"].clear_instant_caches() + + assert tree.get_at_instant(INSTANT).group.x == 0.7 + assert materialise(tree.get_at_instant(INSTANT)) == snapshot(tree, INSTANT) + + +def test_clear_instant_caches_clears_the_subtree_and_the_nodes_above(): + tree = build_tree() + materialise(tree.get_at_instant(INSTANT)) + group = tree.children["group"] + assert tree._at_instant_cache and group._at_instant_cache + + group.clear_instant_caches() + + assert not tree._at_instant_cache + assert not group._at_instant_cache + assert not group.children["inner"]._at_instant_cache + # A sibling's cache holds nothing that depends on the group. + assert tree.children["by_zone"]._at_instant_cache + + +def test_a_system_reads_values_changed_after_a_reform_read_them_during_construction(): + """Reforms applied while a system is built run before the parameter + operations, which change values in place; the system clears its caches + once they have run.""" + + class read_then_change_in_place(Reform): + def apply(self): + assert self.parameters(INSTANT).taxes.income_tax_rate == 0.15 + for value_at_instant in self.parameters.taxes.income_tax_rate.values_list: + value_at_instant.value = 0.4 + + system = CountryTaxBenefitSystem(reform=read_then_change_in_place) + + assert system.parameters(INSTANT).taxes.income_tax_rate == 0.4 + assert system.get_parameters_at_instant(INSTANT).taxes.income_tax_rate == 0.4 + + +def test_a_node_held_across_a_change_keeps_what_it_read(): + """The documented exception: a node requested before the change.""" + tree = build_tree() + held = tree.get_at_instant(INSTANT).group + assert held.x == 0.1 + + tree.children["group"].children["x"].update(period="year:2017:1", value=0.9) + + assert held.x == 0.1 + assert tree.get_at_instant(INSTANT).group.x == 0.9 + + +@pytest.mark.parametrize( + "parent_child", + [{"values": {"2010-01-01": 1}}, {"x": {"values": {"2010-01-01": 1}}}], + ids=["parameter", "node"], +) +def test_a_child_named_parent_does_not_trap_cache_clearing(parent_child): + """``add_child`` sets each child as an attribute, so a child named + ``parent`` replaces the node's ``parent`` and points back at the node + until the node is attached; clearing caches meanwhile must not loop.""" + tree = ParameterNode( + "", + data={ + "group": {"parent": parent_child, "later": {"values": {"2010-01-01": 2}}} + }, + ) + group = tree.children["group"] + assert group.parent is tree + + group.children["later"].update(period="year:2017:1", value=3) + + assert tree.get_at_instant(INSTANT).group.later == 3 + + +def test_add_child_still_rejects_a_name_in_use(): + tree = build_tree() + with pytest.raises(ValueError): + tree.children["group"].add_child( + "x", Parameter("group.x", data={"values": {"2010-01-01": 1}}) + ) + with pytest.raises(ParameterNotFoundError): + tree.get_at_instant(INSTANT).group.missing diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py new file mode 100644 index 000000000..5140b58c5 --- /dev/null +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -0,0 +1,621 @@ +"""``ParameterNodeAtInstant`` resolves its children on first read. + +It used to build every descendant in its constructor, so each instant a +system was asked about cost a copy of the whole parameter tree (about 34,000 +node objects and 16-33 MB per instant for policyengine-us) that the system +then kept. These tests pin that nothing is built before it is read, and that +everything read is what the up-front build held. +""" + +import copy +import gc +import os +import pickle +import select +import signal +import threading +import weakref + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.parameters.parameter_node_at_instant import _STATE +from policyengine_core.tracers import FullTracer +from tests.fixtures.parameter_nodes_at_instant import ( + INSTANTS, + build_tree, + materialise, + ordered, + snapshot, +) + + +@pytest.mark.parametrize("instant", INSTANTS) +def test_fully_read_node_matches_the_up_front_build(instant): + tree = build_tree() + assert ordered(materialise(tree.get_at_instant(instant))) == ordered( + snapshot(tree, instant) + ) + + +@pytest.mark.parametrize( + "instant", ["2000-01-01", "2012-06-01", "2015-01-01", "2017-01-01", "2030-01-01"] +) +def test_country_template_matches_the_up_front_build(instant): + parameters = CountryTaxBenefitSystem().parameters + assert ordered(materialise(parameters.get_at_instant(instant))) == ordered( + snapshot(parameters, instant) + ) + + +def resolved(node_at_instant): + return list(node_at_instant.__dict__[_STATE].resolved) + + +def test_nothing_is_resolved_until_it_is_read(): + at_instant = build_tree().get_at_instant("2017-03-01") + assert resolved(at_instant) == [] + + assert at_instant.group.x == 0.1 + # Only the path that was read: the root's ``group`` and the group's ``x``. + assert resolved(at_instant) == ["group"] + assert resolved(at_instant.group) == ["x"] + + +def test_asking_for_many_instants_builds_one_node_each(): + tree = build_tree() + gc.collect() + before = sum(isinstance(o, ParameterNodeAtInstant) for o in gc.get_objects()) + for year in range(2000, 2100): + tree.get_at_instant(f"{year}-01-01") + after = sum(isinstance(o, ParameterNodeAtInstant) for o in gc.get_objects()) + # The up-front build made 9 nodes per instant for this tree (900 here). + assert after - before == 100 + + +def test_a_read_child_is_resolved_once(): + at_instant = build_tree().get_at_instant("2017-03-01") + # A scale is rebuilt by every ``_get_at_instant`` call, so identity shows + # that the node kept what it resolved. + assert at_instant.scale is at_instant.scale + assert at_instant["scale"] is at_instant.scale + assert at_instant._children["scale"] is at_instant.scale + assert at_instant.group is at_instant["group"] + + +def test_children_are_in_parameter_order_whatever_was_read_first(): + tree = build_tree() + at_instant = tree.get_at_instant("2030-01-01") + at_instant.by_zone + at_instant["late"] + at_instant.group + assert list(at_instant) == list(snapshot(tree, "2030-01-01")) + assert list(at_instant._children) == list(snapshot(tree, "2030-01-01")) + + +def test_child_without_a_value_yet_is_absent(): + at_instant = build_tree().get_at_instant("2017-03-01") + with pytest.raises(ParameterNotFoundError): + at_instant.late + with pytest.raises(ParameterNotFoundError) as error: + at_instant.group.y + assert str(error.value) == ( + "The parameter 'group[y]' was not found in the 2017-03-01 " + "tax and benefit system." + ) + with pytest.raises(KeyError): + at_instant["late"] + assert not hasattr(at_instant, "late") + assert "late" not in at_instant + assert "late" not in at_instant._children + assert at_instant.group.inner._children == {} + + +def test_unknown_names_raise_as_before(): + at_instant = build_tree().get_at_instant("2017-03-01") + for read_first in (False, True): + if read_first: + at_instant._children + with pytest.raises(ParameterNotFoundError) as error: + at_instant.group.nope + assert str(error.value) == ( + "The parameter 'group[nope]' was not found in the 2017-03-01 " + "tax and benefit system." + ) + with pytest.raises(KeyError): + at_instant.group["nope"] + assert not hasattr(at_instant, "nope") + + +def test_vectorial_indexing_reads_the_whole_subtree(): + at_instant = build_tree().get_at_instant("2020-01-01") + zones = np.array(["z1", "z3", "z2"]) + np.testing.assert_array_equal(at_instant.by_zone[zones].amount, [200, 600, 400]) + np.testing.assert_allclose(at_instant.by_zone[zones].rate, [0.1, 0.3, 0.2]) + # Nothing outside the indexed subtree was resolved for it. + assert resolved(at_instant) == ["by_zone"] + + +def test_dir_lists_children_for_completion(): + at_instant = build_tree().get_at_instant("2017-03-01") + names = dir(at_instant) + assert {"flat_rate", "group", "scale", "by_zone"} <= set(names) + assert "late" not in names + + +def test_repr_shows_every_child(): + text = repr(build_tree().get_at_instant("2017-03-01")) + for name in ("flat_rate", "group", "scale", "by_zone", "z3"): + assert name in text + + +def test_add_child_before_and_after_a_full_read(): + at_instant = build_tree().get_at_instant("2017-03-01") + at_instant.add_child("extra", 1) + assert at_instant.extra == 1 and at_instant["extra"] == 1 + assert list(at_instant._children)[-1] == "extra" + at_instant.add_child("later", 2) + assert at_instant.later == 2 and at_instant._children["later"] == 2 + + +def test_a_node_asked_for_after_an_update_reads_the_update(): + tree = build_tree() + assert tree.get_at_instant("2017-03-01").flat_rate == 0.1 + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert tree.get_at_instant("2017-03-01").flat_rate == 0.5 + + +def test_a_child_read_before_an_update_keeps_what_was_read(): + tree = build_tree() + held = tree.get_at_instant("2017-03-01") + assert held.flat_rate == 0.1 + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert held.flat_rate == 0.1 + + +def test_a_child_not_read_before_an_update_reads_the_update(): + """Intended difference from the up-front build, which held 0.1 here: a + node held across a change to its parameters reads, for each child, the + value at the time that child is first read.""" + tree = build_tree() + held = tree.get_at_instant("2017-03-01") + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert held.flat_rate == 0.5 + + +def test_deep_copied_tree_keeps_its_own_nodes_at_instant(): + tree = build_tree() + tree.get_at_instant("2017-03-01").group.x # cache a partly read node + clone = copy.deepcopy(tree) + cached = clone._at_instant_cache["2017-03-01"] + # The copy's cached node is a snapshot: a change to the original tree + # does not reach it. + tree.flat_rate.update(value=0.5, period="year:2010:30") + assert cached.flat_rate == 0.1 + # A node requested from the copy after a change to the copy reads the + # copy, and the original is untouched. + clone.group.y.update(value=0.9, period="year:2010:30") + clone.flat_rate.update(value=0.7, period="year:2010:30") + assert clone.get_at_instant("2017-03-01").flat_rate == 0.7 + assert clone.get_at_instant("2017-03-01").group.y == 0.9 + assert tree.get_at_instant("2017-03-01").flat_rate == 0.5 + assert not hasattr(tree.get_at_instant("2017-03-01").group, "y") + + +def test_a_node_is_built_only_at_a_date(): + """Built up front, a node at "None" (``parameters(None)``) failed on its + first scale and was never cached; ``Simulation._calculate`` relies on + that when it checks an eternal variable's abolition.""" + tree = build_tree() + with pytest.raises(ValueError): + tree(None) + assert "None" not in tree._at_instant_cache + copy.deepcopy(tree) + + +def test_copy_of_a_node_at_instant_reads_the_same_values(): + at_instant = build_tree().get_at_instant("2017-03-01") + at_instant.group + for duplicate in (copy.copy(at_instant), copy.deepcopy(at_instant)): + assert materialise(duplicate) == materialise(at_instant) + + +@pytest.mark.parametrize( + "child_name", + # Not ``add_child``: ParameterNode itself cannot hold a child of that name. + ["_name", "_instant_str", "_vectorial_node", "_node", "_resolved"] + + ["_resolve", "_resolve_all", "_store", "_materialise"], +) +def test_children_named_like_the_nodes_own_attributes(child_name): + """Such a child reads as its value whatever the read order, as when every + child was set up front, and leaves its siblings readable.""" + tree = ParameterNode( + "", + data={ + child_name: {"values": {"2010-01-01": 42}}, + "ok": {"values": {"2010-01-01": 7}}, + }, + ) + for read_name_first in (True, False): + tree._at_instant_cache.clear() + at_instant = tree.get_at_instant("2017-03-01") + if read_name_first: + assert getattr(at_instant, child_name) == 42 + assert at_instant.ok == 7 and at_instant["ok"] == 7 + else: + assert at_instant["ok"] == 7 and at_instant.ok == 7 + assert getattr(at_instant, child_name) == 42 + assert at_instant[child_name] == 42 + assert set(at_instant._children) == {child_name, "ok"} + + +class BlockingChild: + """A parameter child whose first resolution waits until released.""" + + def __init__(self): + self.calls = 0 + self.entered = threading.Event() + self.release = threading.Event() + + def _get_at_instant(self, instant): + self.calls += 1 + if self.calls == 1: + self.entered.set() + assert self.release.wait(10), "never released" + return object() + + +@pytest.mark.parametrize( + "second_read", + [ + lambda node: node.slow, + lambda node: node["slow"], + lambda node: node._children["slow"], + lambda node: [node[name] for name in node][0], + ], +) +def test_concurrent_first_reads_resolve_a_child_once(second_read): + tree = ParameterNode("", data={"other": {"values": {"2010-01-01": 1}}}) + slow = BlockingChild() + tree.children = {"slow": slow, **tree.children} + at_instant = tree.get_at_instant("2017-03-01") + results, errors = [], [] + + def read(reader): + try: + results.append(reader(at_instant)) + except Exception as error: # noqa: BLE001 - recorded for the assertion + errors.append(error) + + first = threading.Thread(target=read, args=(lambda node: node.slow,)) + first.start() + assert slow.entered.wait(10) + second = threading.Thread(target=read, args=(second_read,)) + second.start() + second.join(0.2) # let the second reader reach the child too + slow.release.set() + first.join(10) + second.join(10) + + assert not errors, errors + assert len(results) == 2 and results[0] is results[1] + assert slow.calls == 1 + + +def test_shallow_copies_read_the_same_children(): + held = build_tree().get_at_instant("2017-03-01") + duplicate = copy.copy(held) + original_scale = held.scale + # A scale is rebuilt by each resolution, so identity shows one resolution. + assert duplicate.scale is original_scale + assert duplicate["scale"] is original_scale + assert held["scale"] is original_scale + duplicate._children + assert held._children["scale"] is original_scale + assert held.by_zone is duplicate["by_zone"] + + +@pytest.mark.parametrize( + "duplicate", [copy.deepcopy, lambda n: pickle.loads(pickle.dumps(n))] +) +def test_traced_nodes_copy_and_pickle(duplicate): + tree = build_tree() + tree.trace = True + tree.tracer = FullTracer() + tree.branch_name = "default" + traced = tree.get_at_instant("2017-03-01") + traced.group.x # partly resolved, through the tracing wrapper + raw = traced.parameter_node_at_instant + for node in (raw, traced): + assert duplicate(node).flat_rate == 0.1 + assert duplicate(node).group.x == 0.1 + + +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_a_partly_read_node_pickles_under_every_protocol(protocol): + """Protocols 0 and 1 refuse an object with ``__slots__`` and no + ``__getstate__``, and the lazy state is one.""" + at_instant = build_tree().get_at_instant("2017-03-01") + assert at_instant.group.x == 0.1 + restored = pickle.loads(pickle.dumps(at_instant, protocol=protocol)) + assert restored.group.x == 0.1 + assert restored.flat_rate == 0.1 # not read before pickling + assert ordered(materialise(restored)) == ordered(materialise(at_instant)) + + +@pytest.mark.parametrize( + "duplicate", [copy.deepcopy, lambda node: pickle.loads(pickle.dumps(node))] +) +def test_a_copied_node_is_a_snapshot_without_the_parameter_tree(duplicate): + """Unread children hold parameter nodes, whose parents reach the root: a + copy that kept them would carry, and keep alive, the whole tree.""" + tree = ParameterNode( + "", + data={ + "small": { + "x": {"values": {"2010-01-01": 1}}, + "y": {"values": {"2010-01-01": 2}}, + }, + "large": { + f"parameter_{index}": {"values": {"2010-01-01": index}} + for index in range(2000) + }, + }, + ) + small = tree.get_at_instant("2017-01-01").small + assert small.x == 1 + + copied = duplicate(small) + + assert copied.__dict__[_STATE].source is None + assert (copied.x, copied.y, list(copied)) == (1, 2, ["x", "y"]) + assert len(pickle.dumps(small)) * 20 < len(pickle.dumps(tree)) + tree_reference = weakref.ref(tree) + del tree, small + gc.collect() + assert tree_reference() is None + assert copied.y == 2 + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +@pytest.mark.filterwarnings("ignore::DeprecationWarning") +def test_a_forked_child_resolves_while_another_thread_was_resolving(): + """A thread holding the resolution lock at the fork does not exist in the + child, so the child needs a lock of its own.""" + resolving, release = threading.Event(), threading.Event() + + class SlowChild: + def _get_at_instant(self, instant): + resolving.set() + release.wait(10) + return 1 + + tree = ParameterNode("", data={"other": {"values": {"2010-01-01": 7}}}) + tree.children = {"slow": SlowChild(), **tree.children} + held = tree.get_at_instant("2017-01-01") + reader = threading.Thread(target=lambda: held.slow) + reader.start() + assert resolving.wait(10) + unrelated = ParameterNode("", data={"x": {"values": {"2010-01-01": 3}}}) + + read_end, write_end = os.pipe() + child = os.fork() + if child == 0: # pragma: no cover - runs in the forked child + try: + value = unrelated.get_at_instant("2017-01-01").x + os.write(write_end, str(value).encode()) + finally: + os._exit(0) + os.close(write_end) + try: + ready, _, _ = select.select([read_end], [], [], 10) + received = os.read(read_end, 16) if ready else b"" + finally: + release.set() + reader.join() + if not received: + os.kill(child, signal.SIGKILL) + os.waitpid(child, 0) + os.close(read_end) + assert received == b"3" + + +def test_a_child_without_a_value_is_resolved_once(): + tree = ParameterNode("", data={"late": {"values": {"2030-01-01": 1}}}) + held = tree.get_at_instant("2020-01-01") + calls = [] + original = tree.children["late"]._get_at_instant + + def counting(instant): + calls.append(instant) + return original(instant) + + tree.children["late"]._get_at_instant = counting + for _ in range(3): + with pytest.raises(ParameterNotFoundError): + held.late + with pytest.raises(KeyError): + held["late"] + assert list(held) == [] + assert len(calls) == 1 + + # A child already read keeps what that read found, value or none. + tree.children["late"].update(value=2, period="year:2020:1") + with pytest.raises(KeyError): + held["late"] + assert tree.get_at_instant("2020-01-01").late == 2 + + +def test_a_node_lets_go_of_its_parameter_node_once_every_child_is_read(): + tree = ParameterNode( + "", + data={ + "group": { + "x": {"values": {"2010-01-01": 1}}, + "late": {"values": {"2030-01-01": 2}}, + }, + "unused": {"values": {"2010-01-01": 3}}, + }, + ) + group = tree.get_at_instant("2017-03-01").group + tree_ref = weakref.ref(tree) + del tree + gc.collect() + # Partly read: the node still needs the tree to resolve the rest. + assert tree_ref() is not None + assert group.x == 1 + with pytest.raises(ParameterNotFoundError): + group.late + gc.collect() + assert tree_ref() is None + assert group.x == 1 and list(group) == ["x"] + + +def plain_tree(**values): + return ParameterNode( + "", data={name: {"values": {"2010-01-01": v}} for name, v in values.items()} + ) + + +def test_the_reserved_state_name_is_rejected(): + with pytest.raises(ValueError): + plain_tree(_ParameterNodeAtInstant__state=1).get_at_instant("2017-01-01") + at_instant = plain_tree(x=1).get_at_instant("2017-01-01") + with pytest.raises(ValueError): + at_instant.add_child("_ParameterNodeAtInstant__state", 2) + + +@pytest.mark.parametrize( + "child_name", + [ + "_ParameterNodeAtInstant__resolved", + "_ParameterNodeAtInstant__node", + "_ParameterNodeAtInstant__business", + "__custom__", + "__iter__", + "__dir__", + ], +) +def test_other_unusual_names_are_plain_children(child_name): + at_instant = plain_tree(**{child_name: 42, "ok": 7}).get_at_instant("2017-01-01") + assert getattr(at_instant, child_name) == 42 + assert at_instant[child_name] == 42 + assert at_instant.ok == 7 and at_instant["ok"] == 7 + assert child_name in dir(at_instant) + + +def test_many_added_children_do_not_break_the_first_read(): + names = {f"p{i}": i for i in range(2000)} + at_instant = plain_tree(**names).get_at_instant("2017-01-01") + for i in range(2000): + at_instant.add_child(f"extra{i}", i) + assert at_instant.p0 == 0 + assert len(at_instant._children) == 4000 + assert list(at_instant._children)[:2] == ["p0", "p1"] + + +def test_add_child_on_a_copy_changes_only_shared_items_as_before(): + tree = ParameterNode( + "", + data={ + "x": {"values": {"2010-01-01": 10}}, + "late": {"values": {"2030-01-01": 5}}, + }, + ) + held = tree.get_at_instant("2017-01-01") + with pytest.raises(ParameterNotFoundError): + held.late + duplicate = copy.copy(held) + duplicate.add_child("x", 99) + duplicate.add_child("late", 88) + duplicate.add_child("extra", 77) + # Items are one mapping shared by copies; attributes belong to each copy. + assert held["x"] == 99 and held.x == 10 + assert held["late"] == 88 + with pytest.raises(ParameterNotFoundError): + held.late + assert held["extra"] == 77 + with pytest.raises(ParameterNotFoundError): + held.extra + assert duplicate.x == 99 and duplicate.late == 88 and duplicate.extra == 77 + assert list(held._children) == ["x", "late", "extra"] + + +def test_a_copy_installs_what_it_reads_as_attributes(): + held = build_tree().get_at_instant("2017-03-01") + duplicate = copy.copy(held) + scale = held.scale + assert duplicate.scale is scale + assert duplicate.__dict__["scale"] is scale + + +class _MethodCollision(ParameterNodeAtInstant): + def policy(self): + return "method" + + +class _DoublingAddChild(ParameterNodeAtInstant): + def add_child(self, child_name, child_at_instant): + super().add_child(child_name, child_at_instant * 2) + + +def test_subclasses_keep_their_collisions_and_add_child_hooks(): + tree = plain_tree(policy=42, x=5) + collided = _MethodCollision("", tree, "2017-01-01") + assert collided.policy == 42 and collided["policy"] == 42 + doubled = _DoublingAddChild("", tree, "2017-01-01") + assert doubled.x == 10 and doubled["x"] == 10 + assert dict(doubled._children) == {"policy": 84, "x": 10} + + +class ReenteringChild: + """A child whose first resolution reads the same child again.""" + + def __init__(self): + self.node_at_instant = None + self.calls = 0 + + def _get_at_instant(self, instant): + self.calls += 1 + result = object() + if self.calls == 1: + self.inner = self.node_at_instant["again"] + return result + + +def test_a_reentrant_read_keeps_the_first_recorded_child(): + tree = plain_tree(other=1) + child = ReenteringChild() + tree.children = {"again": child, **tree.children} + at_instant = tree.get_at_instant("2017-01-01") + child.node_at_instant = at_instant + outer = at_instant["again"] + assert outer is child.inner + assert at_instant["again"] is outer and at_instant.again is outer + + +def test_children_and_their_order_are_fixed_when_the_node_is_built(): + tree = plain_tree(x=10, y=20, z=30) + held = tree.get_at_instant("2017-01-01") + assert held.x == 10 + del tree.children["x"] + tree.children["w"] = ParameterNode("", data={"v": {"values": {"2010-01-01": 1}}}) + assert list(held) == ["x", "y", "z"] + assert held.x == 10 + assert "w" not in held._children + + +def test_a_node_pickled_when_every_child_was_built_up_front_loads(): + """State as the up-front build pickled it: everything in ``_children``.""" + restored = ParameterNodeAtInstant.__new__(ParameterNodeAtInstant) + restored.__setstate__( + {"_name": "", "_instant_str": "2017-01-01", "_children": {"x": 10}, "x": 10} + ) + assert restored.x == 10 and restored["x"] == 10 and list(restored) == ["x"] + restored.add_child("y", 2) + assert restored["y"] == 2 and list(restored) == ["x", "y"] + with pytest.raises(KeyError): + restored["nope"] diff --git a/tests/core/parameters/test_parameter_node_at_instant_property.py b/tests/core/parameters/test_parameter_node_at_instant_property.py new file mode 100644 index 000000000..57d96ee53 --- /dev/null +++ b/tests/core/parameters/test_parameter_node_at_instant_property.py @@ -0,0 +1,153 @@ +"""Reads of a lazily resolved ``ParameterNodeAtInstant`` in any order. + +For any instant and any sequence of reads (attribute, item, iteration, +``_children``, ``dir``, ``repr``, vectorial indexing, names that do not +exist), every read returns what the up-front build held, and the node read in +full afterwards equals that build. ``test_parameter_node_at_instant_lazy.py`` +pins the same behaviour with examples. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import Parameter, ParameterNodeAtInstant +from tests.fixtures.parameter_nodes_at_instant import ( + INSTANTS, + build_tree, + describe, + materialise, + ordered, + snapshot, +) + +# Node paths in the fixture tree, and names to ask each for (present at some +# instants, absent at others, or never there). +PATHS = [(), ("group",), ("group", "inner"), ("by_zone",), ("by_zone", "z2")] +NAMES = ["flat_rate", "late", "group", "scale", "by_zone", "x", "y", "inner"] +NAMES += ["1", "2", "z1", "z2", "z3", "amount", "rate", "nope"] +KINDS = ["getattr", "getitem", "iter", "children", "dir", "repr", "vector"] + +reads = st.tuples( + st.sampled_from(PATHS), st.sampled_from(KINDS), st.sampled_from(NAMES) +) + + +def _comparable(value): + if isinstance(value, ParameterNodeAtInstant): + return materialise(value) + return describe(value) + + +@hypothesis.given( + instant=st.sampled_from(INSTANTS), sequence=st.lists(reads, max_size=30) +) +@hypothesis.settings(max_examples=400, deadline=None) +def test_reads_in_any_order_match_the_up_front_build(instant, sequence): + tree = build_tree() + expected_root = snapshot(tree, instant) + root = tree.get_at_instant(instant) + + for path, kind, name in sequence: + node, expected = root, expected_root + for step in path: + node, expected = node[step], expected[step] + + if kind == "getattr": + if name in expected: + assert _comparable(getattr(node, name)) == expected[name] + else: + with pytest.raises(ParameterNotFoundError): + getattr(node, name) + assert not hasattr(node, name) + elif kind == "getitem": + if name in expected: + assert _comparable(node[name]) == expected[name] + else: + with pytest.raises(KeyError): + node[name] + elif kind == "iter": + assert list(node) == list(expected) + assert (name in node) == (name in expected) + elif kind == "children": + assert ordered(materialise(node)) == ordered(expected) + elif kind == "dir": + assert set(expected) <= set(dir(node)) + elif kind == "repr": + assert all(child in repr(node) for child in expected) + elif kind == "vector" and path == ("by_zone",) and expected["z1"]: + zones = np.array(["z3", "z1"]) + np.testing.assert_array_equal( + node[zones].amount, + [expected["z3"]["amount"], expected["z1"]["amount"]], + ) + + assert ordered(materialise(root)) == ordered(expected_root) + + +# Parameters a change can update: plain values, a value deep in the tree, a +# zone of a vectorisable node, and a scale's brackets. +UPDATABLE = [ + lambda tree: tree.children["flat_rate"], + lambda tree: tree.children["group"].children["x"], + lambda tree: tree.children["group"].children["inner"].children["2"], + lambda tree: tree.children["by_zone"].children["z2"].children["amount"], + lambda tree: tree.children["scale"].brackets[0].children["threshold"], + lambda tree: tree.children["scale"].brackets[1].children["rate"], +] +# Nodes a change can add a child to. +EXTENDABLE = [(), ("group",), ("group", "inner")] + +steps = st.one_of( + st.tuples(st.just("read"), st.sampled_from(PATHS), st.sampled_from(NAMES)), + st.tuples( + st.just("update"), + st.integers(0, len(UPDATABLE) - 1), + st.sampled_from([0.5, 7, 2000]), + ), + st.tuples(st.just("add_child"), st.sampled_from(EXTENDABLE), st.just(None)), +) + + +@hypothesis.given( + instant=st.sampled_from(INSTANTS), sequence=st.lists(steps, max_size=25) +) +@hypothesis.settings(max_examples=300, deadline=None) +def test_a_node_requested_after_changes_reads_the_tree_as_it_is(instant, sequence): + """Reads, ``Parameter.update`` (scale brackets included) and + ``ParameterNode.add_child`` in any order: a node requested after the last + change reads exactly what the tree holds, whatever earlier reads resolved.""" + tree = build_tree() + added = 0 + for step, target, argument in sequence: + if step == "read": + node = tree.get_at_instant(instant) + expected = snapshot(tree, instant) + for name in target: + if name not in expected: + break + node, expected = node[name], expected[name] + else: + if argument in expected: + assert _comparable(getattr(node, argument)) == expected[argument] + elif step == "update": + UPDATABLE[target](tree).update(period="year:2000:100", value=argument) + else: + parent = tree + for name in target: + parent = parent.children[name] + added += 1 + parent.add_child( + f"added_{added}", + Parameter(f"added_{added}", data={"values": {"2000-01-01": added}}), + ) + + assert ordered(materialise(tree.get_at_instant(instant))) == ordered( + snapshot(tree, instant) + ) diff --git a/tests/core/parameters/test_path_copy.py b/tests/core/parameters/test_path_copy.py new file mode 100644 index 000000000..67a30b007 --- /dev/null +++ b/tests/core/parameters/test_path_copy.py @@ -0,0 +1,179 @@ +"""``copy_paths`` copies only the nodes on the paths it is given.""" + +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import Parameter +from policyengine_core.parameters.path_copy import copy_paths +from policyengine_core.reforms import set_parameter + +SYSTEM = CountryTaxBenefitSystem() +INSTANTS = ["2012-06-01", "2015-01-01", "2017-01-01", "2020-01-01", "2035-01-01"] + + +def _leaf_values(tree, instant): + """Every parameter's value at ``instant``, by name.""" + return { + descendant.name: descendant(instant) + for descendant in tree.get_descendants() + if isinstance(descendant, Parameter) + } + + +def _set(tree, path, value): + return set_parameter(path, value, period="year:2000:40", return_modifier=True)(tree) + + +def _leaf_paths(tree): + """A path to every parameter of ``tree`` in ``set_parameter``'s syntax.""" + paths = [] + + def visit(node, prefix): + if isinstance(node, Parameter): + paths.append(prefix) + elif hasattr(node, "brackets"): + for index, bracket in enumerate(node.brackets): + visit(bracket, f"{prefix}[{index}]") + else: + for name, child in node.children.items(): + visit(child, f"{prefix}.{name}" if prefix else name) + + visit(tree, "") + return [path for path in paths if not path.startswith("baseline.")] + + +LEAF_PATHS = _leaf_paths(SYSTEM.parameters) + + +def test_only_the_nodes_on_the_path_are_new(): + original = SYSTEM.parameters.clone() + copied = copy_paths(original, ["taxes.income_tax_rate"]) + + assert copied is not original + assert copied.taxes is not original.taxes + assert copied.taxes.income_tax_rate is not original.taxes.income_tax_rate + assert copied.benefits is original.benefits + assert copied.taxes.housing_tax is original.taxes.housing_tax + assert copied.children["taxes"] is copied.taxes + assert copied.taxes.parent is copied + assert copied.taxes.income_tax_rate.parent is copied.taxes + + +def test_a_bracket_path_copies_the_scale_and_the_bracket(): + original = SYSTEM.parameters.clone() + path = "taxes.social_security_contribution[1].rate" + copied = copy_paths(original, [path]) + + scale = copied.taxes.social_security_contribution + original_scale = original.taxes.social_security_contribution + assert scale is not original_scale + assert scale.brackets[1] is not original_scale.brackets[1] + assert scale.brackets[1].parent is scale + assert scale.brackets[0] is original_scale.brackets[0] + assert scale.brackets[1].threshold is original_scale.brackets[1].threshold + + _set(copied, path, 0.5) + assert copied.get_child(path)("2020-01-01") == 0.5 + assert original.get_child(path)("2020-01-01") == 0.06 + + +def test_updating_the_copy_leaves_the_original_and_its_caches_alone(): + original = SYSTEM.parameters.clone() + warm_node = original.get_at_instant("2017-01-01") + warm_benefits = original.benefits.get_at_instant("2017-01-01") + + copied = copy_paths(original, ["taxes.income_tax_rate"]) + _set(copied, "taxes.income_tax_rate", 0.3) + + assert copied.taxes.income_tax_rate("2017-01-01") == 0.3 + assert copied.get_at_instant("2017-01-01").taxes.income_tax_rate == 0.3 + assert original.taxes.income_tax_rate("2017-01-01") == 0.15 + assert original.get_at_instant("2017-01-01") is warm_node + assert warm_node.taxes.income_tax_rate == 0.15 + # A branch the change did not touch keeps the values it cached. + assert copied.benefits.get_at_instant("2017-01-01") is warm_benefits + + +def test_paths_with_a_common_prefix_share_their_copies(): + original = SYSTEM.parameters.clone() + copied = copy_paths( + original, ["taxes.housing_tax.rate", "taxes.housing_tax.minimal_amount"] + ) + housing_tax = copied.taxes.housing_tax + assert housing_tax is not original.taxes.housing_tax + assert housing_tax.rate.parent is housing_tax + assert housing_tax.minimal_amount.parent is housing_tax + + +def test_an_unknown_path_raises_before_anything_is_copied(): + original = SYSTEM.parameters.clone() + with pytest.raises(ValueError, match="Could not find the parameter"): + copy_paths(original, ["taxes.income_tax_rate", "taxes.no_such_tax"]) + + +@settings(max_examples=60, deadline=None) +@given( + changes=st.lists( + st.tuples( + st.sampled_from(LEAF_PATHS), + st.floats(min_value=0, max_value=1e6, allow_nan=False), + ), + min_size=1, + max_size=4, + unique_by=lambda change: change[0], + ), + warm=st.booleans(), +) +def test_copying_paths_matches_cloning_the_whole_tree(changes, warm): + """For any set of parameter changes, a tree that copies only their paths + reads exactly what a full clone with the same changes reads, and the + original reads what it did before, whether or not its at-instant caches + were warm when it was copied.""" + original = SYSTEM.parameters.clone() + before = {instant: _leaf_values(original, instant) for instant in INSTANTS} + if warm: + for instant in INSTANTS: + original.get_at_instant(instant) + + paths = [path for path, _ in changes] + copied = copy_paths(original, paths) + reference = original.clone() + for path, value in changes: + _set(copied, path, value) + _set(reference, path, value) + + for instant in INSTANTS: + assert _leaf_values(copied, instant) == _leaf_values(reference, instant) + assert _leaf_values(original, instant) == before[instant] + # Read through the at-instant trees too, which is how formulas read + # parameters and where a stale shared cache would show. + assert _values_at_instant(copied, instant) == _values_at_instant( + reference, instant + ) + assert _values_at_instant(original, instant) == { + path: before[instant][original.get_child(path).name] + for path in PLAIN_LEAF_PATHS + } + + +# Paths outside scales, which an at-instant tree reaches by attribute. +PLAIN_LEAF_PATHS = [path for path in LEAF_PATHS if "[" not in path] + + +def _values_at_instant(tree, instant): + """Each plain parameter's value read through ``tree``'s at-instant tree, + or ``None`` where that tree leaves it out (it has no value then).""" + node_at_instant = tree.get_at_instant(instant) + values = {} + for path in PLAIN_LEAF_PATHS: + value = node_at_instant + try: + for name in path.split("."): + value = getattr(value, name) + except ParameterNotFoundError: + value = None + values[path] = value + return values diff --git a/tests/core/parameters/test_vectorial_parameter_node_copy.py b/tests/core/parameters/test_vectorial_parameter_node_copy.py index 98af4f9b2..5372ae7d2 100644 --- a/tests/core/parameters/test_vectorial_parameter_node_copy.py +++ b/tests/core/parameters/test_vectorial_parameter_node_copy.py @@ -9,7 +9,7 @@ ``recarray``. ``ParameterNodeAtInstant`` caches its vectorial node, and -``Reform.modify_parameters`` deep-copies the baseline's parameter tree, so a +``Reform.modify_parameters`` deep-copied the baseline's parameter tree, so a reform found a ``recarray`` where its cached node should be. """ @@ -200,7 +200,7 @@ def test_copied_parameter_tree_keeps_its_cached_vectorial_node(copier): assert_near(rate[families][housings].z2, expected) -def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node( +def test_reform_that_modifies_parameters_reads_the_vectorial_node( isolated_tax_benefit_system, ): system = isolated_tax_benefit_system @@ -210,10 +210,11 @@ def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node( class KeepParameters(Reform): def apply(self): - # Deep-copies the baseline's tree, cache included. + # Given the reform's own copy of the tree, not the baseline's, + # which has cached that node. self.modify_parameters(lambda parameters: parameters) general = KeepParameters(system).get_parameters_at_instant("2017-01-01").general - assert type(general._vectorial_node) is VectorialParameterNodeAtInstant assert_near(general[names], expected) + assert type(general._vectorial_node) is VectorialParameterNodeAtInstant diff --git a/tests/core/test_parameter_tracing_cache.py b/tests/core/test_parameter_tracing_cache.py new file mode 100644 index 000000000..5b6d73b0e --- /dev/null +++ b/tests/core/test_parameter_tracing_cache.py @@ -0,0 +1,109 @@ +"""A traced parameter tree wraps its cached nodes on each read and caches only +the plain nodes, so no tracer outlives its simulation and no cached node goes +untraced. The tax-benefit system's own cache of the tree at each instant +(``TaxBenefitSystem.get_parameters_at_instant``) follows the same rule.""" + +import gc +import weakref + +from policyengine_core.country_template.entities import entities +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant + +DATES = [f"{year}-01-01" for year in range(2011, 2027)] + + +def tree(): + return ParameterNode("", data={"x": {"values": {"2010-01-01": 1}}}) + + +def system(): + tax_benefit_system = TaxBenefitSystem(entities) + tax_benefit_system.parameters = tree() + return tax_benefit_system + + +def trace(node, tracer, branch_name="default"): + node.trace, node.tracer, node.branch_name = True, tracer, branch_name + + +def untrace(node): + node.trace, node.tracer, node.branch_name = False, None, None + + +def test_a_node_cached_before_tracing_is_traced_once_tracing_starts(): + parameters = tree() + plain = parameters.get_at_instant("2017-01-01") + trace(parameters, FullTracer()) + traced = parameters.get_at_instant("2017-01-01") + assert isinstance(traced, TracingParameterNodeAtInstant) + assert traced.parameter_node_at_instant is plain + + +def test_the_cache_never_holds_a_tracer(): + parameters = tree() + for tracer in (FullTracer(), FullTracer()): + trace(parameters, tracer) + assert parameters.get_at_instant("2017-01-01").tracer is tracer + assert parameters.get_at_instant("2018-01-01").tracer is tracer + assert all( + type(node) is ParameterNodeAtInstant + for node in parameters._at_instant_cache.values() + ) + parameters.trace, parameters.tracer = False, None + assert type(parameters.get_at_instant("2017-01-01")) is ParameterNodeAtInstant + + +def test_the_system_cache_never_holds_a_tracer(): + """Reads at distinct dates, each under a tracer of its own, as traced + cases on one cached system make them: the system used to cache one + tracing wrapper per date, each keeping its tracer alive.""" + tax_benefit_system = system() + refs = [] + for date in DATES: + tracer = FullTracer() + refs.append(weakref.ref(tracer)) + trace(tax_benefit_system.parameters, tracer) + assert tax_benefit_system.get_parameters_at_instant(date).tracer is tracer + del tracer + untrace(tax_benefit_system.parameters) + gc.collect() + + cached = tax_benefit_system._parameters_at_instant_cache + assert len(cached) == len(DATES) + assert all(type(node) is ParameterNodeAtInstant for node in cached.values()) + assert [ref for ref in refs if ref() is not None] == [] + + +def test_a_system_node_cached_before_tracing_is_traced_once_tracing_starts(): + tax_benefit_system = system() + plain = tax_benefit_system.get_parameters_at_instant("2017-01-01") + tracer = FullTracer() + trace(tax_benefit_system.parameters, tracer, "branch") + traced = tax_benefit_system.get_parameters_at_instant("2017-01-01") + assert isinstance(traced, TracingParameterNodeAtInstant) + assert traced.parameter_node_at_instant is plain + assert (traced.tracer, traced.branch_name) == (tracer, "branch") + untrace(tax_benefit_system.parameters) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01") is plain + + +def test_each_system_read_records_in_the_tracer_current_at_the_read(): + """A read through the system records in the tree's tracer at the time of + the read, and nowhere once tracing stops: it used to record every later + read in the tracer current when the instant was first read.""" + tax_benefit_system = system() + first, second = FullTracer(), FullTracer() + for tracer in (first, second): + tracer.record_calculation_start("y", "2017") + + trace(tax_benefit_system.parameters, first) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01").x == 1 + trace(tax_benefit_system.parameters, second) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01").x == 1 + untrace(tax_benefit_system.parameters) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01").x == 1 + + assert len(first.trees[0].parameters) == 1 + assert len(second.trees[0].parameters) == 1 diff --git a/tests/core/test_parameter_tracing_cache_property.py b/tests/core/test_parameter_tracing_cache_property.py new file mode 100644 index 000000000..32918fb54 --- /dev/null +++ b/tests/core/test_parameter_tracing_cache_property.py @@ -0,0 +1,147 @@ +"""The parameter-instant caches under tracing, for any sequence of tracing +switches and reads at any dates. + +Two paths read the parameter tree at an instant: the tree's own +(``ParameterNode.get_at_instant``, cached in ``_at_instant_cache``) and the +tax-benefit system's (``TaxBenefitSystem.get_parameters_at_instant``, cached +in ``_parameters_at_instant_cache``). Invariants checked after every step: + +- both caches hold plain ``ParameterNodeAtInstant`` objects only, so neither + keeps a tracer; +- differential: both paths return the same plain node for an instant, traced + exactly when the tree is traced at the time of the read, with the tree's + tracer and branch name at that time; +- a traced read records in that tracer and in no other, and tracing never + changes a value read; +- once tracing stops and the caller drops its tracers, none is alive. + +``test_parameter_tracing_cache.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +import gc +import weakref + +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.country_template.entities import entities +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant + +# The parameter's value from each start date on: the reference the reads are +# checked against. +X_VALUES = {"2010-01-01": 1, "2015-01-01": 2, "2020-01-01": 3} +DATES = ["2012-06-01", "2015-01-01", "2017-01-01", "2020-01-01", "2024-12-31"] +TRACERS = 3 + + +def expected_x(date: str) -> int: + return X_VALUES[max(start for start in X_VALUES if start <= date)] + + +def system() -> TaxBenefitSystem: + tax_benefit_system = TaxBenefitSystem(entities) + tax_benefit_system.parameters = ParameterNode( + "", data={"x": {"values": dict(X_VALUES)}} + ) + return tax_benefit_system + + +steps = st.lists( + st.one_of( + st.tuples( + st.just("trace"), + st.integers(min_value=0, max_value=TRACERS - 1), + st.sampled_from(["default", "branch"]), + ), + st.tuples(st.just("untrace")), + st.tuples(st.sampled_from(["system", "tree"]), st.sampled_from(DATES)), + ), + max_size=30, +) + + +def check_caches(tax_benefit_system: TaxBenefitSystem) -> None: + for cache in ( + tax_benefit_system._parameters_at_instant_cache, + tax_benefit_system.parameters._at_instant_cache, + ): + assert all(type(node) is ParameterNodeAtInstant for node in cache.values()) + + +def plain(node): + if isinstance(node, TracingParameterNodeAtInstant): + return node.parameter_node_at_instant + return node + + +@hypothesis.given(sequence=steps) +# ``too_slow`` times input generation by the wall clock, so it fails on a busy +# host, not on anything this property checks. +@hypothesis.settings( + max_examples=300, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +def test_caches_stay_tracer_free_and_both_paths_agree(sequence): + tax_benefit_system = system() + parameters = tax_benefit_system.parameters + tracers = [FullTracer() for _ in range(TRACERS)] + for tracer in tracers: + tracer.record_calculation_start("y", "2017") + refs = [weakref.ref(tracer) for tracer in tracers] + current = None # (tracer index, branch name) while the tree is traced + reads = [0] * TRACERS + + for step in sequence: + if step[0] == "trace": + _, index, branch_name = step + parameters.trace = True + parameters.tracer = tracers[index] + parameters.branch_name = branch_name + current = (index, branch_name) + elif step[0] == "untrace": + parameters.trace, parameters.tracer, parameters.branch_name = ( + False, + None, + None, + ) + current = None + else: + path, date = step + read = ( + tax_benefit_system.get_parameters_at_instant(date) + if path == "system" + else parameters.get_at_instant(date) + ) + other = ( + parameters.get_at_instant(date) + if path == "system" + else tax_benefit_system.get_parameters_at_instant(date) + ) + for node in (read, other): + if current is None: + assert type(node) is ParameterNodeAtInstant + else: + index, branch_name = current + assert isinstance(node, TracingParameterNodeAtInstant) + assert node.tracer is tracers[index] + assert node.branch_name == branch_name + assert plain(read) is plain(other) + assert read.x == expected_x(date) + if current is not None: + reads[current[0]] += 1 + del read, other, node + check_caches(tax_benefit_system) + assert [len(tracer.trees[0].parameters) for tracer in tracers] == reads + + parameters.trace, parameters.tracer, parameters.branch_name = False, None, None + del tracers, tracer + gc.collect() + assert [ref for ref in refs if ref() is not None] == [] diff --git a/tests/core/test_parameters_at_instant_cache.py b/tests/core/test_parameters_at_instant_cache.py new file mode 100644 index 000000000..ecc63733c --- /dev/null +++ b/tests/core/test_parameters_at_instant_cache.py @@ -0,0 +1,47 @@ +"""``TaxBenefitSystem.get_parameters_at_instant`` follows changes to the tree. + +The system caches the parameters at each instant on top of the tree's own +cache. ``Parameter.update`` clears the tree's caches up to the root but not +the system's, which went on returning the node it cached before the change. +""" + +from policyengine_core.country_template import CountryTaxBenefitSystem + +INSTANT = "2015-01-01" + + +def test_repeated_requests_return_the_cached_node(): + system = CountryTaxBenefitSystem() + first = system.get_parameters_at_instant(INSTANT) + assert system.get_parameters_at_instant(INSTANT) is first + + +def test_a_request_after_an_update_reads_the_update(): + system = CountryTaxBenefitSystem() + before = system.get_parameters_at_instant(INSTANT) + assert before.taxes.income_tax_rate == 0.15 + + taxes = system.parameters.taxes + taxes.income_tax_rate.update(period="year:2015:1", value=0.99) + taxes.housing_tax.rate.update(period="year:2015:1", value=77) + + after = system.get_parameters_at_instant(INSTANT) + assert after is not before + # Read before the update and not: both as the tree now holds them. + assert (after.taxes.income_tax_rate, after.taxes.housing_tax.rate) == (0.99, 77) + assert system.get_parameters_at_instant(INSTANT) is after + + +def test_an_update_refreshes_every_cached_instant(): + """``update`` clears the root's cache at every instant, so the system + rebuilds each instant it had cached, not only the one the update names.""" + system = CountryTaxBenefitSystem() + other_instant = system.get_parameters_at_instant("2016-01-01") + system.get_parameters_at_instant(INSTANT) + + system.parameters.taxes.income_tax_rate.update(period="year:2015:1", value=0.99) + + assert system.get_parameters_at_instant("2016-01-01") is not other_instant + assert system.get_parameters_at_instant("2016-01-01").taxes.income_tax_rate == ( + system.parameters.taxes.income_tax_rate("2016-01-01") + ) diff --git a/tests/core/test_reforms.py b/tests/core/test_reforms.py index 711888b50..ecec9285c 100644 --- a/tests/core/test_reforms.py +++ b/tests/core/test_reforms.py @@ -3,7 +3,11 @@ import pytest from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.country_template.entities import Household, Person +from policyengine_core.country_template.reforms.modify_social_security_taxation import ( + modify_social_security_taxation, +) from policyengine_core.model_api import * from policyengine_core.parameters import ParameterNode from policyengine_core.periods import Instant @@ -299,6 +303,98 @@ def apply(self): assert parameters_at_instant.new_node.new_param is True +def test_modify_parameters_reads_its_changes_after_the_baseline_cached_them(): + """A modifier that changes values or brackets in place is read even where + the baseline had cached its parameters. The modifier used to get a deep + copy of the baseline's tree that kept those cached values.""" + + def contribution_scale(system): + return system.get_parameters_at_instant( + "2017-01-01" + ).taxes.social_security_contribution + + expected = contribution_scale( + modify_social_security_taxation(CountryTaxBenefitSystem()) + ) + read_baseline = CountryTaxBenefitSystem() + baseline_scale = contribution_scale(read_baseline) + assert baseline_scale.rates != expected.rates + + reformed = contribution_scale(modify_social_security_taxation(read_baseline)) + assert reformed.thresholds == expected.thresholds + assert reformed.rates == expected.rates + + +def _set_income_tax_rate_in_place(parameters): + """Changes values without ``Parameter.update``, which drops cached values.""" + for value_at_instant in parameters.taxes.income_tax_rate.values_list: + value_at_instant.value = 0.5 + return parameters + + +def _double_basic_income(parameters): + basic_income = parameters.benefits.basic_income + basic_income.update(period="year:2017:1", value=2 * basic_income("2017-01-01")) + return parameters + + +class set_income_tax_rate_in_place(Reform): + def apply(self): + self.modify_parameters(_set_income_tax_rate_in_place) + + +class update_then_two_modifiers(Reform): + def apply(self): + self.parameters.taxes.income_tax_rate.update(period="year:2017:1", value=0.2) + self.modify_parameters(_double_basic_income) + self.modify_parameters(_double_basic_income) + + +def _income_tax_rate_and_basic_income(system): + parameters = system.get_parameters_at_instant("2017-01-01") + return parameters.taxes.income_tax_rate, parameters.benefits.basic_income + + +def test_an_in_place_reform_reads_its_changes_after_the_system_cached_them(): + """``apply_reform_set`` runs ``apply`` on the system itself, whose tree + may already hold its values at an instant.""" + system = CountryTaxBenefitSystem() + assert _income_tax_rate_and_basic_income(system) == (0.15, 600) + + system.apply_reform_set(set_income_tax_rate_in_place) + + assert _income_tax_rate_and_basic_income(system) == (0.5, 600) + + +def test_modify_parameters_leaves_the_baseline_unchanged(): + baseline = CountryTaxBenefitSystem() + + reform = set_income_tax_rate_in_place(baseline) + + assert _income_tax_rate_and_basic_income(reform) == (0.5, 600) + assert _income_tax_rate_and_basic_income(baseline) == (0.15, 600) + + +def test_modify_parameters_keeps_what_apply_changed_before_it(): + """The modifier used to get a fresh copy of the baseline's tree, which + dropped the update made before it and the first modifier's change.""" + reform = update_then_two_modifiers(CountryTaxBenefitSystem()) + + assert _income_tax_rate_and_basic_income(reform) == (0.2, 2400) + + +def test_a_reform_instance_and_an_in_place_reform_give_the_same_parameters(): + """``Reform(baseline)`` and ``apply_reform_set`` run the same ``apply``.""" + in_place = CountryTaxBenefitSystem() + in_place.apply_reform_set(update_then_two_modifiers) + + reform = update_then_two_modifiers(CountryTaxBenefitSystem()) + + assert _income_tax_rate_and_basic_income( + reform + ) == _income_tax_rate_and_basic_income(in_place) + + def _income_tax_rate(reform, year): return reform.get_parameters_at_instant(Instant((year, 1, 1))).taxes.income_tax_rate diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py index 8c4f5e95a..d727e0f5a 100644 --- a/tests/core/test_tracing_parameter_isolation.py +++ b/tests/core/test_tracing_parameter_isolation.py @@ -16,7 +16,7 @@ - ``ParameterNode.clone`` copies ``trace`` and ``tracer``, so every clone of the system traced as well; - copying a ``TracingParameterNodeAtInstant`` (``copy``, ``deepcopy`` or - ``pickle``, as ``Reform.modify_parameters`` does when it deep-copies the + ``pickle``, as ``Reform.modify_parameters`` did when it deep-copied the tree) recursed without end, because ``__getattr__`` read ``parameter_node_at_instant`` from an instance the copy protocol had not filled in yet. @@ -81,7 +81,8 @@ def test_traced_calculation_leaves_the_system_untraced(isolated_tax_benefit_syst def test_reform_after_a_traced_calculation_can_copy_the_parameters( isolated_tax_benefit_system, ): - """``Reform.modify_parameters`` deep-copies the baseline's parameter tree.""" + """A reform built after a traced calculation reads its own parameters + and leaves the system's unchanged.""" system = isolated_tax_benefit_system simulation = build_simulation(system, trace=True) simulation.calculate("income_tax", JANUARY) @@ -358,23 +359,34 @@ def test_tracing_parameter_node_can_be_copied(isolated_tax_benefit_system, dupli ) -def test_parameters_holding_traced_nodes_can_be_deep_copied( +def test_traced_parameters_can_be_deep_copied_without_caching_wrappers( isolated_tax_benefit_system, ): - """A tree whose ``trace`` flag is set caches traced nodes; copying it works.""" + """A traced tree caches plain nodes and returns a wrapper for each read.""" parameters = isolated_tax_benefit_system.parameters parameters.trace = True parameters.tracer = FullTracer() - parameters(INSTANT) - assert type(parameters._at_instant_cache[INSTANT]) is ( - TracingParameterNodeAtInstant - ) + traced = parameters(INSTANT) + cached = parameters._at_instant_cache[INSTANT] + assert type(cached) is ParameterNodeAtInstant + assert type(traced) is TracingParameterNodeAtInstant + assert traced.parameter_node_at_instant is cached + assert traced.tracer is parameters.tracer + assert parameters(INSTANT) is not traced duplicated = copy.deepcopy(parameters) - - assert duplicated(INSTANT).taxes.income_tax_rate == ( - parameters(INSTANT).taxes.income_tax_rate - ) + copied_read = duplicated(INSTANT) + + assert type(copied_read) is TracingParameterNodeAtInstant + assert copied_read.tracer is duplicated.tracer + assert copied_read.tracer is not parameters.tracer + assert copied_read.parameter_node_at_instant is not cached + assert copied_read.taxes.income_tax_rate == traced.taxes.income_tax_rate + for tree in (parameters, duplicated): + assert all( + type(node) is ParameterNodeAtInstant + for node in tree._at_instant_cache.values() + ) class _SlottedParameterNode(ParameterNode): diff --git a/tests/core/tools/test_runner/test_forked_test_runner.py b/tests/core/tools/test_runner/test_forked_test_runner.py new file mode 100644 index 000000000..b9695d38f --- /dev/null +++ b/tests/core/tools/test_runner/test_forked_test_runner.py @@ -0,0 +1,147 @@ +"""Batches run in children forked from the process that built the system.""" + +import os +import time + +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.tools import forked_test_runner +from policyengine_core.tools.forked_test_runner import ( + PYTEST_INTERNAL_ERROR, + run_batches_in_forks, +) +from policyengine_core.tools.test_runner import _tax_benefit_system_cache +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV + +pytestmark = [ + pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork"), + # The test process has native threads it did not start (the BLAS + # library's pool, which Python counts), so every fork here warns. The + # children only run pytest on the country template. + pytest.mark.filterwarnings( + "ignore:This process .* is multi-threaded:DeprecationWarning" + ), +] + +OPTIONS = {} + + +@pytest.fixture(autouse=True) +def nested_pytest_arguments(monkeypatch): + """Run each child's session with the memory tests' arguments.""" + main = forked_test_runner.run_tests.__globals__["pytest"].main + + def nested_main(args, plugins): + return main( + [*args, "-q", "-p", "no:rerunfailures", "-p", "no:cacheprovider"], + plugins=plugins, + ) + + monkeypatch.setattr( + forked_test_runner.run_tests.__globals__["pytest"], "main", nested_main + ) + + +def _write_case(path, expected_income_tax): + path.write_text( + yaml.safe_dump( + [ + { + "name": path.stem, + "period": "2017-01", + "input": {"salary": 1000}, + "output": {"income_tax": expected_income_tax}, + } + ] + ) + ) + return str(path) + + +def test_each_batch_reports_its_own_exit_code_and_output(tmp_path): + passing = _write_case(tmp_path / "passing.yaml", 150) + failing = _write_case(tmp_path / "failing.yaml", 999) + finished = [] + + results = run_batches_in_forks( + CountryTaxBenefitSystem(), + [[passing], [failing], [passing, failing]], + OPTIONS, + workers=2, + on_finish=lambda index, result: finished.append(index), + ) + + assert [result.exit_code for result in results] == [0, 1, 1] + assert "1 passed" in results[0].output + assert "1 failed" in results[1].output + assert "1 failed, 1 passed" in results[2].output + assert sorted(finished) == [0, 1, 2] + assert all(result.elapsed > 0 for result in results) + assert not any(result.timed_out for result in results) + + +def test_the_reform_free_system_is_built_once_before_forking(tmp_path): + system = CountryTaxBenefitSystem() + run_batches_in_forks(system, [[_write_case(tmp_path / "case.yaml", 150)]]) + # Built in the parent, where every child shares it. + assert _tax_benefit_system_cache[system].reform_free is not None + + +def test_no_more_than_workers_batches_run_at_once(tmp_path, monkeypatch): + def record_while_running(tax_benefit_system, paths, options): + (marker,) = paths + with open(marker, "a") as log: + log.write(f"start {time.monotonic()}\n") + time.sleep(0.5) + with open(marker, "a") as log: + log.write(f"end {time.monotonic()}\n") + return 0 + + monkeypatch.setattr(forked_test_runner, "run_tests", record_while_running) + markers = [str(tmp_path / f"batch-{index}.log") for index in range(5)] + results = run_batches_in_forks( + CountryTaxBenefitSystem(), [[marker] for marker in markers], workers=2 + ) + + assert [result.exit_code for result in results] == [0] * 5 + events = [] + for marker in markers: + with open(marker) as log: + for line in log: + kind, moment = line.split() + events.append((float(moment), 1 if kind == "start" else -1)) + running = most = 0 + for _, change in sorted(events): + running += change + most = max(most, running) + assert most == 2 + + +def test_a_batch_over_its_timeout_is_stopped(monkeypatch): + monkeypatch.setattr( + forked_test_runner, "run_tests", lambda *args: time.sleep(60) or 0 + ) + started = time.monotonic() + (result,) = run_batches_in_forks( + CountryTaxBenefitSystem(), [["anything"]], timeout=0.5 + ) + assert result.timed_out + assert result.exit_code < 0 + assert time.monotonic() - started < 30 + + +def test_an_error_outside_pytest_is_reported_as_an_internal_error(monkeypatch): + def fail(*args): + raise RuntimeError("the batch could not start") + + monkeypatch.setattr(forked_test_runner, "run_tests", fail) + (result,) = run_batches_in_forks(CountryTaxBenefitSystem(), [["anything"]]) + assert result.exit_code == PYTEST_INTERNAL_ERROR + assert "the batch could not start" in result.output + + +def test_workers_must_be_positive(): + with pytest.raises(ValueError, match="workers must be 1 or more"): + run_batches_in_forks(CountryTaxBenefitSystem(), [], workers=0) diff --git a/tests/core/tools/test_runner/test_parameter_inputs.py b/tests/core/tools/test_runner/test_parameter_inputs.py new file mode 100644 index 000000000..833a034e4 --- /dev/null +++ b/tests/core/tools/test_runner/test_parameter_inputs.py @@ -0,0 +1,149 @@ +"""A case's dotted parameter inputs run on a system that copies only the +parameters they set, on top of the system the case's reforms select.""" + +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.reforms import Reform, set_parameter +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.tools.test_runner import ( + OpenFiscaPlugin, + _ParameterInputReform, + _get_tax_benefit_system, + _tax_benefit_system_cache, +) +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV, write_cases + +SCALE_REFORM = ( + "policyengine_core.country_template.reforms." + "modify_social_security_taxation.modify_social_security_taxation" +) +PARAMETER_INPUTS = [ + ("taxes.income_tax_rate", 0.3), + ("taxes.housing_tax.minimal_amount", 350), + ("taxes.social_security_contribution[1].rate", 0.5), +] + + +def _full_clone_system(baseline, reforms, parameter_inputs): + """The system the runner built before: a full clone, the named reforms, + then a reform that sets the inputs on a full copy of the tree.""" + system = baseline.clone() + for reform in reforms: + system = system.apply_reform(reform) + + class inline_reform_class(Reform): + def apply(self): + for key, value in parameter_inputs: + self.parameters = set_parameter( + key, value, return_modifier=True, period="year:2000:40" + )(self.parameters) + + return inline_reform_class(system) + + +def _run(path, baseline): + return pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], + plugins=[OpenFiscaPlugin(baseline, {})], + ) + + +def test_the_system_shares_every_parameter_it_does_not_set(): + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + system = _ParameterInputReform(reform_free, [("taxes.income_tax_rate", 0.3)]) + + assert system.baseline is reform_free + assert system.parameters.taxes.income_tax_rate("2017-01-01") == 0.3 + assert reform_free.parameters.taxes.income_tax_rate("2017-01-01") == 0.15 + assert system.parameters.benefits is reform_free.parameters.benefits + assert system.parameters.taxes.housing_tax is ( + reform_free.parameters.taxes.housing_tax + ) + + +@pytest.mark.parametrize("reforms", [[], [SCALE_REFORM]]) +@pytest.mark.parametrize("instant", ["2015-06-01", "2017-01-01", "2030-01-01"]) +def test_it_reads_what_a_full_clone_with_the_same_inputs_reads(reforms, instant): + """Including a bracket index, which counts the brackets the named reform + left: the inputs apply after the reforms, as they did on a full clone.""" + baseline = CountryTaxBenefitSystem() + structural = _get_tax_benefit_system(baseline, reforms, []) + system = _ParameterInputReform(structural, PARAMETER_INPUTS) + reference = _full_clone_system(baseline, reforms, PARAMETER_INPUTS) + + values = system.get_parameters_at_instant(instant) + expected = reference.get_parameters_at_instant(instant) + assert values.taxes.income_tax_rate == expected.taxes.income_tax_rate + assert ( + values.taxes.housing_tax.minimal_amount + == expected.taxes.housing_tax.minimal_amount + ) + scale = values.taxes.social_security_contribution + expected_scale = expected.taxes.social_security_contribution + assert list(scale.rates) == list(expected_scale.rates) + assert list(scale.thresholds) == list(expected_scale.thresholds) + + +def test_cases_give_the_same_results_as_on_a_full_clone(tmp_path): + baseline = CountryTaxBenefitSystem() + situation = {"salary": 50_000} + inputs = {**situation, **{key: value for key, value in PARAMETER_INPUTS}} + variables = ["income_tax", "social_security_contribution"] + cases = [] + for reforms in [[], [SCALE_REFORM]]: + reference = _full_clone_system(baseline, reforms, PARAMETER_INPUTS) + builder = SimulationBuilder() + builder.set_default_period("2017-01") + simulation = builder.build_from_dict(reference, dict(situation)) + expected = { + variable: float(simulation.calculate(variable, "2017-01")[0]) + for variable in variables + } + assert expected["income_tax"] == pytest.approx(15_000) + for variable in variables: + cases.append( + { + "name": f"{variable} with reforms {reforms}", + "period": "2017-01", + **({"reforms": reforms} if reforms else {}), + "input": inputs, + "output": {variable: expected[variable]}, + "absolute_error_margin": 0.01, + } + ) + # The reform changes the scale, so the two contributions must differ for + # this to show the inputs were applied on top of it. + assert cases[1]["output"] != cases[3]["output"] + path = tmp_path / "cases.yaml" + path.write_text(yaml.safe_dump(cases, sort_keys=False)) + assert _run(path, baseline) == 0 + + +def test_cases_with_parameter_inputs_add_nothing_to_the_reform_cache(tmp_path): + baseline = CountryTaxBenefitSystem() + path = write_cases(tmp_path / "cases.yaml", 30, reform_every=1) + assert _run(path, baseline) == 0 + cached = _tax_benefit_system_cache[baseline] + assert cached.reform_free is not None + assert len(cached.reforms) == 0 + + +def test_an_unknown_parameter_input_fails_its_case(tmp_path, capsys): + path = tmp_path / "unknown.yaml" + path.write_text( + yaml.safe_dump( + [ + { + "name": "unknown", + "period": "2017-01", + "input": {"salary": 1000, "taxes.no_such_tax": 1}, + "output": {"income_tax": 150}, + } + ] + ) + ) + assert _run(path, CountryTaxBenefitSystem()) == 1 + assert "Could not find the parameter taxes.no_such_tax" in capsys.readouterr().out diff --git a/tests/core/tools/test_runner/test_read_only_arrays.py b/tests/core/tools/test_runner/test_read_only_arrays.py new file mode 100644 index 000000000..42eda2c3a --- /dev/null +++ b/tests/core/tools/test_runner/test_read_only_arrays.py @@ -0,0 +1,148 @@ +"""With ``read_only_arrays``, a formula that writes into an array it read from +another variable fails at that line instead of changing that variable.""" + +import numpy as np +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.periods import MONTH, period +from policyengine_core.tools.test_runner import OpenFiscaPlugin, run_tests +from policyengine_core.variables import Variable +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV + +JANUARY = period("2017-01") + + +class salary_plus_one(Variable): + """Writes into the salary array it read: changes ``salary`` itself.""" + + value_type = float + entity = Person + label = "Salary plus one, written into the salary array" + definition_period = MONTH + + def formula(person, period, parameters): + salary = person("salary", period) + salary += 1 + return salary + + +@pytest.fixture +def read_only(monkeypatch): + monkeypatch.setattr(InMemoryStorage, "read_only_arrays", True) + + +def test_stored_arrays_stay_writeable_by_default(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0, 2.0]), JANUARY) + stored = storage.get(JANUARY) + stored += 1 + assert list(storage.get(JANUARY)) == [2.0, 3.0] + + +def test_stored_arrays_are_read_only_and_the_callers_array_is_not(read_only): + storage = InMemoryStorage(is_eternal=False) + given = np.array([1.0, 2.0]) + storage.put(given, JANUARY) + stored = storage.get(JANUARY) + + with pytest.raises(ValueError, match="read-only"): + stored += 1 + assert given.flags.writeable + assert list(storage.get(JANUARY)) == [1.0, 2.0] + + +def test_an_array_a_clone_copies_on_first_read_is_read_only(read_only): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0, 2.0]), JANUARY) + clone = storage.clone(share_arrays=True) + + copied = clone.get(JANUARY) + with pytest.raises(ValueError, match="read-only"): + copied[0] = 5 + assert list(storage.get(JANUARY)) == [1.0, 2.0] + + +def _system_with_a_writing_formula(): + system = CountryTaxBenefitSystem() + system.add_variable(salary_plus_one) + return system + + +def _cases(tmp_path): + path = tmp_path / "writes.yaml" + path.write_text( + yaml.safe_dump( + [ + { + "name": "writes into salary", + "period": "2017-01", + "input": {"salary": 1000}, + "output": {"salary_plus_one": 1001}, + } + ] + ) + ) + return path + + +def test_a_formula_writing_into_a_cached_array_passes_without_the_option(tmp_path): + exit_code = pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(_cases(tmp_path))], + plugins=[OpenFiscaPlugin(_system_with_a_writing_formula(), {})], + ) + assert exit_code == 0 + + +def test_the_option_fails_a_formula_writing_into_a_cached_array( + tmp_path, capsys, monkeypatch +): + monkeypatch.setattr( + "policyengine_core.tools.test_runner.pytest.main", + _main_without_plugins_from_outside(), + ) + exit_code = run_tests( + _system_with_a_writing_formula(), + str(_cases(tmp_path)), + {"read_only_arrays": True}, + ) + assert exit_code == 1 + assert "read-only" in capsys.readouterr().out + # The run restores the setting it found. + assert InMemoryStorage.read_only_arrays is False + + +def _main_without_plugins_from_outside(): + """``pytest.main`` with the arguments the memory tests use, so the nested + session does not load this session's rerun plugin.""" + main = pytest.main + + def nested_main(args, plugins): + return main( + [*args, "-p", "no:rerunfailures", "-p", "no:cacheprovider"], plugins=plugins + ) + + return nested_main + + +def test_command_line_passes_read_only_arrays(monkeypatch): + from policyengine_core.scripts import policyengine_command, run_test + + seen = {} + + def fake_run_tests(tax_benefit_system, paths, options): + seen.update(options) + return 0 + + monkeypatch.setattr(run_test, "run_tests", fake_run_tests) + monkeypatch.setattr(run_test, "build_tax_benefit_system", lambda *args: object()) + for arguments, expected in [([], False), (["--read-only-arrays"], True)]: + monkeypatch.setattr( + "sys.argv", ["policyengine-core", "test", "cases.yaml", *arguments] + ) + with pytest.raises(SystemExit): + run_test.main(policyengine_command.get_parser()) + assert seen["read_only_arrays"] is expected diff --git a/tests/core/tools/test_runner/test_reform_cache_property.py b/tests/core/tools/test_runner/test_reform_cache_property.py new file mode 100644 index 000000000..3ce2cdccd --- /dev/null +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -0,0 +1,108 @@ +"""The YAML runner's reform-system cache against a reference LRU. + +For any sequence of requests, each with its own cache size, the cache holds +at most that call's ``cache_size`` reform systems after the call, returns a cached object exactly when a plain +LRU of that size would, always hands back a system built for the reforms and +extensions requested, and keeps one reform-free system throughout, which no +reform or extension changes although reforms are built on it. +``test_runner_memory.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +from collections import OrderedDict + +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.tools.test_runner import ( + _get_tax_benefit_system, + _tax_benefit_system_cache, +) + + +class StubSystem: + """Records what the runner applied to it, in order.""" + + def __init__(self, applied=()): + self.applied = tuple(applied) + + def clone(self): + return StubSystem(self.applied) + + def apply_reform(self, path): + return StubSystem(self.applied + (("reform", path),)) + + def load_extension(self, extension): + self.applied += (("extension", extension),) + + +requests = st.tuples( + st.lists(st.sampled_from(["a", "b", "c"]), max_size=2), + st.lists(st.sampled_from(["x", "y"]), max_size=2, unique=True), + st.sampled_from(["", "p:1", "p:2"]), + # Each call's bound; runs on one baseline can ask for different ones. + st.integers(min_value=0, max_value=3), +) + + +@hypothesis.given(sequence=st.lists(requests, max_size=40)) +# ``too_slow`` times input generation by the wall clock, so it fails on a busy +# host, not on anything this property checks. +@hypothesis.settings( + max_examples=300, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +def test_cache_matches_reference_lru(sequence): + baseline = StubSystem() + reference = OrderedDict() + last_returned = {} + reform_free = None + + for reforms, extensions, reform_key, cache_size in sequence: + # The reference applies each call's bound first, oldest out first. + while len(reference) > cache_size: + reference.popitem(last=False) + system = _get_tax_benefit_system( + baseline, + list(reforms), + list(extensions), + reform_key=reform_key, + cache_size=cache_size, + ) + assert system is not baseline + # Reforms apply in the order given; extension order is not part of + # the cache key, so only the set is fixed. + applied_reforms = [name for kind, name in system.applied if kind == "reform"] + applied_extensions = { + name for kind, name in system.applied if kind == "extension" + } + assert applied_reforms == list(reforms) + assert applied_extensions == set(extensions) + + cached = _tax_benefit_system_cache[baseline] + if reform_free is None: + reform_free = cached.reform_free + assert cached.reform_free is reform_free + assert reform_free.applied == () + if not reforms and not extensions and not reform_key: + assert system is reform_free + else: + key = (tuple(reforms), reform_key, frozenset(extensions)) + if key in reference: + reference.move_to_end(key) + assert system is last_returned[key] + else: + assert system is not last_returned.get(key) + if cache_size: + reference[key] = True + while len(reference) > cache_size: + reference.popitem(last=False) + last_returned[key] = system + + assert len(cached.reforms) <= cache_size + assert list(cached.reforms) == list(reference) diff --git a/tests/core/tools/test_runner/test_reform_system_size.py b/tests/core/tools/test_runner/test_reform_system_size.py new file mode 100644 index 000000000..f3843e9ba --- /dev/null +++ b/tests/core/tools/test_runner/test_reform_system_size.py @@ -0,0 +1,60 @@ +"""A reform system holds one copy of the parameter tree, and the cache never +holds more reform systems than its size, even while it builds one.""" + +import gc +import weakref + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.reforms import Reform +from policyengine_core.tools.test_runner import _get_tax_benefit_system + +NEW_TAX = "policyengine_core.country_template.reforms.add_new_tax.add_new_tax" + + +def test_a_named_reform_is_built_on_the_reform_free_system_and_leaves_it_alone(): + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + variables = dict(reform_free.variables) + parameters = reform_free.parameters + entities = list(reform_free.entities) + + system = _get_tax_benefit_system(baseline, [NEW_TAX], []) + + assert system.baseline is reform_free + assert system.parameters is not parameters + assert "new_tax" in system.variables + assert "new_tax" not in reform_free.variables + assert reform_free.variables == variables + assert reform_free.parameters is parameters + assert reform_free.entities == entities + assert all( + entity._tax_benefit_system is reform_free for entity in reform_free.entities + ) + + +def test_a_reform_function_is_applied_to_a_clone(): + """Only reforms that copy what they change are built on the shared system.""" + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + received = [] + + def reform_function(system): + received.append(system) + return system + + _get_tax_benefit_system(baseline, [reform_function], [], reform_key="function") + assert received and received[0] is not reform_free + + +def test_the_cache_releases_a_system_before_building_the_one_replacing_it(): + baseline = CountryTaxBenefitSystem() + first = weakref.ref(_get_tax_benefit_system(baseline, [NEW_TAX], [], cache_size=1)) + alive_while_building = [] + + class probe(Reform): + def apply(self): + gc.collect() + alive_while_building.append(first() is not None) + + _get_tax_benefit_system(baseline, [probe], [], cache_size=1) + assert alive_while_building == [False] diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py new file mode 100644 index 000000000..fd4795505 --- /dev/null +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -0,0 +1,372 @@ +"""The YAML runner must not keep per-case state alive across a run. + +pytest holds every collected item until the session ends. When each +``YamlItem`` kept its simulation, and the runner kept every reform system it +ever built, memory grew with every case: a policyengine-us run of ~1,500 +files in one process reached 118 GB. +""" + +import gc +import weakref + +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.reforms import Reform, set_parameter +from policyengine_core.tools import test_runner +from policyengine_core.tools.test_runner import ( + DEFAULT_REFORM_CACHE_SIZE, + _get_tax_benefit_system, + _tax_benefit_system_cache, +) +from policyengine_core.tools.test_runner import OpenFiscaPlugin +from tests.fixtures.yaml_runner_memory import ( + RUNNER_ARGV, + mark_parameter_trees_traced, + record_results, + run_with_probe, + system_reading_its_parameters, + write_cases, + write_cases_at_distinct_dates, +) +from tests.fixtures.tracing import build_simulation + + +def _rate_reform(rate): + """An inline reform class like the one the runner builds for a dotted + parameter input, and the ``reform_key`` it would get.""" + modifier = set_parameter( + "taxes.income_tax_rate", rate, return_modifier=True, period="year:2000:40" + ) + + class inline_reform_class(Reform): + def apply(self): + self.parameters = modifier(self.parameters) + + return [inline_reform_class], f"taxes.income_tax_rate:{rate}" + + +def _system_for_rate(baseline, rate, cache_size=None): + reforms, key = _rate_reform(rate) + return _get_tax_benefit_system( + baseline, reforms, [], reform_key=key, cache_size=cache_size + ) + + +def _rate(system): + return system.parameters.taxes.income_tax_rate("2017-01-01") + + +def test_reform_cache_keeps_at_most_cache_size_reform_systems(): + baseline = CountryTaxBenefitSystem() + refs = [weakref.ref(_system_for_rate(baseline, 0.2 + i / 100)) for i in range(8)] + gc.collect() + + cached = _tax_benefit_system_cache[baseline].reforms + assert len(cached) == DEFAULT_REFORM_CACHE_SIZE + # Evicted systems are freed, not just dropped from the cache. + alive = [ref() for ref in refs if ref() is not None] + assert len(alive) == DEFAULT_REFORM_CACHE_SIZE + assert set(map(id, alive)) == set(map(id, cached.values())) + + +def test_reform_cache_evicts_least_recently_used(): + baseline = CountryTaxBenefitSystem() + a = _system_for_rate(baseline, 0.21, cache_size=2) + b = _system_for_rate(baseline, 0.22, cache_size=2) + assert _system_for_rate(baseline, 0.21, cache_size=2) is a # hit; A is newest + _system_for_rate(baseline, 0.23, cache_size=2) # evicts B, not A + assert _system_for_rate(baseline, 0.21, cache_size=2) is a + rebuilt_b = _system_for_rate(baseline, 0.22, cache_size=2) + assert rebuilt_b is not b + assert _rate(rebuilt_b) == pytest.approx(0.22) + + +def test_reform_free_system_is_kept_through_evictions(): + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], [], reform_key="") + for i in range(6): + _system_for_rate(baseline, 0.3 + i / 100, cache_size=1) + assert _get_tax_benefit_system(baseline, [], [], reform_key="") is reform_free + assert _get_tax_benefit_system(baseline, [], []) is reform_free + assert reform_free is not baseline + assert _rate(reform_free) == pytest.approx(0.15) + assert _rate(baseline) == pytest.approx(0.15) + + +def test_cache_size_zero_builds_a_fresh_reform_system_each_time(): + baseline = CountryTaxBenefitSystem() + first = _system_for_rate(baseline, 0.25, cache_size=0) + second = _system_for_rate(baseline, 0.25, cache_size=0) + assert first is not second + assert _rate(first) == _rate(second) == pytest.approx(0.25) + assert len(_tax_benefit_system_cache[baseline].reforms) == 0 + + +def test_a_smaller_cache_size_applies_to_what_is_already_cached(): + baseline = CountryTaxBenefitSystem() + first = _system_for_rate(baseline, 0.21, cache_size=3) + _system_for_rate(baseline, 0.22, cache_size=3) + _system_for_rate(baseline, 0.23, cache_size=3) + # A cache hit under a smaller bound still trims, least recently used first. + assert _system_for_rate(baseline, 0.23, cache_size=1) is not None + assert len(_tax_benefit_system_cache[baseline].reforms) == 1 + # Any request applies its bound, a reform-free one too; 0 empties it. + _get_tax_benefit_system(baseline, [], [], cache_size=0) + assert len(_tax_benefit_system_cache[baseline].reforms) == 0 + assert _system_for_rate(baseline, 0.21, cache_size=0) is not first + + +def test_class_reforms_without_a_key_are_cached_apart(): + baseline = CountryTaxBenefitSystem() + (reform_a,), _ = _rate_reform(0.21) + (reform_b,), _ = _rate_reform(0.22) + system_a = _get_tax_benefit_system(baseline, [reform_a], []) + system_b = _get_tax_benefit_system(baseline, [reform_b], []) + assert _rate(system_a) == pytest.approx(0.21) + assert _rate(system_b) == pytest.approx(0.22) + assert _get_tax_benefit_system(baseline, [reform_a], []) is system_a + + +def test_negative_cache_size_is_rejected(): + with pytest.raises(ValueError): + _system_for_rate(CountryTaxBenefitSystem(), 0.25, cache_size=-1) + + +def test_finished_cases_release_their_simulations(tmp_path): + path = write_cases(tmp_path / "cases.yaml", 60) + probe = run_with_probe(CountryTaxBenefitSystem(), path) + assert len(probe.outcomes) == 60 + # Every item is still referenced by the session here; none may keep the + # simulation it ran. + assert probe.live_simulations == 0 + + +def test_live_systems_do_not_grow_with_the_number_of_cases(tmp_path): + short = run_with_probe( + CountryTaxBenefitSystem(), write_cases(tmp_path / "short.yaml", 12) + ) + long = run_with_probe( + CountryTaxBenefitSystem(), write_cases(tmp_path / "long.yaml", 90) + ) + # 4 vs 30 distinct reform systems requested; the live count must not move. + assert long.live_systems == short.live_systems + + +def test_traced_cases_release_their_tracers(tmp_path): + """Verbose (traced) cases at different dates keep no tracer: the cached + system's parameters used to cache one tracing wrapper, holding that case's + tracer and results, per date.""" + cases = [ + { + "name": f"case {i}", + "period": f"20{15 + i // 12}-{i % 12 + 1:02d}", + "input": {"salary": 1000}, + "output": {"income_tax": 150}, + } + for i in range(16) + ] + path = tmp_path / "traced.yaml" + path.write_text(yaml.safe_dump(cases)) + system = CountryTaxBenefitSystem() + probe = run_with_probe(system, path, {"verbose": True}) + assert probe.live_tracers == 0 + assert probe.live_simulations == 0 + reform_free = _tax_benefit_system_cache[system].reform_free + assert reform_free.parameters.trace is False + assert reform_free.parameters.tracer is None + + +def _run_cases_reading_the_systems_parameters(tmp_path, monkeypatch): + """Run 16 verbose cases at distinct dates whose formula reads + ``TaxBenefitSystem.get_parameters_at_instant``; return the probe, the + weak references to every result their tracers recorded, and the system + they ran on.""" + results = record_results(monkeypatch) + system = system_reading_its_parameters() + path = write_cases_at_distinct_dates(tmp_path / "traced.yaml", 16) + probe = run_with_probe(system, path, {"verbose": True}) + assert len(probe.outcomes) == 16 + assert results + return probe, results, _tax_benefit_system_cache[system].reform_free + + +def test_traced_cases_reading_the_systems_parameters_release_their_tracers( + tmp_path, monkeypatch +): + """The system caches the parameter tree at each instant it is asked for. + While core's simulations marked the tree traced, that cache kept a + tracing wrapper, with its case's tracer and results, for every date (16 + tracers and 32 result arrays after 16 dates).""" + probe, results, reform_free = _run_cases_reading_the_systems_parameters( + tmp_path, monkeypatch + ) + assert probe.live_tracers == 0 + assert [ref for ref in results if ref() is not None] == [] + assert probe.live_simulations == 0 + assert len(reform_free._parameters_at_instant_cache) == 16 + + +def test_cases_on_a_tree_marked_traced_release_their_tracers(tmp_path, monkeypatch): + """The same cases, with each simulation marking its system's parameter + tree traced, as code outside core can: the system's cache must still + keep only plain nodes, and teardown unmarks the tree.""" + mark_parameter_trees_traced(monkeypatch) + probe, results, reform_free = _run_cases_reading_the_systems_parameters( + tmp_path, monkeypatch + ) + assert probe.live_tracers == 0 + assert [ref for ref in results if ref() is not None] == [] + assert probe.live_simulations == 0 + assert reform_free.parameters.trace is False + assert reform_free.parameters.tracer is None + + +@pytest.mark.parametrize( + "invalid_input", + [ + pytest.param({"salary": "not a number"}, id="variable-value"), + pytest.param( + {"persons": {"person": {"salary": {"2017-01": "not a number"}}}}, + id="entity-value", + ), + pytest.param({"unknown_variable": 1}, id="unknown-variable"), + ], +) +@pytest.mark.parametrize("existing_simulation", [False, True]) +def test_cases_whose_simulation_fails_to_build_release_it( + tmp_path, invalid_input, existing_simulation +): + """Both builder paths release failed cases and restore any prior backlink.""" + cases = [ + { + "name": f"bad input {i}", + "period": "2017-01", + "input": invalid_input, + "output": {"income_tax": 0}, + } + for i in range(8) + ] + path = tmp_path / "bad.yaml" + path.write_text(yaml.safe_dump(cases)) + system = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(system, [], []) + previous = build_simulation(reform_free) if existing_simulation else None + probe = run_with_probe(system, path, outcome="failed") + assert len(probe.outcomes) == len(probe.items) == 8 + assert probe.live_simulations == 0 + assert getattr(reform_free, "simulation", None) is previous + for item in probe.items: + assert item.simulation is None + assert item.tax_benefit_system is None + assert item._system_simulation is None + + +def test_successful_cases_restore_the_systems_previous_simulation(tmp_path): + """Successful and failed builds have the same ownership after teardown.""" + system = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(system, [], []) + previous = build_simulation(reform_free) + path = write_cases(tmp_path / "good.yaml", 3, reform_every=0) + + probe = run_with_probe(system, path) + + assert len(probe.outcomes) == len(probe.items) == 3 + assert probe.live_simulations == 0 + assert reform_free.simulation is previous + for item in probe.items: + assert item.simulation is None + assert item.tax_benefit_system is None + assert item._system_simulation is None + + +def test_inline_parameter_values_key_by_type(tmp_path): + """A number and the same digits as a string are different inputs: the + string must not reuse the system built for the number.""" + base = {"period": "2017-01", "input": {"salary": 1000}} + cases = [ + { + **base, + "name": "number", + "input": {**base["input"], "taxes.income_tax_rate": 0.2}, + "output": {"income_tax": 200}, + }, + { + **base, + "name": "string", + "input": {**base["input"], "taxes.income_tax_rate": "0.2"}, + "output": {"income_tax": 200}, + }, + ] + path = tmp_path / "typed.yaml" + path.write_text(yaml.safe_dump(cases)) + exit_code = pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], + plugins=[OpenFiscaPlugin(CountryTaxBenefitSystem(), {})], + ) + # With a value-only key the string case passed on the number's system. + assert exit_code == 1 + + +@pytest.mark.parametrize("reform_every", [0, 3]) +def test_traced_memory_stays_flat_over_hundreds_of_cases(tmp_path, reform_every): + """A few hundred cases in one process: memory after the last case is + within a small bound of memory after the first hundred.""" + path = write_cases(tmp_path / "cases.yaml", 300, reform_every=reform_every) + probe = run_with_probe(CountryTaxBenefitSystem(), path) + assert len(probe.outcomes) == 300 + # One sample every 10 cases: index 9 is after case 100. + growth = max(probe.traced[10:]) - probe.traced[9] + # Measured over these 200 cases: 0.5 MB with the fix (pytest's own + # per-case reports); 4.0 MB (plain cases) and 18.4 MB (every third case a + # reform) without it. + assert growth < 2_000_000, f"traced memory grew {growth / 1e6:.1f} MB" + + +def test_reform_cache_size_option_reaches_the_cache(tmp_path, monkeypatch): + seen = [] + original = test_runner._get_tax_benefit_system + + def spy(*args, **kwargs): + seen.append(kwargs.get("cache_size")) + return original(*args, **kwargs) + + monkeypatch.setattr(test_runner, "_get_tax_benefit_system", spy) + path = write_cases(tmp_path / "cases.yaml", 3) + run_with_probe(CountryTaxBenefitSystem(), path, {"reform_cache_size": 0}) + assert seen == [0, 0, 0] + + +def test_command_line_rejects_a_negative_reform_cache_size(capsys): + from policyengine_core.scripts import policyengine_command + + with pytest.raises(SystemExit) as exit_info: + policyengine_command.get_parser().parse_args( + ["test", "cases.yaml", "--reform-cache-size", "-1"] + ) + assert exit_info.value.code == 2 + assert "must be 0 or more" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "arguments, expected", [([], None), (["--reform-cache-size", "3"], 3)] +) +def test_command_line_passes_reform_cache_size(monkeypatch, arguments, expected): + from policyengine_core.scripts import policyengine_command, run_test + + seen = {} + + def fake_run_tests(tax_benefit_system, paths, options): + seen.update(options) + return 0 + + monkeypatch.setattr(run_test, "run_tests", fake_run_tests) + monkeypatch.setattr(run_test, "build_tax_benefit_system", lambda *args: object()) + monkeypatch.setattr( + "sys.argv", ["policyengine-core", "test", "cases.yaml", *arguments] + ) + with pytest.raises(SystemExit) as exit_info: + run_test.main(policyengine_command.get_parser()) + assert exit_info.value.code == 0 + assert seen["reform_cache_size"] == expected diff --git a/tests/fixtures/parameter_nodes_at_instant.py b/tests/fixtures/parameter_nodes_at_instant.py new file mode 100644 index 000000000..7b991041b --- /dev/null +++ b/tests/fixtures/parameter_nodes_at_instant.py @@ -0,0 +1,131 @@ +"""A parameter tree and an up-front reference for the at-instant tests. + +``snapshot`` computes what a ``ParameterNodeAtInstant`` held when it built +every descendant in its constructor, straight from the parameter tree and +without using ``ParameterNodeAtInstant``. ``materialise`` reads the same +thing out of a ``ParameterNodeAtInstant``. +""" + +import numpy as np + +from policyengine_core.parameters import ( + Parameter, + ParameterNode, + ParameterNodeAtInstant, + ParameterScale, +) + +INSTANTS = [ + "2009-12-31", + "2010-01-01", + "2014-06-30", + "2015-01-01", + "2017-03-01", + "2018-01-01", + "2020-01-01", + "2029-12-31", + "2030-01-01", + "2035-07-01", +] + + +def _zone(amount, rate): + return { + "amount": {"values": {"2010-01-01": amount, "2020-01-01": amount * 2}}, + "rate": {"values": {"2010-01-01": rate}}, + } + + +TREE = { + "flat_rate": {"values": {"2015-01-01": 0.1, "2020-01-01": 0.2}}, + # Not defined before 2030, so absent from earlier instants. + "late": {"values": {"2030-01-01": 5}}, + "group": { + "x": {"values": {"2010-01-01": 0.1}}, + "y": {"values": {"2018-01-01": 0.2}}, + # Every child undefined before 2018: an empty node, not an absent one. + "inner": { + "1": {"values": {"2018-01-01": 10}}, + "2": {"values": {"2018-01-01": 20, "2020-01-01": 30}}, + }, + }, + "scale": { + "brackets": [ + { + "threshold": {"2010-01-01": {"value": 0}}, + "rate": {"2010-01-01": {"value": 0.1}}, + }, + { + "threshold": {"2010-01-01": {"value": 1000}}, + "rate": { + "2010-01-01": {"value": 0.2}, + "2020-01-01": {"value": 0.3}, + }, + }, + ] + }, + # Homogeneous children: vectorisable with ``node[array]``. + "by_zone": {"z1": _zone(100, 0.1), "z2": _zone(200, 0.2), "z3": _zone(300, 0.3)}, +} + + +def build_tree() -> ParameterNode: + return ParameterNode("", data=TREE) + + +def describe(value): + """A comparable form of one at-instant value.""" + if isinstance(value, (int, float, bool, str)) or value is None: + return value + if isinstance(value, (list, tuple)): + return tuple(describe(item) for item in value) + if isinstance(value, dict): + return {key: describe(item) for key, item in value.items()} + if hasattr(value, "thresholds"): + # Tax scales refuse ``==``. + return ( + type(value).__name__, + tuple(value.thresholds), + tuple(getattr(value, "rates", getattr(value, "amounts", ()))), + ) + if isinstance(value, np.ndarray): + return ("ndarray", str(value.dtype), value.tolist()) + if isinstance(value, np.generic): + return value.item() + return (type(value).__name__, repr(value)) + + +def snapshot(node: ParameterNode, instant_str: str) -> dict: + """Every descendant's value at the instant, in the node's own order, + leaving out children that have no value yet.""" + result = {} + for name, child in node.children.items(): + if isinstance(child, ParameterScale): + value = describe(child._get_at_instant(instant_str)) + elif isinstance(child, Parameter): + value = child._get_at_instant(instant_str) + else: + value = snapshot(child, instant_str) + if value is not None: + result[name] = value + return result + + +def materialise(node_at_instant: ParameterNodeAtInstant) -> dict: + """The same nested mapping, read from a ``ParameterNodeAtInstant``.""" + return { + name: ( + materialise(child) + if isinstance(child, ParameterNodeAtInstant) + else describe(child) + ) + for name, child in node_at_instant._children.items() + } + + +def ordered(mapping: dict): + """A form that compares key order as well as content.""" + return [ + (key, ordered(value) if isinstance(value, dict) else value) + for key, value in mapping.items() + ] diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py new file mode 100644 index 000000000..fb6264a0d --- /dev/null +++ b/tests/fixtures/yaml_runner_memory.py @@ -0,0 +1,203 @@ +"""Helpers for the YAML runner memory tests: generated case files and a pytest +plugin that measures what a run keeps alive.""" + +import gc +import sys +import tracemalloc +import weakref +from pathlib import Path + +import numpy as np +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.variables.taxes import ( + income_tax as template_income_tax, +) +from policyengine_core.simulations import Simulation +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tracers import FullTracer +from policyengine_core.tools.test_runner import OpenFiscaPlugin + +# pytest-rerunfailures 16.7 stores failed ExceptionInfo objects on collected +# items even when no reruns are requested, retaining the builder's simulation +# through traceback locals. Measure the runner's ownership without this dev +# plugin; the outer test session can still rerun these tests normally. +RUNNER_ARGV = [ + "--capture", + "no", + "--maxfail", + "0", + "--tb", + "short", + "-p", + "no:rerunfailures", +] + + +def write_cases(path: Path, count: int, reform_every: int = 3) -> Path: + """Write ``count`` country-template income tax cases to ``path``. + + Every ``reform_every``-th case sets the income tax rate as a dotted + parameter input with a value no other case uses, so it needs a reform + system of its own; the others run on the reform-free system. ``0`` + writes no reform cases. + """ + cases = [] + for i in range(count): + salary = 1000 + i + case = { + "name": f"case {i}", + "period": "2017-01", + "input": {"salary": salary}, + "absolute_error_margin": 0.01, + } + if reform_every and i % reform_every == 0: + rate = round(0.10 + i / 10_000, 4) + case["input"]["taxes.income_tax_rate"] = rate + case["output"] = {"income_tax": salary * rate} + else: + case["output"] = {"income_tax": salary * 0.15} + cases.append(case) + path.write_text(yaml.safe_dump(cases, sort_keys=False)) + return path + + +def write_cases_at_distinct_dates(path: Path, count: int, people: int = 256) -> Path: + """Write ``count`` income tax cases for ``people`` people each, every one + in a month no other case uses, so each reads the parameters at an instant + of its own.""" + cases = [ + { + "name": f"case {i}", + "period": f"{2017 + i // 12}-{1 + i % 12:02d}", + "input": {"salary": [1000] * people}, + "output": {"income_tax": [150] * people}, + } + for i in range(count) + ] + path.write_text(yaml.safe_dump(cases)) + return path + + +def system_reading_its_parameters() -> CountryTaxBenefitSystem: + """A country-template system whose income tax formula reads the rate + through the public ``TaxBenefitSystem.get_parameters_at_instant``, not + through the formula's ``parameters`` argument.""" + system = CountryTaxBenefitSystem() + + class income_tax(template_income_tax): + def formula(person, period, parameters): + system_parameters = ( + person.simulation.tax_benefit_system.get_parameters_at_instant( + period.start + ) + ) + return person("salary", period) * system_parameters.taxes.income_tax_rate + + system.update_variable(income_tax) + return system + + +def mark_parameter_trees_traced(monkeypatch) -> None: + """Have every simulation that switches tracing on mark its system's + parameter tree traced with its own tracer and branch name. + + Core's simulations trace per call and leave the tree untraced; code + outside core still marks it (policyengine-us's + ``isolate_parameter_tracing`` marks the root of a tree no one else reads + where it stands). + """ + original = Simulation.trace + + def set_trace(simulation, trace): + original.fset(simulation, trace) + if trace: + parameters = simulation.tax_benefit_system.parameters + parameters.trace = True + parameters.tracer = simulation.tracer + parameters.branch_name = simulation.branch_name + + monkeypatch.setattr(Simulation, "trace", property(original.fget, set_trace)) + + +def record_results(monkeypatch) -> list: + """Weak references to every array a ``FullTracer`` records as a + calculation result from now on.""" + refs = [] + original = FullTracer.record_calculation_result + + def record(tracer, value): + if isinstance(value, np.ndarray): + refs.append(weakref.ref(value)) + original(tracer, value) + + monkeypatch.setattr(FullTracer, "record_calculation_result", record) + return refs + + +def live(cls) -> weakref.WeakSet: + """The live instances of ``cls`` after a full collection.""" + gc.collect() + return weakref.WeakSet(o for o in gc.get_objects() if isinstance(o, cls)) + + +class MemoryProbe: + """pytest plugin recording, after every ``sample_every``-th case, the bytes + Python has allocated (``tracemalloc``, after a full collection) and, at the + end of the session, how many simulations and tax-benefit systems are + alive that the run created (not ones other tests left behind).""" + + def __init__(self, sample_every: int = 10): + self.sample_every = sample_every + self.items = [] + self.traced = [] + self.outcomes = [] + self.live_simulations = None + self.live_systems = None + self.live_tracers = None + + def pytest_sessionstart(self, session): + self.simulations_before = live(Simulation) + self.systems_before = live(TaxBenefitSystem) + self.tracers_before = live(FullTracer) + tracemalloc.start() + + def pytest_runtest_logreport(self, report): + if report.when == "call": + self.outcomes.append(report.outcome) + + def pytest_runtest_logfinish(self, nodeid, location): + if len(self.outcomes) % self.sample_every == 0: + gc.collect() + self.traced.append(tracemalloc.get_traced_memory()[0]) + + def pytest_sessionfinish(self, session, exitstatus): + tracemalloc.stop() + # Keep the collected items alive even after pytest.main returns, so + # collecting the session itself cannot hide a runner-owned reference. + self.items = list(session.items) + # pytest keeps the last failure's traceback in ``sys.last_*``; that one + # root is pytest's, not the runner's. + for name in ("last_type", "last_value", "last_traceback", "last_exc"): + if hasattr(sys, name): + delattr(sys, name) + self.live_simulations = len(live(Simulation) - self.simulations_before) + self.live_systems = len(live(TaxBenefitSystem) - self.systems_before) + self.live_tracers = len(live(FullTracer) - self.tracers_before) + + +def run_with_probe( + tax_benefit_system, path: Path, options=None, outcome="passed" +) -> MemoryProbe: + """Run ``path`` the way ``run_tests`` does, with a ``MemoryProbe``; every + case must end with ``outcome``.""" + probe = MemoryProbe() + exit_code = pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], + plugins=[OpenFiscaPlugin(tax_benefit_system, options or {}), probe], + ) + assert exit_code == (0 if outcome == "passed" else 1), exit_code + assert probe.outcomes and set(probe.outcomes) == {outcome} + return probe