fix(multimodal): bound MmKwargsNixlSender.cleanup and always release buffers - #12759
Conversation
GitHub Actions
Inactive
GitHub Actions
Inactive
|
👋 Hi yifjiang! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
Bot
added
the
external-contribution
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
WalkthroughNIXL sender cleanup now retains operation objects, waits for completion with a configurable 60-second default timeout, logs failures, and releases registered memory regions in all cases. Tests cover timeout, normal completion, and empty input. ChangesNIXL cleanup lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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: 3
🧹 Nitpick comments (1)
components/src/dynamo/common/multimodal/mm_kwargs_transfer.py (1)
251-254: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAlign the cleanup-item type contract.
_encode_item()declaresAwaitable[None]but returns an operation that provideswait_for_completion()and__exit__(). Define a cleanup-handleProtocoland use it in the sender contracts. Cleanup items are passed only tocleanup()and are not awaited directly.🤖 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 `@components/src/dynamo/common/multimodal/mm_kwargs_transfer.py` around lines 251 - 254, Define a cleanup-handle Protocol exposing the operation methods used by cleanup, including wait_for_completion() and __exit__(), then update _encode_item() and the sender cleanup-item type contracts to use this Protocol instead of Awaitable[None]. Preserve returning the operation itself from _encode_item() and ensure cleanup items are handled only through cleanup().
🤖 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 `@components/src/dynamo/common/multimodal/mm_kwargs_transfer.py`:
- Around line 305-314: Update the operation-release loop to collect exceptions
from each op.__exit__ call instead of only logging them at debug level. Continue
attempting release for every operation, then re-raise the retained release error
after the loop so cleanup failures propagate to callers.
- Around line 290-303: The cleanup await in the NIXL transfer completion flow
must wait for every operation before buffers are released. Update the
asyncio.gather call around op.wait_for_completion() to use
return_exceptions=True, then inspect gathered results and re-raise completion
errors only after all operations finish; preserve the existing timeout handling,
and add a regression test covering one failing operation alongside one blocking
operation.
In `@components/src/dynamo/common/tests/multimodal/test_mm_kwargs_transfer.py`:
- Around line 146-161: Add the pytest timeout marker to
test_cleanup_is_bounded_and_releases_when_never_read while retaining its local
asyncio.wait_for assertion and existing cleanup behavior.
---
Nitpick comments:
In `@components/src/dynamo/common/multimodal/mm_kwargs_transfer.py`:
- Around line 251-254: Define a cleanup-handle Protocol exposing the operation
methods used by cleanup, including wait_for_completion() and __exit__(), then
update _encode_item() and the sender cleanup-item type contracts to use this
Protocol instead of Awaitable[None]. Preserve returning the operation itself
from _encode_item() and ensure cleanup items are handled only through cleanup().
🪄 Autofix
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: f26923f2-58ce-4434-bdcc-19d76104a816
📒 Files selected for processing (2)
components/src/dynamo/common/multimodal/mm_kwargs_transfer.pycomponents/src/dynamo/common/tests/multimodal/test_mm_kwargs_transfer.py
|
/ok to test a18e47a |
This comment has been minimized.
This comment has been minimized.
|
/ok to test a49b80a |
GitHub Actions
Inactive
|
Tagging reviewers with a suggested split, since this touches two areas:
Context on urgency: this is a live leak. On a multimodal deployment the frontend grew to ~62 GiB (97% of its limit) and OOMKilled repeatedly; disabling the NIXL multimodal path holds it flat at ~2 GiB with zero restarts. Anyone running CI is green (56 passed, 0 failures), including |
|
/ok to test 9bb7745 |
GitHub Actions
Inactive
|
Thanks — all three findings were valid and are addressed in 9bb7745. 1. Premature release (@coderabbitai, Major) — real bug, fixed. Confirmed against the previous revision with a slow operation alongside a raising one:
Added the regression test you asked for (one failing op + one blocking op). 2. Swallowed release errors (@devin-ai-integration) — fixed. You're right that this contradicted 3. Late-read race after timeout (@devin-ai-integration) — confirmed, and now documented rather than changed. Your reading is correct: once the timeout elapses the descriptor is deregistered, so a backend that reads at, say, 61 s gets a NIXL transfer error rather than data. That is the deliberate trade-off — the alternative is the unbounded growth this PR exists to stop — and it is now stated explicitly in the On the 60 s default: it is meant to be comfortably above worst-case backend queueing, and it is tunable via |
|
Tip For best results, initiate chat on the files or code changes.
The The documented timeout trade-off is clear. The configurable timeout is appropriate because backend queueing latency is deployment-specific. 🐇 ✅ ✏️ Learnings added
You are interacting with an AI system. |
End-to-end verification: the leak is reproduced, and this fix stops itUntil now this PR only had unit-level evidence ( Reproduction. Drives Results (RSS growth after the loop, gc settled):
Unpatched growth scales linearly with request count — unbounded. Patched growth is constant at 59 MB across both loads, i.e. working set rather than a leak. Not one unpatched cleanup task ever completed; every patched one did. Both fixes are required — worth knowing before this landsRunning the same test without
The two changes address different halves and neither is sufficient alone:
That explains something that was previously unclear: pinning the allocator threshold alone was deployed first and was only ever a partial fix — it roughly doubled frontend lifetime but the process still OOMKilled, because the buffers were never actually released. Anyone applying only one of the two should expect exactly that. I nearly reported this fix as ineffective on the strength of the first run. Worth stating plainly so a reviewer reproducing it doesn't reach the same wrong conclusion. |
|
/ok to test d19f7a0 |
GitHub Actions
Inactive
|
All review findings are now addressed in d19f7a0.
One deliberate narrowing on the re-raise: only Re-verified end to end after these changesThe behaviour changed, so the leak measurement was re-run on real NIXL (GB200 arm64) rather than assumed:
No regression: growth stays constant with load rather than scaling with it. Also re-checked: |
d19f7a0 to
a2eb20a
Compare
GitHub Actions
Inactive
|
/ok to test a2eb20a |
1 similar comment
|
/ok to test a2eb20a |
krishung5
left a comment
krishung5
left a comment
There was a problem hiding this comment.
Thanks for the fix! Left some comments.
| # ReadableOperation.__exit__ -> _release() -> deregister. | ||
| # _release() skips descriptors that are already deregistered, | ||
| # so this stays safe alongside __del__. | ||
| op.__exit__(None, None, None) |
There was a problem hiding this comment.
_release_all loops N synchronous native releases here: op.__exit__() → _release() → Descriptor.deregister_with_connector() → connection._nixl.deregister_memory(), which is a blocking native call. This runs inside async cleanup() on the single frontend, once per request. If deregister_memory isn't cheap, this stalls the event loop proportional to the buffer count. Could we confirm the perf impact, or maybe offloading via run_in_executor?
There was a problem hiding this comment.
Measured it rather than guessing — you were right to ask, and the answer is that it is cheap.
Timed op.__exit__() -> _release() -> deregister_memory() per buffer, real NIXL on GB200:
| buffer size | p50 | p95 | max |
|---|---|---|---|
| 20 MB | 0.005 ms | 0.016 ms | 0.039 ms |
| 8 MB | 0.005 ms | 0.013 ms | 0.038 ms |
So ~5 us per release: ~0.02 ms of event-loop time for a 4-image request, 0.2 ms for 20 buffers. That is below the cost of a run_in_executor hand-off, so I have left it inline and recorded the numbers in the _release_all docstring so the question does not have to be re-derived later.
Leaving this thread open rather than resolving it — if you would still prefer the executor defensively (e.g. you expect pathological buffer counts per request that I have not exercised), say so and I will switch it.
|
/ok to test 82bcb2b |
GitHub Actions
Inactive
|
/ok to test 82bcb2b |
…buffers
MmKwargsNixlSender.cleanup() awaited transfer completion with no timeout:
await asyncio.gather(*items)
ReadableOperation.wait_for_completion() only resolves once the backend
actually reads the registered buffer. A request that is cancelled,
rejected before the read, or routed to a worker that dies leaves it
pending forever. Because the pending coroutine holds a reference to the
operation, _release() never runs and the NIXL registration is never
dropped, so the frontend retains the full payload of every un-read
transfer for the lifetime of the process.
On a production multimodal deployment this presented as frontend RSS
climbing steadily under image traffic until the container hit its memory
limit and was OOMKilled, repeatedly, while workers were unaffected.
Two changes, both needed:
1. _encode_item() now returns the ReadableOperation instead of only its
completion awaitable. cleanup_item is documented as an opaque handle
passed straight back to cleanup(), so this does not alter any public
signature.
2. cleanup() bounds the wait (DYN_MM_NIXL_CLEANUP_TIMEOUT_S, default 60)
and releases the operations in a finally block.
The release is the part that fixes the leak. A timeout alone would stop
the hang while still pinning the memory, since cleanup() previously had
no handle on which to call _release(). _release() skips descriptors that
are already deregistered, so the explicit __exit__ remains safe
alongside __del__.
Verified by loading the module directly (the package __init__ chain
requires compiled extensions):
unfixed: cleanup never returns; 0 buffers released
fixed: returns in 0.20s; all buffers released
Adds regression tests covering the never-read timeout path, normal
completion, the empty-items no-op, and release when the await raises.
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
mypy: _encode_item still declared tuple[TensorTransferSpec, Awaitable[None]] while it now returns the ReadableOperation. Annotates the second element as Any, matching the base hook's opaque cleanup_item contract and avoiding a dependency on dynamo.nixl_connect, which is imported lazily so the module stays importable where NIXL is unavailable. Drops the now-unused Awaitable import. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
… errors
Addresses review findings on this PR.
1. Premature release (real bug). asyncio.gather propagates the first
completion error while its siblings are still running, so the release
that follows could deregister a buffer whose backend read was still in
flight. Adds return_exceptions=True so every operation is awaited (or
cancelled by the timeout) before anything is released, and inspects
the results afterwards, logging failures at warning level.
Confirmed against the previous revision with a slow operation
alongside a raising one:
before slow op finished=False released mid-read=True
after slow op finished=True released mid-read=False
2. Release errors were logged at debug and swallowed, against
.ai/python-guidelines.md ("if you must catch broadly ... always
re-raise after logging"). The broad except stays, because releasing
is best-effort across all items and aborting early would leak exactly
the buffers this code frees -- but it now logs at warning, collects
the first failure, and re-raises it once every item has been
attempted. The raise sits outside the finally so it can never mask an
exception that was already propagating.
3. Documents the deliberate trade-off raised in review: after the
timeout the buffer is deregistered, so a backend that reads later
sees a NIXL transfer error rather than data. The alternative is
unbounded growth; the timeout is generous and tunable via
DYN_MM_NIXL_CLEANUP_TIMEOUT_S.
Adds the regression test asked for (one failing operation alongside one
blocking operation) plus one asserting a failing release does not skip
the remaining buffers.
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
Remaining review items. Completion errors were logged but swallowed. They are now re-raised, after every sibling has been awaited and every buffer released, so a failed transfer is visible rather than silent. A release failure takes precedence over a completion failure, because it means a buffer is still registered -- the condition this method exists to prevent. Deliberately narrowed to Exception: a cancelled sibling captured by return_exceptions=True is not a fault worth converting into one here. Adds @pytest.mark.timeout(10) to the cleanup tests. The local asyncio.wait_for assertions stay; the marker additionally stops the test if the event loop itself stalls. pytest-timeout is already a dependency and the marker is registered in pyproject.toml. Updates the failing-op regression test to expect the re-raise while still asserting no sibling was released mid-read. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
Reverts the re-raise added earlier. Review is right that it breaks the
calling contract, and both points check out in the source:
- MmKwargsShmSender.cleanup() logs and continues, so raising here made
the two senders diverge on failure
- the only caller (vllm_processor._generator_inner) awaits this from a
bare finally with no guard, so a raise surfaces after the stream has
completed and, on a client cancel, replaces the CancelledError
The release in the finally is what fixes the leak, and it is unchanged.
Re-measured after this revert: 40 iterations of un-read 20 MB transfers
grow 59 MB with the fix versus 841 MB without, 40/40 cleanups complete
-- identical to the version that re-raised.
The broad except is now justified inline, since it is a deliberate
exception to the "re-raise after logging" guideline.
Also records the measured cost of the synchronous native release, which
review asked about: p50 0.005 ms, p95 0.016 ms per buffer for 8-20 MB on
GB200, so ~0.02 ms of event-loop time for a 4-image request. That is
below the cost of an executor hand-off, so it stays inline rather than
moving to run_in_executor.
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com>
82bcb2b to
06f2e15
Compare
GitHub Actions
Inactive
|
/ok to test 06f2e15 |
1 similar comment
|
/ok to test 06f2e15 |
|
Thanks @krishung5 — both points check out, and I've taken both. Pushed in 06f2e15. Re-raise removed (
|
| growth | cleanups completed | |
|---|---|---|
| unpatched | 841 MB | 0/40 |
| with re-raise | 59 MB | 40/40 |
| best-effort (this revision) | 59 MB | 40/40 |
Release cost measured (:295)
Good question — I measured it rather than guessing. Timed op.__exit__() → _release() → deregister_memory() per buffer, real NIXL on GB200:
| buffer size | p50 | p95 | max |
|---|---|---|---|
| 20 MB | 0.005 ms | 0.016 ms | 0.039 ms |
| 8 MB | 0.005 ms | 0.013 ms | 0.038 ms |
~5 µs per release, so ~0.02 ms of event-loop time for a 4-image request, and 0.2 ms for 20 buffers. That is below the cost of a run_in_executor hand-off, so I've left it inline and recorded the numbers in the docstring so the question doesn't have to be re-asked. Happy to switch to an executor if you'd still prefer it defensively, but on these numbers it looks like it would cost more than it saves.
4 participants
Problem
MmKwargsNixlSender.cleanup()awaits transfer completion with no bound:ReadableOperation.wait_for_completion()resolves only once the backend actually reads the registered buffer. A request that is cancelled, rejected before the read, or routed to a worker that dies leaves it pending forever. The pending coroutine holds a reference to the operation, so_release()never runs, the NIXL registration is never dropped, and the frontend retains the payload of every un-read transfer for the life of the process.Observed on a multimodal deployment as frontend RSS climbing under image traffic until the container hit its memory limit and was OOMKilled, repeatedly, over days. Workers were unaffected. Disabling the NIXL multimodal path stopped the growth: frontends have since held flat at ~2 GiB where they previously reached ~62 GiB.
cleanup()receives only the completion awaitables, so it has no handle to release with — a timeout alone would stop the hang while still pinning the memory. Both halves below are required.Fix
_encode_item()returns theReadableOperationinstead of onlyreadable_op.wait_for_completion(). The abstract hook documents this value as "an opaque handle the caller passes tocleanup()", so no public signature changes.cleanup()bounds the wait withasyncio.wait_for(DYN_MM_NIXL_CLEANUP_TIMEOUT_S, default 60) and releases every operation in afinallyblock.ReadableOperation.__exit__calls_release(), which skips descriptors that are already deregistered, so the explicit release stays safe alongside__del__.Test
The package
__init__chain imports compiled extensions (dynamo._core,dynamo.llm) that I cannot build locally, so I loaded the module directly and drovecleanup()with a stand-in operation whose completion never resolves:cleanup()returnsAdds
TestMmKwargsNixlSenderCleanupcovering the never-read timeout path, normal completion, the empty-items no-op, and release when the await raises.pytest components/src/dynamo/common/tests/multimodal/test_mm_kwargs_transfer.py -k CleanupNotes
conftest.pyimportsdynamo._core, which needs the Rust extension built. CI should be the judge.mm_kwargs_transfer.pyhas not been modified since it was added in feat(multimodal): move MM routing into vLLM frontend processor #8065 (2026-04-24), so this behaviour is present in every release since.__del__would eventually release even without the explicit__exit__. The explicit call makes it deterministic rather than GC-timing dependent.Summary by CodeRabbit