Skip to content

Resolve parameter nodes at instant on first read - #570

Draft
MaxGhenis wants to merge 4 commits into
masterfrom
lazy-parameter-node-at-instant
Draft

MaxGhenis wants to merge 4 commits into
masterfrom
lazy-parameter-node-at-instant

Conversation

@MaxGhenis

@MaxGhenis MaxGhenis commented Oct 2, 2026 •

Copy link
Copy Markdown
Contributor

Draft: measurements on policyengine-us and policyengine-uk are running; this description gets the before/after numbers when they finish.

Problem

ParameterNodeAtInstant.__init__ built every descendant up front. Every instant a tax-benefit system is asked about therefore materialised the whole parameter tree, and ParameterNode._at_instant_cache (plus TaxBenefitSystem._parameters_at_instant_cache) keeps it for the system's lifetime.

Measured on policyengine-us (132,396 parameter descendants) on master b78b0ba:

  • each instant costs about 34,000 node objects, 16 to 33 MB, and 4 to 5 s on a loaded host, on every system;
  • a household on a reform system added 349 MB, most of it the reform system's own at-instant trees.

A YAML run touches every year it tests and every month of a monthly variable, so the shared system and each reform system all pay this many times over. It was the third cause of the 118 GB policyengine-core test runs on 2026-10-02, next to the test-runner causes fixed in #569; the two PRs are independent.

Change

A node at instant resolves a child the first time it is read (__getattr__, __getitem__) and keeps it. It resolves all of its direct children only when something needs all of them: _children, iteration, repr, dir, vectorial indexing. A resolved child is a plain instance attribute, so hot reads cost what they did.

Invariants (tested)

  • For any instant and any order of reads (attribute, item, iteration, _children, dir, repr, vectorial indexing, names that do not exist), every read returns what the up-front build held, and a node read in full equals the up-front build, including child order. test_parameter_node_at_instant_property.py (Hypothesis) checks this against a reference that walks the ParameterNode tree directly, and it passes on master too.
  • A child resolves once; its identity is stable.
  • Unknown names raise the same ParameterNotFoundError / KeyError as before.
  • Intended difference: a node at instant held across an in-place change to its parameter tree reads, for a child not yet read, the value at the time of that first read; the old node held the value at construction. A node requested after the change behaves as before. (test_a_child_not_read_before_an_update_reads_the_update)

Tests

  • uv run pytest tests/core/parameters tests/core/parameters_fancy_indexing tests/core/test_parameters.py tests/core/test_tracers.py tests/core/test_reforms.py
  • The whole core suite passes locally.

🤖 Generated with Claude Code

Building every descendant in the constructor made each instant a system is
asked about cost a copy of the whole parameter tree, kept for the system's
lifetime. For policyengine-us that is about 34,000 node objects, 16-33 MB and
several seconds per instant, per tax-benefit system.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MaxGhenis and others added 2 commits October 2, 2026 09:52
…-coherent

Review of 404a3a3 found five regressions against the up-front build:
- concurrent first reads could raise KeyError/AttributeError or resolve a
  child twice: first resolution now runs under a module RLock;
- children named like the new state or helpers (_node, _resolved, _resolve,
  ...) read as internal objects: state is name-mangled, helpers are module
  functions, and children named like the node's own attributes resolve at
  construction, as before;
- a shallow copy and its original could resolve a scale twice: both now read
  one grow-only resolved map;
- a child with no value at the instant was re-resolved on every read: its
  absence is cached;
- deepcopy/pickle of a traced tree's raw node recursed through
  TracingParameterNodeAtInstant.__getattr__, which now raises AttributeError
  for its own unset fields.
A node also drops its parameter node once every child has been read.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Measurements of this PR combined with #569 on the two policyengine-us CI batches that are failing on 16 GB runners: see #569 (comment). Together they take the NY/OH/OK batch from 16.3 GB (3.32.11) to 7.0 GB and the USDA batch from 17.1 GB to 6.4 GB peak footprint.

…ases

Round-2 review of a90b524 found reads that still differed from the up-front
build. A node now keeps all of its own state in one _LazyChildren object under
a single reserved name (a parameter of that name is rejected):
- which children a node has, and their order, are fixed at construction
  (a child deleted from the tree later keeps its place);
- completion counts only those children, so add_child extras can no longer
  send the first read into recursive completion;
- resolution results and add_child entries are kept apart: items are shared
  by shallow copies, attributes belong to each copy, as before;
- a copy installs what it reads as attributes (no lock on warm reads);
- children named like attributes of the node's runtime type (dunders and
  subclass methods included) resolve at construction; a subclass overriding
  add_child gets the up-front build through it;
- a reentrant read keeps the first recorded child;
- nodes pickled by the up-front build load (__setstate__).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant