lmcache--lmcache/tests/v1/test_device_id_race.py

# SPDX-License-Identifier: Apache-2.0
"""Unit tests for NixlDynamicStorageBackend._alloc_device_ids.

These tests verify the fix for the NIXL OBJ devIdToObjKey_ race condition
(https://github.com/LMCache/LMCache/issues/2983).  When nixl_async_put is
enabled, async PUT cleanup and sync GET run concurrently.  Both previously
used device_id = 0, 1, 2, ... for every call, causing the PUT's
deregister to erase the GET registration in NIXL's flat map.

The fix uses a monotonically increasing counter so each register/deregister
cycle gets globally unique device_ids.  These tests validate correctness
and thread safety of that counter without requiring NIXL or CUDA hardware.
"""

# Standard
import threading

# Third Party
import pytest


# ---------------------------------------------------------------------------
# Helper: build a minimal NixlDynamicStorageBackend with only the fields
# needed by _alloc_device_ids, bypassing __init__'s heavy dependencies.
# ---------------------------------------------------------------------------
def _make_stub_backend():
    """Return a NixlDynamicStorageBackend instance with __init__ bypassed."""
    # Import inside function so the test file can be collected even when
    # nixl is not installed (the importorskip below handles the skip).
    # First Party
    from lmcache.v1.storage_backend.nixl_storage_backend import (
        NixlDynamicStorageBackend,
    )

    obj = object.__new__(NixlDynamicStorageBackend)
    obj._device_id_counter = 0
    obj._device_id_lock = threading.Lock()
    return obj


# Skip the entire module if nixl is not importable (mirrors existing tests)
pytest.importorskip("nixl", reason="nixl package is required for nixl tests")


# ---- Basic correctness ----------------------------------------------------


class TestAllocDeviceIds:
    """Tests for _alloc_device_ids correctness."""

    def test_single_alloc_returns_sequential_range(self):
        backend = _make_stub_backend()
        ids = backend._alloc_device_ids(5)
        assert ids == [0, 1, 2, 3, 4]

    def test_successive_allocs_are_non_overlapping(self):
        backend = _make_stub_backend()
        first = backend._alloc_device_ids(3)
        second = backend._alloc_device_ids(4)
        third = backend._alloc_device_ids(2)
        assert first == [0, 1, 2]
        assert second == [3, 4, 5, 6]
        assert third == [7, 8]

    def test_alloc_zero_returns_empty(self):
        backend = _make_stub_backend()
        ids = backend._alloc_device_ids(0)
        assert ids == []
        # Counter should not advance
        assert backend._device_id_counter == 0

    def test_alloc_one_returns_single_element(self):
        backend = _make_stub_backend()
        ids = backend._alloc_device_ids(1)
        assert ids == [0]
        assert backend._device_id_counter == 1

    def test_counter_advances_correctly(self):
        backend = _make_stub_backend()
        backend._alloc_device_ids(10)
        assert backend._device_id_counter == 10
        backend._alloc_device_ids(5)
        assert backend._device_id_counter == 15

    def test_all_ids_globally_unique_across_many_allocs(self):
        backend = _make_stub_backend()
        all_ids = []
        for n in [1, 3, 7, 2, 10, 5]:
            all_ids.extend(backend._alloc_device_ids(n))
        assert len(all_ids) == len(set(all_ids)), "duplicate device_ids found"
        assert all_ids == list(range(28))


# ---- Thread safety ---------------------------------------------------------


class TestAllocDeviceIdsThreadSafety:
    """Verify no duplicate IDs under concurrent access."""

    def test_concurrent_allocs_produce_unique_ids(self):
        """Simulate the race: multiple threads calling _alloc_device_ids
        concurrently, as happens when async PUT cleanup and sync GET
        overlap."""
        backend = _make_stub_backend()
        num_threads = 16
        allocs_per_thread = 200
        batch_size = 5  # typical: one device_id per key in a batch
        results: list[list[int]] = [[] for _ in range(num_threads)]

        barrier = threading.Barrier(num_threads)

        def worker(thread_idx):
            barrier.wait()  # maximize contention
            for _ in range(allocs_per_thread):
                ids = backend._alloc_device_ids(batch_size)
                results[thread_idx].extend(ids)

        threads = [
            threading.Thread(target=worker, args=(i,)) for i in range(num_threads)
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        all_ids = []
        for r in results:
            all_ids.extend(r)

        expected_total = num_threads * allocs_per_thread * batch_size
        assert len(all_ids) == expected_total
        assert len(set(all_ids)) == expected_total, (
            f"found {expected_total - len(set(all_ids))} duplicate device_ids "
            f"under concurrent access"
        )

    def test_interleaved_put_get_simulation(self):
        """Simulate interleaved PUT (async cleanup) and GET (sync worker)
        operations to verify the fix prevents the exact race from #44."""
        backend = _make_stub_backend()
        put_ids: list[list[int]] = []
        get_ids: list[list[int]] = []
        barrier = threading.Barrier(2)

        def put_worker():
            """Simulates async PUT: register → transfer → deregister."""
            barrier.wait()
            for _ in range(100):
                ids = backend._alloc_device_ids(3)
                put_ids.append(ids)

        def get_worker():
            """Simulates sync GET: register → prepXfer → postXfer → deregister."""
            barrier.wait()
            for _ in range(100):
                ids = backend._alloc_device_ids(3)
                get_ids.append(ids)

        t_put = threading.Thread(target=put_worker)
        t_get = threading.Thread(target=get_worker)
        t_put.start()
        t_get.start()
        t_put.join()
        t_get.join()

        # Flatten and check uniqueness
        all_put = [id_ for batch in put_ids for id_ in batch]
        all_get = [id_ for batch in get_ids for id_ in batch]

        # No overlap between PUT and GET id ranges
        overlap = set(all_put) & set(all_get)
        assert len(overlap) == 0, (
            f"PUT and GET device_ids overlap: {sorted(list(overlap))[:10]}..."
        )

        # All IDs globally unique
        combined = all_put + all_get
        assert len(combined) == len(set(combined))