diff --git a/docs/design/shm_tensor_arena.md b/docs/design/shm_tensor_arena.md new file mode 100644 index 000000000000..9d5952ecd5cc --- /dev/null +++ b/docs/design/shm_tensor_arena.md @@ -0,0 +1,207 @@ +# Zero-Copy Shared-Memory Tensor Arena for `MessageQueue` + +This document describes a performance problem in vLLM's engine→worker IPC path that +severely penalizes multimodal serving under tensor parallelism, and the fix implemented +on this branch: a zero-copy shared-memory tensor arena inside +`vllm/distributed/device_communicators/shm_broadcast.py`. + +## TL;DR + +With TP > 1, every `execute_model` call is broadcast from the EngineCore process to the +TP worker processes through `MessageQueue.enqueue`, which serializes the whole payload +with `pickle` — **including raw multimodal `pixel_values` tensors**. Torch tensors do +not participate in pickle's out-of-band buffer protocol, so a large image is byte-copied +into the pickle stream on the engine side and byte-copied back out **once per TP rank** +on the reader side. For a ~200 MB pixel tensor on a TP=4 worker that is ~1.8 GB of host +byte-copying, most of it serial and **inside the engine's step loop** — the scheduler +stops, and all GPUs on the worker go idle for **0.9–1.3 s per large image**, including +decode for unrelated in-flight requests (head-of-line blocking). The vision-encoder +compute for the same images is 50–90 ms; transport dominates compute ~15:1. + +The fix routes large CPU tensors around the pickle stream entirely, through a slotted +shared-memory arena: one memcpy into a slot on the writer, a `torch.frombuffer` view +(zero bytes copied) on each reader, with the mapping `cudaHostRegister`-pinned so the +subsequent H2D copy is a true DMA. Measured end to end on an interactive multimodal +workload (uncapped input images, identical seeded request sequence): TTFT p99 +1321 ms → **862 ms**, TTFT max 2139 ms → **1375 ms**, requests over 1.5 s: 9 → **0**, +TTFT p50 unchanged (no overhead on the small-request path). + +## 1. The problem + +### Where it bites + +Any multimodal model served with `--tensor-parallel-size > 1` through the multiproc +executor. The `rpc_broadcast_mq` `MessageQueue` (`v1/executor/multiproc_executor.py`) +carries scheduler output — with the multimodal kwargs, i.e. the raw preprocessed pixel +tensors — from the EngineCore to every TP worker on each step. Image inputs at native +resolution easily reach tens to hundreds of MB of `pixel_values` (a ~4k×4k image is +~200 MB as bf16 patches). Latency-sensitive (interactive) serving makes the resulting +tail visible in TTFT and, through head-of-line blocking, in the end-to-end latency of +*other* requests. + +### The mechanism, at source level + +`MessageQueue.enqueue` (before this branch): + +```python +serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback) # ← the cost lives here +if total_bytes >= self.buffer.max_chunk_bytes: # 16 MB default + # write a 1-byte overflow marker into the shm ring... + self.local_socket.send_multipart(all_buffers) # ...and ZMQ to every local reader +else: + # copy the bytes into a shm ring chunk +``` + +Three compounding costs: + +1. **Serialize.** Torch tensors do not implement pickle protocol-5 out-of-band + buffers, so `pickle.dumps` drives `THPStorage_writeFileRaw`, which byte-copies the + entire tensor into the pickle stream through Python-driven small writes. Measured at + ~6.5 ms/MB (~155 MB/s) on the critical path — ~1.3 s for a 200 MB tensor. +2. **It blocks the engine.** The serialize runs inside the EngineCore step loop + (`collective_rpc → enqueue`), so no new GPU work is scheduled anywhere on the worker + until it finishes. CUDA-level profiling shows all GPUs of the TP group completely + silent for the duration — decode graph launches for other requests included. +3. **Deserialize, per rank.** Payloads above `max_chunk_bytes` bypass the ring and go + through a local ZMQ socket to each reader; every rank then runs `pickle.loads` → + `THPStorage_readFileRaw` — another full byte-copy per rank — before the (cheap, + ~20 ms) H2D copy. + +Per 200 MB image on TP=4: ~200 MB × (1 serialize + 4 socket transfers + 4 deserializes) +≈ **1.8 GB of host data motion**, where the model itself only needs one 200 MB H2D. + +### What does *not* fix it + +Raising `VLLM_MQ_MAX_CHUNK_BYTES_MB` so the payload takes the shm-ring path instead of +the ZMQ socket was tested first (same seeded workload, per-image A/B): **no +improvement**. The path switch happens (zmq frames disappear from the stall +callchains), but `pickle.dumps` runs *before* the ring-vs-socket branch — the serialize +and the ×N deserialize copies are common to both paths and are the actual cost. This +negative result pins the fix target: the pickle byte-copies themselves must go. + +## 2. The fix + +All changes are contained in `shm_broadcast.py`. + +### 2.1 `ShmTensorArena` + +A second shared-memory region created alongside the existing `ShmRingBuffer`: + +- **N slots × slot_bytes** (default 8 × 256 MB, env-tunable), plus per-slot metadata + `[written_flag, reader0_done … readerN_done]`. +- Concurrency uses the **same lock-free single-writer/N-reader protocol as + `ShmRingBuffer`** (memory fences, per-reader done flags), so the model is one the + codebase already trusts. +- Created by the queue **writer** in `MessageQueue.__init__` when all readers are + node-local; readers attach via a new `tensor_arena_handle` field on the queue + `Handle`. Queues with remote readers get no arena and keep today's behavior. +- Pages are allocated lazily by the kernel on first write, so arenas on queues that + never carry big tensors (e.g. worker→engine response queues) cost approximately + nothing. + +### 2.2 Writer path — `_ArenaPickler.reducer_override` + +```python +class _ArenaPickler(pickle.Pickler): + def reducer_override(self, obj): + if (isinstance(obj, torch.Tensor) and obj.device.type == "cpu" + and obj.layout is torch.strided and obj.is_contiguous() + and obj.numel() * obj.element_size() >= MIN_BYTES): + idx = self.arena.write_tensor(obj) # ONE memcpy into a free slot + if idx is not None: + return (_rebuild_arena_tensor, + (arena_name, idx, nbytes, dtype_str, shape)) + return NotImplemented # everything else: default pickling +``` + +`reducer_override` is consulted for every object before its `__reduce_ex__`, so a +diverted tensor never reaches `THPStorage_writeFileRaw`: the pickle stream carries a +~100-byte rebuild stub instead of the tensor bytes. The one remaining copy — +`torch.frombuffer(slot).copy_(t.view(torch.uint8))` — is a multithreaded memcpy at +memory bandwidth (~50–70 ms for 200 MB, vs ~1.3 s for the pickle serialize). + +**The writer never blocks.** If no slot is free or the tensor exceeds the slot size, +`write_tensor` returns `None` and the pickler falls through to the default path — the +worst case is exactly today's behavior, and deadlock is structurally impossible. + +### 2.3 Reader path — zero copies + +The stub unpickles through a module-level rebuild function: + +```python +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + arena = _TENSOR_ARENAS[arena_name] # this process's attached arena + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) + # get_tensor: torch.frombuffer over the mapped slot — zero bytes copied +``` + +The rebuilt tensor *is* the shared memory — no deserialize on any rank. + +**Slot lifecycle.** The rebuilt tensor is consumed (H2D'd) while the worker executes +that step, so the reader must not release the slot at unpickle time. Releases are +deferred and flushed at the reader's **next `dequeue`**: the worker loop is strictly +sequential (dequeue step N → execute → dequeue step N+1), so by the next dequeue the +previous step's H2D has completed. The writer requires all readers' done flags before +reusing a slot. + +### 2.4 Pinning — the zero-copy trap + +The first version of the patch only halved the stall. Profiling showed the residual +entirely inside `cudaMemcpyAsync`: the zero-copy view references tmpfs pages the reader +process has never touched, so the H2D pays **first-touch page faults on ~50k pages plus +pageable-copy staging**. (The old deserialize, for all its waste, incidentally left the +bytes in hot process-local heap pages — which is why the baseline H2D was fast.) + +Fix: each reader lazily `cudaHostRegister`s the whole arena mapping once (~1 s, first +use), after which every H2D from the arena is a **pinned-memory DMA** (~10 ms for +192 MB). Processes without a CUDA context skip registration silently. + +## 3. Configuration + +| Env var | Default | Meaning | +|---|---|---| +| `VLLM_SHM_TENSOR_ARENA` | `1` | Enable the arena (`0` disables; behavior reverts to stock). | +| `VLLM_SHM_TENSOR_ARENA_SLOTS` | `8` | Number of slots. | +| `VLLM_SHM_TENSOR_ARENA_SLOT_MB` | `256` | Slot size; tensors larger than this fall back to pickle. | +| `VLLM_SHM_TENSOR_ARENA_MIN_MB` | `8` | Minimum tensor size to divert; smaller tensors pickle as before. | + +## 4. Validation + +1. **Unit test**: a 199 MB bf16 tensor pushed through a real `MessageQueue` to two + forked reader processes — byte-exact checksums, arena path confirmed, deferred + release, slot reuse, and exhaustion/oversize fallbacks exercised. Enqueue of the + 199 MB payload: **66.7 ms** (vs ~1275 ms observed in production for the same size). +2. **Same-seed A/B**, interactive multimodal workload, uncapped input images, TP=4 ×2 + workers on one 8-GPU node, low qps so individual images decompose cleanly + (~1.2k aligned requests): + + | TTFT (ms) | p50 | p90 | p99 | max | >1 s | >1.5 s | + |---|---:|---:|---:|---:|---:|---:| + | stock | 93 | 401 | 1321 | 2139 | 30 | 9 | + | bigger MQ chunk (config lever) | 91 | 368 | 1287 | 2027 | 22 | 5 | + | **arena + pinning** | **89** | **241** | **862** | **1375** | **7** | **0** | + + p50 unchanged → `reducer_override` adds no measurable overhead when no large tensor + is present. +3. **Per-image stall** (worst offenders, CUDA-level all-GPU idle gap around the pixel + H2D): 192 MB image 1318 ms → 681 ms (arena only) → **351 ms class** (arena + + pinning); vision-encoder compute for the same image is ~80 ms. +4. **Load scaling**: on the same workload the stock build failed a 1.5 s e2e-p99 target + even at very low arrival rates (the stall is a fixed per-image cost); with the patch + the passing arrival rate is bounded by prefill/decode compute interference instead — + transport no longer appears in the profile at the tail. + +## 5. Limitations and future work + +- **Slot release granularity**: releases are flushed at the reader's next dequeue. A + CUDA-event-based release would close the theoretical reuse window under extreme + multi-image bursts deeper than the slot count (today such bursts safely fall back to + pickle when the arena is exhausted). +- **Fallback observability**: arena exhaustion / oversize fallbacks are rate-limited + log lines today; a counter metric would be better. +- **Scope**: the arena activates only when every queue reader is node-local. Remote + readers (multi-node PP/TP) keep the existing socket path. +- **Generality**: the problem is generic to any multimodal model on TP > 1 with the + multiproc executor; upstreaming the approach (or an equivalent out-of-band tensor + channel) is worth discussing. diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index dc7e6d151a48..ef27e6682c7f 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import io +import os import pickle import sys import threading @@ -344,6 +346,266 @@ def get_metadata(self, current_idx: int): yield buf +# --------------------------------------------------------------------------- +# Zero-copy shared-memory arena for large CPU tensors (multimodal fast path). +# +# Problem: `MessageQueue.enqueue` pickles the whole payload. For multimodal +# models the payload contains raw `pixel_values` tensors (up to hundreds of +# MB), and torch's storage pickling copies every byte into the pickle stream +# on the writer (THPStorage_writeFileRaw) and back out on EVERY local reader +# (THPStorage_readFileRaw). On a TP=4 engine this serialize/deserialize chain +# blocks the EngineCore step loop for ~1s per large image with all GPUs idle. +# +# Fast path: intercept large contiguous CPU tensors during pickling +# (`_ArenaPickler.reducer_override`), memcpy them ONCE into a free slot of a +# shared-memory arena, and pickle only a tiny (slot, nbytes, dtype, shape) +# stub. Readers rebuild the tensor as a zero-copy view of the mapped slot +# (`torch.frombuffer`) — no byte-copy on either side. +# +# Slot lifecycle: single writer, n_reader readers, per-slot metadata +# [written_flag, reader0_done, ..., readerN_done] (same protocol as +# ShmRingBuffer). A reader does NOT release the slot when the tensor is +# rebuilt — the tensor is still consumed (HtoD copy) while the worker +# executes that step. Instead releases are deferred and flushed at the +# reader's NEXT `dequeue` call: the worker loop is sequential, so by then the +# previous step (and its pageable-source HtoD, which completes staging before +# `cudaMemcpyAsync` returns) has finished. +# +# The writer NEVER blocks on the arena: if no slot is free (or the tensor is +# larger than a slot), it falls back to the default pickle path for that +# tensor. Worst case is the status-quo behavior, and deadlock is impossible. +# +# Physical memory: slots are allocated lazily by the kernel (pages are backed +# on first write), so arenas on queues that never carry big tensors cost ~0. +# --------------------------------------------------------------------------- + +VLLM_SHM_TENSOR_ARENA = os.getenv("VLLM_SHM_TENSOR_ARENA", "1") != "0" +VLLM_SHM_TENSOR_ARENA_SLOTS = int(os.getenv("VLLM_SHM_TENSOR_ARENA_SLOTS", "8")) +VLLM_SHM_TENSOR_ARENA_SLOT_MB = int(os.getenv("VLLM_SHM_TENSOR_ARENA_SLOT_MB", "256")) +VLLM_SHM_TENSOR_ARENA_MIN_MB = int(os.getenv("VLLM_SHM_TENSOR_ARENA_MIN_MB", "8")) + +# Reader-side registry: arena shm name -> attached ShmTensorArena. Populated +# by MessageQueue.create_from_handle; consumed by _rebuild_arena_tensor when +# unpickling a tensor stub. +_TENSOR_ARENAS: dict[str, "ShmTensorArena"] = {} + + +class ShmTensorArena: + """Slotted shared-memory arena: one writer, n_reader zero-copy readers. + + Memory layout: [slot0 | slot1 | ... | slotN-1 | meta0 | meta1 | ... ] + where each meta is (1 + n_reader) bytes: [written_flag, reader_done...]. + Slot states mirror ShmRingBuffer: + written=0 -> free (never written / being written) + written=1, some reader_done=0 -> in use, cannot reuse + written=1, all reader_done=1 -> consumed, can reuse + """ + + def __init__( + self, + n_reader: int, + slot_bytes: int, + n_slots: int, + name: str | None = None, + reader_rank: int = -1, + ): + self.n_reader = n_reader + self.slot_bytes = slot_bytes + self.n_slots = n_slots + self.reader_rank = reader_rank # -1 for the writer + self.metadata_size = 1 + n_reader + self.metadata_offset = slot_bytes * n_slots + self.total_bytes = (slot_bytes + self.metadata_size) * n_slots + self._next_slot = 0 + self._fallbacks = 0 + self._pin_attempted = False + self._pinned = False + # slots whose tensors were handed out by THIS reader and not yet + # released (flushed at the next dequeue). + self._pending_release: list[int] = [] + + if name is None: + self.is_creator = True + self.shared_memory = shared_memory.SharedMemory( + create=True, size=self.total_bytes + ) + assert self.shared_memory.buf is not None + with self.shared_memory.buf[self.metadata_offset :] as meta: + torch.frombuffer(meta, dtype=torch.uint8).fill_(0) + else: + self.is_creator = False + # same resource_tracker workaround as ShmRingBuffer + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + try: + self.shared_memory = shared_memory.SharedMemory(name=name) + assert self.shared_memory.size >= self.total_bytes + except FileNotFoundError: + # deserialized on a different node; arena unused there + pass + + def handle(self): + return (self.n_reader, self.slot_bytes, self.n_slots, self.shared_memory.name) + + def _meta(self, idx: int) -> memoryview: + start = self.metadata_offset + idx * self.metadata_size + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + self.metadata_size] + + def _slot(self, idx: int, nbytes: int) -> memoryview: + start = idx * self.slot_bytes + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + nbytes] + + # ---- writer side ---- + + def write_tensor(self, t: torch.Tensor) -> int | None: + """Copy tensor bytes into a free slot; return slot idx or None + (caller must then fall back to the default pickle path).""" + nbytes = t.numel() * t.element_size() + if nbytes > self.slot_bytes: + return None + memory_fence() + for probe in range(self.n_slots): + idx = (self._next_slot + probe) % self.n_slots + with self._meta(idx) as meta: + free = meta[0] == 0 or sum(meta[1:]) == self.n_reader + if not free: + continue + meta[0] = 0 # claim + src = t.detach().reshape(-1).view(torch.uint8) + slot_mv = self._slot(idx, nbytes) + try: + dst = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) + dst.copy_(src) + finally: + slot_mv.release() + with self._meta(idx) as meta: + for i in range(1, self.n_reader + 1): + meta[i] = 0 + memory_fence() + meta[0] = 1 + memory_fence() + self._next_slot = (idx + 1) % self.n_slots + return idx + self._fallbacks += 1 + if self._fallbacks == 1 or self._fallbacks % 100 == 0: + logger.info( + "ShmTensorArena: no free slot (%d bytes, %d fallbacks so far); " + "falling back to in-band pickling. Consider raising " + "VLLM_SHM_TENSOR_ARENA_SLOTS/SLOT_MB.", + nbytes, + self._fallbacks, + ) + return None + + # ---- reader side ---- + + def _ensure_pinned(self): + """cudaHostRegister the whole arena mapping in THIS process (lazy, + once). Without it the HtoD of a zero-copy tensor pays first-touch + page faults on the tmpfs mapping plus pageable staging (~hundreds of + ms for 200MB); registration allocates+pins the pages once, making + every later HtoD a true DMA. Failure (no CUDA in this process, etc.) + is fine — the copy still works, just slower.""" + if self._pin_attempted: + return + self._pin_attempted = True + try: + if not torch.cuda.is_available(): + return + import ctypes + + buf = self.shared_memory.buf + assert buf is not None + ptr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) + ret = torch.cuda.cudart().cudaHostRegister(ptr, self.total_bytes, 0) + self._pinned = int(ret) == 0 + logger.info( + "ShmTensorArena: cudaHostRegister(%d MB) -> %s", + self.total_bytes >> 20, + "pinned" if self._pinned else f"error {int(ret)}", + ) + except Exception as e: + logger.info("ShmTensorArena: host-register skipped: %s", e) + + def get_tensor( + self, idx: int, nbytes: int, dtype: torch.dtype, shape: tuple[int, ...] + ) -> torch.Tensor: + """Zero-copy view of a slot as a tensor. The slot is released at this + reader's next flush_releases() (called from MessageQueue.dequeue).""" + self._ensure_pinned() + # NOT a context-manager view: the tensor must keep the mapping alive. + slot_mv = self._slot(idx, nbytes) + t8 = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) + t = t8.view(dtype).view(shape) + self._pending_release.append(idx) + return t + + def flush_releases(self): + if not self._pending_release: + return + for idx in self._pending_release: + with self._meta(idx) as meta: + meta[1 + self.reader_rank] = 1 + memory_fence() + self._pending_release.clear() + + def __del__(self): + if hasattr(self, "shared_memory"): + try: + self.shared_memory.close() + if self.is_creator: + self.shared_memory.unlink() + except BufferError: + # zero-copy tensor views may still hold exported pointers at + # interpreter shutdown; the mapping dies with the process. + pass + + +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook: rebuild a tensor as a zero-copy view of an arena slot.""" + arena = _TENSOR_ARENAS[arena_name] + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) + + +class _ArenaPickler(pickle.Pickler): + """Pickler that diverts large contiguous CPU tensors into the arena.""" + + def __init__(self, file, arena: ShmTensorArena, buffer_callback=None): + super().__init__( + file, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=buffer_callback, + ) + self.arena = arena + + def reducer_override(self, obj): + if ( + isinstance(obj, torch.Tensor) + and obj.device.type == "cpu" + and obj.layout == torch.strided + and obj.is_contiguous() + and obj.numel() * obj.element_size() + >= VLLM_SHM_TENSOR_ARENA_MIN_MB * 1024 * 1024 + ): + idx = self.arena.write_tensor(obj) + if idx is not None: + return ( + _rebuild_arena_tensor, + ( + self.arena.shared_memory.name, + idx, + obj.numel() * obj.element_size(), + str(obj.dtype).removeprefix("torch."), + tuple(obj.shape), + ), + ) + return NotImplemented + + @dataclass class Handle: local_reader_ranks: list[int] = field(default_factory=list) @@ -353,6 +615,7 @@ class Handle: local_notify_addr: str | None = None remote_subscribe_addr: str | None = None remote_addr_ipv6: bool = False + tensor_arena_handle: tuple[int, int, int, str] | None = None class MessageQueue: @@ -377,12 +640,24 @@ def __init__( self.shutting_down = False context = Context() + self.tensor_arena: ShmTensorArena | None = None if n_local_reader > 0: # for local readers, we will: # 1. create a shared memory ring buffer to communicate small data # 2. create a publish-subscribe socket to communicate large data self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks) + # Zero-copy arena for large CPU tensors (see ShmTensorArena). + # Local readers only: remote readers receive the pickled bytes + # over a socket and cannot map the arena, so the substitution + # would break them. + if VLLM_SHM_TENSOR_ARENA and n_remote_reader == 0: + self.tensor_arena = ShmTensorArena( + n_local_reader, + VLLM_SHM_TENSOR_ARENA_SLOT_MB * 1024 * 1024, + VLLM_SHM_TENSOR_ARENA_SLOTS, + ) + # XPUB is very similar to PUB, # except that it can receive subscription messages # to confirm the number of subscribers @@ -443,6 +718,9 @@ def __init__( local_notify_addr=local_notify_addr, remote_subscribe_addr=remote_subscribe_addr, remote_addr_ipv6=remote_addr_ipv6, + tensor_arena_handle=( + self.tensor_arena.handle() if self.tensor_arena is not None else None + ), ) logger.debug("vLLM message queue communication handle: %s", self.handle) @@ -458,6 +736,7 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": context = Context() + self.tensor_arena = None if rank in handle.local_reader_ranks: assert handle.buffer_handle is not None self.buffer = ShmRingBuffer(*handle.buffer_handle) @@ -466,6 +745,15 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": self._is_local_reader = True self._is_remote_reader = False + arena_handle = getattr(handle, "tensor_arena_handle", None) + if arena_handle is not None: + self.tensor_arena = ShmTensorArena( + *arena_handle, reader_rank=self.local_reader_rank + ) + _TENSOR_ARENAS[self.tensor_arena.shared_memory.name] = ( + self.tensor_arena + ) + self.local_socket = context.socket(SUB) self.local_socket.setsockopt_string(SUBSCRIBE, "") socket_addr = handle.local_subscribe_addr @@ -733,9 +1021,15 @@ def oob_callback(buf: PickleBuffer) -> bool: total_bytes += len(raw_buf) + 4 return False - all_buffers[0] = pickle.dumps( - obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback - ) + arena = getattr(self, "tensor_arena", None) + if arena is not None: + bio = io.BytesIO() + _ArenaPickler(bio, arena, buffer_callback=oob_callback).dump(obj) + all_buffers[0] = bio.getbuffer() + else: + all_buffers[0] = pickle.dumps( + obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback + ) if self.n_local_reader > 0: if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: with self.acquire_write(timeout) as buf: @@ -769,6 +1063,12 @@ def dequeue( ): """Read from message queue with optional timeout (in seconds)""" if self._is_local_reader: + # Release arena slots consumed by the PREVIOUS message: the worker + # loop is sequential, so the previous step (incl. its HtoD of any + # zero-copy tensors) has completed by the time we dequeue again. + arena = getattr(self, "tensor_arena", None) + if arena is not None: + arena.flush_releases() with self.acquire_read(timeout, indefinite) as buf: overflow = buf[0] == 1 if not overflow: