From ef07ed6c4e364d633a9cd60fcaf6543e2b0696ce Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 13:37:11 -0400 Subject: [PATCH 1/8] Compile each variable file once per process add_variables_from_file compiled every variable file again for each new system, so each system got new code objects for the same source. Coverage's sys.monitoring core keeps every code object alive, so it retained every system a process built. Each file is now compiled once per version (path, mtime, size) and the code object runs in each system's own module, so systems keep their own modules, classes, functions and globals. Fixes #607 Co-Authored-By: Claude Opus 5.5 --- changelog.d/share-variable-code.fixed.md | 1 + .../taxbenefitsystems/tax_benefit_system.py | 22 +++- tests/core/test_variable_code_sharing.py | 104 ++++++++++++++++++ 3 files changed, 126 insertions(+), 1 deletion(-) create mode 100644 changelog.d/share-variable-code.fixed.md create mode 100644 tests/core/test_variable_code_sharing.py diff --git a/changelog.d/share-variable-code.fixed.md b/changelog.d/share-variable-code.fixed.md new file mode 100644 index 000000000..87e03b133 --- /dev/null +++ b/changelog.d/share-variable-code.fixed.md @@ -0,0 +1 @@ +Compile each variable file once per process and run that code object in every system's own module, so tools that keep code objects alive, such as coverage's sys.monitoring core, no longer retain every system built (#607). diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 212de058d..c416dc549 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -10,6 +10,7 @@ from pathlib import Path import sys import traceback +import types import typing from typing import ( TYPE_CHECKING, @@ -55,6 +56,24 @@ log = logging.getLogger(__name__) +# Code compiled from each variable file, keyed by the file's absolute path, +# modification time and size. A file's code object does not depend on the +# system that runs it, so every system loaded from the same file runs the same +# code object, in a module of its own. Compiling again for each system would +# make new code objects every time, which tools that key their records by code +# object, such as coverage's sys.monitoring core, keep alive forever. +_variable_file_code: dict[tuple[str, int, int], types.CodeType] = {} + + +def _compiled_variable_file(loader, module_name: str, file_path: str): + """The code object for ``file_path``, compiled once per version of the file.""" + stat = os.stat(file_path) + key = (os.path.abspath(file_path), stat.st_mtime_ns, stat.st_size) + code = _variable_file_code.get(key) + if code is None: + code = _variable_file_code[key] = loader.get_code(module_name) + return code + class TaxBenefitSystem: """ @@ -349,7 +368,8 @@ def add_variables_from_file(self, file_path: str) -> None: spec = importlib.util.spec_from_file_location(module_name, file_path) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module - spec.loader.exec_module(module) + code = _compiled_variable_file(spec.loader, module_name, file_path) + exec(code, module.__dict__) except NameError as e: logging.error( str(e) diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py new file mode 100644 index 000000000..764d05274 --- /dev/null +++ b/tests/core/test_variable_code_sharing.py @@ -0,0 +1,104 @@ +"""Systems loaded from the same variable files run the same code objects. + +Each system still gets its own modules, classes and functions; only the code +compiled from a file is shared, and a changed file is compiled again. +""" + +import os +import sys +import textwrap + +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.taxbenefitsystems import tax_benefit_system as tbs_module + +VARIABLE_FILE = textwrap.dedent( + """ + from policyengine_core.country_template.entities import Person + from policyengine_core.periods import MONTH + from policyengine_core.variables import Variable + + + class shared_code_probe(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Shared code probe" + + def formula(person, period, parameters): + return person("salary", period) * 0 + {value} + """ +) + + +def _formula(system, name): + return system.variables[name].get_formula() + + +def _system_with_probe(path): + system = CountryTaxBenefitSystem() + system.add_variables_from_file(str(path)) + return system + + +def _value(system): + from policyengine_core.simulations import SimulationBuilder + + simulation = SimulationBuilder().build_from_entities( + system, {"persons": {"a": {"salary": {"2024-01": 100}}}} + ) + return simulation.calculate("shared_code_probe", "2024-01")[0] + + +def test_systems_share_code_but_not_modules_classes_or_functions(): + first = CountryTaxBenefitSystem() + second = CountryTaxBenefitSystem() + first_class = type(first.variables["disposable_income"]) + second_class = type(second.variables["disposable_income"]) + first_formula = _formula(first, "disposable_income") + second_formula = _formula(second, "disposable_income") + + assert first_class is not second_class + assert ( + sys.modules[first_class.__module__] + is not (sys.modules[second_class.__module__]) + ) + assert first_formula is not second_formula + assert first_formula.__globals__ is not second_formula.__globals__ + assert first_formula.__code__ is second_formula.__code__ + + +def test_shared_code_is_compiled_from_the_current_source(): + CountryTaxBenefitSystem() + assert tbs_module._variable_file_code + for (path, mtime_ns, size), code in tbs_module._variable_file_code.items(): + if not os.path.exists(path): + continue # a temporary file from another test + stat = os.stat(path) + if (stat.st_mtime_ns, stat.st_size) != (mtime_ns, size): + continue # an older version of a file that changed + with open(path, encoding="utf-8") as f: + source = f.read() + assert code == compile(source, path, "exec", dont_inherit=True), path + + +def test_a_changed_file_is_compiled_again(tmp_path): + path = tmp_path / "probe.py" + path.write_text(VARIABLE_FILE.format(value=1)) + first = _system_with_probe(path) + second = _system_with_probe(path) + path.write_text(VARIABLE_FILE.format(value=20)) + stat = os.stat(path) + os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000_000)) + third = _system_with_probe(path) + + assert _formula(first, "shared_code_probe").__code__ is ( + _formula(second, "shared_code_probe").__code__ + ) + assert _formula(third, "shared_code_probe").__code__ is not ( + _formula(first, "shared_code_probe").__code__ + ) + assert _value(first) == pytest.approx(1) + assert _value(second) == pytest.approx(1) + assert _value(third) == pytest.approx(20) From db9182b0c258857de0dcf86b3381885fff786cf3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 16:12:48 -0400 Subject: [PATCH 2/8] Key shared variable code by the source's SHA-256, keep one version per file Review of ef07ed6c4: a cache keyed by path, mtime and size served stale code after an edit that kept both (or a symlink retargeted to a file with the same metadata); two threads missing at once could publish different code objects; and old versions were never evicted. The cache now maps each absolute path to the SHA-256 of the source it compiled and the code, filled under a lock and compiled from the bytes it hashed (about 0.07 s of CPU for all 1,110 policyengine-uk variable files, once per process). The changelog no longer implies the variable modules stop being retained (#520). Co-Authored-By: Claude Opus 5.5 --- changelog.d/share-variable-code.fixed.md | 2 +- .../taxbenefitsystems/tax_benefit_system.py | 44 ++++++++++------- tests/core/test_variable_code_sharing.py | 47 +++++++++++++++---- 3 files changed, 65 insertions(+), 28 deletions(-) diff --git a/changelog.d/share-variable-code.fixed.md b/changelog.d/share-variable-code.fixed.md index 87e03b133..03368ba72 100644 --- a/changelog.d/share-variable-code.fixed.md +++ b/changelog.d/share-variable-code.fixed.md @@ -1 +1 @@ -Compile each variable file once per process and run that code object in every system's own module, so tools that keep code objects alive, such as coverage's sys.monitoring core, no longer retain every system built (#607). +Compile each variable file once per process for each version of its source, and run that code object in every system's own module, so tools that keep code objects alive, such as coverage's sys.monitoring core, no longer hold a new copy of every variable file's code for each system built (#607). The variable modules themselves still stay in sys.modules (#520). diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index c416dc549..f94a911c4 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -2,6 +2,7 @@ import copy import glob +import hashlib import importlib import inspect import logging @@ -9,6 +10,7 @@ import yaml from pathlib import Path import sys +import threading import traceback import types import typing @@ -56,23 +58,29 @@ log = logging.getLogger(__name__) -# Code compiled from each variable file, keyed by the file's absolute path, -# modification time and size. A file's code object does not depend on the -# system that runs it, so every system loaded from the same file runs the same -# code object, in a module of its own. Compiling again for each system would -# make new code objects every time, which tools that key their records by code -# object, such as coverage's sys.monitoring core, keep alive forever. -_variable_file_code: dict[tuple[str, int, int], types.CodeType] = {} - - -def _compiled_variable_file(loader, module_name: str, file_path: str): - """The code object for ``file_path``, compiled once per version of the file.""" - stat = os.stat(file_path) - key = (os.path.abspath(file_path), stat.st_mtime_ns, stat.st_size) - code = _variable_file_code.get(key) - if code is None: - code = _variable_file_code[key] = loader.get_code(module_name) - return code +# Code compiled from each variable file, by the file's absolute path, with the +# SHA-256 of the source it was compiled from. A file's code object does not +# depend on the system that runs it, so every system loaded from the same +# source runs the same code object, in a module of its own. Compiling again for +# each system would make new code objects every time, and tools that keep code +# objects alive, such as coverage's sys.monitoring core, would then hold a copy +# of every variable file's code for each system built. Only the latest version +# of each file is kept. +_variable_file_code: dict[str, tuple[bytes, types.CodeType]] = {} +_variable_file_code_lock = threading.Lock() + + +def _compiled_variable_file(loader, file_path: str) -> types.CodeType: + """The code object for ``file_path``'s current source, compiled once.""" + source = loader.get_data(file_path) + digest = hashlib.sha256(source).digest() + path = os.path.abspath(file_path) + with _variable_file_code_lock: + cached = _variable_file_code.get(path) + if cached is None or cached[0] != digest: + cached = (digest, loader.source_to_code(source, file_path)) + _variable_file_code[path] = cached + return cached[1] class TaxBenefitSystem: @@ -368,7 +376,7 @@ def add_variables_from_file(self, file_path: str) -> None: spec = importlib.util.spec_from_file_location(module_name, file_path) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module - code = _compiled_variable_file(spec.loader, module_name, file_path) + code = _compiled_variable_file(spec.loader, file_path) exec(code, module.__dict__) except NameError as e: logging.error( diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index 764d05274..5ddbc696a 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -1,12 +1,17 @@ """Systems loaded from the same variable files run the same code objects. Each system still gets its own modules, classes and functions; only the code -compiled from a file is shared, and a changed file is compiled again. +compiled from a file's source is shared. Any change to the source, even one +that keeps the file's size and modification time, is compiled again, and only +the latest version of each file is kept. """ +import hashlib import os import sys import textwrap +import threading +from importlib.machinery import SourceFileLoader import pytest @@ -72,14 +77,13 @@ def test_systems_share_code_but_not_modules_classes_or_functions(): def test_shared_code_is_compiled_from_the_current_source(): CountryTaxBenefitSystem() assert tbs_module._variable_file_code - for (path, mtime_ns, size), code in tbs_module._variable_file_code.items(): + for path, (digest, code) in tbs_module._variable_file_code.items(): if not os.path.exists(path): continue # a temporary file from another test - stat = os.stat(path) - if (stat.st_mtime_ns, stat.st_size) != (mtime_ns, size): - continue # an older version of a file that changed - with open(path, encoding="utf-8") as f: + with open(path, "rb") as f: source = f.read() + if hashlib.sha256(source).digest() != digest: + continue # a file another test changed after loading it assert code == compile(source, path, "exec", dont_inherit=True), path @@ -88,9 +92,11 @@ def test_a_changed_file_is_compiled_again(tmp_path): path.write_text(VARIABLE_FILE.format(value=1)) first = _system_with_probe(path) second = _system_with_probe(path) - path.write_text(VARIABLE_FILE.format(value=20)) + # The same size and modification time: only the content tells them apart. stat = os.stat(path) - os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000_000)) + path.write_text(VARIABLE_FILE.format(value=2)) + os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns)) + assert os.stat(path).st_size == stat.st_size third = _system_with_probe(path) assert _formula(first, "shared_code_probe").__code__ is ( @@ -101,4 +107,27 @@ def test_a_changed_file_is_compiled_again(tmp_path): ) assert _value(first) == pytest.approx(1) assert _value(second) == pytest.approx(1) - assert _value(third) == pytest.approx(20) + assert _value(third) == pytest.approx(2) + cached = [key for key in tbs_module._variable_file_code if key == str(path)] + assert cached == [str(path)] + + +def test_concurrent_loads_share_one_code_object(tmp_path): + path = tmp_path / "probe.py" + path.write_text(VARIABLE_FILE.format(value=3)) + loader = SourceFileLoader("probe", str(path)) + threads = 8 + barrier = threading.Barrier(threads) + codes = [None] * threads + + def load(i): + barrier.wait() + codes[i] = tbs_module._compiled_variable_file(loader, str(path)) + + workers = [threading.Thread(target=load, args=(i,)) for i in range(threads)] + for worker in workers: + worker.start() + for worker in workers: + worker.join() + + assert all(code is codes[0] for code in codes) From ba633db813c255104bed7a784cac7ffe467e8dc0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 17:04:35 -0400 Subject: [PATCH 3/8] Compile under the loader's absolute path, survive fork, tighten the tests Opus delta review of db9182b0c: - Code compiled with the caller's spelling of the path, so a relative path gave relative co_filename and error file names (exec_module uses the loader's absolute path), and the cache kept whichever spelling came first. The read, key and compile now all use the loader's absolute path. - A child forked while another thread held the lock would deadlock; the lock is re-created after fork. - The thread test passed when every load raised and could not see a missing lock under the GIL. It now uses a loader that counts and sleeps in each compile, collects thread errors, and asserts one compilation. - The eviction assertion could not fail; it now checks the stored digest and that the old code object is gone. - New Hypothesis test (own file, importorskip like the repo's other property tests): any sequence of writes and loads returns the current content's code under the absolute path, the same object until the next write, and only the latest version is cached. Co-Authored-By: Claude Opus 5.5 --- .../taxbenefitsystems/tax_benefit_system.py | 29 ++++++++--- tests/core/test_variable_code_sharing.py | 49 ++++++++++++++++--- .../test_variable_code_sharing_property.py | 49 +++++++++++++++++++ 3 files changed, 113 insertions(+), 14 deletions(-) create mode 100644 tests/core/test_variable_code_sharing_property.py diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index f94a911c4..8d6dbd329 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -65,20 +65,35 @@ # each system would make new code objects every time, and tools that keep code # objects alive, such as coverage's sys.monitoring core, would then hold a copy # of every variable file's code for each system built. Only the latest version -# of each file is kept. +# of each file is kept, for the life of the process (entries for files deleted +# since are not removed). Variable files' .pyc files are not used: the code is +# compiled from the bytes that were hashed, once per process. _variable_file_code: dict[str, tuple[bytes, types.CodeType]] = {} _variable_file_code_lock = threading.Lock() -def _compiled_variable_file(loader, file_path: str) -> types.CodeType: - """The code object for ``file_path``'s current source, compiled once.""" - source = loader.get_data(file_path) +def _reset_variable_file_code_lock() -> None: + # A child forked while another thread held the lock would otherwise + # deadlock on its first variable file. + global _variable_file_code_lock + _variable_file_code_lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_variable_file_code_lock) + + +def _compiled_variable_file(loader) -> types.CodeType: + """The code object for the loader's file's current source, compiled once.""" + # The loader's path, as exec_module would use: absolute, whatever spelling + # the caller passed, so the key and every code object's co_filename agree. + path = os.path.abspath(loader.get_filename()) + source = loader.get_data(path) digest = hashlib.sha256(source).digest() - path = os.path.abspath(file_path) with _variable_file_code_lock: cached = _variable_file_code.get(path) if cached is None or cached[0] != digest: - cached = (digest, loader.source_to_code(source, file_path)) + cached = (digest, loader.source_to_code(source, path)) _variable_file_code[path] = cached return cached[1] @@ -376,7 +391,7 @@ def add_variables_from_file(self, file_path: str) -> None: spec = importlib.util.spec_from_file_location(module_name, file_path) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module - code = _compiled_variable_file(spec.loader, file_path) + code = _compiled_variable_file(spec.loader) exec(code, module.__dict__) except NameError as e: logging.error( diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index 5ddbc696a..3ce251f99 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -11,6 +11,8 @@ import sys import textwrap import threading +import time +import types from importlib.machinery import SourceFileLoader import pytest @@ -85,6 +87,19 @@ def test_shared_code_is_compiled_from_the_current_source(): if hashlib.sha256(source).digest() != digest: continue # a file another test changed after loading it assert code == compile(source, path, "exec", dont_inherit=True), path + # Code equality ignores the file name, so check it separately. + assert code.co_filename == path + + +def test_a_relative_path_compiles_under_its_absolute_path(tmp_path, monkeypatch): + (tmp_path / "relative_probe.py").write_text(VARIABLE_FILE.format(value=4)) + monkeypatch.chdir(tmp_path) + by_relative = _system_with_probe("relative_probe.py") + by_absolute = _system_with_probe(tmp_path / "relative_probe.py") + code = _formula(by_relative, "shared_code_probe").__code__ + + assert code.co_filename == str(tmp_path / "relative_probe.py") + assert _formula(by_absolute, "shared_code_probe").__code__ is code def test_a_changed_file_is_compiled_again(tmp_path): @@ -92,6 +107,7 @@ def test_a_changed_file_is_compiled_again(tmp_path): path.write_text(VARIABLE_FILE.format(value=1)) first = _system_with_probe(path) second = _system_with_probe(path) + old = tbs_module._variable_file_code[str(path)][1] # The same size and modification time: only the content tells them apart. stat = os.stat(path) path.write_text(VARIABLE_FILE.format(value=2)) @@ -108,21 +124,37 @@ def test_a_changed_file_is_compiled_again(tmp_path): assert _value(first) == pytest.approx(1) assert _value(second) == pytest.approx(1) assert _value(third) == pytest.approx(2) - cached = [key for key in tbs_module._variable_file_code if key == str(path)] - assert cached == [str(path)] + # Only the new version is kept. + digest, _ = tbs_module._variable_file_code[str(path)] + assert digest == hashlib.sha256(path.read_bytes()).digest() + assert not any(code is old for _, code in tbs_module._variable_file_code.values()) + + +class _CountingLoader(SourceFileLoader): + """Counts compilations and sleeps in each, so that threads overlap.""" + + compilations = 0 + + def source_to_code(self, data, path, *args, **kwargs): + type(self).compilations += 1 + time.sleep(0.05) # releases the GIL + return super().source_to_code(data, path, *args, **kwargs) -def test_concurrent_loads_share_one_code_object(tmp_path): +def test_concurrent_loads_compile_once_and_share_the_code(tmp_path): path = tmp_path / "probe.py" path.write_text(VARIABLE_FILE.format(value=3)) - loader = SourceFileLoader("probe", str(path)) + loader = _CountingLoader("probe", str(path)) threads = 8 barrier = threading.Barrier(threads) - codes = [None] * threads + codes, errors = [None] * threads, [] def load(i): - barrier.wait() - codes[i] = tbs_module._compiled_variable_file(loader, str(path)) + try: + barrier.wait() + codes[i] = tbs_module._compiled_variable_file(loader) + except BaseException as error: # reported below, not lost in a thread + errors.append(error) workers = [threading.Thread(target=load, args=(i,)) for i in range(threads)] for worker in workers: @@ -130,4 +162,7 @@ def load(i): for worker in workers: worker.join() + assert not errors + assert _CountingLoader.compilations == 1 + assert all(isinstance(code, types.CodeType) for code in codes) assert all(code is codes[0] for code in codes) diff --git a/tests/core/test_variable_code_sharing_property.py b/tests/core/test_variable_code_sharing_property.py new file mode 100644 index 000000000..9c43855d9 --- /dev/null +++ b/tests/core/test_variable_code_sharing_property.py @@ -0,0 +1,49 @@ +"""Property test for the shared variable code cache (see test_variable_code_sharing.py).""" + +import hashlib +import tempfile +from importlib.machinery import SourceFileLoader +from pathlib import Path + +import pytest + +pytest.importorskip("hypothesis") +from hypothesis import given, settings, strategies as st + +from policyengine_core.taxbenefitsystems import tax_benefit_system as tbs_module +from tests.core.test_variable_code_sharing import VARIABLE_FILE + + +# Any sequence of writes and loads of one file: every load returns the code of +# the file's current content under its absolute path, a load with no write +# since the last one returns the same object, and the cache holds only the +# latest version. +@settings(max_examples=50, deadline=None) +@given( + steps=st.lists( + st.one_of(st.integers(0, 3).map(lambda v: ("write", v)), st.just(("load", 0))), + min_size=1, + max_size=12, + ) +) +def test_any_writes_and_loads_give_the_current_code(steps): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory).resolve() / "probe.py" + path.write_text(VARIABLE_FILE.format(value=0)) + loader = SourceFileLoader("probe", str(path)) + previous = None + for action, value in steps: + if action == "write": + path.write_text(VARIABLE_FILE.format(value=value)) + previous = None + continue + code = tbs_module._compiled_variable_file(loader) + source = path.read_bytes() + assert code == compile(source, str(path), "exec", dont_inherit=True) + assert code.co_filename == str(path) + if previous is not None: + assert code is previous + previous = code + digest, cached = tbs_module._variable_file_code[str(path)] + assert digest == hashlib.sha256(source).digest() and cached is code + tbs_module._variable_file_code.pop(str(path), None) From f52aefaa3b9356c2a2b718c338a4960466130175 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 17:22:08 -0400 Subject: [PATCH 4/8] Reset the thread test's counter, clean up the property test, test the fork reset Opus confirmation review of ba633db81 (approved) left nits: the counting loader's class counter wasn't reset, so a rerun of the thread test failed; the property test leaked cache entries when an example failed; nothing tested the lock's reset after fork (a POSIX test now forks while holding it and checks the child's lock is free); and the comment now says the path is absolute and normalized rather than as exec_module spells it. Co-Authored-By: Claude Opus 5.5 --- .../taxbenefitsystems/tax_benefit_system.py | 4 +-- tests/core/test_variable_code_sharing.py | 15 +++++++- .../test_variable_code_sharing_property.py | 34 ++++++++++--------- 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 8d6dbd329..431f90923 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -85,8 +85,8 @@ def _reset_variable_file_code_lock() -> None: def _compiled_variable_file(loader) -> types.CodeType: """The code object for the loader's file's current source, compiled once.""" - # The loader's path, as exec_module would use: absolute, whatever spelling - # the caller passed, so the key and every code object's co_filename agree. + # The loader's path, absolute and normalized whatever spelling the caller + # passed, so the key and every code object's co_filename agree. path = os.path.abspath(loader.get_filename()) source = loader.get_data(path) digest = hashlib.sha256(source).digest() diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index 3ce251f99..8fcc09812 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -141,9 +141,10 @@ def source_to_code(self, data, path, *args, **kwargs): return super().source_to_code(data, path, *args, **kwargs) -def test_concurrent_loads_compile_once_and_share_the_code(tmp_path): +def test_concurrent_loads_compile_once_and_share_the_code(tmp_path, monkeypatch): path = tmp_path / "probe.py" path.write_text(VARIABLE_FILE.format(value=3)) + monkeypatch.setattr(_CountingLoader, "compilations", 0) loader = _CountingLoader("probe", str(path)) threads = 8 barrier = threading.Barrier(threads) @@ -166,3 +167,15 @@ def load(i): assert _CountingLoader.compilations == 1 assert all(isinstance(code, types.CodeType) for code in codes) assert all(code is codes[0] for code in codes) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_forked_child_gets_a_free_lock(): + # A child forked while another thread held the lock would deadlock on its + # first variable file unless the lock is re-created after fork. + with tbs_module._variable_file_code_lock: + pid = os.fork() + if pid == 0: + os._exit(1 if tbs_module._variable_file_code_lock.locked() else 0) + _, status = os.waitpid(pid, 0) + assert os.waitstatus_to_exitcode(status) == 0 diff --git a/tests/core/test_variable_code_sharing_property.py b/tests/core/test_variable_code_sharing_property.py index 9c43855d9..497373f28 100644 --- a/tests/core/test_variable_code_sharing_property.py +++ b/tests/core/test_variable_code_sharing_property.py @@ -31,19 +31,21 @@ def test_any_writes_and_loads_give_the_current_code(steps): path = Path(directory).resolve() / "probe.py" path.write_text(VARIABLE_FILE.format(value=0)) loader = SourceFileLoader("probe", str(path)) - previous = None - for action, value in steps: - if action == "write": - path.write_text(VARIABLE_FILE.format(value=value)) - previous = None - continue - code = tbs_module._compiled_variable_file(loader) - source = path.read_bytes() - assert code == compile(source, str(path), "exec", dont_inherit=True) - assert code.co_filename == str(path) - if previous is not None: - assert code is previous - previous = code - digest, cached = tbs_module._variable_file_code[str(path)] - assert digest == hashlib.sha256(source).digest() and cached is code - tbs_module._variable_file_code.pop(str(path), None) + try: + previous = None + for action, value in steps: + if action == "write": + path.write_text(VARIABLE_FILE.format(value=value)) + previous = None + continue + code = tbs_module._compiled_variable_file(loader) + source = path.read_bytes() + assert code == compile(source, str(path), "exec", dont_inherit=True) + assert code.co_filename == str(path) + if previous is not None: + assert code is previous + previous = code + digest, cached = tbs_module._variable_file_code[str(path)] + assert digest == hashlib.sha256(source).digest() and cached is code + finally: + tbs_module._variable_file_code.pop(str(path), None) From 847fe46822918e0957b25fe83d9ad8b06af6e364 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 17:30:43 -0400 Subject: [PATCH 5/8] Exit the forked child in a finally, as test_data_storage_dir does So an error in the child can never return it to pytest to run the rest of the suite while the parent waits. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_variable_code_sharing.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index 8fcc09812..cb3416731 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -176,6 +176,10 @@ def test_a_forked_child_gets_a_free_lock(): with tbs_module._variable_file_code_lock: pid = os.fork() if pid == 0: - os._exit(1 if tbs_module._variable_file_code_lock.locked() else 0) + status = 2 + try: + status = 1 if tbs_module._variable_file_code_lock.locked() else 0 + finally: + os._exit(status) _, status = os.waitpid(pid, 0) assert os.waitstatus_to_exitcode(status) == 0 From 1ff8941b347b47f43a6b8bb0f1b1d166373f00d8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 21:09:42 -0400 Subject: [PATCH 6/8] Preserve symlink parent traversal when loading variable source --- .../taxbenefitsystems/tax_benefit_system.py | 12 +++--- tests/core/test_variable_code_sharing.py | 42 +++++++++++++++++++ 2 files changed, 49 insertions(+), 5 deletions(-) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 431f90923..63cb2eb94 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -58,7 +58,7 @@ log = logging.getLogger(__name__) -# Code compiled from each variable file, by the file's absolute path, with the +# Code compiled from each variable file, by its absolute loader path, with the # SHA-256 of the source it was compiled from. A file's code object does not # depend on the system that runs it, so every system loaded from the same # source runs the same code object, in a module of its own. Compiling again for @@ -85,10 +85,12 @@ def _reset_variable_file_code_lock() -> None: def _compiled_variable_file(loader) -> types.CodeType: """The code object for the loader's file's current source, compiled once.""" - # The loader's path, absolute and normalized whatever spelling the caller - # passed, so the key and every code object's co_filename agree. - path = os.path.abspath(loader.get_filename()) - source = loader.get_data(path) + # Read the loader's exact filename: normalizing '..' before following a + # symlink can select a different file. Preserve that spelling in the + # absolute cache key and co_filename too, so different files cannot collide. + filename = loader.get_filename() + path = filename if os.path.isabs(filename) else os.path.join(os.getcwd(), filename) + source = loader.get_data(filename) digest = hashlib.sha256(source).digest() with _variable_file_code_lock: cached = _variable_file_code.get(path) diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index cb3416731..21a3ca5b8 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -102,6 +102,48 @@ def test_a_relative_path_compiles_under_its_absolute_path(tmp_path, monkeypatch) assert _formula(by_absolute, "shared_code_probe").__code__ is code +@pytest.mark.parametrize("relative", [False, True]) +def test_a_symlink_then_parent_path_reads_the_os_resolved_file( + tmp_path, monkeypatch, relative +): + repo = tmp_path / "repo" + other = tmp_path / "other" + repo.mkdir() + (other / "subdir").mkdir(parents=True) + link = repo / "link" + try: + link.symlink_to(other / "subdir", target_is_directory=True) + except (OSError, NotImplementedError) as error: + pytest.skip(f"Cannot create directory symlinks: {error}") + intended = other / "probe.py" + intended.write_text(VARIABLE_FILE.format(value=1)) + decoy = repo / "probe.py" + decoy.write_text(VARIABLE_FILE.format(value=9)) + linked_path = link / ".." / "probe.py" + if relative: + monkeypatch.chdir(tmp_path) + linked_path = linked_path.relative_to(tmp_path) + + first = _system_with_probe(linked_path) + assert _value(first) == pytest.approx(1) + intended.write_text(VARIABLE_FILE.format(value=2)) + second = _system_with_probe(linked_path) + assert _value(second) == pytest.approx(2) + assert _value(first) == pytest.approx(1) + assert _value(_system_with_probe(decoy)) == pytest.approx(9) + assert _formula(first, "shared_code_probe").__code__ is not ( + _formula(second, "shared_code_probe").__code__ + ) + # Even identical sources at these different paths need separate entries: + # collapsing '..' would give both files the decoy's cache key. + decoy.write_text(VARIABLE_FILE.format(value=2)) + by_decoy = _system_with_probe(decoy) + assert _value(by_decoy) == pytest.approx(2) + assert _formula(by_decoy, "shared_code_probe").__code__ is not ( + _formula(second, "shared_code_probe").__code__ + ) + + def test_a_changed_file_is_compiled_again(tmp_path): path = tmp_path / "probe.py" path.write_text(VARIABLE_FILE.format(value=1)) From 5d8187921e2e2d646a252055af00839601a65f89 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 21:52:00 -0400 Subject: [PATCH 7/8] Make concurrent variable compilation test deterministic --- tests/core/test_variable_code_sharing.py | 67 ++++++++++++++++++------ 1 file changed, 50 insertions(+), 17 deletions(-) diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index 21a3ca5b8..6675a0f50 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -11,7 +11,6 @@ import sys import textwrap import threading -import time import types from importlib.machinery import SourceFileLoader @@ -172,29 +171,61 @@ def test_a_changed_file_is_compiled_again(tmp_path): assert not any(code is old for _, code in tbs_module._variable_file_code.values()) -class _CountingLoader(SourceFileLoader): - """Counts compilations and sleeps in each, so that threads overlap.""" - - compilations = 0 - - def source_to_code(self, data, path, *args, **kwargs): - type(self).compilations += 1 - time.sleep(0.05) # releases the GIL - return super().source_to_code(data, path, *args, **kwargs) - - def test_concurrent_loads_compile_once_and_share_the_code(tmp_path, monkeypatch): path = tmp_path / "probe.py" path.write_text(VARIABLE_FILE.format(value=3)) - monkeypatch.setattr(_CountingLoader, "compilations", 0) - loader = _CountingLoader("probe", str(path)) + loader = SourceFileLoader("probe", str(path)) threads = 8 barrier = threading.Barrier(threads) + all_attempted = threading.Event() + observations_lock = threading.Lock() + attempted = set() + compilations = 0 codes, errors = [None] * threads, [] + def record_attempt(): + with observations_lock: + attempted.add(threading.get_ident()) + if len(attempted) == threads: + all_attempted.set() + + real_lock = tbs_module._variable_file_code_lock + + class ObservedLock: + def __enter__(self): + # Count followers before they block on the first compiler's lock. + record_attempt() + return real_lock.__enter__() + + def __exit__(self, *args): + return real_lock.__exit__(*args) + + class ObservedCache(dict): + def get(self, key, default=None): + cached = super().get(key, default) + # Without the lock, every thread must capture its miss before the + # first compiler is allowed to publish, including the last thread. + record_attempt() + return cached + + source_to_code = loader.source_to_code + + def compile_when_all_attempted(*args, **kwargs): + nonlocal compilations + with observations_lock: + compilations += 1 + # The timeout only diagnoses a deadlock; overlap is controlled by the + # event, irrespective of the speed or scheduling of the workers. + assert all_attempted.wait(10), "not all workers reached the cache" + return source_to_code(*args, **kwargs) + + monkeypatch.setattr(tbs_module, "_variable_file_code_lock", ObservedLock()) + monkeypatch.setattr(tbs_module, "_variable_file_code", ObservedCache()) + monkeypatch.setattr(loader, "source_to_code", compile_when_all_attempted) + def load(i): try: - barrier.wait() + barrier.wait(timeout=10) codes[i] = tbs_module._compiled_variable_file(loader) except BaseException as error: # reported below, not lost in a thread errors.append(error) @@ -203,10 +234,12 @@ def load(i): for worker in workers: worker.start() for worker in workers: - worker.join() + worker.join(timeout=15) + assert not any(worker.is_alive() for worker in workers) assert not errors - assert _CountingLoader.compilations == 1 + assert len(attempted) == threads + assert compilations == 1 assert all(isinstance(code, types.CodeType) for code in codes) assert all(code is codes[0] for code in codes) From 3f0590789789fcaa4259bf699331cb5441b4d2d1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 12 Oct 2026 06:45:37 -0400 Subject: [PATCH 8/8] Preserve bytecode variable loading with source code sharing --- changelog.d/share-variable-code.fixed.md | 2 +- .../taxbenefitsystems/tax_benefit_system.py | 12 ++++++++---- tests/core/test_variable_code_sharing.py | 15 +++++++++++++++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/changelog.d/share-variable-code.fixed.md b/changelog.d/share-variable-code.fixed.md index 03368ba72..61102ef5f 100644 --- a/changelog.d/share-variable-code.fixed.md +++ b/changelog.d/share-variable-code.fixed.md @@ -1 +1 @@ -Compile each variable file once per process for each version of its source, and run that code object in every system's own module, so tools that keep code objects alive, such as coverage's sys.monitoring core, no longer hold a new copy of every variable file's code for each system built (#607). The variable modules themselves still stay in sys.modules (#520). +Compile each variable source file once per process for each source version, reusing that code object in each system's own module so tools such as coverage's sys.monitoring core retain fewer copies (#607), while keeping variable modules in sys.modules (#520) and preserving the original execution path for bytecode-only and other loaders. diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 63cb2eb94..6766dbc7c 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -4,6 +4,7 @@ import glob import hashlib import importlib +from importlib.machinery import SourceFileLoader import inspect import logging import os @@ -66,8 +67,8 @@ # objects alive, such as coverage's sys.monitoring core, would then hold a copy # of every variable file's code for each system built. Only the latest version # of each file is kept, for the life of the process (entries for files deleted -# since are not removed). Variable files' .pyc files are not used: the code is -# compiled from the bytes that were hashed, once per process. +# since are not removed). For source loaders, .pyc files are not used: the code +# is compiled from the bytes that were hashed, once per process. _variable_file_code: dict[str, tuple[bytes, types.CodeType]] = {} _variable_file_code_lock = threading.Lock() @@ -393,8 +394,11 @@ def add_variables_from_file(self, file_path: str) -> None: spec = importlib.util.spec_from_file_location(module_name, file_path) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module - code = _compiled_variable_file(spec.loader) - exec(code, module.__dict__) + if isinstance(spec.loader, SourceFileLoader): + code = _compiled_variable_file(spec.loader) + exec(code, module.__dict__) + else: + spec.loader.exec_module(module) except NameError as e: logging.error( str(e) diff --git a/tests/core/test_variable_code_sharing.py b/tests/core/test_variable_code_sharing.py index 6675a0f50..256565f35 100644 --- a/tests/core/test_variable_code_sharing.py +++ b/tests/core/test_variable_code_sharing.py @@ -8,6 +8,7 @@ import hashlib import os +import py_compile import sys import textwrap import threading @@ -90,6 +91,20 @@ def test_shared_code_is_compiled_from_the_current_source(): assert code.co_filename == path +def test_sourceless_bytecode_variable_file_loads_and_runs_formula(tmp_path): + source = tmp_path / "probe.py" + bytecode = source.with_suffix(".pyc") + source.write_text(VARIABLE_FILE.format(value=4)) + try: + py_compile.compile(str(source), cfile=str(bytecode), doraise=True) + except (OSError, NotImplementedError) as error: + pytest.skip(f"Cannot write bytecode: {error}") + source.unlink() + + system = _system_with_probe(bytecode) + assert _value(system) == pytest.approx(4) + + def test_a_relative_path_compiles_under_its_absolute_path(tmp_path, monkeypatch): (tmp_path / "relative_probe.py").write_text(VARIABLE_FILE.format(value=4)) monkeypatch.chdir(tmp_path)