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

@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.

@anth-volk

Copy link
Copy Markdown
Collaborator

Rebased onto master (b909d5a6) and force-pushed: de47cde0 → 089761c5. Beyond the rebase:

Details are under "October 14 rebase" in the description.

anth-volk added a commit that referenced this pull request Oct 14, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
anth-volk added a commit that referenced this pull request Oct 14, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
anth-volk added a commit that referenced this pull request Oct 14, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
anth-volk added a commit that referenced this pull request Oct 14, 2026
#570 now provides the walk, so modify_parameters uses it instead of its own.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
anth-volk added a commit that referenced this pull request Oct 15, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MaxGhenis and others added 4 commits October 16, 2026 00:20
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>
…-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>
…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>
@anth-volk

Copy link
Copy Markdown
Collaborator

I've dropped the commits I added to this branch and restored Max's description.

  • Dropped:
    • 089761c5, which let lazy nodes pickle under protocols 0 and 1;
    • 62d63079 and b26aa0a7, cache-consistency fixes;
    • 5719353e, a merge of master;
    • four review fixes that were never pushed.
  • Now on the branch: only Max's 4 commits, rebased onto master (4adfea70) and force-pushed: 5719353e → 31469a4d. The October 14 rebase's one change to Max's code remains: the _OWN_ATTRIBUTES guard in the tracing wrapper is gone, because master's _wrapped covers that case.
  • Expected CI failures: without 089761c5, four of master's copy-and-pickle tests from Let vectorial parameter nodes be copied and pickled #584 fail on this branch, the protocol 0 and 1 cases of:
    • test_tracing_wrapper_of_a_slotted_node_can_be_pickled;
    • test_copied_parameter_tree_keeps_its_cached_vectorial_node.

Review kept finding edge cases in how Core's parameter caches are filled, invalidated, copied and shared. Some predate this PR; others come from nodes changing their own state as they're read. #621 proposes rules for those caches, which this PR should wait for.

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.

2 participants