DeepInfer/tests/ve/test_shared_pool_groups.py

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()