Skip to content

Fix flaky test_siblings_without_eval - #3621

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:flaky-leak-test
Jun 12, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:flaky-leak-test

Conversation

@zcbenz

@zcbenz zcbenz commented Jun 4, 2026

Copy link
Copy Markdown
Member

Instead of measuring the RAM used by the process, which is very unreliable, change the test to directly verify we are not leaking ArrayDesc references.

@zcbenz
zcbenz force-pushed the flaky-leak-test branch from 1b99f39 to 104fdc6 Compare June 4, 2026 03:24

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🙏

@zcbenz
zcbenz merged commit 337f736 into ml-explore:main Jun 12, 2026
16 checks passed
@zcbenz
zcbenz deleted the flaky-leak-test branch June 12, 2026 00:31
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.
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