Merge e6afeb54aac21c1865653700b88023d7a9ad22b3 into 553741557a57df44a46fb3ab72b6aafeaf945337

deepinfer/ve/multimodal_processor.py

@@ -298,8 +298,10 @@ class MultimodalProcessor:
                 for key, (handle, shape, dtype) in multi_modal_data["handle"].items():
                     multi_modal_data[key] = get_multi_modal_pool().get_tensor(handle, shape, dtype)
 
-            self.multi_modal_dict[id].clear()
-            self.multi_modal_dict[id].update(multi_modal_data)
+            # The sequence keeps this dict alive; the processor only tracks pending work.
+            request_data = self.multi_modal_dict.pop(id)
+            request_data.clear()
+            request_data.update(multi_modal_data)
             self.scheduler.set_schedulable(id)
 
         self._check_worker_hang()

tests/ve/test_multimodal_processor.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import gc
+import queue
+import weakref
+from unittest.mock import Mock, call
+
+import pytest
+import torch
+
+from deepinfer.ve import multimodal_processor, schema
+from deepinfer.ve.multimodal_processor import MultimodalProcessor
+
+pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
+
+
+@pytest.fixture
+def processor(monkeypatch: pytest.MonkeyPatch) -> MultimodalProcessor:
+    monkeypatch.setattr(multimodal_processor.envs, "DEEPINFER_USE_MULTIMODAL_POOL", False)
+    processor = object.__new__(MultimodalProcessor)
+    processor.input_queue = queue.Queue()
+    processor.output_queue = queue.Queue()
+    processor.multi_modal_dict = {}
+    processor.to_abort_ids = []
+    processor._pending = set()
+    processor._last_output_ts = 0.0
+    processor._worker_timeout = 0.0
+    processor.scheduler = Mock()
+    return processor
+
+
+def test_fetch_releases_completed_results(processor: MultimodalProcessor) -> None:
+    for i in range(4):
+        request_id = str(i)
+        sequence_data = {}
+        processor.add(request_id, sequence_data)
+        processor.input_queue.get_nowait()  # The worker has consumed the input.
+
+        images = torch.zeros(16, dtype=torch.uint8, device="cpu")
+        images_ref = weakref.ref(images)
+        processor.output_queue.put((request_id, {"images_raw": images}))
+        del images
+
+        processor.fetch()
+
+        assert sequence_data["images_raw"] is images_ref()
+        processor.scheduler.set_schedulable.assert_called_with(request_id)
+        assert not processor.has_pending()
+
+        del sequence_data
+        gc.collect()
+        assert images_ref() is None
+        assert not processor.multi_modal_dict
+
+
+def test_fetch_keeps_unfinished_requests(processor: MultimodalProcessor) -> None:
+    ready_data, pending_data = {}, {}
+    processor.add("ready", ready_data)
+    processor.add("pending", pending_data)
+    processor.input_queue.get_nowait()
+    processor.input_queue.get_nowait()
+    processor.output_queue.put(("ready", {"images_raw": b"ready"}))
+
+    processor.fetch()
+
+    assert ready_data == {"images_raw": b"ready"}
+    assert processor.multi_modal_dict == {"pending": pending_data}
+    assert processor.has_pending()
+    processor.scheduler.set_schedulable.assert_called_once_with("ready")
+
+    processor.output_queue.put(("pending", {"images_raw": b"pending"}))
+    processor.fetch()
+
+    assert pending_data == {"images_raw": b"pending"}
+    assert not processor.multi_modal_dict
+    assert not processor.has_pending()
+    assert processor.scheduler.set_schedulable.call_args_list == [call("ready"), call("pending")]
+
+
+@pytest.mark.parametrize("abort_before_fetch", [False, True])
+def test_fetch_pool_handle_ownership(processor: MultimodalProcessor, monkeypatch: pytest.MonkeyPatch, abort_before_fetch: bool) -> None:
+    monkeypatch.setattr(multimodal_processor.envs, "DEEPINFER_USE_MULTIMODAL_POOL", True)
+    pool = Mock()
+    monkeypatch.setattr(multimodal_processor, "get_multi_modal_pool", lambda: pool)
+    monkeypatch.setattr(schema, "get_multi_modal_pool", lambda: pool)
+    images = torch.zeros(16, dtype=torch.uint8, device="cpu")
+    pool.get_tensor.return_value = images
+    handle = object()
+    handles = {"images_raw": (handle, images.shape, images.dtype)}
+    sequence_data = {}
+    processor.add("request", sequence_data)
+    processor.input_queue.get_nowait()
+    if abort_before_fetch:
+        processor.abort(["request"])
+    processor.output_queue.put(("request", {"images_raw": None, "handle": handles}))
+
+    processor.fetch()
+
+    assert not processor.multi_modal_dict
+    assert not processor.has_pending()
+    if abort_before_fetch:
+        pool.get_tensor.assert_not_called()
+        pool.free.assert_called_once_with(handle)
+        processor.scheduler.set_schedulable.assert_not_called()
+    else:
+        assert sequence_data["images_raw"] is images
+        assert sequence_data["handle"] is handles
+        pool.get_tensor.assert_called_once_with(handle, images.shape, images.dtype)
+        pool.free.assert_not_called()
+        processor.scheduler.set_schedulable.assert_called_once_with("request")
+        schema.free_multi_modal_pool_handles(sequence_data, pool=pool)
+        pool.free.assert_called_once_with(handle)
+
+
+def test_fetch_reports_worker_failure(processor: MultimodalProcessor) -> None:
+    sequence_data = {}
+    processor.add("failed", sequence_data)
+    processor.input_queue.get_nowait()
+    processor.output_queue.put(("failed", None))
+
+    processor.fetch()
+
+    assert processor.fetch_to_abort_requests() == ["failed"]
+    assert processor.fetch_to_abort_requests() == []
+    assert not processor.has_pending()
+    processor.scheduler.set_schedulable.assert_not_called()
+    processor.abort(["failed"])
+    assert not processor.multi_modal_dict