DeepInfer/tests/ve/test_multimodal_processor.py

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