From a6e50ac0e534b1e7ad260e4d12a7c7a8fdaeb0b8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 21:16:22 -0400 Subject: [PATCH 1/3] Cut the behavioural-response copy tests' CI time (WIP) Raise the cycle collector's first threshold while this module copies policy trees, and run each pinned copy case for all three wage levels in one three-household simulation, replacing three derandomized generated examples that re-ran pinned paths at other wages. Co-Authored-By: Claude Opus 5.5 --- ...d-behavioral-measurement-copies.changed.md | 1 + ..._behavioral_response_measurement_copies.py | 88 +++++++++++++------ 2 files changed, 63 insertions(+), 26 deletions(-) create mode 100644 changelog.d/ci-speed-behavioral-measurement-copies.changed.md 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..6b69b28062f --- /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 each pinned copy case for all three wage levels 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..e7811b984ad 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,25 @@ }, } 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: millions of new objects, nearly all still in use. A + collection starts each time allocations exceed deallocations by the first + threshold (2,000 by default), so one copy ran thousands of collections + over live objects, about two thirds of its time (2026-10-23). A higher + threshold reclaims the same garbage in fewer passes; this module's peak + memory stayed the same. + """ + thresholds = gc.get_threshold() + gc.set_threshold(GC_THRESHOLD, *thresholds[1:]) + yield + gc.set_threshold(*thresholds) def household(wages): @@ -70,6 +91,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,12 +125,14 @@ 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' side by + side. Values only, so sharing them across tests shares no simulation state. """ key = (wages, period) if key not in _fresh: + situation = households(wages) if isinstance(wages, tuple) else household(wages) _fresh[key] = response( - make_simulation(situation=household(wages), reform=LSR_REFORM), period + make_simulation(situation=situation, reform=LSR_REFORM), period ) return _fresh[key] @@ -107,33 +145,31 @@ 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) +# Every household takes another's wages, so each pin below runs all three +# levels in one simulation: one household goes from 50,000 to 100,000. +SWAPPED_WAGES = (50_000, 100_000, 20_000) + + +@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), + 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 -) -@example( - kind="clone", after_source_measured=False, source_wages=50_000, copy_wages=50_000 -) -@example( - kind="clone", after_source_measured=True, source_wages=50_000, copy_wages=100_000 -) -@example( - kind="branch", after_source_measured=True, source_wages=50_000, copy_wages=100_000 -) +@example(kind="branch", after_source_measured=False, copy_wages=WAGES) +@example(kind="clone", after_source_measured=False, copy_wages=WAGES) +@example(kind="clone", after_source_measured=True, copy_wages=SWAPPED_WAGES) +@example(kind="branch", after_source_measured=True, copy_wages=SWAPPED_WAGES) def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( - kind, after_source_measured, source_wages, copy_wages + kind, after_source_measured, copy_wages ): - source = make_simulation(situation=household(source_wages), reform=LSR_REFORM) + source = make_simulation(situation=households(WAGES), reform=LSR_REFORM) if after_source_measured: source_response = response(source) measured = { @@ -143,14 +179,16 @@ def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( ].items() } copy = source.clone() if kind == "clone" else source.get_branch("copy") - if copy_wages != source_wages: + if copy_wages != WAGES: set_wages(copy, copy_wages) copy_response = response(copy) if kind == "branch" and after_source_measured: # 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, WAGES), + atol=1e-3, ) copied_measurements = copy.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR][str(YEAR)] assert copied_measurements.keys() == measured.keys() @@ -159,9 +197,7 @@ def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( else: np.testing.assert_allclose(copy_response, fresh_response(copy_wages), atol=1e-3) if after_source_measured: - np.testing.assert_allclose( - response(source), fresh_response(source_wages), atol=1e-3 - ) + np.testing.assert_allclose(response(source), fresh_response(WAGES), atol=1e-3) def test_migrated_input_records_are_private_and_explicit_overrides_survive(): From 59dd4ac7dcae7bd8f5136d65b1a4413380dad5d5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 21:51:57 -0400 Subject: [PATCH 2/3] Measure each pinned household against its own one-household reference Expected responses come from a new one-household reform simulation per wage level, as before; the three are shared with the fully-calculated clone test, so the three-household pins add no reference runs. State the GC fixture's measured effect, and why only the pins run. Co-Authored-By: Claude Opus 5.5 --- ..._behavioral_response_measurement_copies.py | 29 ++++++++++++------- 1 file changed, 18 insertions(+), 11 deletions(-) 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 e7811b984ad..a3f8f925533 100644 --- a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py +++ b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py @@ -67,12 +67,13 @@ 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: millions of new objects, nearly all still in use. A + parameter tree: about five million objects, nearly all still in use. A collection starts each time allocations exceed deallocations by the first - threshold (2,000 by default), so one copy ran thousands of collections - over live objects, about two thirds of its time (2026-10-23). A higher - threshold reclaims the same garbage in fewer passes; this module's peak - memory stayed the same. + threshold, and at the default of 2,000 there were one or two full + collections, each traversing every live tree, per copy: half the time of + building, measuring and cloning a reform simulation (2026-10-23). At + 100,000 there were about a quarter as many, this module ran about a + quarter faster, and its peak memory rose about 2%. """ thresholds = gc.get_threshold() gc.set_threshold(GC_THRESHOLD, *thresholds[1:]) @@ -125,14 +126,16 @@ def response(simulation, period=YEAR): def fresh_response(wages, period=YEAR): """The response a new reform simulation with these wages reports. - ``wages`` is one household's, or a tuple of several households' side by - side. 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: - situation = households(wages) if isinstance(wages, tuple) else household(wages) _fresh[key] = response( - make_simulation(situation=situation, reform=LSR_REFORM), period + make_simulation(situation=household(wages), reform=LSR_REFORM), period ) return _fresh[key] @@ -150,11 +153,15 @@ def set_wages(simulation, wages, period=YEAR): ) -# Every household takes another's wages, so each pin below runs all three -# levels in one simulation: one household goes from 50,000 to 100,000. +# The source holds one household at each of WAGES, so each pin 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 generated examples were fixed, +# and each reran a pinned path at other wages, which the households now cover. +# Drop ``phases`` to explore the remaining combinations. @settings(phases=[Phase.explicit], deadline=None) @given( kind=st.sampled_from(["clone", "branch"]), From 2ce982ce2a98e3ca05d9a2608ce0106f05f5ee33 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 24 Oct 2026 07:21:46 -0400 Subject: [PATCH 3/3] Pin the all-zero source case and keep only measured facts in the GC note Review (GPT-6.1 Sol) found that the old derandomized example (clone after measuring a lone 20,000 household, then 50,000) was the only source whose measured response was zero for every household; every three-household source has nonzero households. Pin it again as a fifth example, with the source's wages a parameter once more (a tuple). The fixture docstring now states only saved measurements: 6.4 million tracked objects per tree copy, none garbage after it; 29 against 8 full collections; and the CI effect on the tests this change does not otherwise touch. Co-Authored-By: Claude Opus 5.5 --- ...d-behavioral-measurement-copies.changed.md | 2 +- ..._behavioral_response_measurement_copies.py | 66 +++++++++++++------ 2 files changed, 47 insertions(+), 21 deletions(-) diff --git a/changelog.d/ci-speed-behavioral-measurement-copies.changed.md b/changelog.d/ci-speed-behavioral-measurement-copies.changed.md index 6b69b28062f..70498d030d0 100644 --- a/changelog.d/ci-speed-behavioral-measurement-copies.changed.md +++ b/changelog.d/ci-speed-behavioral-measurement-copies.changed.md @@ -1 +1 @@ -Cut the CI time of the behavioural-response measurement copy tests: collect garbage less often while they copy policy trees, and run each pinned copy case for all three wage levels in one simulation instead of three fixed generated examples. +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 a3f8f925533..466c161011d 100644 --- a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py +++ b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py @@ -67,13 +67,14 @@ 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: about five million objects, nearly all still in use. A - collection starts each time allocations exceed deallocations by the first - threshold, and at the default of 2,000 there were one or two full - collections, each traversing every live tree, per copy: half the time of - building, measuring and cloning a reform simulation (2026-10-23). At - 100,000 there were about a quarter as many, this module ran about a - quarter faster, and its peak memory rose about 2%. + 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:]) @@ -153,30 +154,53 @@ def set_wages(simulation, wages, period=YEAR): ) -# The source holds one household at each of WAGES, so each pin covers every +# 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 generated examples were fixed, -# and each reran a pinned path at other wages, which the households now cover. -# Drop ``phases`` to explore the remaining combinations. +# 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.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, copy_wages=WAGES) -@example(kind="clone", after_source_measured=False, copy_wages=WAGES) -@example(kind="clone", after_source_measured=True, copy_wages=SWAPPED_WAGES) -@example(kind="branch", after_source_measured=True, copy_wages=SWAPPED_WAGES) +@example( + 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=True, + source_wages=WAGES, + copy_wages=SWAPPED_WAGES, +) +@example( + 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="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, copy_wages + kind, after_source_measured, source_wages, copy_wages ): - source = make_simulation(situation=households(WAGES), reform=LSR_REFORM) + source = make_simulation(situation=households(source_wages), reform=LSR_REFORM) if after_source_measured: source_response = response(source) measured = { @@ -186,7 +210,7 @@ def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( ].items() } copy = source.clone() if kind == "clone" else source.get_branch("copy") - if copy_wages != WAGES: + if copy_wages != source_wages: set_wages(copy, copy_wages) copy_response = response(copy) if kind == "branch" and after_source_measured: @@ -194,7 +218,7 @@ def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( # With this reform's uniform elasticities, the response scales with wages. np.testing.assert_allclose( copy_response, - source_response * np.divide(copy_wages, WAGES), + source_response * np.divide(copy_wages, source_wages), atol=1e-3, ) copied_measurements = copy.__dict__[BEHAVIORAL_RESPONSE_CACHE_ATTR][str(YEAR)] @@ -204,7 +228,9 @@ def test_a_copy_reports_the_response_for_its_inputs_and_measurement_state( else: np.testing.assert_allclose(copy_response, fresh_response(copy_wages), atol=1e-3) if after_source_measured: - np.testing.assert_allclose(response(source), fresh_response(WAGES), atol=1e-3) + np.testing.assert_allclose( + response(source), fresh_response(source_wages), atol=1e-3 + ) def test_migrated_input_records_are_private_and_explicit_overrides_survive():