Compile each variable file once per process

changelog.d/share-variable-code.fixed.md

@@ -0,0 +1 @@
+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.

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -2,14 +2,18 @@
 
 import copy
 import glob
+import hashlib
 import importlib
+from importlib.machinery import SourceFileLoader
 import inspect
 import logging
 import os
 import yaml
 from pathlib import Path
 import sys
+import threading
 import traceback
+import types
 import typing
 from typing import (
     TYPE_CHECKING,
@@ -55,6 +59,47 @@
 
 log = logging.getLogger(__name__)
 
+# 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
+# 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, for the life of the process (entries for files deleted
+# 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()
+
+
+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."""
+    # 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)
+        if cached is None or cached[0] != digest:
+            cached = (digest, loader.source_to_code(source, path))
+            _variable_file_code[path] = cached
+    return cached[1]
+
 
 class TaxBenefitSystem:
     """
@@ -349,7 +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
-                spec.loader.exec_module(module)
+                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)

tests/core/test_variable_code_sharing.py

@@ -0,0 +1,275 @@
+"""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'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 py_compile
+import sys
+import textwrap
+import threading
+import types
+from importlib.machinery import SourceFileLoader
+
+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, (digest, code) in tbs_module._variable_file_code.items():
+        if not os.path.exists(path):
+            continue  # a temporary file from another test
+        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
+        # Code equality ignores the file name, so check it separately.
+        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)
+    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
+
+
+@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))
+    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))
+    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 (
+        _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(2)
+    # 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())
+
+
+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))
+    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(timeout=10)
+            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:
+        worker.start()
+    for worker in workers:
+        worker.join(timeout=15)
+
+    assert not any(worker.is_alive() for worker in workers)
+    assert not errors
+    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)
+
+
+@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:
+            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