612 lines
21 KiB
Python
612 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from unittest.mock import MagicMock, patch
|
|
import asyncio
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import threading
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.utils import CacheEngineKey, DiskCacheMetadata
|
|
from lmcache.v1.config import LMCacheEngineConfig
|
|
from lmcache.v1.config_base import _parse_local_disk
|
|
from lmcache.v1.memory_management import MemoryFormat, MemoryObj
|
|
from lmcache.v1.metadata import LMCacheMetadata
|
|
from lmcache.v1.storage_backend.local_cpu_backend import LocalCPUBackend
|
|
from lmcache.v1.storage_backend.local_disk_backend import LocalDiskBackend
|
|
|
|
|
|
class MockLookupServer:
|
|
def __init__(self):
|
|
self.removed_keys = []
|
|
self.inserted_keys = []
|
|
|
|
def batched_remove(self, keys):
|
|
self.removed_keys.extend(keys)
|
|
|
|
def batched_insert(self, keys):
|
|
self.inserted_keys.extend(keys)
|
|
|
|
|
|
class MockLMCacheWorker:
|
|
def __init__(self):
|
|
self.messages = []
|
|
|
|
def put_msg(self, msg):
|
|
self.messages.append(msg)
|
|
|
|
|
|
def create_test_config(
|
|
disk_path: str,
|
|
max_disk_size: float = 1.0,
|
|
local_disk_path_sharding: str = "by_gpu",
|
|
):
|
|
"""Create a test configuration for LocalDiskBackend."""
|
|
config = LMCacheEngineConfig.from_defaults(
|
|
chunk_size=256,
|
|
local_disk=disk_path,
|
|
local_disk_path_sharding=local_disk_path_sharding,
|
|
max_local_disk_size=max_disk_size,
|
|
lmcache_instance_id="test_instance",
|
|
)
|
|
return config
|
|
|
|
|
|
def create_test_metadata():
|
|
"""Create a test metadata for LMCacheMetadata."""
|
|
return LMCacheMetadata(
|
|
model_name="test_model",
|
|
world_size=1,
|
|
local_world_size=1,
|
|
worker_id=0,
|
|
local_worker_id=0,
|
|
kv_dtype=torch.bfloat16,
|
|
kv_shape=(28, 2, 256, 8, 128),
|
|
)
|
|
|
|
|
|
def create_test_key(key_id: int = 0) -> CacheEngineKey:
|
|
"""Create a test CacheEngineKey."""
|
|
return CacheEngineKey(
|
|
model_name="test_model",
|
|
world_size=3,
|
|
worker_id=1,
|
|
chunk_hash=hash(key_id),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_disk_path():
|
|
"""Create a temporary directory for disk storage tests."""
|
|
temp_dir = tempfile.mkdtemp()
|
|
yield temp_dir
|
|
# Cleanup
|
|
if os.path.exists(temp_dir):
|
|
shutil.rmtree(temp_dir)
|
|
|
|
|
|
@pytest.fixture
|
|
def async_loop():
|
|
"""Create an asyncio event loop for testing."""
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
yield loop
|
|
loop.close()
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def local_cpu_backend(memory_allocator):
|
|
"""Create a LocalCPUBackend for testing."""
|
|
config = LMCacheEngineConfig.from_legacy(chunk_size=256)
|
|
return LocalCPUBackend(config, memory_allocator=memory_allocator)
|
|
|
|
|
|
@pytest.fixture
|
|
def local_disk_backend(temp_disk_path, async_loop, local_cpu_backend):
|
|
"""Create a LocalDiskBackend for testing."""
|
|
config = create_test_config(temp_disk_path)
|
|
return LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
|
|
|
|
class TestLocalDiskBackend:
|
|
"""Test cases for LocalDiskBackend."""
|
|
|
|
def test_init(self, temp_disk_path, async_loop, local_cpu_backend):
|
|
"""Test LocalDiskBackend initialization."""
|
|
config = create_test_config(temp_disk_path)
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
|
|
assert backend.dst_device == "cuda:0"
|
|
assert backend.local_cpu_backend == local_cpu_backend
|
|
assert backend.path == temp_disk_path
|
|
assert os.path.exists(temp_disk_path)
|
|
assert backend.lmcache_worker is None
|
|
assert backend.instance_id == "test_instance"
|
|
assert backend.usage == 0
|
|
assert len(backend.dict) == 0
|
|
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_init_with_lookup_server_and_worker(
|
|
self, temp_disk_path, async_loop, local_cpu_backend
|
|
):
|
|
"""Test LocalDiskBackend initialization with lookup server and worker."""
|
|
config = create_test_config(temp_disk_path)
|
|
lmcache_worker = MockLMCacheWorker()
|
|
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
lmcache_worker=lmcache_worker,
|
|
)
|
|
|
|
assert backend.lmcache_worker == lmcache_worker
|
|
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_str(self, local_disk_backend):
|
|
"""Test string representation."""
|
|
assert str(local_disk_backend) == "LocalDiskBackend"
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_key_to_path(self, local_disk_backend):
|
|
"""Test key to path conversion."""
|
|
key = create_test_key(1)
|
|
path = local_disk_backend._key_to_path(key)
|
|
|
|
expected_filename = key.to_string().replace("/", "-") + ".pt"
|
|
assert path == os.path.join(local_disk_backend.path, expected_filename)
|
|
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_contains_key_not_exists(self, local_disk_backend):
|
|
"""Test contains() when key doesn't exist."""
|
|
key = create_test_key(2)
|
|
assert not local_disk_backend.contains(key)
|
|
assert not local_disk_backend.contains(key, pin=True)
|
|
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_get_blocking_key_not_exists(self, local_disk_backend):
|
|
"""Test get_blocking() when key doesn't exist."""
|
|
key = create_test_key(2)
|
|
result = local_disk_backend.get_blocking(key)
|
|
|
|
assert result is None
|
|
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
|
|
class TestMultiPathDiskBackend:
|
|
"""Test cases for multi-path (multi-device) LocalDiskBackend."""
|
|
|
|
def test_init_multi_path(self, async_loop, local_cpu_backend):
|
|
"""Test initialisation with comma-separated paths."""
|
|
dir_a = tempfile.mkdtemp()
|
|
dir_b = tempfile.mkdtemp()
|
|
try:
|
|
combined = f"{dir_a},{dir_b}"
|
|
config = create_test_config(combined)
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
|
|
# Path selected by device_id (0 % 2 = 0 -> dir_a)
|
|
assert backend.path == dir_a
|
|
# Both directories should exist
|
|
assert os.path.isdir(dir_a)
|
|
assert os.path.isdir(dir_b)
|
|
# Block size is a plain int for the selected path
|
|
assert isinstance(backend.os_disk_bs, int)
|
|
finally:
|
|
shutil.rmtree(dir_a, ignore_errors=True)
|
|
shutil.rmtree(dir_b, ignore_errors=True)
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_gpu_affinity_selects_path(self, async_loop, local_cpu_backend):
|
|
"""Different cuda devices select different paths via modulo."""
|
|
dir_a = tempfile.mkdtemp()
|
|
dir_b = tempfile.mkdtemp()
|
|
try:
|
|
combined = f"{dir_a},{dir_b}"
|
|
config = create_test_config(combined)
|
|
|
|
dirs_by_gpu = {}
|
|
for device in ("cuda:0", "cuda:1"):
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device=device,
|
|
)
|
|
dirs_by_gpu[device] = backend.path
|
|
|
|
assert dirs_by_gpu["cuda:0"] == dir_a
|
|
assert dirs_by_gpu["cuda:1"] == dir_b
|
|
finally:
|
|
shutil.rmtree(dir_a, ignore_errors=True)
|
|
shutil.rmtree(dir_b, ignore_errors=True)
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_all_directories_created(self, async_loop, local_cpu_backend):
|
|
"""All paths in the list get their directories created."""
|
|
base = tempfile.mkdtemp()
|
|
try:
|
|
paths = [os.path.join(base, f"nvme{i}") for i in range(3)]
|
|
combined = ",".join(paths)
|
|
config = create_test_config(combined)
|
|
LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
for p in paths:
|
|
assert os.path.isdir(p), f"{p} should exist"
|
|
finally:
|
|
shutil.rmtree(base, ignore_errors=True)
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_single_path_backward_compat(
|
|
self, temp_disk_path, async_loop, local_cpu_backend
|
|
):
|
|
"""A single path (no commas) works exactly as before."""
|
|
config = create_test_config(temp_disk_path)
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
assert backend.path == temp_disk_path
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_path_sharding_default(self, temp_disk_path, async_loop, local_cpu_backend):
|
|
"""Default local_disk_path_sharding is 'by_gpu' (backend inits OK)."""
|
|
config = create_test_config(temp_disk_path)
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
assert backend.path == temp_disk_path
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_path_sharding_explicit_by_gpu(
|
|
self, temp_disk_path, async_loop, local_cpu_backend
|
|
):
|
|
"""Explicitly setting local_disk_path_sharding='by_gpu' works."""
|
|
config = create_test_config(temp_disk_path, local_disk_path_sharding="by_gpu")
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
assert backend.path == temp_disk_path
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_path_sharding_unsupported_raises(
|
|
self, temp_disk_path, async_loop, local_cpu_backend
|
|
):
|
|
"""Unsupported local_disk_path_sharding raises ValueError."""
|
|
config = create_test_config(
|
|
temp_disk_path, local_disk_path_sharding="round_robin"
|
|
)
|
|
with pytest.raises(ValueError, match="Unsupported path sharding strategy"):
|
|
LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cuda:0",
|
|
)
|
|
|
|
def test_cpu_dst_device_defaults_to_first_path(self, async_loop, local_cpu_backend):
|
|
"""dst_device='cpu' should fall back to device_id=0."""
|
|
dir_a = tempfile.mkdtemp()
|
|
dir_b = tempfile.mkdtemp()
|
|
try:
|
|
combined = f"{dir_a},{dir_b}"
|
|
config = create_test_config(combined)
|
|
backend = LocalDiskBackend(
|
|
config=config,
|
|
loop=async_loop,
|
|
local_cpu_backend=local_cpu_backend,
|
|
dst_device="cpu",
|
|
)
|
|
# device_id=0 -> 0 % 2 = 0 -> dir_a
|
|
assert backend.path == dir_a
|
|
finally:
|
|
shutil.rmtree(dir_a, ignore_errors=True)
|
|
shutil.rmtree(dir_b, ignore_errors=True)
|
|
local_cpu_backend.memory_allocator.close()
|
|
|
|
|
|
class TestParseLocalDisk:
|
|
"""Unit tests for the _parse_local_disk config parser."""
|
|
|
|
def test_none(self):
|
|
assert _parse_local_disk(None) is None
|
|
|
|
def test_single_raw_path(self):
|
|
assert _parse_local_disk("/mnt/nvme0/cache/") == "/mnt/nvme0/cache/"
|
|
|
|
def test_single_file_uri(self):
|
|
assert _parse_local_disk("file:///mnt/nvme0/cache/") == "/mnt/nvme0/cache/"
|
|
|
|
def test_single_file_uri_no_trailing_slash(self):
|
|
assert _parse_local_disk("file:///mnt/nvme0/cache") == "/mnt/nvme0/cache"
|
|
|
|
def test_comma_separated_raw(self):
|
|
result = _parse_local_disk("/mnt/nvme0/,/mnt/nvme1/")
|
|
assert result == "/mnt/nvme0/,/mnt/nvme1/"
|
|
|
|
def test_comma_separated_file_uris(self):
|
|
result = _parse_local_disk("file:///mnt/nvme0/,file:///mnt/nvme1/")
|
|
assert result == "/mnt/nvme0/,/mnt/nvme1/"
|
|
|
|
def test_mixed_uri_and_raw(self):
|
|
result = _parse_local_disk("file:///mnt/nvme0/,/mnt/nvme1/")
|
|
assert result == "/mnt/nvme0/,/mnt/nvme1/"
|
|
|
|
def test_whitespace_around_paths(self):
|
|
result = _parse_local_disk(" /mnt/nvme0/ , /mnt/nvme1/ ")
|
|
assert result == "/mnt/nvme0/,/mnt/nvme1/"
|
|
|
|
def test_empty_string(self):
|
|
assert _parse_local_disk("") is None
|
|
|
|
|
|
class TestGetBlockingCachePolicyUpdate:
|
|
"""Regression tests for phantom cache hit in get_blocking() (issue #3015).
|
|
|
|
``get_blocking()`` must call ``cache_policy.update_on_hit()`` only when
|
|
``load_bytes_from_disk()`` returns a valid ``MemoryObj``. Calling it
|
|
before confirming load success records a phantom hit that skews future
|
|
eviction decisions.
|
|
"""
|
|
|
|
def _inject_key(
|
|
self,
|
|
backend: LocalDiskBackend,
|
|
key: CacheEngineKey,
|
|
shape: torch.Size,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
"""Insert a key into backend.dict without writing anything to disk."""
|
|
meta = DiskCacheMetadata(
|
|
path="/nonexistent/path.pt",
|
|
size=0,
|
|
shape=shape,
|
|
dtype=dtype,
|
|
cached_positions=None,
|
|
fmt=MemoryFormat.KV_2LTD,
|
|
pin_count=0,
|
|
)
|
|
with backend.disk_lock:
|
|
backend.dict[key] = meta
|
|
backend.cache_policy.update_on_put(key)
|
|
|
|
def test_no_phantom_hit_when_load_fails(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""update_on_hit must NOT be called when load_bytes_from_disk returns None."""
|
|
key = create_test_key(101)
|
|
shape = torch.Size([28, 2, 256, 8, 128])
|
|
self._inject_key(local_disk_backend, key, shape, torch.bfloat16)
|
|
|
|
with patch.object(
|
|
local_disk_backend, "load_bytes_from_disk", return_value=None
|
|
):
|
|
with patch.object(
|
|
local_disk_backend.cache_policy, "update_on_hit"
|
|
) as mock_update:
|
|
result = local_disk_backend.get_blocking(key)
|
|
|
|
assert result is None
|
|
mock_update.assert_not_called()
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_updates_cache_policy_on_successful_load(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""update_on_hit must be called exactly once when the load succeeds."""
|
|
key = create_test_key(102)
|
|
shape = torch.Size([28, 2, 256, 8, 128])
|
|
self._inject_key(local_disk_backend, key, shape, torch.bfloat16)
|
|
|
|
fake_memory_obj = MagicMock(spec=MemoryObj)
|
|
with patch.object(
|
|
local_disk_backend, "load_bytes_from_disk", return_value=fake_memory_obj
|
|
):
|
|
with patch.object(
|
|
local_disk_backend.cache_policy, "update_on_hit"
|
|
) as mock_update:
|
|
result = local_disk_backend.get_blocking(key)
|
|
|
|
assert result is fake_memory_obj
|
|
mock_update.assert_called_once_with(key, local_disk_backend.dict)
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_key_absent_returns_none_without_policy_update(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""get_blocking must return None immediately when the key is not cached."""
|
|
key = create_test_key(103)
|
|
|
|
with patch.object(
|
|
local_disk_backend.cache_policy, "update_on_hit"
|
|
) as mock_update:
|
|
result = local_disk_backend.get_blocking(key)
|
|
|
|
assert result is None
|
|
mock_update.assert_not_called()
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
|
|
class TestBatchedGetBlocking:
|
|
"""Tests for the concurrent batched_get_blocking() override."""
|
|
|
|
_SHAPE = torch.Size([28, 2, 256, 8, 128])
|
|
_DTYPE = torch.bfloat16
|
|
|
|
def _write_key(
|
|
self,
|
|
backend: LocalDiskBackend,
|
|
key: CacheEngineKey,
|
|
) -> bytes:
|
|
"""Write random data to disk for *key* and register it in the backend.
|
|
|
|
Returns the raw bytes that were written so callers can verify reads.
|
|
"""
|
|
path = backend._key_to_path(key)
|
|
nbytes = 1
|
|
for s in self._SHAPE:
|
|
nbytes *= s
|
|
nbytes *= self._DTYPE.itemsize
|
|
data = os.urandom(nbytes)
|
|
with open(path, "wb") as f:
|
|
f.write(data)
|
|
backend.insert_key(
|
|
key,
|
|
size=nbytes,
|
|
shape=self._SHAPE,
|
|
dtype=self._DTYPE,
|
|
fmt=MemoryFormat.KV_2LTD,
|
|
)
|
|
return data
|
|
|
|
def test_all_keys_missing(self, local_disk_backend: LocalDiskBackend) -> None:
|
|
"""batched_get_blocking returns [None, …] when no keys are cached."""
|
|
keys = [create_test_key(i) for i in range(200, 204)]
|
|
results = local_disk_backend.batched_get_blocking(keys)
|
|
|
|
assert len(results) == len(keys)
|
|
assert all(r is None for r in results)
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_reads_match_written_data(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""Data loaded by batched_get_blocking matches what was written."""
|
|
keys = [create_test_key(i) for i in range(210, 214)]
|
|
expected_data = {}
|
|
for key in keys:
|
|
expected_data[key] = self._write_key(local_disk_backend, key)
|
|
|
|
results = local_disk_backend.batched_get_blocking(keys)
|
|
|
|
assert len(results) == len(keys)
|
|
for key, mem_obj in zip(keys, results, strict=True):
|
|
assert mem_obj is not None, f"Expected data for {key}"
|
|
actual = bytes(mem_obj.byte_array)
|
|
assert actual == expected_data[key]
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_mixed_hit_and_miss(self, local_disk_backend: LocalDiskBackend) -> None:
|
|
"""Handles a mix of cached and missing keys correctly."""
|
|
present_key = create_test_key(220)
|
|
missing_key = create_test_key(221)
|
|
expected = self._write_key(local_disk_backend, present_key)
|
|
|
|
results = local_disk_backend.batched_get_blocking([present_key, missing_key])
|
|
|
|
assert len(results) == 2
|
|
assert results[0] is not None
|
|
assert bytes(results[0].byte_array) == expected
|
|
assert results[1] is None
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_cache_policy_updated_only_for_hits(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""update_on_hit is called only for keys that were successfully loaded."""
|
|
hit_key = create_test_key(230)
|
|
miss_key = create_test_key(231)
|
|
self._write_key(local_disk_backend, hit_key)
|
|
|
|
with patch.object(
|
|
local_disk_backend.cache_policy, "update_on_hit"
|
|
) as mock_update:
|
|
results = local_disk_backend.batched_get_blocking([hit_key, miss_key])
|
|
|
|
assert results[0] is not None
|
|
assert results[1] is None
|
|
mock_update.assert_called_once_with(hit_key, local_disk_backend.dict)
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_single_key_delegates_to_get_blocking(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""A single-key batch falls through to get_blocking (fast path)."""
|
|
key = create_test_key(240)
|
|
sentinel = MagicMock(spec=MemoryObj)
|
|
with patch.object(
|
|
local_disk_backend, "get_blocking", return_value=sentinel
|
|
) as mock_get:
|
|
results = local_disk_backend.batched_get_blocking([key])
|
|
|
|
assert results == [sentinel]
|
|
mock_get.assert_called_once_with(key)
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_concurrent_reads_use_multiple_threads(
|
|
self, local_disk_backend: LocalDiskBackend
|
|
) -> None:
|
|
"""Verify that reads actually fan out across threads."""
|
|
keys = [create_test_key(i) for i in range(250, 254)]
|
|
for key in keys:
|
|
self._write_key(local_disk_backend, key)
|
|
|
|
thread_ids: list[int] = []
|
|
lock = threading.Lock()
|
|
original_read_file = local_disk_backend.read_file
|
|
|
|
def tracking_read_file(key, buffer, path):
|
|
with lock:
|
|
thread_ids.append(threading.current_thread().ident)
|
|
return original_read_file(key, buffer, path)
|
|
|
|
with patch.object(
|
|
local_disk_backend, "read_file", side_effect=tracking_read_file
|
|
):
|
|
results = local_disk_backend.batched_get_blocking(keys)
|
|
|
|
assert all(r is not None for r in results)
|
|
# With 4 keys and 4 threads, we should see more than 1 unique thread.
|
|
assert len(set(thread_ids)) > 1, (
|
|
f"Expected multiple threads, got {set(thread_ids)}"
|
|
)
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|
|
|
|
def test_empty_keys_list(self, local_disk_backend: LocalDiskBackend) -> None:
|
|
"""An empty key list returns an empty result list."""
|
|
results = local_disk_backend.batched_get_blocking([])
|
|
assert results == []
|
|
local_disk_backend.local_cpu_backend.memory_allocator.close()
|