[Bugfix] Don't reuse engine core payload buffer while zmq is sending it - #50053
Merged
njhill merged 3 commits intoJul 28, 2026
Merged
Conversation
`EngineCoreProc.process_output_sockets` recycles the msgpack payload
bytearray across messages via `MsgpackEncoder.encode_into`, gated on the
tracker returned by `send_multipart(copy=False, track=True)`.
That gate never engages. `Socket.send_multipart()` returns a tracker for
the *last* frame only, and pyzmq copies any frame smaller than
`zmq.COPY_THRESHOLD` (64KiB), returning the always-done
`_FINISHED_TRACKER` for it. The last frame is a small tensor buffer, so
the tracker is always done and `pending` is dead code, while the payload
frame is genuinely zero-copy once it exceeds 64KiB. The next
`encode_into` then overwrites it mid-flight, and the client decodes a
newer payload alongside the older message's tensor frames:
RuntimeError: shape '[29, 2]' is invalid for input of size 60
msgspec.ValidationError: cannot unpack non-iterable int object
- at `$[1][69][3][0]`
Send the payload frame separately so the tracker covers the buffer we
actually reuse. pyzmq's copy threshold then works in our favour: a small
payload is copied and reused immediately, a large one is tracked
properly. Also cap `reuse_buffers` when reclaiming from `pending`, which
previously could not grow because that path was unreachable.
Seen intermittently in CI as v1/sample/test_logprobs_e2e.py failures,
where lm_eval's large prompt-logprobs batches push the payload frame
either side of the 64KiB threshold.
Also drop the retention of objects whose buffers were extracted for
zero-copy send, both here and in `MPClient.pending_messages`. For a
zero-copy frame pyzmq registers the buffer in its own gc registry
(`zmq/backend/cython/_zmq.py`, `_gc.store(data, ...)`) and releases it
only once libzmq is finished, and frames below `zmq.COPY_THRESHOLD` are
copied outright. `tensor_data()` hands zmq a memoryview whose base chain
reaches the source tensor, so refcounting - not timing - keeps the memory
from being freed and reused. The retained references never covered device
tensors anyway, since `tensor_data()` copies those to a temporary host
tensor that the retained object does not reference. Neither client path
reuses a send buffer, so that bookkeeping had no effect beyond looking
like protection it did not provide.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
This was referenced Jul 27, 2026
Member
Author
|
CI failure unrelated, should be fixed by #50060 |
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.
EngineCoreProc.process_output_socketsrecycles the msgpack payload bytearray across messages viaMsgpackEncoder.encode_into, gated on the tracker returned bysend_multipart(copy=False, track=True).That gate never engages.
Socket.send_multipart()returns a tracker for the last frame only, and pyzmq copies any frame smaller thanzmq.COPY_THRESHOLD(64KiB), returning the always-done_FINISHED_TRACKERfor it. The last frame is a small tensor buffer, so the tracker is always done andpendingis dead code, while the payload frame is genuinely zero-copy once it exceeds 64KiB. The nextencode_intothen overwrites it mid-flight, and the client decodes a newer payload alongside the older message's tensor frames:Send the payload frame separately so the tracker covers the buffer we actually reuse. pyzmq's copy threshold then works in our favour: a small payload is copied and reused immediately, a large one is tracked properly. Also cap
reuse_bufferswhen reclaiming frompending, which previously could not grow because that path was unreachable.Seen intermittently in CI as v1/sample/test_logprobs_e2e.py failures, where lm_eval's large prompt-logprobs batches push the payload frame either side of the 64KiB threshold.
Also drop the retention of objects whose buffers were extracted for zero-copy send, both here and in
MPClient.pending_messages. For a zero-copy frame pyzmq registers the buffer in its own gc registry (zmq/backend/cython/_zmq.py,_gc.store(data, ...)) and releases it only once libzmq is finished, and frames belowzmq.COPY_THRESHOLDare copied outright.tensor_data()hands zmq a memoryview whose base chain reaches the source tensor, so refcounting - not timing - keeps the memory from being freed and reused. The retained references never covered device tensors anyway, sincetensor_data()copies those to a temporary host tensor that the retained object does not reference. Neither client path reuses a send buffer, so that bookkeeping had no effect beyond looking like protection it did not provide.