diff --git a/changelog.d/randomness-guard-module-removed.removed.md b/changelog.d/randomness-guard-module-removed.removed.md new file mode 100644 index 000000000..c6630cf31 --- /dev/null +++ b/changelog.d/randomness-guard-module-removed.removed.md @@ -0,0 +1 @@ +The internal `policyengine_core.simulations.randomness_guard` module (the `forbid_randomness` runtime guard added in 3.27.0) has been removed. Formula-time randomness is now rejected statically at variable registration, and `NonDeterministicFormulaError` now lives in `policyengine_core.errors` (also re-exported from `policyengine_core.variables.formula_randomness`). diff --git a/changelog.d/static-formula-randomness-check.fixed.md b/changelog.d/static-formula-randomness-check.fixed.md new file mode 100644 index 000000000..017c50670 --- /dev/null +++ b/changelog.d/static-formula-randomness-check.fixed.md @@ -0,0 +1 @@ +Formula-time randomness is now rejected by a static check when a variable is registered, instead of a runtime guard that mutated the process-global `numpy.random`. This fixes intermittent, misattributed `NonDeterministicFormulaError`s raised under concurrent simulations, where one thread's running formula collided with another thread's legitimate setup seeding in `Simulation.__init__` (policyengine-core#518). The check fails fast at load, names the offending variable, and additionally catches `from numpy.random import ...` drawing functions and module-scope generators that the runtime guard could not. diff --git a/policyengine_core/errors/__init__.py b/policyengine_core/errors/__init__.py index f22917e0d..cd7cb6a09 100644 --- a/policyengine_core/errors/__init__.py +++ b/policyengine_core/errors/__init__.py @@ -1,6 +1,7 @@ from .cycle_error import CycleError from .empty_argument_error import EmptyArgumentError from .nan_creation_error import NaNCreationError +from .non_deterministic_formula_error import NonDeterministicFormulaError from .parameter_not_found_error import ParameterNotFoundError from .parameter_parsing_error import ParameterParsingError from .period_mismatch_error import PeriodMismatchError diff --git a/policyengine_core/errors/non_deterministic_formula_error.py b/policyengine_core/errors/non_deterministic_formula_error.py new file mode 100644 index 000000000..99723c795 --- /dev/null +++ b/policyengine_core/errors/non_deterministic_formula_error.py @@ -0,0 +1,9 @@ +class NonDeterministicFormulaError(RuntimeError): + """A variable's formula references a random number generator. + + Rules-engine formulas must be deterministic functions of their inputs. + Raised at variable registration by + ``policyengine_core.variables.formula_randomness.check_formula_determinism``. + """ + + pass diff --git a/policyengine_core/simulations/randomness_guard.py b/policyengine_core/simulations/randomness_guard.py deleted file mode 100644 index 3514b1842..000000000 --- a/policyengine_core/simulations/randomness_guard.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Forbid non-deterministic randomness inside variable formulas. - -A rules engine must be a pure function of its inputs: identical inputs must -always produce identical outputs. Calling a random number generator inside a -formula breaks that contract — the same household can get different results on -different runs — and makes whole datasets non-reproducible. - -While a formula is executing, this guard replaces the callables exposed by -``numpy.random`` and the standard library ``random`` module with functions that -raise :class:`NonDeterministicFormulaError`. This includes the RNG constructors -``np.random.default_rng``/``Generator``/``RandomState``, so *building* a -generator inside a formula — even a seeded one — also raises. Seeding does not -make randomness acceptable in a formula: stochastic inputs belong in the dataset -(computed once, deterministically, and stored), not in the formula. The guard is -re-entrant, so nested formula evaluation is handled correctly, and it restores -the originals once the outermost guarded formula returns. - -Known limitations (pinned by tests in ``test_randomness_guard`` so they cannot -drift silently): - -* A ``Generator``/``RandomState`` instance built *before* the formula runs - (e.g. ``rng = np.random.default_rng(0)`` at module scope) is not intercepted - when its methods are called inside the formula. ``numpy``'s generator classes - are immutable C extension types, so their bound methods cannot be patched. -* A *bare drawing function* bound into another module's namespace before the - guard installs (e.g. ``from numpy.random import random``) is not intercepted, - because the guard patches the ``numpy.random`` module attribute, not every - rebinding of it. - -Both require deliberately hoisting randomness out of the ``np.random.`` form; -use ``np.random.(...)`` or construct the generator inside the formula (both -caught) rather than importing or pre-building drawing callables. -""" - -from __future__ import annotations - -import random as _stdlib_random -from typing import Callable - -import numpy as np - - -class NonDeterministicFormulaError(RuntimeError): - """Raised when a formula invokes a random number generator.""" - - -# Public callables exposed by each randomness namespace, captured once at -# import so the per-formula swap is a cheap dict iteration rather than a -# fresh ``dir()`` scan. -def _public_callables(module) -> dict[str, Callable]: - return { - name: getattr(module, name) - for name in dir(module) - if not name.startswith("_") and callable(getattr(module, name)) - } - - -_GUARDED_NAMESPACES = ( - ("numpy.random", np.random, _public_callables(np.random)), - ("random", _stdlib_random, _public_callables(_stdlib_random)), -) - -# Re-entrancy bookkeeping: only the outermost guarded formula installs and -# removes the patches; the active variable name is tracked as a stack so the -# error message always names the formula that actually made the call. -_depth = 0 -_variable_stack: list[str] = [] - - -def _make_raiser(namespace: str, attribute: str) -> Callable: - qualified = f"{namespace}.{attribute}" - - def _raise(*args, **kwargs): - variable = _variable_stack[-1] if _variable_stack else "" - raise NonDeterministicFormulaError( - f"The formula for '{variable}' called {qualified}(), but rules-engine " - f"formulas must be deterministic functions of their inputs. Remove the " - f"random call. If you need a stochastic input, compute it once when " - f"building the dataset (with a seeded generator) and store it as an " - f"input variable instead." - ) - - return _raise - - -# Pre-build the raisers once per (namespace, attribute). -_RAISERS = { - id(module): {name: _make_raiser(namespace, name) for name in originals} - for namespace, module, originals in _GUARDED_NAMESPACES -} - - -def _install() -> None: - for _namespace, module, originals in _GUARDED_NAMESPACES: - raisers = _RAISERS[id(module)] - for name in originals: - setattr(module, name, raisers[name]) - - -def _restore() -> None: - for _namespace, module, originals in _GUARDED_NAMESPACES: - for name, original in originals.items(): - setattr(module, name, original) - - -class forbid_randomness: - """Context manager that bans RNG use while a formula runs. - - Re-entrant: nested formulas reuse the single installed patch set and only - the outermost context restores the originals. - """ - - __slots__ = ("variable_name",) - - def __init__(self, variable_name: str): - self.variable_name = variable_name - - def __enter__(self) -> "forbid_randomness": - global _depth - if _depth == 0: - _install() - _depth += 1 - _variable_stack.append(self.variable_name) - return self - - def __exit__(self, exc_type, exc_value, traceback) -> bool: - global _depth - _variable_stack.pop() - _depth -= 1 - if _depth == 0: - _restore() - return False diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 3864a5dc6..b7f52591e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -14,7 +14,6 @@ from policyengine_core.entities.entity import Entity from policyengine_core.enums import Enum, EnumArray from policyengine_core.errors import CycleError, SpiralError -from policyengine_core.simulations.randomness_guard import forbid_randomness from policyengine_core.holders.holder import Holder from policyengine_core.periods import Period from policyengine_core.periods.config import ETERNITY, MONTH, YEAR @@ -593,8 +592,9 @@ def calculate( self.tracer.record_calculation_start(variable_name, period, self.branch_name) # No per-variable RNG seeding: formulas may not use randomness at all - # (enforced by forbid_randomness in _run_formula), so there is nothing - # to make reproducible here. + # (enforced statically at variable registration by + # check_formula_determinism), so there is nothing to make reproducible + # here. try: result = self._calculate(variable_name, period) @@ -1106,13 +1106,12 @@ def _run_formula( parameters_at = self.tax_benefit_system.parameters # A rules-engine formula must be a pure, deterministic function of its - # inputs. Forbid any random number generation while it runs so the same - # inputs always produce the same outputs. - with forbid_randomness(variable.name): - if formula.__code__.co_argcount == 2: - array = formula(population, period) - else: - array = formula(population, period, parameters_at) + # inputs. Randomness is forbidden statically at variable registration + # (check_formula_determinism), so no runtime guard is needed here. + if formula.__code__.co_argcount == 2: + array = formula(population, period) + else: + array = formula(population, period, parameters_at) return array diff --git a/policyengine_core/variables/formula_randomness.py b/policyengine_core/variables/formula_randomness.py new file mode 100644 index 000000000..a12ee104e --- /dev/null +++ b/policyengine_core/variables/formula_randomness.py @@ -0,0 +1,168 @@ +"""Static detection of formula-time randomness (policyengine-core#518). + +A rules-engine formula must be a pure, deterministic function of its inputs. +Rather than forbidding randomness with a runtime monkeypatch of the +process-global ``numpy.random`` -- which is not safe under concurrent +simulations (see #518: the guard installed by one thread's formula leaks onto +another thread's legitimate setup ``np.random.seed(0)``) -- this module scans a +formula's bytecode *once, when the variable is constructed* and raises if it +references numpy's or the standard library's random facilities. + +Detection is by **identity**, not by name: a local variable, a string argument, +or an attribute that merely happens to be spelled ``random`` is never flagged. +Only genuine *global* and *closure* loads are resolved against the formula's +namespace (attribute loads such as ``foo.random`` are not), which keeps the +check free of name-collision false positives. + +Because it is static, this also catches forms the old runtime guard explicitly +could not: ``from numpy.random import default_rng`` (a drawing callable imported +by name) and a generator hoisted to module scope or captured through a closure +(``rng = np.random.default_rng(0)``) and used inside the formula. +""" + +from __future__ import annotations + +import dis +import random as _stdlib_random +from types import ModuleType +from typing import Optional + +import numpy + +from policyengine_core.errors import NonDeterministicFormulaError + +_RANDOM_MODULES = (numpy.random, _stdlib_random) +_RNG_INSTANCE_TYPES = (numpy.random.Generator, numpy.random.RandomState) + +# Scanning is memoised by code object: each unique formula is inspected once, +# no matter how many tax-benefit systems instantiate the variable. +_cache: dict = {} + + +def _is_random_module(obj: object) -> bool: + """True if ``obj`` is the ``numpy.random`` or stdlib ``random`` module.""" + return any(obj is module for module in _RANDOM_MODULES) + + +def _is_random_member(obj: object) -> bool: + """True if ``obj`` is a callable / generator that belongs to a random module. + + Catches ``from numpy.random import default_rng`` / ``from random import + random`` (imported drawing callables), pre-built ``Generator`` / + ``RandomState`` instances hoisted to module scope, and the retired + ``commons.formulas.random`` helper (kept as a raiser since #494). + """ + module = getattr(obj, "__module__", None) + if isinstance(module, str) and ( + module == "random" or module.startswith("numpy.random") + ): + return True + if isinstance(obj, _RNG_INSTANCE_TYPES): + return True + if ( + module == "policyengine_core.commons.formulas" + and getattr(obj, "__name__", None) == "random" + ): + return True + return False + + +def _describe(obj: object) -> str: + """A short, human-readable name for the offending reference.""" + if isinstance(obj, ModuleType): + return obj.__name__ + module = getattr(obj, "__module__", None) + name = getattr(obj, "__name__", None) + if module and name: + return f"{module}.{name}" + if isinstance(obj, _RNG_INSTANCE_TYPES): + return f"numpy.random.{type(obj).__name__} instance" + return repr(obj) + + +def _freevars(formula) -> dict: + """Map each of a formula's closure variable names to its current value.""" + result: dict = {} + if formula.__closure__: + for name, cell in zip(formula.__code__.co_freevars, formula.__closure__): + try: + result[name] = cell.cell_contents + except ValueError: + continue # empty cell + return result + + +def _resolve_load(instruction, global_namespace: dict, freevars: dict) -> object: + """Resolve a name-load instruction to the object it loads, or ``None``. + + Only global and closure loads are resolved -- never ``LOAD_ATTR`` -- so an + attribute merely spelled ``random`` is not mistaken for the module, even + when ``random`` is also imported at module scope. + """ + if instruction.opname in ("LOAD_GLOBAL", "LOAD_NAME"): + return global_namespace.get(instruction.argval) + if instruction.opname in ("LOAD_DEREF", "LOAD_CLASSDEREF"): + return freevars.get(instruction.argval) + return None + + +def _referenced_randomness(formula) -> Optional[str]: + """Describe the randomness a formula references, or ``None``. + + A single pass over the bytecode: each global/closure load is checked by + identity, and a numpy top-level module load followed by a ``.random`` + attribute access (e.g. ``np.random.seed``) is flagged too. + """ + global_namespace = formula.__globals__ + freevars = _freevars(formula) + previous = None + for instruction in dis.get_instructions(formula): + loaded = _resolve_load(instruction, global_namespace, freevars) + if loaded is not None and ( + _is_random_module(loaded) or _is_random_member(loaded) + ): + return _describe(loaded) + if ( + instruction.opname in ("LOAD_ATTR", "LOAD_METHOD") + and instruction.argval == "random" + and previous is not None + and _resolve_load(previous, global_namespace, freevars) is numpy + ): + return "numpy.random" + previous = instruction + return None + + +def formula_uses_randomness(formula) -> Optional[str]: + """Return a description of the randomness a formula references, or ``None``. + + ``formula`` is a plain function (a ``def formula(...)`` collected from a + ``Variable`` subclass). Objects without a code object (e.g. ``None``) return + ``None``. Memoised by code object. + """ + code = getattr(formula, "__code__", None) + if code is None: + return None + if code not in _cache: + _cache[code] = _referenced_randomness(formula) + return _cache[code] + + +def check_formula_determinism(variable) -> None: + """Raise ``NonDeterministicFormulaError`` if any of a variable's formulas + references randomness. + + Intended to be called from ``Variable.__init__`` alongside + ``check_computation_modes`` so a randomness-using formula fails fast, at + variable construction, with the offending variable named. + """ + for formula in variable.formulas.values(): + offending = formula_uses_randomness(formula) + if offending is not None: + raise NonDeterministicFormulaError( + f"The formula for '{variable.name}' references {offending}, but " + f"rules-engine formulas must be deterministic functions of their " + f"inputs. Remove the random call. If you need a stochastic input, " + f"compute it once when building the dataset (with a seeded " + f"generator) and store it as an input variable instead." + ) diff --git a/policyengine_core/variables/variable.py b/policyengine_core/variables/variable.py index 7f55c3ea9..7cb91b47b 100644 --- a/policyengine_core/variables/variable.py +++ b/policyengine_core/variables/variable.py @@ -18,6 +18,7 @@ from policyengine_core.periods import DAY, ETERNITY from . import config, helpers +from .formula_randomness import check_formula_determinism class QuantityType: @@ -325,6 +326,7 @@ def __init__(self, baseline_variable=None): self.formulas = self.set_formulas(formulas_attr) self.check_computation_modes() + check_formula_determinism(self) if unexpected_attrs: raise ValueError( diff --git a/tests/core/simulations/test_formula_randomness_concurrency.py b/tests/core/simulations/test_formula_randomness_concurrency.py new file mode 100644 index 000000000..3205d1418 --- /dev/null +++ b/tests/core/simulations/test_formula_randomness_concurrency.py @@ -0,0 +1,62 @@ +"""Concurrent simulations must not spuriously raise (policyengine-core#518). + +The removed runtime guard mutated the process-global ``numpy.random`` while a +formula ran, so one thread's formula could make another thread's legitimate +``Simulation.__init__`` ``np.random.seed(0)`` raise +``NonDeterministicFormulaError`` -- intermittently, and misattributed to an +innocent formula. Now that enforcement is static (at variable registration), +building and calculating simulations concurrently is safe. + +This test reproduces that race: several threads each repeatedly build a default +simulation (each build calls the legitimate setup ``np.random.seed(0)``) and +compute a formula variable, with a barrier to maximise overlap. It passes with +the static check and fails against the reinstated runtime guard. +""" + +import threading + +import numpy as np + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables.formula_randomness import ( + NonDeterministicFormulaError, +) + +PERIOD = "2026-01" +THREADS = 12 +ITERATIONS = 60 + + +def test_concurrent_build_and_calculate_is_safe(): + barrier = threading.Barrier(THREADS) + errors: list = [] + + def worker(): + # A per-thread system isolates unrelated state; the #518 race is on the + # shared global numpy.random, which per-thread systems still exercise. + system = CountryTaxBenefitSystem() + barrier.wait() # release all threads together to maximise overlap + try: + for _ in range(ITERATIONS): + simulation = SimulationBuilder().build_default_simulation(system) + simulation.calculate("disposable_income", PERIOD) + except Exception as exc: # noqa: BLE001 - captured for the assertion + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(THREADS)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + nondeterministic = [ + e for e in errors if isinstance(e, NonDeterministicFormulaError) + ] + assert not nondeterministic, ( + f"concurrent simulations raised spurious randomness errors: " + f"{nondeterministic!r}" + ) + assert not errors, f"concurrent simulations raised: {errors!r}" + # numpy randomness still works normally in the main thread afterwards. + assert isinstance(float(np.random.random()), float) diff --git a/tests/core/test_randomness_guard.py b/tests/core/test_randomness_guard.py deleted file mode 100644 index f96443a23..000000000 --- a/tests/core/test_randomness_guard.py +++ /dev/null @@ -1,189 +0,0 @@ -"""A variable formula must not invoke a random number generator. - -Rules-engine formulas have to be deterministic functions of their inputs, so -calling ``numpy.random`` or the stdlib ``random`` module inside a formula raises -:class:`NonDeterministicFormulaError`. These tests pin that behaviour and verify -the guard restores the randomness namespaces afterwards. -""" - -import random -from numpy.random import random as _bare_random_import - -import numpy as np -import pytest - -# A generator hoisted to module scope, built before any formula runs. Its bound -# methods cannot be patched (numpy generator classes are immutable C types), so -# using it inside a formula is a known, pinned gap in the guard. -_PREBUILT_RNG = np.random.default_rng(0) - -from policyengine_core import periods -from policyengine_core.country_template import CountryTaxBenefitSystem, entities -from policyengine_core.simulations import SimulationBuilder -from policyengine_core.simulations.randomness_guard import ( - NonDeterministicFormulaError, - forbid_randomness, -) -from policyengine_core.variables import Variable - -PERIOD = "2013-01" - - -def _simulation_with(*variable_classes): - system = CountryTaxBenefitSystem() - system.add_variables(*variable_classes) - return SimulationBuilder().build_default_simulation(system) - - -class uses_numpy_random(Variable): - value_type = float - entity = entities.Person - definition_period = periods.MONTH - label = "formula that draws from numpy.random" - - def formula(person, period): - return np.random.random(person.count) - - -class uses_stdlib_random(Variable): - value_type = float - entity = entities.Person - definition_period = periods.MONTH - label = "formula that draws from the random module" - - def formula(person, period): - return random.random() - - -class uses_seeded_generator(Variable): - value_type = float - entity = entities.Person - definition_period = periods.MONTH - label = "formula that builds a seeded generator" - - def formula(person, period): - # Seeding does not make randomness acceptable inside a formula. - return np.random.default_rng(0).random(person.count) - - -class deterministic(Variable): - value_type = int - entity = entities.Person - definition_period = periods.MONTH - label = "deterministic formula" - - def formula(person, period): - return person.count - - -class raises_value_error(Variable): - value_type = float - entity = entities.Person - definition_period = periods.MONTH - label = "formula that raises a non-RNG exception" - - def formula(person, period): - raise ValueError("boom") - - -class uses_prebuilt_generator(Variable): - value_type = float - entity = entities.Person - definition_period = periods.MONTH - label = "formula that uses a module-scope generator (known gap)" - - def formula(person, period): - return _PREBUILT_RNG.random(person.count) - - -class uses_bare_imported_function(Variable): - value_type = float - entity = entities.Person - definition_period = periods.MONTH - label = "formula that uses a by-name imported drawing function (known gap)" - - def formula(person, period): - return _bare_random_import() - - -def test_numpy_random_in_formula_raises(): - simulation = _simulation_with(uses_numpy_random) - with pytest.raises(NonDeterministicFormulaError, match="uses_numpy_random"): - simulation.calculate("uses_numpy_random", PERIOD) - - -def test_stdlib_random_in_formula_raises(): - simulation = _simulation_with(uses_stdlib_random) - with pytest.raises(NonDeterministicFormulaError, match="random"): - simulation.calculate("uses_stdlib_random", PERIOD) - - -def test_seeded_generator_in_formula_still_raises(): - simulation = _simulation_with(uses_seeded_generator) - with pytest.raises(NonDeterministicFormulaError): - simulation.calculate("uses_seeded_generator", PERIOD) - - -def test_deterministic_formula_is_unaffected(): - simulation = _simulation_with(deterministic) - result = simulation.calculate("deterministic", PERIOD) - assert (result == 1).all() - - -def test_randomness_restored_after_guarded_formula(): - simulation = _simulation_with(uses_numpy_random) - with pytest.raises(NonDeterministicFormulaError): - simulation.calculate("uses_numpy_random", PERIOD) - # Outside any formula, numpy and stdlib randomness work normally again. - assert isinstance(float(np.random.random()), float) - assert isinstance(random.random(), float) - - -def test_randomness_restored_after_non_rng_exception_in_formula(): - # If a formula raises a normal exception, _run_formula's guarded block must - # still restore the randomness namespaces (no leak of the patched state). - simulation = _simulation_with(raises_value_error) - with pytest.raises(ValueError, match="boom"): - simulation.calculate("raises_value_error", PERIOD) - assert isinstance(float(np.random.random()), float) - assert isinstance(random.random(), float) - - -def test_constructing_generator_inside_formula_is_caught(): - # Building a generator inside the formula hits the patched np.random - # constructor, so even this seeded form raises (covered by - # uses_seeded_generator too; kept explicit for the boundary). - simulation = _simulation_with(uses_seeded_generator) - with pytest.raises(NonDeterministicFormulaError): - simulation.calculate("uses_seeded_generator", PERIOD) - - -def test_prebuilt_generator_is_a_known_gap(): - # Documented limitation: a generator built before the formula runs cannot be - # intercepted (numpy generator classes are immutable). Pin the behaviour so a - # future change to it is noticed. - simulation = _simulation_with(uses_prebuilt_generator) - result = simulation.calculate("uses_prebuilt_generator", PERIOD) - assert result is not None - - -def test_by_name_imported_function_is_a_known_gap(): - # Documented limitation: a drawing function imported by name before the guard - # installs is not intercepted. Pin the behaviour. - simulation = _simulation_with(uses_bare_imported_function) - result = simulation.calculate("uses_bare_imported_function", PERIOD) - assert result is not None - - -def test_guard_is_reentrant(): - # Entering twice and leaving the inner context must not restore the - # originals while the outer context is still active. - with forbid_randomness("outer"): - with forbid_randomness("inner"): - with pytest.raises(NonDeterministicFormulaError, match="inner"): - np.random.random() - # Still guarded: the outer context owns the patch. - with pytest.raises(NonDeterministicFormulaError, match="outer"): - np.random.random() - # Fully restored once the outermost context exits. - assert isinstance(float(np.random.random()), float) diff --git a/tests/core/variables/test_formula_randomness.py b/tests/core/variables/test_formula_randomness.py new file mode 100644 index 000000000..ff85a5439 --- /dev/null +++ b/tests/core/variables/test_formula_randomness.py @@ -0,0 +1,231 @@ +"""Unit tests for the static formula-randomness detector (policyengine-core#518). + +These exercise ``formula_uses_randomness`` / ``check_formula_determinism`` +directly on hand-written functions -- no tax-benefit system is built here. The +detector is wired into ``Variable.__init__`` in a later commit; these tests pin +its contract independently of that wiring. +""" + +import random # stdlib, imported at module scope on purpose (collision bait) + +import numpy +import numpy as np +import numpy.random as npr +import pytest +from numpy.random import default_rng +from numpy.random import random as bare_np_random + +from policyengine_core.commons.formulas import random as commons_random +from policyengine_core.variables.formula_randomness import ( + NonDeterministicFormulaError, + check_formula_determinism, + formula_uses_randomness, +) + +# A generator hoisted to module scope -- the "pre-built generator" the runtime +# guard could not catch. +MODULE_RNG = np.random.default_rng(0) + + +# --- positives: every form of formula-time randomness must be detected -------- + + +def _np_random_seed(person, period, parameters): + np.random.seed(0) + return person("salary", period) + + +def _np_random_draw(person, period, parameters): + return np.random.random(len(person("salary", period))) + + +def _numpy_random_fullname(person, period, parameters): + return numpy.random.choice([0, 1], size=3) + + +def _stdlib_random(person, period, parameters): + return random.random() + + +def _aliased_numpy_random_module(person, period, parameters): + npr.seed(0) + return person("salary", period) + + +def _imported_default_rng(person, period, parameters): + return default_rng(0).random(3) + + +def _bare_imported_drawing_function(person, period, parameters): + return bare_np_random(3) + + +def _prebuilt_module_generator(person, period, parameters): + return MODULE_RNG.random(3) + + +def _retired_commons_helper(person, period, parameters): + return commons_random(person) + + +POSITIVES = [ + _np_random_seed, + _np_random_draw, + _numpy_random_fullname, + _stdlib_random, + _aliased_numpy_random_module, + _imported_default_rng, + _bare_imported_drawing_function, + _prebuilt_module_generator, + _retired_commons_helper, +] + + +@pytest.mark.parametrize("formula", POSITIVES, ids=lambda f: f.__name__) +def test_randomness_is_detected(formula): + assert formula_uses_randomness(formula) is not None + + +def test_bare_imported_drawing_function_is_now_caught(): + # Former runtime-guard gap: `from numpy.random import random` imported by name. + assert formula_uses_randomness(_bare_imported_drawing_function) is not None + + +def test_prebuilt_module_generator_is_now_caught(): + # Former runtime-guard gap: a Generator built before the formula runs. + assert formula_uses_randomness(_prebuilt_module_generator) is not None + + +# --- negatives: deterministic numpy use and name collisions must NOT flag ----- + + +def _deterministic_numpy(person, period, parameters): + salary = person("salary", period) + return np.where(salary > 0, np.maximum(salary, 1), 0) + + +def _attribute_named_random(person, period, parameters): + # `.random` here is an attribute access on a non-random object; `random` is + # also imported at module scope -- the detector must not conflate the two. + return person.random + + +def _local_named_random(person, period, parameters): + random = 5 # a local, not the module + return random + + +def _string_argument_random(person, period, parameters): + # "random_draw" is a precomputed input variable, read like any other. + return person("random_draw", period) + + +NEGATIVES = [ + _deterministic_numpy, + _attribute_named_random, + _local_named_random, + _string_argument_random, +] + + +@pytest.mark.parametrize("formula", NEGATIVES, ids=lambda f: f.__name__) +def test_deterministic_formula_is_not_flagged(formula): + assert formula_uses_randomness(formula) is None + + +def test_non_function_returns_none(): + assert formula_uses_randomness(None) is None + assert formula_uses_randomness("not a formula") is None + + +# --- closures: randomness captured through a closure is resolved too ---------- + + +def _closure_over_generator(): + rng = np.random.default_rng(0) + + def formula(person, period, parameters): + return rng.random(3) + + return formula + + +def _closure_over_drawing_function(): + draw = np.random.random + + def formula(person, period, parameters): + return draw(3) + + return formula + + +def _closure_over_numpy_module(): + local_numpy = np + + def formula(person, period, parameters): + return local_numpy.random.random(3) + + return formula + + +def _closure_over_deterministic_numpy(): + op = np.maximum + + def formula(person, period, parameters): + return op(person("salary", period), 0) + + return formula + + +@pytest.mark.parametrize( + "factory", + [ + _closure_over_generator, + _closure_over_drawing_function, + _closure_over_numpy_module, + ], + ids=lambda f: f.__name__, +) +def test_randomness_captured_through_closure_is_detected(factory): + assert formula_uses_randomness(factory()) is not None + + +def test_deterministic_closure_is_not_flagged(): + assert formula_uses_randomness(_closure_over_deterministic_numpy()) is None + + +# --- check_formula_determinism (the variable-level entry point) ---------------- + + +class _FakeVariable: + def __init__(self, name, formulas): + self.name = name + self.formulas = formulas + + +def test_check_raises_naming_the_variable(): + variable = _FakeVariable("uses_randomness", {"0001-01-01": _np_random_seed}) + with pytest.raises(NonDeterministicFormulaError, match="uses_randomness"): + check_formula_determinism(variable) + + +def test_check_passes_for_deterministic_formulas(): + variable = _FakeVariable( + "clean", + {"0001-01-01": _deterministic_numpy, "2020-01-01": _string_argument_random}, + ) + check_formula_determinism(variable) # must not raise + + +def test_check_passes_for_formula_less_variable(): + check_formula_determinism(_FakeVariable("input_only", {})) # must not raise + + +def test_exception_canonical_home_is_errors_package(): + # NonDeterministicFormulaError lives in policyengine_core.errors; the module + # that raises it re-exports the same class. + from policyengine_core.errors import ( + NonDeterministicFormulaError as FromErrors, + ) + + assert FromErrors is NonDeterministicFormulaError diff --git a/tests/core/variables/test_variable_formula_determinism.py b/tests/core/variables/test_variable_formula_determinism.py new file mode 100644 index 000000000..294919e93 --- /dev/null +++ b/tests/core/variables/test_variable_formula_determinism.py @@ -0,0 +1,157 @@ +"""Registration-time enforcement of formula determinism (policyengine-core#518). + +A variable whose formula references randomness must be rejected when it is +*registered* with a tax-benefit system -- i.e. when it is instantiated by +``add_variables`` -- so a bad formula fails fast at load, with the variable +named, instead of intermittently at calculation time. A deterministic variable +must still register and compute normally. + +These integration tests prove the wiring into ``Variable.__init__`` (the pure +detector is unit-tested in ``test_formula_randomness``). +""" + +import random +from numpy.random import random as _bare_random_import + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.reforms import Reform +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable +from policyengine_core.variables.formula_randomness import ( + NonDeterministicFormulaError, +) + +# A generator hoisted to module scope, built before any formula runs. +_PREBUILT_RNG = np.random.default_rng(0) + +PERIOD = "2013-01" + + +class uses_numpy_random(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "formula that draws from numpy.random" + + def formula(person, period): + return np.random.random(person.count) + + +class uses_stdlib_random(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "formula that draws from the random module" + + def formula(person, period): + return random.random() + + +class uses_seeded_generator(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "formula that builds a seeded generator inside the formula" + + def formula(person, period): + return np.random.default_rng(0).random(person.count) + + +class uses_prebuilt_generator(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "formula that uses a module-scope generator" + + def formula(person, period): + return _PREBUILT_RNG.random(person.count) + + +class uses_bare_imported_function(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "formula that uses a by-name imported drawing function" + + def formula(person, period): + return _bare_random_import() + + +class deterministic(Variable): + value_type = int + entity = entities.Person + definition_period = periods.MONTH + label = "deterministic formula" + + def formula(person, period): + return person.count + + +class uses_random_in_dated_formula(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "clean base formula but randomness in a later dated formula" + + def formula(person, period): + return person.count + + def formula_2020(person, period): + return np.random.random(person.count) + + +RANDOMNESS_VARIABLES = [ + uses_numpy_random, + uses_stdlib_random, + uses_seeded_generator, + uses_prebuilt_generator, + uses_bare_imported_function, +] + + +def _register(*variable_classes): + system = CountryTaxBenefitSystem() + system.add_variables(*variable_classes) + return system + + +@pytest.mark.parametrize( + "variable_class", RANDOMNESS_VARIABLES, ids=lambda c: c.__name__ +) +def test_randomness_variable_rejected_at_registration(variable_class): + with pytest.raises(NonDeterministicFormulaError, match=variable_class.__name__): + _register(variable_class) + + +def test_deterministic_variable_registers_and_computes(): + system = _register(deterministic) + simulation = SimulationBuilder().build_default_simulation(system) + result = simulation.calculate("deterministic", PERIOD) + assert (result == 1).all() + + +def test_randomness_in_a_dated_formula_is_rejected(): + # check_formula_determinism scans every dated formula, not just the first. + with pytest.raises( + NonDeterministicFormulaError, match="uses_random_in_dated_formula" + ): + _register(uses_random_in_dated_formula) + + +def test_reform_update_variable_with_randomness_is_rejected(): + # The reform / update_variable path re-runs Variable.__init__, so a reform + # that introduces a randomness-using formula is rejected when applied. + class disposable_income(Variable): + def formula(person, period): + return np.random.random(person.count) + + class reform(Reform): + def apply(self): + self.update_variable(disposable_income) + + with pytest.raises(NonDeterministicFormulaError, match="disposable_income"): + reform(CountryTaxBenefitSystem())