wehub-resource-sync
922 lines
31 KiB
Python
922 lines
31 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for the DAX MP L2 adapter.
|
|
"""
|
|
|
|
# Standard
|
|
from typing import cast
|
|
import select
|
|
import threading
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.native_storage_ops import Bitmap
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.config import (
|
|
EvictionConfig,
|
|
L1ManagerConfig,
|
|
L1MemoryManagerConfig,
|
|
StorageManagerConfig,
|
|
)
|
|
from lmcache.v1.distributed.error import L1Error
|
|
from lmcache.v1.distributed.internal_api import L2AdapterListener
|
|
from lmcache.v1.distributed.l2_adapters.base import L2AdapterInterface
|
|
from lmcache.v1.distributed.l2_adapters.config import (
|
|
L2AdaptersConfig,
|
|
get_registered_l2_adapter_types,
|
|
)
|
|
from lmcache.v1.distributed.l2_adapters.dax_l2_adapter import (
|
|
DaxDeviceConfig,
|
|
DaxL2Adapter,
|
|
DaxL2AdapterConfig,
|
|
)
|
|
from lmcache.v1.distributed.l2_adapters.reconfiguration import (
|
|
L2ReconfigurableAdapter,
|
|
L2ReconfigureError,
|
|
)
|
|
from lmcache.v1.distributed.storage_manager import StorageManager
|
|
from lmcache.v1.memory_allocators.ad_hoc_memory_allocator import AdHocMemoryAllocator
|
|
from lmcache.v1.memory_management import (
|
|
MemoryFormat,
|
|
MemoryObj,
|
|
)
|
|
from lmcache.v1.platform import consume_fd
|
|
|
|
_EMPTY_LAYOUT = MemoryLayoutDesc(shapes=[], dtypes=[])
|
|
|
|
_DEFAULT_SHAPE = torch.Size([2, 4, 8])
|
|
|
|
|
|
class _RecordingListener(L2AdapterListener):
|
|
def __init__(self):
|
|
self.stored: list[list[ObjectKey]] = []
|
|
self.accessed: list[list[ObjectKey]] = []
|
|
self.deleted: list[list[ObjectKey]] = []
|
|
|
|
def on_l2_keys_stored(self, keys: list[ObjectKey], sizes: list[int]):
|
|
self.stored.append(list(keys))
|
|
|
|
def on_l2_keys_accessed(self, keys: list[ObjectKey]):
|
|
self.accessed.append(list(keys))
|
|
|
|
def on_l2_keys_deleted(self, keys: list[ObjectKey]):
|
|
self.deleted.append(list(keys))
|
|
|
|
|
|
def create_object_key(
|
|
chunk_id: int,
|
|
model_name: str = "test_model",
|
|
cache_salt: str = "",
|
|
) -> ObjectKey:
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
|
|
model_name=model_name,
|
|
kv_rank=0,
|
|
cache_salt=cache_salt,
|
|
)
|
|
|
|
|
|
def create_memory_obj(
|
|
*,
|
|
shape: torch.Size = _DEFAULT_SHAPE,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
fill_value: float = 0,
|
|
fmt: MemoryFormat = MemoryFormat.KV_2LTD,
|
|
) -> MemoryObj:
|
|
allocator = AdHocMemoryAllocator(device="cpu")
|
|
obj = allocator.allocate([shape], [dtype], fmt=fmt)
|
|
assert obj is not None
|
|
assert obj.tensor is not None
|
|
obj.tensor.fill_(fill_value)
|
|
return obj
|
|
|
|
|
|
def wait_for_event_fd(event_fd: int, timeout: float = 5.0) -> bool:
|
|
poll = select.poll()
|
|
poll.register(event_fd, select.POLLIN)
|
|
events = poll.poll(timeout * 1000)
|
|
if not events:
|
|
return False
|
|
consume_fd(event_fd)
|
|
return True
|
|
|
|
|
|
def wait_for_condition(predicate, timeout: float = 5.0, interval: float = 0.02) -> bool:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(interval)
|
|
return False
|
|
|
|
|
|
def make_adapter(
|
|
tmp_path,
|
|
*,
|
|
slot_bytes: int = 4096,
|
|
max_slots: int = 2,
|
|
num_store_workers: int = 1,
|
|
num_lookup_workers: int = 1,
|
|
num_load_workers: int = 1,
|
|
) -> tuple[DaxL2Adapter, DaxL2AdapterConfig]:
|
|
device_path = tmp_path / "dax.bin"
|
|
arena_bytes = slot_bytes * max_slots
|
|
with open(device_path, "wb") as fout:
|
|
fout.truncate(arena_bytes)
|
|
|
|
config = DaxL2AdapterConfig(
|
|
devices=[
|
|
DaxDeviceConfig(
|
|
device_path=str(device_path),
|
|
max_dax_size_gb=arena_bytes / (1024**3),
|
|
)
|
|
],
|
|
slot_bytes=slot_bytes,
|
|
num_store_workers=num_store_workers,
|
|
num_lookup_workers=num_lookup_workers,
|
|
num_load_workers=num_load_workers,
|
|
)
|
|
return DaxL2Adapter(config), config
|
|
|
|
|
|
def make_hotplug_adapter(tmp_path, *, slot_bytes: int = 2048) -> DaxL2Adapter:
|
|
devices = []
|
|
for i in range(2):
|
|
device_path = tmp_path / f"dax_hotplug_{i}.bin"
|
|
with open(device_path, "wb") as fout:
|
|
fout.truncate(slot_bytes * 2)
|
|
devices.append(
|
|
DaxDeviceConfig(
|
|
device_path=str(device_path),
|
|
max_dax_size_gb=(slot_bytes * 2) / (1024**3),
|
|
)
|
|
)
|
|
|
|
return DaxL2Adapter(
|
|
DaxL2AdapterConfig(
|
|
devices=devices,
|
|
hotplug_enabled=True,
|
|
slot_bytes=slot_bytes,
|
|
num_store_workers=1,
|
|
num_lookup_workers=1,
|
|
num_load_workers=1,
|
|
)
|
|
)
|
|
|
|
|
|
def store_and_wait(adapter: DaxL2Adapter, key: ObjectKey, obj: MemoryObj) -> None:
|
|
task_id = adapter.submit_store_task([key], [obj])
|
|
assert wait_for_event_fd(adapter.get_store_event_fd())
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert completed[task_id].is_successful()
|
|
|
|
|
|
def bitmap_to_bools(bitmap: Bitmap, size: int) -> list[bool]:
|
|
return [bitmap.test(i) for i in range(size)]
|
|
|
|
|
|
def lookup_and_wait(adapter: DaxL2Adapter, keys: list[ObjectKey]) -> list[bool]:
|
|
task_id = adapter.submit_lookup_and_lock_task(keys, _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
|
|
bitmap = adapter.query_lookup_and_lock_result(task_id)
|
|
assert bitmap is not None
|
|
return bitmap_to_bools(bitmap, len(keys))
|
|
|
|
|
|
def load_and_wait(
|
|
adapter: DaxL2Adapter,
|
|
keys: list[ObjectKey],
|
|
objs: list[MemoryObj],
|
|
) -> list[bool]:
|
|
task_id = adapter.submit_load_task(keys, objs)
|
|
assert wait_for_event_fd(adapter.get_load_event_fd())
|
|
bitmap = adapter.query_load_result(task_id)
|
|
assert bitmap is not None
|
|
return bitmap_to_bools(bitmap, len(keys))
|
|
|
|
|
|
def test_dax_adapter_registers_and_has_distinct_eventfds(tmp_path):
|
|
adapter, _ = make_adapter(tmp_path)
|
|
try:
|
|
assert "dax" in get_registered_l2_adapter_types()
|
|
store_fd = adapter.get_store_event_fd()
|
|
lookup_fd = adapter.get_lookup_and_lock_event_fd()
|
|
load_fd = adapter.get_load_event_fd()
|
|
assert len({store_fd, lookup_fd, load_fd}) == 3
|
|
finally:
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_hotplug_remove_migrate_preserves_loadability(tmp_path):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
obj = create_memory_obj(fill_value=9)
|
|
target = create_memory_obj(fill_value=0)
|
|
try:
|
|
key = create_object_key(80)
|
|
store_and_wait(adapter, key, obj)
|
|
source_path = adapter.hotplug_status()["devices"][0]["device_path"]
|
|
|
|
result = adapter.reconfigure(
|
|
"remove",
|
|
{
|
|
"device_path": source_path,
|
|
"mode": "migrate",
|
|
},
|
|
)
|
|
|
|
assert result["state"] == "removed"
|
|
assert result["moved_keys"] == 1
|
|
assert result["deleted_keys"] == 0
|
|
assert result["source_slots_freed"] == 1
|
|
assert listener.deleted == []
|
|
assert lookup_and_wait(adapter, [key]) == [True]
|
|
assert load_and_wait(adapter, [key], [target]) == [True]
|
|
assert target.tensor is not None
|
|
assert torch.all(target.tensor == 9)
|
|
adapter.submit_unlock([key])
|
|
finally:
|
|
obj.ref_count_down()
|
|
target.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_implements_generic_reconfigure_status(tmp_path):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
try:
|
|
assert isinstance(adapter, L2ReconfigurableAdapter)
|
|
status = adapter.reconfigure("status", {})
|
|
assert status == {
|
|
"backend": "dax",
|
|
"supported_operations": ["status", "add", "remove", "resize"],
|
|
"status": adapter.hotplug_status(),
|
|
}
|
|
finally:
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_hotplug_drain_stops_new_writes_to_source_device(tmp_path):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
obj0 = create_memory_obj(fill_value=1)
|
|
obj1 = create_memory_obj(fill_value=2)
|
|
try:
|
|
store_and_wait(adapter, create_object_key(81), obj0)
|
|
devices_before = adapter.hotplug_status()["devices"]
|
|
source_path = devices_before[0]["device_path"]
|
|
source_live_slots = devices_before[0]["live_slot_count"]
|
|
|
|
result = adapter.hotplug_remove_device(source_path, "drain")
|
|
|
|
assert result["status"] == "ok"
|
|
assert result["state"] == "draining"
|
|
store_and_wait(adapter, create_object_key(82), obj1)
|
|
devices = adapter.hotplug_status()["devices"]
|
|
assert devices[0]["live_slot_count"] == source_live_slots
|
|
assert devices[1]["live_slot_count"] == 1
|
|
finally:
|
|
obj0.ref_count_down()
|
|
obj1.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_hotplug_remove_blocked_restores_active_state(tmp_path):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
obj = create_memory_obj(fill_value=5)
|
|
try:
|
|
key = create_object_key(84)
|
|
store_and_wait(adapter, key, obj)
|
|
source_path = adapter.hotplug_status()["devices"][0]["device_path"]
|
|
assert lookup_and_wait(adapter, [key]) == [True]
|
|
|
|
with pytest.raises(L2ReconfigureError) as exc_info:
|
|
adapter.hotplug_remove_device(source_path, "migrate")
|
|
|
|
assert exc_info.value.status_code == 409
|
|
assert adapter.hotplug_status()["devices"][0]["state"] == "active"
|
|
adapter.submit_unlock([key])
|
|
finally:
|
|
obj.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_hotplug_remove_migrate_rejects_duplicate_destination_key(tmp_path):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
source_obj = create_memory_obj(fill_value=1)
|
|
duplicate_obj = create_memory_obj(fill_value=9)
|
|
target = create_memory_obj(fill_value=0)
|
|
try:
|
|
key = create_object_key(83)
|
|
store_and_wait(adapter, key, source_obj)
|
|
source_path = adapter.hotplug_status()["devices"][0]["device_path"]
|
|
destination = adapter._devices[1]
|
|
assert destination.core.put_many([key], [duplicate_obj]) == [True]
|
|
|
|
with pytest.raises(L2ReconfigureError) as exc_info:
|
|
adapter.hotplug_remove_device(source_path, "migrate")
|
|
|
|
assert exc_info.value.status_code == 409
|
|
assert adapter.hotplug_status()["devices"][0]["state"] == "active"
|
|
assert lookup_and_wait(adapter, [key]) == [True]
|
|
assert load_and_wait(adapter, [key], [target]) == [True]
|
|
assert target.tensor is not None
|
|
assert torch.all(target.tensor == 1)
|
|
adapter.submit_unlock([key])
|
|
finally:
|
|
source_obj.ref_count_down()
|
|
duplicate_obj.ref_count_down()
|
|
target.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_hotplug_remove_evict_notifies_logical_delete(tmp_path):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
obj = create_memory_obj(fill_value=1)
|
|
try:
|
|
key = create_object_key(85)
|
|
store_and_wait(adapter, key, obj)
|
|
source_path = adapter.hotplug_status()["devices"][0]["device_path"]
|
|
|
|
result = adapter.hotplug_remove_device(source_path, "evict")
|
|
|
|
assert result["deleted_keys"] == 1
|
|
assert result["source_slots_freed"] == 1
|
|
assert listener.deleted == [[key]]
|
|
finally:
|
|
obj.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_hotplug_add_sanitizes_mapping_errors(tmp_path):
|
|
adapter = DaxL2Adapter(
|
|
DaxL2AdapterConfig(
|
|
devices=[],
|
|
hotplug_enabled=True,
|
|
slot_bytes=2048,
|
|
num_store_workers=1,
|
|
num_lookup_workers=1,
|
|
num_load_workers=1,
|
|
)
|
|
)
|
|
missing_path = str(tmp_path / "missing_dax.bin")
|
|
try:
|
|
with pytest.raises(L2ReconfigureError) as exc_info:
|
|
adapter.hotplug_add_device(missing_path, 2048)
|
|
|
|
assert exc_info.value.status_code == 400
|
|
assert exc_info.value.payload == {"error": "failed to map DAX device"}
|
|
assert missing_path not in str(exc_info.value.payload)
|
|
finally:
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_lookup_batches_unmapped_keys_by_device(tmp_path, monkeypatch):
|
|
adapter = make_hotplug_adapter(tmp_path)
|
|
objs = [create_memory_obj(fill_value=i) for i in range(4)]
|
|
keys = [create_object_key(500 + i) for i in range(4)]
|
|
calls: list[tuple[int, int, bool]] = []
|
|
|
|
try:
|
|
task_id = adapter.submit_store_task(keys, objs)
|
|
assert wait_for_event_fd(adapter.get_store_event_fd())
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert completed[task_id].is_successful()
|
|
|
|
with adapter._device_lock:
|
|
adapter._key_to_device.clear()
|
|
for entry in adapter._devices:
|
|
original = entry.core.exists_many
|
|
|
|
def wrapped_exists_many(
|
|
lookup_keys,
|
|
lock=False,
|
|
*,
|
|
device_id=entry.device_id,
|
|
original=original,
|
|
):
|
|
calls.append((device_id, len(lookup_keys), lock))
|
|
return original(lookup_keys, lock=lock)
|
|
|
|
monkeypatch.setattr(entry.core, "exists_many", wrapped_exists_many)
|
|
|
|
assert lookup_and_wait(adapter, keys) == [True, True, True, True]
|
|
|
|
assert [count for _, count, _ in calls] == [4, 2]
|
|
assert all(lock for _, _, lock in calls)
|
|
finally:
|
|
adapter.submit_unlock(keys)
|
|
adapter.close()
|
|
|
|
|
|
class _FakeReconfigurableAdapter:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, dict[str, object]]] = []
|
|
|
|
def reconfigure_status(self) -> dict:
|
|
return {
|
|
"backend": "fake",
|
|
"supported_operations": ["flip"],
|
|
"status": {"ready": True},
|
|
}
|
|
|
|
def reconfigure(self, operation: str, payload: dict[str, object]) -> dict:
|
|
self.calls.append((operation, payload))
|
|
return {"status": "ok", "operation": operation, "payload": payload}
|
|
|
|
|
|
class _SerdeLikeWrapper:
|
|
def __init__(self, inner_adapter: _FakeReconfigurableAdapter) -> None:
|
|
self.inner_adapter = inner_adapter
|
|
|
|
|
|
class _FakeAdapterDescriptor:
|
|
def __init__(self, type_name: str) -> None:
|
|
self.type_name = type_name
|
|
|
|
|
|
def test_storage_manager_routes_generic_l2_reconfigure_to_adapter():
|
|
sm = StorageManager.__new__(StorageManager)
|
|
adapter = _FakeReconfigurableAdapter()
|
|
sm._adapters_lock = threading.Lock()
|
|
sm._l2_adapters = {0: cast(L2AdapterInterface, adapter)}
|
|
|
|
result = sm.reconfigure_l2_adapter(0, "flip", {"enabled": True})
|
|
|
|
assert result == {
|
|
"status": "ok",
|
|
"operation": "flip",
|
|
"payload": {"enabled": True},
|
|
"adapter_index": 0,
|
|
}
|
|
assert adapter.calls == [("flip", {"enabled": True})]
|
|
|
|
|
|
def test_storage_manager_finds_serde_wrapped_reconfigurable_adapter():
|
|
sm = StorageManager.__new__(StorageManager)
|
|
sm._adapters_lock = threading.Lock()
|
|
sm._l2_adapters = {
|
|
0: cast(L2AdapterInterface, _SerdeLikeWrapper(_FakeReconfigurableAdapter()))
|
|
}
|
|
sm._adapter_descriptors = {0: _FakeAdapterDescriptor("configured_fake")}
|
|
|
|
status = sm.get_l2_adapter_reconfigure_status()
|
|
|
|
assert status["enabled"] is True
|
|
assert status["num_adapters"] == 1
|
|
assert status["adapters"][0]["backend"] == "configured_fake"
|
|
assert status["adapters"][0]["adapter_index"] == 0
|
|
assert status["adapters"][0]["l2_adapter_index"] == 0
|
|
|
|
|
|
def test_dax_adapter_store_lookup_load_and_one_shot_results(tmp_path):
|
|
adapter, _ = make_adapter(tmp_path, slot_bytes=2048, max_slots=4)
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
obj0 = create_memory_obj(fill_value=3)
|
|
obj2 = create_memory_obj(fill_value=7)
|
|
miss_target = create_memory_obj(fill_value=0)
|
|
load_target0 = create_memory_obj(fill_value=0)
|
|
load_target2 = create_memory_obj(fill_value=0)
|
|
try:
|
|
key0 = create_object_key(10)
|
|
key1 = create_object_key(11)
|
|
key2 = create_object_key(12)
|
|
|
|
store_task = adapter.submit_store_task([key0, key2], [obj0, obj2])
|
|
assert wait_for_event_fd(adapter.get_store_event_fd())
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert completed[store_task].is_successful()
|
|
assert adapter.pop_completed_store_tasks() == {}
|
|
assert listener.stored == [[key0, key2]]
|
|
|
|
lookup_task = adapter.submit_lookup_and_lock_task(
|
|
[key0, key1, key2], _EMPTY_LAYOUT
|
|
)
|
|
assert wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
|
|
lookup_bitmap = adapter.query_lookup_and_lock_result(lookup_task)
|
|
assert lookup_bitmap is not None
|
|
assert bitmap_to_bools(lookup_bitmap, 3) == [True, False, True]
|
|
assert adapter.query_lookup_and_lock_result(lookup_task) is None
|
|
|
|
load_task = adapter.submit_load_task(
|
|
[key0, key1, key2],
|
|
[load_target0, miss_target, load_target2],
|
|
)
|
|
assert wait_for_event_fd(adapter.get_load_event_fd())
|
|
load_bitmap = adapter.query_load_result(load_task)
|
|
assert load_bitmap is not None
|
|
assert bitmap_to_bools(load_bitmap, 3) == [True, False, True]
|
|
assert adapter.query_load_result(load_task) is None
|
|
|
|
assert load_target0.tensor is not None
|
|
assert load_target2.tensor is not None
|
|
assert miss_target.tensor is not None
|
|
assert torch.all(load_target0.tensor == 3)
|
|
assert torch.all(load_target2.tensor == 7)
|
|
assert torch.all(miss_target.tensor == 0)
|
|
assert listener.accessed == [[key0, key2]]
|
|
finally:
|
|
obj0.ref_count_down()
|
|
obj2.ref_count_down()
|
|
miss_target.ref_count_down()
|
|
load_target0.ref_count_down()
|
|
load_target2.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_unlock_refcount_and_delete_skips_locked_keys(tmp_path):
|
|
adapter, _ = make_adapter(tmp_path, slot_bytes=2048, max_slots=2)
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
obj = create_memory_obj(fill_value=5)
|
|
try:
|
|
key = create_object_key(20)
|
|
store_and_wait(adapter, key, obj)
|
|
|
|
first_lookup = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
|
|
assert adapter.query_lookup_and_lock_result(first_lookup) is not None
|
|
|
|
second_lookup = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
|
|
assert adapter.query_lookup_and_lock_result(second_lookup) is not None
|
|
|
|
adapter.delete([key])
|
|
assert listener.deleted == []
|
|
assert adapter.report_status()["locked_key_count"] == 1
|
|
|
|
adapter.submit_unlock([key])
|
|
adapter.delete([key])
|
|
assert listener.deleted == []
|
|
|
|
adapter.submit_unlock([key])
|
|
adapter.delete([key])
|
|
assert listener.deleted == [[key]]
|
|
finally:
|
|
obj.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_usage_and_status_track_pending_free_slots(tmp_path):
|
|
adapter, _ = make_adapter(tmp_path, slot_bytes=2048, max_slots=2)
|
|
obj = create_memory_obj(fill_value=9)
|
|
try:
|
|
assert adapter.supports_global_eviction is True
|
|
key = create_object_key(30)
|
|
store_and_wait(adapter, key, obj)
|
|
|
|
reservations, hits = adapter._core.reserve_reads([key], prefix_only=False)
|
|
assert hits == [True]
|
|
|
|
adapter.delete([key])
|
|
usage = adapter.get_usage()
|
|
assert usage.usage_fraction == pytest.approx(0.5)
|
|
assert usage.total_bytes_used == 2048
|
|
|
|
status = adapter.report_status()
|
|
assert status["live_slot_count"] == 1
|
|
assert status["borrowed_slot_count"] == 1
|
|
assert status["locked_key_count"] == 0
|
|
assert status["inflight_store_tasks"] == 0
|
|
assert status["inflight_lookup_tasks"] == 0
|
|
assert status["inflight_load_tasks"] == 0
|
|
_, usage_after_eviction = adapter._core.usage()
|
|
assert usage_after_eviction == pytest.approx(0.0)
|
|
assert status["supports_restart_recovery"] is False
|
|
|
|
adapter._core.finalize_reads(reservations, set())
|
|
assert adapter.get_usage().usage_fraction == pytest.approx(0.0)
|
|
assert adapter._core.usage()[1] == pytest.approx(0.0)
|
|
finally:
|
|
obj.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_full_arena_does_not_evict_internally(tmp_path):
|
|
adapter, _ = make_adapter(tmp_path, slot_bytes=2048, max_slots=2)
|
|
listener = _RecordingListener()
|
|
adapter.register_listener(listener)
|
|
|
|
obj0 = create_memory_obj(fill_value=1)
|
|
obj1 = create_memory_obj(fill_value=2)
|
|
obj2 = create_memory_obj(fill_value=3)
|
|
try:
|
|
key0 = create_object_key(31)
|
|
key1 = create_object_key(32)
|
|
key2 = create_object_key(33)
|
|
|
|
store_and_wait(adapter, key0, obj0)
|
|
store_and_wait(adapter, key1, obj1)
|
|
|
|
task_id = adapter.submit_store_task([key2], [obj2])
|
|
assert wait_for_event_fd(adapter.get_store_event_fd())
|
|
completed = adapter.pop_completed_store_tasks()
|
|
assert not completed[task_id].is_successful()
|
|
assert listener.stored == [[key0], [key1]]
|
|
assert adapter.get_usage().usage_fraction == pytest.approx(1.0)
|
|
|
|
lookup_task = adapter.submit_lookup_and_lock_task(
|
|
[key0, key1, key2], _EMPTY_LAYOUT
|
|
)
|
|
assert wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
|
|
bitmap = adapter.query_lookup_and_lock_result(lookup_task)
|
|
assert bitmap is not None
|
|
assert bitmap_to_bools(bitmap, 3) == [True, True, False]
|
|
adapter.submit_unlock([key0, key1])
|
|
finally:
|
|
obj0.ref_count_down()
|
|
obj1.ref_count_down()
|
|
obj2.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_usage_tracks_cache_salt_by_slot(tmp_path):
|
|
adapter, _ = make_adapter(tmp_path, slot_bytes=2048, max_slots=4)
|
|
|
|
obj0 = create_memory_obj(fill_value=1)
|
|
obj1 = create_memory_obj(fill_value=2)
|
|
obj2 = create_memory_obj(fill_value=3)
|
|
try:
|
|
alice0 = create_object_key(34, cache_salt="alice")
|
|
bob = create_object_key(35, cache_salt="bob")
|
|
alice1 = create_object_key(36, cache_salt="alice")
|
|
|
|
store_and_wait(adapter, alice0, obj0)
|
|
store_and_wait(adapter, bob, obj1)
|
|
store_and_wait(adapter, alice1, obj2)
|
|
|
|
usage = adapter.get_usage()
|
|
assert usage.total_bytes_used == 2048 * 3
|
|
assert dict(usage.bytes_by_cache_salt) == {
|
|
"alice": 2048 * 2,
|
|
"bob": 2048,
|
|
}
|
|
|
|
adapter.delete([alice0])
|
|
usage = adapter.get_usage()
|
|
assert usage.total_bytes_used == 2048 * 2
|
|
assert dict(usage.bytes_by_cache_salt) == {
|
|
"alice": 2048,
|
|
"bob": 2048,
|
|
}
|
|
|
|
adapter.delete([alice1])
|
|
usage = adapter.get_usage()
|
|
assert usage.total_bytes_used == 2048
|
|
assert dict(usage.bytes_by_cache_salt) == {"bob": 2048}
|
|
finally:
|
|
obj0.ref_count_down()
|
|
obj1.ref_count_down()
|
|
obj2.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_close_waits_for_inflight_tasks(tmp_path, monkeypatch):
|
|
adapter, _ = make_adapter(tmp_path, slot_bytes=2048, max_slots=2)
|
|
obj = create_memory_obj(fill_value=2)
|
|
target = create_memory_obj(fill_value=0)
|
|
try:
|
|
key = create_object_key(40)
|
|
store_and_wait(adapter, key, obj)
|
|
|
|
load_started = threading.Event()
|
|
allow_load = threading.Event()
|
|
close_returned = threading.Event()
|
|
original_load = adapter._core.load_many_into
|
|
|
|
def _blocking_load_many_into(keys, objs):
|
|
load_started.set()
|
|
assert allow_load.wait(timeout=2)
|
|
return original_load(keys, objs)
|
|
|
|
monkeypatch.setattr(adapter._core, "load_many_into", _blocking_load_many_into)
|
|
|
|
adapter.submit_load_task([key], [target])
|
|
assert load_started.wait(timeout=2)
|
|
|
|
closer = threading.Thread(
|
|
target=lambda: (adapter.close(), close_returned.set())
|
|
)
|
|
closer.start()
|
|
time.sleep(0.05)
|
|
assert not close_returned.is_set()
|
|
|
|
allow_load.set()
|
|
closer.join(timeout=2)
|
|
assert not closer.is_alive()
|
|
assert close_returned.is_set()
|
|
finally:
|
|
obj.ref_count_down()
|
|
target.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_dax_adapter_restart_is_volatile_only(tmp_path):
|
|
adapter, config = make_adapter(tmp_path, slot_bytes=2048, max_slots=2)
|
|
obj = create_memory_obj(fill_value=4)
|
|
try:
|
|
key = create_object_key(50)
|
|
store_and_wait(adapter, key, obj)
|
|
adapter.close()
|
|
|
|
reopened = DaxL2Adapter(config)
|
|
try:
|
|
lookup_task = reopened.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
|
|
assert wait_for_event_fd(reopened.get_lookup_and_lock_event_fd())
|
|
bitmap = reopened.query_lookup_and_lock_result(lookup_task)
|
|
assert bitmap is not None
|
|
assert bitmap_to_bools(bitmap, 1) == [False]
|
|
finally:
|
|
reopened.close()
|
|
finally:
|
|
obj.ref_count_down()
|
|
adapter.close()
|
|
|
|
|
|
def test_storage_manager_dax_adapter_roundtrip(tmp_path):
|
|
shape = torch.Size([2, 4, 8])
|
|
dtype = torch.bfloat16
|
|
slot_bytes = shape.numel() * dtype.itemsize
|
|
|
|
device_path = tmp_path / "sm_dax.bin"
|
|
with open(device_path, "wb") as fout:
|
|
fout.truncate(slot_bytes * 4)
|
|
|
|
storage_config = StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=1 << 20,
|
|
use_lazy=False,
|
|
init_size_in_bytes=1 << 20,
|
|
align_bytes=0x1000,
|
|
),
|
|
write_ttl_seconds=60,
|
|
read_ttl_seconds=60,
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="noop"),
|
|
l2_adapter_config=L2AdaptersConfig(
|
|
[
|
|
DaxL2AdapterConfig(
|
|
devices=[
|
|
DaxDeviceConfig(
|
|
device_path=str(device_path),
|
|
max_dax_size_gb=(slot_bytes * 4) / (1024**3),
|
|
)
|
|
],
|
|
slot_bytes=slot_bytes,
|
|
)
|
|
]
|
|
),
|
|
store_policy="skip_l1",
|
|
)
|
|
|
|
layout = MemoryLayoutDesc(shapes=[shape], dtypes=[dtype])
|
|
key = create_object_key(60)
|
|
sm = StorageManager(storage_config)
|
|
|
|
try:
|
|
adapter = sm._l2_adapters[0]
|
|
assert isinstance(adapter, DaxL2Adapter)
|
|
|
|
reserved = sm.reserve_write([key], layout, mode="new")
|
|
assert key in reserved
|
|
assert reserved[key].tensor is not None
|
|
reserved[key].tensor.fill_(11)
|
|
sm.finish_write([key])
|
|
|
|
assert wait_for_condition(
|
|
lambda: (
|
|
sm.report_status()["l1_manager"]["total_object_count"] == 0
|
|
and adapter.get_usage().usage_fraction > 0
|
|
),
|
|
timeout=5.0,
|
|
)
|
|
|
|
handle = sm.submit_prefetch_task([key], layout)
|
|
assert wait_for_condition(
|
|
lambda: sm.query_prefetch_lookup_hits(handle) is not None,
|
|
timeout=5.0,
|
|
)
|
|
lookup_hits = sm.query_prefetch_lookup_hits(handle)
|
|
assert lookup_hits == 1
|
|
|
|
final_result: dict[str, int | None] = {"value": None}
|
|
|
|
def _capture_prefetch_result() -> bool:
|
|
result = sm.query_prefetch_status(handle)
|
|
if result is None:
|
|
return False
|
|
final_result["value"] = result.count_leading_ones()
|
|
return True
|
|
|
|
assert wait_for_condition(_capture_prefetch_result, timeout=5.0)
|
|
final_hits = final_result["value"]
|
|
assert final_hits == 1
|
|
|
|
with sm.read_prefetched_results([key]) as results:
|
|
assert results is not None
|
|
assert len(results) == 1
|
|
assert results[0].tensor is not None
|
|
assert torch.all(results[0].tensor == 11)
|
|
|
|
assert adapter.report_status()["locked_key_count"] == 0
|
|
assert sm._l1_manager.delete([key])[key] == L1Error.KEY_IS_LOCKED
|
|
|
|
sm.finish_read_prefetched([key])
|
|
assert sm._l1_manager.delete([key])[key] == L1Error.KEY_NOT_EXIST
|
|
finally:
|
|
sm.close()
|
|
|
|
|
|
def test_storage_manager_dax_adapter_uses_global_l2_eviction(tmp_path):
|
|
shape = torch.Size([2, 4, 8])
|
|
dtype = torch.bfloat16
|
|
slot_bytes = shape.numel() * dtype.itemsize
|
|
|
|
device_path = tmp_path / "sm_dax_eviction.bin"
|
|
with open(device_path, "wb") as fout:
|
|
fout.truncate(slot_bytes * 2)
|
|
|
|
dax_config = DaxL2AdapterConfig(
|
|
devices=[
|
|
DaxDeviceConfig(
|
|
device_path=str(device_path),
|
|
max_dax_size_gb=(slot_bytes * 2) / (1024**3),
|
|
)
|
|
],
|
|
slot_bytes=slot_bytes,
|
|
)
|
|
dax_config.eviction_config = EvictionConfig(
|
|
eviction_policy="LRU",
|
|
trigger_watermark=0.5,
|
|
eviction_ratio=0.5,
|
|
)
|
|
|
|
storage_config = StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=1 << 20,
|
|
use_lazy=False,
|
|
init_size_in_bytes=1 << 20,
|
|
align_bytes=0x1000,
|
|
),
|
|
write_ttl_seconds=60,
|
|
read_ttl_seconds=60,
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="noop"),
|
|
l2_adapter_config=L2AdaptersConfig([dax_config]),
|
|
store_policy="skip_l1",
|
|
)
|
|
|
|
layout = MemoryLayoutDesc(shapes=[shape], dtypes=[dtype])
|
|
sm = StorageManager(storage_config)
|
|
|
|
try:
|
|
sm._l2_eviction_controller.stop()
|
|
adapter = sm._l2_adapters[0]
|
|
assert isinstance(adapter, DaxL2Adapter)
|
|
assert adapter.supports_global_eviction is True
|
|
assert len(sm._l2_eviction_controller._adapter_states) == 1
|
|
|
|
key0 = create_object_key(70)
|
|
key1 = create_object_key(71)
|
|
key2 = create_object_key(72)
|
|
|
|
def _write_key(key: ObjectKey, fill_value: int, usage_fraction: float) -> None:
|
|
reserved = sm.reserve_write([key], layout, mode="new")
|
|
assert key in reserved
|
|
assert reserved[key].tensor is not None
|
|
reserved[key].tensor.fill_(fill_value)
|
|
sm.finish_write([key])
|
|
assert wait_for_condition(
|
|
lambda: (
|
|
adapter.get_usage().usage_fraction == pytest.approx(usage_fraction)
|
|
),
|
|
timeout=5.0,
|
|
)
|
|
|
|
_write_key(key0, 1, 0.5)
|
|
_write_key(key1, 2, 1.0)
|
|
assert adapter.get_usage().usage_fraction == pytest.approx(1.0)
|
|
|
|
eviction_state = sm._l2_eviction_controller._adapter_states[0]
|
|
sm._l2_eviction_controller._check_and_evict(eviction_state)
|
|
assert adapter.get_usage().usage_fraction == pytest.approx(0.5)
|
|
|
|
hits = adapter._core.exists_many([key0, key1], lock=True)
|
|
assert hits == [False, True]
|
|
adapter.submit_unlock([key1])
|
|
|
|
_write_key(key2, 3, 1.0)
|
|
assert adapter.get_usage().usage_fraction == pytest.approx(1.0)
|
|
finally:
|
|
sm.close()
|