Fix stale compile cache entry when released on another thread - #4096
Fix stale compile cache entry when released on another thread#4096yentur wants to merge 3 commits into
Conversation
The compile cache is thread local, but compile_erase runs on whichever thread drops the last reference to a compiled function. When that is not the thread that traced it, the erase hits the wrong cache and the entry stays behind. It does not only leak. fun_id is the address of the callable and the cache matches on it plus shapes, dtypes, stream and constants, so once the address is reused a later compile of an unrelated function can match the stale entry and get the dead function's tape. Register the caches process wide so an erase can reach them all. Threads other than the caller are handed the id instead of having their cache touched from the outside, and they apply it the next time they use the cache. The check for pending ids is a single atomic load, so compiled calls stay lock free.
|
thanks @yentur, and thanks for the credit. good to see this hit independently, and the process-wide registry is the right shape. one thing worth checking before it lands, as a race rather than a certain break: the lazy drain introduces a lifetime hazard on the owning thread. find() still returns a bare the outer id gets into that thread's pending_erases_ when a second thread drops a handle to the same callable (fun_id is the reused callable address). full sequence:
narrow, needs that three-way overlap in one window, but each piece is legitimate (per-thread thread_local caches under an inference server hit it). to be clear the lock-order side looks correct: swapping the queue under caches_mutex and doing the destructive it's fixable inside your mechanism: make entries re your question: no need to close #4096 on our account. the goal is just a correct fix landed, so we'd rather converge on your registry and contribute the test plus the shared_ptr lifetime change than run a competing PR. |
find() handed out a bare reference into the vector inside cache_, and the caller holds it across compile_trace(), which runs user code. A traced body that calls an already compiled function on constant inputs reaches find() again, and the drain there can erase the outer in-flight fun_id and free the entry the caller is still writing to. Store the entries as shared_ptr and let the caller keep one for the whole call, so an erase during a trace, re-entrant or from another thread, only drops the cache slot while the entry in use stays alive.
|
Confirmed, it reproduces. I built the branch under ASAN with a test running your exact sequence: thread B calls The part I wanted to check before agreeing was reachability of the nested Worth adding: this is not only a consequence of the registry. The same shape is already on main. A traced body that drops a second wrapper over the same callable calls So the lifetime change is worth having on its own, independent of the erase routing. Pushed as cf6e3ac: entries are Please send the C++ test as |
zcbenz
added
the
await verification
|
here is the erase-during-trace coverage as two cases in verified against your branch: both pass under ASAN on cf6e3ac. reverting only cf6e3ac (back to happy to send this as a PR against your branch instead if that is easier than pasting. TEST_CASE("test compile erase while a trace is in flight") {
// An entry handed to a caller must outlive an erase that lands while the
// caller is still filling it in. Here a second thread erases the outer id
// while the worker is inside its trace, and the traced body then invokes an
// already-compiled helper on a fresh (non-tracer) constant. That helper's
// find() drains the pending erase on the worker's own cache, which frees the
// outer entry unless the caller is holding it.
auto x = zeros({1}, float32);
eval(x);
constexpr std::uintptr_t nested_id = 0xf00d;
constexpr std::uintptr_t outer_id = 0xbeef;
auto nested = detail::compile(
[](const std::vector<array>& in) {
return std::vector<array>{in[0] + 2.0f};
},
nested_id);
eval(nested({zeros({1}, float32)}));
std::mutex mtx;
std::condition_variable cv;
int stage = 0;
auto traced = [&](const std::vector<array>& inputs) {
{
std::lock_guard<std::mutex> lk(mtx);
stage = 1;
}
cv.notify_one();
{
std::unique_lock<std::mutex> lk(mtx);
cv.wait(lk, [&stage] { return stage == 2; });
}
// Non-tracer input, so the compiled helper runs find() here and drains the
// erase queued for the outer id above.
auto tmp = nested({zeros({1}, float32)});
return std::vector<array>{inputs[0] + tmp[0]};
};
float result = 0.0f;
std::thread worker([&]() {
auto compiled = detail::compile(traced, outer_id);
auto outputs = compiled({x});
eval(outputs);
result = outputs[0].item<float>();
});
{
std::unique_lock<std::mutex> lk(mtx);
cv.wait(lk, [&stage] { return stage == 1; });
}
detail::compile_erase(outer_id);
{
std::lock_guard<std::mutex> lk(mtx);
stage = 2;
}
cv.notify_one();
worker.join();
CHECK_EQ(result, 2.0f);
}
TEST_CASE("test compile erase from another thread") {
// The compile cache is thread local, so an erase that reaches only the
// calling thread leaves behind the entry of a function traced elsewhere.
// Since |fun_id| is a reused address, the next compile on that id would then
// be handed the dead function's tape. Trace on a worker thread which stays
// alive, erase from this one, then reuse the id for a different function.
auto x = zeros({1}, float32);
eval(x);
constexpr std::uintptr_t fun_id = 0xc0ffee;
auto add_one = [](const std::vector<array>& inputs) {
return std::vector<array>{inputs[0] + 1.0f};
};
auto add_two = [](const std::vector<array>& inputs) {
return std::vector<array>{inputs[0] + 2.0f};
};
std::mutex mtx;
std::condition_variable cv;
int stage = 0;
float before = 0.0f;
float after = 0.0f;
std::thread worker([&]() {
{
auto compiled = detail::compile(add_one, fun_id);
auto outputs = compiled({x});
eval(outputs);
before = outputs[0].item<float>();
}
{
std::lock_guard<std::mutex> lk(mtx);
stage = 1;
}
cv.notify_one();
{
std::unique_lock<std::mutex> lk(mtx);
cv.wait(lk, [&stage] { return stage == 2; });
}
// Same id, different function. Without the erase reaching this thread the
// cached tape still adds one.
auto compiled = detail::compile(add_two, fun_id);
auto outputs = compiled({x});
eval(outputs);
after = outputs[0].item<float>();
});
{
std::unique_lock<std::mutex> lk(mtx);
cv.wait(lk, [&stage] { return stage == 1; });
}
detail::compile_erase(fun_id);
{
std::lock_guard<std::mutex> lk(mtx);
stage = 2;
}
cv.notify_one();
worker.join();
CHECK_EQ(before, 1.0f);
CHECK_EQ(after, 2.0f);
}these need |
Two cases in compile_tests.cpp. The first erases the outer id from a second thread while a worker is inside its trace, then has the traced body call an already compiled helper on a constant so the nested find() drains on the worker's own cache. The second traces on a worker that stays alive, erases from another thread, then reuses the id for a different function. Both tests were written by @sashko-zakharchuk and are added here with his agreement in ml-explore#4096.
|
Applied both cases as 6242ddb, with I re-derived the three states here rather than taking the report as given, CPU-only build on an M4 Pro:
Since both are threaded and the PR is green right now, I ran each case 30 times on the normal build and 30 times under ASAN, 120 runs, no failures. Full suite is 249 cases and 3329 assertions passing on both builds, and I ran the ASAN suite three times with both cases in place with no reports. Credited you in the commit body. |
3 participants
Proposed changes
Fixes #3940.
compile_eraseruns from~PyCompiledFun, and from theshared_ptrdeleter in the C++compile(), so it runs on whichever thread drops the last reference to a compiledfunction.
compiler_cache()has beenthread_localsince #3280, so when the releasehappens off the thread that traced the function the erase goes to the wrong cache and the
entry stays behind.
That is not only a leak.
fun_idis the address of the callable, andCompilerCache::findmatches onfun_idplus shapes, dtypes, stream and constants, soonce the address is reused a later
mx.compileof an unrelated function matches thestale entry and gets the dead function's tape back.
The caches stay thread local and are now also registered process wide so an erase can
reach all of them. The calling thread erases from its own cache directly, the others are
handed the id and apply it the next time they use their cache, so a cache is still only
ever mutated by the thread that owns it.
findchecks for pending ids with a singleacquire load, so compiled calls stay lock free.
Repro from the issue, on 8c28c38:
With the patch, three runs each:
test_compile_release_on_another_threadcovers the deterministic half of that. On mainit fails with
and it passed 20 runs in a row with the patch. Full runs:
python/tests811 tests OK(75 skipped),
tests/tests247 cases and 3326 assertions passing.Two things worth flagging. I built this CPU only on an M4 Pro because there is no Metal
toolchain on the machine, so the Metal path is untested here, though nothing in the
change is backend specific. And a thread that keeps a compile cache but never touches it
again now accumulates 8 bytes per queued id until it next uses the cache or exits, which
is far below the tape that is leaked today but is not nothing.
The reporter mentioned in the issue that they had a patch ready. If theirs is further
along, close this one.
Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes