From 7cd04027633afda9653d8daf3f5ced17b8401446 Mon Sep 17 00:00:00 2001
From: zhujiale520 <175239653+zhujiale520@users.noreply.github.com>
Date: Tue, 20 Oct 2026 19:02:20 +0800
Subject: [PATCH] fix: allocate multimodal tensor groups without deadlocks or
 leaks

---
 deepinfer/shared_memory.py           |  27 +++++
 deepinfer/ve/multimodal_processor.py |  11 +-
 tests/ve/test_shared_pool_groups.py  | 162 +++++++++++++++++++++++++++
 3 files changed, 193 insertions(+), 7 deletions(-)
 create mode 100644 tests/ve/test_shared_pool_groups.py

diff --git a/deepinfer/shared_memory.py b/deepinfer/shared_memory.py
index 4d256e9..2234efa 100644
--- a/deepinfer/shared_memory.py
+++ b/deepinfer/shared_memory.py
@@ -133,6 +133,8 @@ class SharedTorchMemoryPool:
 
         # 锁：多进程同步
         self.lock = mp.Lock()
+        # Serialize complete tensor groups without blocking consumers that free blocks.
+        self._tensor_group_lock = mp.Lock()
 
         self._init_metadata()
 
@@ -309,6 +311,31 @@ class SharedTorchMemoryPool:
             self.free(handle)
             raise
 
+    def copy_tensors(self, tensors: dict[str, torch.Tensor]) -> dict[str, tuple]:
+        """Copy a complete request into the pool, rolling back on failure.
+
+        Serialize groups so concurrent producers cannot hold partial allocations
+        while waiting for one another. Consumers can still free earlier groups.
+        """
+        tensors = {key: value for key, value in tensors.items() if value.numel() > 0}
+        required = sum(self._align_up(value.numel() * value.element_size()) for value in tensors.values())
+        if required > self.total_bytes:
+            raise MemoryError(f"tensor group requires {required} aligned bytes, but the pool capacity is {self.total_bytes} bytes")
+
+        handles = {}
+        with self._tensor_group_lock:
+            try:
+                for key, value in tensors.items():
+                    handle, shared_data = self.alloc_tensor(value.shape, dtype=value.dtype)
+                    # Record ownership before copy_ so its exceptions also roll back.
+                    handles[key] = (handle, value.shape, value.dtype)
+                    shared_data.copy_(value)
+            except Exception:
+                for handle, _shape, _dtype in handles.values():
+                    self.free(handle)
+                raise
+        return handles
+
     def _validate_handle(self, handle: BlockHandle) -> None:
         idx = int(handle.block_id)
         if idx < 0 or idx >= self.max_blocks:
diff --git a/deepinfer/ve/multimodal_processor.py b/deepinfer/ve/multimodal_processor.py
index 05e0077..1347bd5 100644
--- a/deepinfer/ve/multimodal_processor.py
+++ b/deepinfer/ve/multimodal_processor.py
@@ -202,13 +202,10 @@ class MultimodalProcessor:
                 if envs.DEEPINFER_USE_MULTIMODAL_POOL:
                     if not isinstance(multi_modal_data, dict):
                         raise NotImplementedError(f"multi_modal_data type {type(multi_modal_data)} not supported")
-                    multi_modal_data["handle"] = handles = {}
-                    for key, value in multi_modal_data.items():
-                        if isinstance(value, torch.Tensor) and value.numel() > 0:
-                            handle, shared_data = multi_modal_pool.alloc_tensor(value.shape, dtype=value.dtype)
-                            shared_data.copy_(value)
-                            handles[key] = (handle, value.shape, value.dtype)
-                            multi_modal_data[key] = None
+                    tensors = {key: value for key, value in multi_modal_data.items() if isinstance(value, torch.Tensor)}
+                    multi_modal_data["handle"] = handles = multi_modal_pool.copy_tensors(tensors)
+                    for key in handles:
+                        multi_modal_data[key] = None
                 self.output_queue.put((id, multi_modal_data))
             except Exception:
                 logger.exception("Multi-modal request %s failed during multi_modal_mapper", id)
