Skip to content

Simulations pickled in one process cannot be unpickled in another #568

Description

@MaxGhenis

Summary

A simulation, tax-benefit system or variable pickled in one Python process cannot be unpickled in another. Unpickling raises ModuleNotFoundError for a module named like 4688005354000_-7719298932528201769_income.

This is separate from the RecursionError that stopped every simulation from unpickling, even in the same process (fixed in #567). With that fix, a round trip inside one process works. Crossing a process boundary still fails.

Cause

TaxBenefitSystem.add_variables_from_file (policyengine_core/taxbenefitsystems/tax_benefit_system.py) loads every variable file under a generated module name and registers it in sys.modules:

module_name = f"{id(self)}_{hash(os.path.abspath(file_path))}_{file_name}"
spec = importlib.util.spec_from_file_location(module_name, file_path)
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module

Pickle stores classes and functions by __module__ and __qualname__. Every Variable subclass and formula therefore points at a module that exists only in the process that loaded it: id(self) is the system's address, and hash() of a string is salted per process (PYTHONHASHSEED). A second process has no such module, so loading the pickle fails. The same applies to enums defined in variable files. In #567, an EnumArray unpickled in another process comes back with possible_values = None for this reason instead of failing.

Reproduction

With #567 applied:

# process 1
import pickle
from policyengine_core.country_template import CountryTaxBenefitSystem
from policyengine_core.simulations import SimulationBuilder

s = SimulationBuilder().build_from_entities(
    CountryTaxBenefitSystem(),
    {"persons": {"a": {"salary": {"2025-01": 1000}}}, "households": {"h": {"parents": ["a"]}}},
)
s.calculate("income_tax", "2025-01")
open("sim.pkl", "wb").write(pickle.dumps(s))

# process 2
import pickle
pickle.loads(open("sim.pkl", "rb").read())
# ModuleNotFoundError: No module named '<id>_<hash>_income'

Why it matters

It blocks any use that ships a simulation or system to another process: multiprocessing with the spawn start method (the default on macOS and Windows), ProcessPoolExecutor, joblib/loky, Modal function arguments, and notebook or disk checkpoints. policyengine-us already expects systems to pickle: test_shared_policy_class_has_an_import_path publishes its runtime classes so "anything that pickles a system (a worker pool, joblib, a notebook checkpoint)" can find them, but the variables inside the system still cannot be found. Current downstream code avoids the problem: policyengine-us-data's worker pools pass file paths and build one Microsimulation per worker.

Options

  1. Stable module names. Use a name that every process derives the same way, such as the package-qualified dotted path when the file sits inside an importable package (policyengine_us.variables.gov.irs....), and keep a per-system suffix only when it is needed. The id(self) part exists so the same file loaded by two systems gives two sets of classes; any change has to keep that property where it matters (reforms that replace variables, test runners that build several systems).
  2. Pickle variables by name. Give Variable (and the system) a __reduce__ that stores the country package, system class and variable name, and rebuilds from a freshly loaded system in the receiving process. This avoids renaming modules, but a variable a reform changed at runtime would come back as the unreformed one unless the reform is replayed.
  3. Document it. State that simulations pickle only within one process, and that worker pools should pass inputs, not simulations.

Option 1 is the smallest if the collision concern can be met. It needs a decision before anyone builds it.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions