Repository navigation
Fix flaky test_siblings_without_eval - #3621
Merged
Merged
Conversation
PhilipJohnBasile
added a commit
to PhilipJohnBasile/mlx
that referenced
this pull request
Aug 1, 2026
) compile_dfs deep-copies the traced graph so the cache owns a private mutable copy, then redirects every consumer to the copy and leaves the originals to be freed by reference counting. Single-output originals are. Multi-output originals are not: each output holds its siblings in ArrayDesc::siblings, so the group refers to itself. array::~array() breaks that cycle, but the redirection goes through array::operator=, which only reseats array_desc_ and never runs the destructor, and the vector teardown happens while the not-yet-remapped parents map still holds copies so the destructor's count test fails. The orphaned group then keeps its whole input cone alive -- including any arrays the compiled function captured -- for the lifetime of the process. Trace nodes hold no buffers, so the visible cost is the captured constants: 65,536 B per re-trace in the issue's repro. The trigger is a new trace (fresh closure identity, or new shapes), not an eviction, which is why a plain function is unaffected -- its stable fun_id means it traces once. Detaching the replaced originals where they are created also fixes mx.export_function, the only other caller of compile_dfs, which leaks identically with no compile cache involved at all. The guard fires only for nodes actually replaced: graph inputs, graph outputs and primitive-less nodes all map to themselves. Sibling groups are all-or-nothing because io_set is seeded with o.siblings(), so `a, b = split(x); del a; eval(b)` never enters the branch. The Primitive survives because the copies co-own it via primitive_ptr(). Regression test uses the std::weak_ptr<array::Data> probe from ml-explore#1590/ml-explore#3621 rather than an RSS or get_active_memory assertion. Verified CPU-only (MLX_BUILD_METAL=OFF): the test fails with this hunk removed and passes with it, and the full C++ suite is 245/245 cases / 3240/3240 assertions.
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
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.
Instead of measuring the RAM used by the process, which is very unreliable, change the test to directly verify we are not leaking
ArrayDescreferences.