diff --git a/changelog.d/ci-speed-behavioral-measurement-copies.changed.md b/changelog.d/ci-speed-behavioral-measurement-copies.changed.md new file mode 100644 index 00000000000..70498d030d0 --- /dev/null +++ b/changelog.d/ci-speed-behavioral-measurement-copies.changed.md @@ -0,0 +1 @@ +Cut the CI time of the behavioural-response measurement copy tests: collect garbage less often while they copy policy trees, and run the pinned copy cases over three households in one simulation instead of three fixed generated examples. diff --git a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py index 58d6c5f2bf7..466c161011d 100644 --- a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py +++ b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py @@ -27,9 +27,11 @@ 4. Invalidating a simulation's cached formula output drops its measurements. """ +import gc + import numpy as np import pytest -from hypothesis import example, given, settings +from hypothesis import Phase, example, given, settings from hypothesis import strategies as st from policyengine_core.periods import ETERNITY, YEAR as YEAR_PERIOD from policyengine_core.periods import period as to_period @@ -57,6 +59,27 @@ }, } WAGES = (20_000, 50_000, 100_000) +GC_THRESHOLD = 100_000 + + +@pytest.fixture(autouse=True, scope="module") +def collect_cycles_less_often(): + """Run the cycle collector less often while these tests copy policy. + + Every reform simulation, measurement and clone here deep-copies the + parameter tree: 6.4 million objects the collector tracks, none of them + garbage while the copy is in use. A collection starts each time + allocations exceed deallocations by the first threshold. At the default of + 2,000, three rounds of building, measuring, cloning and branching a reform + simulation ran 29 full collections, each traversing every live tree; at + 100,000 they ran 8 (2026-10-23). In CI the tests this module did not + otherwise change then ran a fifth to a quarter faster (run 40052943537 + against 40038526568 and 40050303812). + """ + thresholds = gc.get_threshold() + gc.set_threshold(GC_THRESHOLD, *thresholds[1:]) + yield + gc.set_threshold(*thresholds) def household(wages): @@ -70,6 +93,21 @@ def household(wages): } +def households(wages): + """``household(w)`` for each of ``wages``, side by side in one situation.""" + situation = {} + for index, wage in enumerate(wages): + for group, units in household(wage).items(): + for name, unit in units.items(): + if "members" in unit: + unit = { + **unit, + "members": [f"{member}_{index}" for member in unit["members"]], + } + situation.setdefault(group, {})[f"{name}_{index}"] = unit + return situation + + def make_simulation(**kwargs): """Reuse read-only policy; each reform gets a private model and inputs.""" kwargs.setdefault("tax_benefit_system", system) @@ -89,8 +127,12 @@ def response(simulation, period=YEAR): def fresh_response(wages, period=YEAR): """The response a new reform simulation with these wages reports. - Values only, so sharing them across tests shares no simulation state. + ``wages`` is one household's, or a tuple of several households', each + measured in a simulation of its own. Values only, so sharing them across + tests shares no simulation state. """ + if isinstance(wages, tuple): + return np.concatenate([fresh_response(wage, period) for wage in wages]) key = (wages, period) if key not in _fresh: _fresh[key] = response( @@ -107,33 +149,58 @@ def set_wages(simulation, wages, period=YEAR): everything but inputs first. """ drop_inherited_values(simulation) - simulation.set_input("employment_income_before_lsr", period, np.array([wages])) + simulation.set_input( + "employment_income_before_lsr", period, np.array(wages, ndmin=1) + ) -@settings(max_examples=3, deadline=None, derandomize=True) +# Most pins' source holds one household at each of WAGES, so each covers every +# level. A copy given other wages moves each household to another's: 20,000 to +# 50,000, 50,000 to 100,000 (the original pins) and 100,000 to 20,000. +SWAPPED_WAGES = (50_000, 100_000, 20_000) + + +# Only the pins run. Under ``derandomize`` the three generated examples were the +# same on every run, and each reran a pinned path at other wages, which the +# households and the last pin now cover. Drop ``phases`` to generate examples +# from these strategies as well. +@settings(phases=[Phase.explicit], deadline=None) @given( kind=st.sampled_from(["clone", "branch"]), after_source_measured=st.booleans(), - source_wages=st.sampled_from(WAGES), - copy_wages=st.sampled_from(WAGES), + source_wages=st.just(WAGES), + copy_wages=st.permutations(WAGES).map(tuple), ) # Each of these failed before: zero response, or the source's relative changes. @example( - kind="branch", after_source_measured=False, source_wages=50_000, copy_wages=50_000 + kind="branch", after_source_measured=False, source_wages=WAGES, copy_wages=WAGES +) +@example( + kind="clone", after_source_measured=False, source_wages=WAGES, copy_wages=WAGES ) @example( - kind="clone", after_source_measured=False, source_wages=50_000, copy_wages=50_000 + kind="clone", + after_source_measured=True, + source_wages=WAGES, + copy_wages=SWAPPED_WAGES, ) @example( - kind="clone", after_source_measured=True, source_wages=50_000, copy_wages=100_000 + kind="branch", + after_source_measured=True, + source_wages=WAGES, + copy_wages=SWAPPED_WAGES, ) +# A source whose every measured response is zero (at 20,000 under this reform). @example( - kind="branch", after_source_measured=True, source_wages=50_000, copy_wages=100_000 + kind="clone", + after_source_measured=True, + source_wages=(20_000,), + copy_wages=(50_000,), ) def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( kind, after_source_measured, source_wages, copy_wages ): - source = make_simulation(situation=household(source_wages), reform=LSR_REFORM) + source = make_simulation(situation=households(source_wages), reform=LSR_REFORM) if after_source_measured: source_response = response(source) measured = { @@ -150,7 +217,9 @@ def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( # d1021: perturbed wages use the household's measured relative changes. # With this reform's uniform elasticities, the response scales with wages. np.testing.assert_allclose( - copy_response, source_response * copy_wages / source_wages, atol=1e-3 + copy_response, + source_response * np.divide(copy_wages, source_wages), + atol=1e-3, ) copied_measurements = copy.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR][str(YEAR)] assert copied_measurements.keys() == measured.keys()