diff --git a/deepinfer/ve/multimodal_processor.py b/deepinfer/ve/multimodal_processor.py index 05e0077..949cc3d 100644 --- a/deepinfer/ve/multimodal_processor.py +++ b/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() diff --git a/tests/ve/test_multimodal_processor.py b/tests/ve/test_multimodal_processor.py new file mode 100644 index 0000000..7f84030 --- /dev/null +++ b/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