[None][fix] Drain in-flight requests before clearing the KV cache reuse state - #17163
[None][fix] Drain in-flight requests before clearing the KV cache reuse state#17163lowsfer wants to merge 3 commits into
Conversation
|
/bot run --disable-fail-fast |
WalkthroughKV cache cleanup now rejects operations while caches remain open. Native and Python destructors report shutdown failures. Prefix-cache reset now drains active requests before invalidating the cache. ChangesKV cache lifecycle safety
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp (1)
147-149: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy liftAdd regression tests for both lifecycle backends.
Test cleanup with one open cache and with a closed cache. Assert that
shutdown()andclearReusableBlocks()reject open caches, preserve storage, and succeed after close. Test destructor paths separately so guard errors are logged or warned without propagating.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp` around lines 147 - 149, Add regression coverage for both lifecycle backends around KvCacheManager::shutdown() and clearReusableBlocks(). Verify open caches are rejected while storage remains intact, closed caches allow both operations to succeed, and repeat each scenario with one open and one closed cache. Exercise destructor cleanup separately, asserting guard failures are logged or warned without propagating exceptions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp`:
- Line 168: Update the argument passed to _checkNoLivingKvCaches in
clearReusableBlocks() to report the native C++ method name,
clearReusableBlocks(), instead of the Python-style clear_reusable_blocks().
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py`:
- Around line 297-300: Update the warnings.warn call in the cleanup exception
handler around self.shutdown() to pass stacklevel=2, so the warning points to
the cleanup trigger instead of __del__.
---
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp`:
- Around line 147-149: Add regression coverage for both lifecycle backends
around KvCacheManager::shutdown() and clearReusableBlocks(). Verify open caches
are rejected while storage remains intact, closed caches allow both operations
to succeed, and repeat each scenario with one open and one closed cache.
Exercise destructor cleanup separately, asserting guard failures are logged or
warned without propagating exceptions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2fbb04da-0a9c-46d8-8826-550c56b53ce9
📒 Files selected for processing (4)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.htensorrt_llm/llmapi/rlhf_utils.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
|
PR_Github #63227 [ run ] triggered by Bot. Commit: |
|
PR_Github #63227 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63235 [ run ] triggered by Bot. Commit: |
|
PR_Github #63235 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63237 [ run ] triggered by Bot. Commit: |
|
PR_Github #63237 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63240 [ run ] triggered by Bot. Commit: |
|
PR_Github #63240 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63252 [ run ] triggered by Bot. Commit: |
|
PR_Github #63252 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63256 [ run ] triggered by Bot. Commit: |
|
PR_Github #63256 [ run ] completed with state |
1b2c042 to
13114f7
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 992-1019: Strengthen
test_clear_reusable_blocks_rejects_open_kv_cache and
test_shutdown_rejects_open_kv_cache by creating committed reusable radix-tree
content and allocated work before invoking the cleanup API with an open cache.
After each expected GuardError, verify the content remains reusable and the
manager can continue allocating or resume work, then close the cache and perform
the successful cleanup.
- Around line 1030-1041: Update the cleanup in the test around
clear_reusable_blocks so caches[0] is closed regardless of whether the first
assertRaises/assertIn sequence succeeds. Ensure the finally block closes every
cache, while preserving the intermediate close and assertions used to verify the
open-cache counts.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a22a3f13-aa3a-4c87-93e8-5da904c8cf58
📒 Files selected for processing (5)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.htensorrt_llm/llmapi/rlhf_utils.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
🚧 Files skipped from review as they are similar to previous changes (3)
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp
- tensorrt_llm/llmapi/rlhf_utils.py
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
The control_action_decorator on reset_prefix_cache() is clearly right — reset_reuse_state() forwards straight to impl.clear_reusable_blocks() without closing anything, so the drain is genuinely load-bearing. Two things before this lands:
1. The shutdown() guard turns a completed teardown into a leaked pool. shutdown() checks first, then clears and destroys storage — so when the check throws, _storage.destroy() / the C++ teardown never runs, and __del__ / ~KvCacheManager swallow it. The only in-tree caller, KVCacheManagerV2.shutdown(), closes every kv_cache_map entry before calling impl.shutdown(), so the guard should be unreachable there. That leaves the question the PR does not answer: can anything hold an open KvCache outside kv_cache_map (fatal-error path, warmup, transient prefetch caches, a direct binding user)? If no, the guard is dead code on this entry point and a warning would be safer. If yes, the new failure mode is a silent GPU-pool leak in a long-lived process instead of a loud error — worse to diagnose than what it replaces. Consider having shutdown() close the living caches itself (mirroring the wrapper) rather than refuse, and keeping the hard throw only on clear_reusable_blocks() where the caller can actually recover.
2. The description contradicts the diff. "Test Coverage: No new test. ... a regression test would only pass on one backend" — but the diff adds TestLivingKvCacheGuard with three cross-backend tests. Worth fixing; a reviewer who reads "no new test" on a KV-cache lifecycle change will stop there.
|
/bot run --disable-fail-fast |
|
PR_Github #63264 [ run ] triggered by Bot. Commit: |
|
PR_Github #63264 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63271 [ run ] triggered by Bot. Commit: |
|
PR_Github #63271 [ run ] completed with state |
|
@BowenFu Thanks — point 2 was right, the description was stale and is now updated (the tests landed after it was written). On point 1, I don't think the premise holds. On the leak, the trigger can't arise. On closing the living caches inside On recoverability: the check is the first statement in |
…se state
WorkerExtension.reset_prefix_cache() was missing @control_action_decorator,
unlike both of its neighbours: update_weights(), whose docstring calls out that
it "uses the control_action_decorator to ensure all active requests are
finished", and wait_for_engine_idle(). It is a public Ray worker-extension
method, so _collective_rpc("reset_prefix_cache") reached
KvCacheManager::clearReusableBlocks() with requests still in flight.
Clearing detaches the whole radix tree, but a live request keeps owning the tree
blocks it matched via SeqBlock::treeBlock, and goes on using them as the parent
for the block its close() commits. resume() deliberately un-commits a partial
trailing block so it is re-committed at close, so this is the normal path, not a
corner case. The resulting block is grafted onto a detached subtree: unreachable
from any root, so the work is silently discarded. On C++ it is worse than
silent, because addOrGetExistingBlock() derives the block size via
prev->tokensPerBlock(), which reads the now-null grandparent link and segfaults
in release builds.
Add the decorator so the reuse state is only cleared once the engine has
drained. Also guard both entry points -- clearReusableBlocks() and shutdown(),
which frees the storage those pages live in -- so misuse reports itself at the
offending call rather than surfacing later as discarded work or a null
dereference. Both backends name the API in the message, and the check counts
sequences that are still open rather than objects that are still referenced,
since KvCache::close() / _KVCache.close() drop the entry.
shutdown() runs from ~KvCacheManager and __del__, which must not propagate, so
both report and skip teardown instead -- leaving the storage alive is safer than
freeing it under live pages.
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
…ross backends The guards added with the drain fix had no regression coverage: every existing clear_reusable_blocks() call site closes its caches first, so nothing exercised the rejection path, and nothing pinned the promise that a rejected call leaves the storage intact. Add TestLivingKvCacheGuard: both entry points reject an open cache, the operation succeeds once it is closed, and the count tracks close() rather than object liveness -- the last one matters because the guard counts sequences that are still open, not objects that are still referenced. Writing it surfaced a divergence the fix had claimed to remove: C++ said "KvCache(s)" where Python said "KV cache(s)", so the two backends reported the same mistake differently. Use the Python wording in C++ so the message is identical either way, and assert the shared text so the two cannot drift apart again unnoticed. The exception types still differ -- Python raises LogicError while TLLM_CHECK_WITH_INFO surfaces as RuntimeError, since only RequestSpecificException is translated -- so the test accepts either rather than pretending they agree. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
The guards were exercised against an empty radix tree: the caches were created but never resumed or committed, so "the rejected call left the storage intact" had nothing to observe. An implementation that cleared the tree and *then* raised passed just as well as the correct one. Seed committed, reusable content first, and assert it is still reusable after the rejection -- then assert it is gone after the call that is allowed through, so the first check cannot pass vacuously. Verified by mutation: reordering clear_reusable_blocks() to clear before checking now fails with 0 != 64, and neutering the clear fails the follow-up assertion. For shutdown(), which frees the storage rather than just the tree, also allocate and resume a fresh sequence after the rejection to show the pool itself survived. Close every cache in the counting test's cleanup rather than all but the first. The first was closed inside the try block, so an assertion failing ahead of it left it open, and tearDown's shutdown() then raised the guard error over the top of the real failure. close() is idempotent, so closing all of them is enough. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
13114f7 to
3681f7f
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py (1)
315-321: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new lifecycle precondition.
Add Google-style docstrings to
shutdown()andclear_reusable_blocks(). State that all KV caches must be closed and thatLogicErroris raised otherwise.The PR objective identifies these methods as public direct-binding APIs. As per coding guidelines, prefer docstrings for external interfaces and document public function arguments.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py` around lines 315 - 321, Add Google-style docstrings to the public methods shutdown() and clear_reusable_blocks(), documenting that all KV caches must be closed before invocation and that LogicError is raised when living caches remain; include argument documentation where applicable while preserving the existing lifecycle checks and behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py`:
- Around line 315-321: Add Google-style docstrings to the public methods
shutdown() and clear_reusable_blocks(), documenting that all KV caches must be
closed before invocation and that LogicError is raised when living caches
remain; include argument documentation where applicable while preserving the
existing lifecycle checks and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c5c5d48a-2a76-44a3-bb02-27ceb69ffea8
📒 Files selected for processing (5)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.htensorrt_llm/llmapi/rlhf_utils.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
🚧 Files skipped from review as they are similar to previous changes (4)
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h
- tensorrt_llm/llmapi/rlhf_utils.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
|
/bot run --disable-fail-fast |
|
PR_Github #63281 [ run ] triggered by Bot. Commit: |
|
PR_Github #63281 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #63290 [ run ] triggered by Bot. Commit: |
|
PR_Github #63290 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63307 [ run ] triggered by Bot. Commit: |
|
PR_Github #63307 [ run ] completed with state
|
|
Read One thing I'd like your read on before approving: the rejected
Two shapes that avoid it: reject only the explicit API call and let the destructor force the teardown, or have Minor: |
3 participants
Dev Engineer Review
WorkerExtension.reset_prefix_cache()now drains active requests before it clears prefix-cache state.shutdown()andclearReusableBlocks()reject cleanup while KV caches remain open in both backends.QA Engineer Review
TestLivingKvCacheGuardcoverage for both backends.tests/integration/test_lists/coverage references were changed.Description
WorkerExtension.reset_prefix_cache()was missing@control_action_decorator, unlike both of itsneighbours in the same class:
update_weights(), whose docstring calls out that it "uses thecontrol_action_decorator to ensure all active requests are finished", and
wait_for_engine_idle().It is a public Ray worker-extension method, so
_collective_rpc("reset_prefix_cache")reachedKvCacheManager::clearReusableBlocks()with requests still in flight.Clearing detaches the whole radix tree, but a live request keeps owning the tree blocks it matched
via
SeqBlock::treeBlock, and goes on using them as the parent for the block itsclose()commits.resume()deliberately un-commits a partial trailing block so it is re-committed at close, so thisis the normal path rather than a corner case. The resulting block is grafted onto a detached
subtree: unreachable from any root, so the work is silently discarded. On C++ it is worse than
silent, because
addOrGetExistingBlock()derives the block size viaprev->tokensPerBlock(), whichreads the now-null grandparent link and segfaults in release builds.
This PR:
actual fix.
clearReusableBlocks()andshutdown(), which additionally frees thestorage those pages live in — so misuse reports itself at the offending call rather than
surfacing later as discarded work or a null dereference.
Both backends name the API in the message, and the check counts sequences that are still open
rather than objects that are still referenced, since
KvCache::close()/_KVCache.close()dropthe entry.
shutdown()runs from~KvCacheManagerand__del__, which must not propagate, so both report andskip teardown instead — leaving the storage alive is safer than freeing it under live pages.
Test Coverage
TestLivingKvCacheGuardintests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py,three tests that run on both backends:
clear_reusable_blocks()andshutdown()each reject an open KV cache, naming the API the calleractually invoked.
call succeeds once the cache is closed. This is what makes the error recoverable rather than fatal
— close the sequences and retry.
close()rather than object liveness: with three caches open, closing one makesthe message report two, and the call succeeds with all three references still held.
The two backends raise different types — Python raises
LogicError, whileTLLM_CHECK_WITH_INFOsurfaces as
RuntimeErrorsince onlyRequestSpecificExceptionis translated — so the test acceptseither rather than pretending they agree.
Writing the test surfaced a divergence this PR had claimed to remove: C++ reported
KvCache(s)wherePython reported
KV cache(s), so the two backends described the same mistake differently. C++ nowuses the Python wording, and the shared text is asserted so they cannot drift apart unnoticed.
Full suite, both backends:
TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp pytest tests/unittest/kv_cache_manager_v2_tests/— 177 passed, 12 skippedTLLM_KV_CACHE_MANAGER_V2_BACKEND=python pytest tests/unittest/kv_cache_manager_v2_tests/— 169 passed, 20 skippedThat is +3 on each backend versus the pre-change counts (174 / 166), matching the three new tests and
confirming they execute under both.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.🤖 Generated with Claude Code