fix(jax): structural defense against cached_property pytree/dict leaks#1302
Merged
Conversation
PR #1300 fixed a specific leak where AbstractPriorModel.parameterization (a `@functools.cached_property` added in commit 4564ae9) leaked its cached string into every ModelInstance via Collection._instance_for_arguments. That broke 38 JAX jit(fit_from) calls and the autofit_workspace overview_1 smoke (clusters C1+C4). The minimal fix renamed the cache key to `_parameterization_cache` so the existing `_`-prefix filter at each `__dict__` iterator skipped it. The structural problem remained: every walker uses an opt-out filter (blacklist + underscore prefix), so any future cached_property declared on a model class silently reproduces the same class of bug. This PR closes the class: - New classmethod `AbstractModel._cached_property_names(cls)` delegates to `autoconf.tools.decorators.cached_property_names` (PyAutoConf #111), returning a frozenset of every functools.cached_property and autoconf CachedProperty descriptor name in the MRO. - Extend the filter at every `__dict__` iteration site to union the pre-existing exclusion with this frozenset: autofit/mapper/model.py (__getstate__, ModelInstance.dict) autofit/mapper/model_object.py (ModelObject._dict — feeds Collection.items) autofit/mapper/prior_model/abstract.py (AbstractModel.items) autofit/mapper/prior_model/collection.py (Collection._instance_for_arguments) autofit/mapper/prior_model/prior_model.py (Model._instance_for_arguments) Identifier-hash stability verified: the unique_identifier walker at `autofit/mapper/identifier.py` does NOT call any of these 6 sites — it walks `__dict__` independently with its own `_`-prefix filter. Three representative model shapes (simple Collection, nested Collection, Model with tuple arg) all produce byte-identical identifier hashes pre- and post-defense: simple: f7f19073a8fb19b3d11231fb6eef7e3b ✓ nested: 04e1328c84a1e4c3a81a9d3544dd19f5 ✓ with_tuple: 36084d2c3fec27e0b7aa504add0bd898 ✓ Tests: - `test_cached_property_names_classmethod_walks_mro`: confirms the classmethod surfaces MRO-declared descriptors and memoises per-class. - `test_cached_property_excluded_from_all_dict_walks`: ships a synthetic GuardedCollection with a cached_property returning a string; asserts the value never appears in instance.__dict__, instance.dict, model.items(), tree_flatten() leaves, __getstate__, or pickle round-trip. - 1415/1415 PyAutoFit tests pass (1413 prior + 2 new). Depends on: PyAutoLabs/PyAutoConf#111. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up to #1300. The previous PR fixed a specific leak; this PR fixes the class of bug: every
__dict__walker on the model side uses an opt-out filter (underscore prefix + hardcoded blacklist), so any future@cached_propertyon a model class silently reproduces the same 38-script JAX failure mode.Six sites in
autofit/mapper/are extended with a frozenset union of the class's cached_property names (sourced from PyAutoLabs/PyAutoConf#111's newcached_property_nameshelper):autofit/mapper/model.py:86AbstractModel.__getstate__autofit/mapper/model.py:467ModelInstance.dictautofit/mapper/model_object.py:332ModelObject._dict(feedsCollection.items)autofit/mapper/prior_model/abstract.py:1280AbstractModel.itemsautofit/mapper/prior_model/collection.py:289Collection._instance_for_argumentsautofit/mapper/prior_model/prior_model.py:497Model._instance_for_argumentsA new classmethod
AbstractModel._cached_property_names()delegates to the autoconf helper so the filter sites all read from a single source of truth.Identifier-hash stability — verified
Critical for release: the
unique_identifierhash (search output directory name, fit resume key, DB lookup) must not shift. Verified:autofit/mapper/identifier.py:66-159) walks__dict__independently — it does NOT call any of the 6 sites modified here._-prefix keys at line 137.@cached_propertyexists on any AbstractModel descendant in PyAutoFit/PyAutoArray/PyAutoGalaxy/PyAutoLens. The defense is purely forward-compat.Collection(gaussian=Model(Gaussian))→f7f19073...✓Collection(a=Model, b=Collection(nested=Model))→04e1328c...✓Collection(Model(Gaussian, centre=(0.0, 0.0)))→36084d2c...✓Test plan
pytest test_autofit/mapper/test_parameterization.py -v— 13/13 pass (2 new tests + 11 existing).pytest test_autofit -q— 1415/1415 pass (1413 prior + 2 new), 1 skipped (no regressions).Dependency
Depends on PyAutoLabs/PyAutoConf#111 (must merge first — provides the
cached_property_namesMRO walker).A paired PyAutoArray PR ships the same defense on the pytree flatten side: PyAutoLabs/PyAutoArray#TBD.
🤖 Generated with Claude Code