diff --git a/tests/ve/test_shared_pool_groups.py b/tests/ve/test_shared_pool_groups.py
new file mode 100644
index 0000000..0d7eb31
--- /dev/null
+++ b/tests/ve/test_shared_pool_groups.py
@@ -0,0 +1,162 @@
+from collections.abc import Generator
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+
+import pytest
+import torch
+
+from deepinfer import shared_memory
+from deepinfer.shared_memory import SharedTorchMemoryPool
+
+pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
+
+
+@pytest.fixture
+def pool(monkeypatch: pytest.MonkeyPatch) -> Generator[SharedTorchMemoryPool, None, None]:
+    # Exercise real shared storage and allocator metadata without an accelerator.
+    monkeypatch.setattr(shared_memory.PLATFORM, "pin_host_memory", lambda _: None)
+    monkeypatch.setattr(shared_memory.PLATFORM, "unpin_host_memory", lambda _: None)
+    pool = SharedTorchMemoryPool(256, max_blocks=32)
+    yield pool
+    pool.close()
+
+
+def free_group(pool: SharedTorchMemoryPool, handles: dict[str, tuple]) -> None:
+    for handle, _shape, _dtype in handles.values():
+        pool.free(handle)
+
+
+def free_bytes(pool: SharedTorchMemoryPool) -> int:
+    return sum(size for _, _, size in pool.dump_free_list())
+
+
+@pytest.mark.parametrize("sizes", [(65,), (33, 17, 1), (32, 32, 1)])
+def test_rejects_impossible_group_before_allocating(
+    pool: SharedTorchMemoryPool, monkeypatch: pytest.MonkeyPatch, sizes: tuple[int, ...]
+) -> None:
+    calls = []
+    original = pool.alloc_tensor
+
+    def allocate(*args: object, **kwargs: object) -> tuple:
+        calls.append(args)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(pool, "alloc_tensor", allocate)
+    with pytest.raises(MemoryError, match="tensor group requires.*pool capacity"):
+        pool.copy_tensors({str(i): torch.ones(n) for i, n in enumerate(sizes)})
+    assert calls == []
+    assert free_bytes(pool) == 256
+
+
+def test_exact_capacity_and_empty_tensors(pool: SharedTorchMemoryPool) -> None:
+    values = {"first": torch.arange(32), "empty": torch.empty(0)}
+    handles = pool.copy_tensors(values)
+    assert set(handles) == {"first"}
+    handle, shape, dtype = handles["first"]
+    assert torch.equal(pool.get_tensor(handle, shape, dtype), values["first"])
+    assert free_bytes(pool) == 0
+    free_group(pool, handles)
+    assert free_bytes(pool) == 256
+    assert pool.copy_tensors({"empty": torch.empty(0)}) == {}
+
+
+def test_copy_failure_returns_all_allocations(pool: SharedTorchMemoryPool) -> None:
+    sparse = torch.sparse_coo_tensor(torch.tensor([[0], [0]]), torch.tensor([1.0]), (2, 2), check_invariants=True)
+    for _ in range(3):
+        with pytest.raises(RuntimeError):
+            pool.copy_tensors({"valid": torch.ones(16), "invalid": sparse})
+        assert free_bytes(pool) == 256
+    handles = pool.copy_tensors({"retry": torch.ones(64)})
+    free_group(pool, handles)
+    assert free_bytes(pool) == 256
+
+
+def test_allocation_failure_returns_earlier_blocks(pool: SharedTorchMemoryPool, monkeypatch: pytest.MonkeyPatch) -> None:
+    original = pool.alloc_tensor
+    calls = 0
+
+    def allocate(*args: object, **kwargs: object) -> tuple:
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise MemoryError("metadata exhausted")
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(pool, "alloc_tensor", allocate)
+    with pytest.raises(MemoryError, match="metadata exhausted"):
+        pool.copy_tensors({"first": torch.ones(16), "second": torch.ones(16)})
+    assert free_bytes(pool) == 256
+
+
+def test_producers_serialize_groups_while_consumers_can_free(pool: SharedTorchMemoryPool, monkeypatch: pytest.MonkeyPatch) -> None:
+    first_allocated = Event()
+    continue_first = Event()
+    second_started = Event()
+    second_allocated = Event()
+    original = pool.alloc_tensor
+
+    def allocate(shape: tuple[int, ...], dtype: torch.dtype = torch.float32) -> tuple:
+        result = original(shape, dtype)
+        if not first_allocated.is_set():
+            first_allocated.set()
+            assert continue_first.wait(5)
+        elif not continue_first.is_set():
+            second_allocated.set()
+        return result
+
+    monkeypatch.setattr(pool, "alloc_tensor", allocate)
+    values = {"first": torch.ones(32), "second": torch.ones(16)}
+
+    def second_producer() -> dict[str, tuple]:
+        second_started.set()
+        return pool.copy_tensors(values)
+
+    # Timeouts turn regressions into failures rather than hanging the suite.
+    executor = ThreadPoolExecutor(2)
+    try:
+        first = executor.submit(pool.copy_tensors, values)
+        assert first_allocated.wait(5)
+        second = executor.submit(second_producer)
+        assert second_started.wait(5)
+        assert not second_allocated.wait(0.1)
+        continue_first.set()
+        first_handles = first.result(timeout=5)
+        # The second producer may be waiting for memory while holding the group
+        # lock. A consumer must still be able to release the first group.
+        free_group(pool, first_handles)
+        second_handles = second.result(timeout=5)
+        free_group(pool, second_handles)
+        assert free_bytes(pool) == 256
+    finally:
+        continue_first.set()
+        executor.shutdown(wait=False, cancel_futures=True)
+
+
+def copy_group_worker(pool: SharedTorchMemoryPool, output: object) -> None:
+    handles = pool.copy_tensors({"first": torch.ones(32), "second": torch.ones(16)})
+    output.put(handles)
+
+
+def test_group_lock_is_shared_between_processes(pool: SharedTorchMemoryPool) -> None:
+    import multiprocessing
+
+    context = multiprocessing.get_context("fork")
+    output = context.Queue()
+    processes = [context.Process(target=copy_group_worker, args=(pool, output)) for _ in range(2)]
+    try:
+        for process in processes:
+            process.start()
+        for _ in processes:
+            handles = output.get(timeout=10)
+            free_group(pool, handles)
+        for process in processes:
+            process.join(timeout=5)
+            assert process.exitcode == 0
+        assert free_bytes(pool) == 256
+    finally:
+        for process in processes:
+            if process.is_alive():
+                process.terminate()
+            process.join(timeout=5)
+        output.close()
+        output.join_thread()