wehub-resource-sync
557 lines
18 KiB
Python
557 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unit tests for L1MemoryManager.
|
|
|
|
These tests verify the behavior of L1MemoryManager as described in the
|
|
interface docstrings. The tests focus on:
|
|
|
|
1. allocate() - Thread-safe allocation returning (error_code, memory_objs)
|
|
- Returns SUCCESS and non-empty list on successful allocation
|
|
- Returns OUT_OF_MEMORY and empty list when allocation fails
|
|
|
|
2. free() - Thread-safe deallocation returning error_code
|
|
- Returns SUCCESS when operation succeeds
|
|
|
|
3. get_l1_memory_desc() - Returns L1MemoryDesc with pointer, size, and
|
|
alignment of the L1 buffer
|
|
"""
|
|
|
|
# Standard
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import threading
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc
|
|
from lmcache.v1.distributed.config import L1MemoryManagerConfig
|
|
from lmcache.v1.distributed.error import L1Error
|
|
|
|
try:
|
|
# First Party
|
|
from lmcache.v1.distributed.internal_api import L1MemoryDesc
|
|
from lmcache.v1.distributed.memory_manager import (
|
|
L1MemoryManager,
|
|
)
|
|
except ImportError:
|
|
# Skip the tests if the L1MemoryManager cannot be imported
|
|
pytest.skip("L1MemoryManager could not be imported", allow_module_level=True)
|
|
|
|
# Skip all tests in this module if CUDA is not available
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason="CUDA is not available"
|
|
)
|
|
|
|
|
|
def should_use_lazy_alloc() -> bool:
|
|
"""Determine if lazy allocation should be used based on CUDA availability."""
|
|
return torch.cuda.is_available()
|
|
|
|
|
|
# =============================================================================
|
|
# Fixtures
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.fixture
|
|
def basic_config():
|
|
"""Create a basic L1MemoryManagerConfig for testing."""
|
|
return L1MemoryManagerConfig(
|
|
size_in_bytes=128 * 1024 * 1024, # 128MB
|
|
use_lazy=should_use_lazy_alloc(),
|
|
init_size_in_bytes=64 * 1024 * 1024, # 64MB
|
|
align_bytes=0x1000, # 4KB
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def small_config():
|
|
"""Create a small L1MemoryManagerConfig to test memory exhaustion.
|
|
|
|
Note: Minimum size is 64MB due to LazyMemoryAllocator's PIN_CHUNK_SIZE.
|
|
"""
|
|
return L1MemoryManagerConfig(
|
|
size_in_bytes=64 * 1024 * 1024, # 64MB
|
|
use_lazy=should_use_lazy_alloc(),
|
|
init_size_in_bytes=64
|
|
* 1024
|
|
* 1024, # 64MB (same as final to prevent expansion)
|
|
align_bytes=0x1000,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def non_lazy_config():
|
|
"""Create a config that does not use lazy allocation."""
|
|
return L1MemoryManagerConfig(
|
|
size_in_bytes=128 * 1024 * 1024, # 128MB
|
|
use_lazy=False,
|
|
init_size_in_bytes=64 * 1024 * 1024,
|
|
align_bytes=0x1000,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def basic_layout():
|
|
"""Create a basic MemoryLayoutDesc for testing."""
|
|
return MemoryLayoutDesc(
|
|
shapes=[torch.Size([100, 2, 512])],
|
|
dtypes=[torch.bfloat16],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def multi_tensor_layout():
|
|
"""Create a MemoryLayoutDesc with multiple tensor shapes."""
|
|
return MemoryLayoutDesc(
|
|
shapes=[torch.Size([100, 2, 512]), torch.Size([100, 2, 512])],
|
|
dtypes=[torch.bfloat16, torch.bfloat16],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def large_layout():
|
|
"""Create a large MemoryLayoutDesc that will exhaust small memory.
|
|
|
|
Each allocation is 8MB (2M elements * 4 bytes).
|
|
"""
|
|
return MemoryLayoutDesc(
|
|
shapes=[torch.Size([2048, 1024])], # 2M elements * 4 bytes = 8MB
|
|
dtypes=[torch.float32],
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for L1MemoryManager.allocate()
|
|
# =============================================================================
|
|
|
|
|
|
class TestAllocate:
|
|
"""
|
|
Tests for L1MemoryManager.allocate() method.
|
|
|
|
Per the docstring:
|
|
- This function should be thread-safe
|
|
- Returns tuple[L1Error, list[MemoryObj]]
|
|
- Error code is OUT_OF_MEMORY if allocation fails, otherwise SUCCESS
|
|
- If allocation fails, the memory object list will be empty
|
|
"""
|
|
|
|
def test_allocate_returns_success_and_memory_objs(self, basic_config, basic_layout):
|
|
"""Test that allocate returns SUCCESS and valid memory objects."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
error, mem_objs = manager.allocate(basic_layout, count=1)
|
|
|
|
assert error == L1Error.SUCCESS
|
|
assert isinstance(mem_objs, list)
|
|
assert len(mem_objs) == 1
|
|
for obj in mem_objs:
|
|
assert obj is not None
|
|
assert obj.is_valid()
|
|
|
|
manager.close()
|
|
|
|
def test_allocate_returns_correct_count(self, basic_config, basic_layout):
|
|
"""Test that allocate returns the requested number of memory objects."""
|
|
manager = L1MemoryManager(basic_config)
|
|
count = 5
|
|
|
|
error, mem_objs = manager.allocate(basic_layout, count=count)
|
|
|
|
assert error == L1Error.SUCCESS
|
|
assert len(mem_objs) == count
|
|
|
|
manager.close()
|
|
|
|
def test_allocate_with_multi_tensor_layout(self, basic_config, multi_tensor_layout):
|
|
"""Test allocation with multiple tensor shapes in the layout."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
error, mem_objs = manager.allocate(multi_tensor_layout, count=2)
|
|
|
|
assert error == L1Error.SUCCESS
|
|
assert len(mem_objs) == 2
|
|
|
|
manager.close()
|
|
|
|
def test_allocate_returns_out_of_memory_when_exhausted(
|
|
self, small_config, large_layout
|
|
):
|
|
"""
|
|
Test that allocate returns OUT_OF_MEMORY and empty list when memory
|
|
is exhausted.
|
|
"""
|
|
manager = L1MemoryManager(small_config)
|
|
|
|
# Request more memory than available (8MB * 10 = 80MB > 64MB)
|
|
error, mem_objs = manager.allocate(large_layout, count=10)
|
|
|
|
assert error == L1Error.OUT_OF_MEMORY
|
|
assert isinstance(mem_objs, list)
|
|
assert len(mem_objs) == 0
|
|
|
|
manager.close()
|
|
|
|
def test_allocate_returns_empty_list_on_failure(self, small_config, large_layout):
|
|
"""Test that the memory object list is empty when allocation fails."""
|
|
manager = L1MemoryManager(small_config)
|
|
|
|
# Request more memory than available (8MB * 10 = 80MB > 64MB)
|
|
error, mem_objs = manager.allocate(large_layout, count=10)
|
|
|
|
# Verify the docstring contract: "If the allocation fails, the memory
|
|
# object list will be empty."
|
|
assert error == L1Error.OUT_OF_MEMORY
|
|
assert mem_objs == []
|
|
|
|
manager.close()
|
|
|
|
def test_allocate_is_thread_safe(self, basic_config, basic_layout):
|
|
"""Test that allocate is thread-safe with concurrent allocations."""
|
|
manager = L1MemoryManager(basic_config)
|
|
num_threads = 10
|
|
allocations_per_thread = 5
|
|
results = []
|
|
errors = []
|
|
|
|
def allocate_task():
|
|
try:
|
|
for _ in range(allocations_per_thread):
|
|
error, mem_objs = manager.allocate(basic_layout, count=1)
|
|
results.append((error, mem_objs))
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=allocate_task) for _ in range(num_threads)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# No exceptions should have been raised
|
|
assert len(errors) == 0, f"Thread-safety errors: {errors}"
|
|
|
|
# All results should have valid error codes
|
|
for error, mem_objs in results:
|
|
assert error in (
|
|
L1Error.SUCCESS,
|
|
L1Error.OUT_OF_MEMORY,
|
|
)
|
|
if error == L1Error.SUCCESS:
|
|
assert len(mem_objs) == 1
|
|
else:
|
|
assert len(mem_objs) == 0
|
|
|
|
manager.close()
|
|
|
|
def test_allocate_with_non_lazy_config(self, non_lazy_config, basic_layout):
|
|
"""Test allocation with non-lazy (MixedMemoryAllocator) configuration."""
|
|
manager = L1MemoryManager(non_lazy_config)
|
|
|
|
error, mem_objs = manager.allocate(basic_layout, count=2)
|
|
|
|
assert error == L1Error.SUCCESS
|
|
assert len(mem_objs) == 2
|
|
|
|
manager.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for L1MemoryManager.free()
|
|
# =============================================================================
|
|
|
|
|
|
class TestFree:
|
|
"""
|
|
Tests for L1MemoryManager.free() method.
|
|
|
|
Per the docstring:
|
|
- This function should be thread-safe
|
|
- Returns L1Error indicating the result
|
|
- Returns SUCCESS if operation succeeds
|
|
"""
|
|
|
|
def test_free_returns_success(self, basic_config, basic_layout):
|
|
"""Test that free returns SUCCESS for valid memory objects."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
error, mem_objs = manager.allocate(basic_layout, count=2)
|
|
assert error == L1Error.SUCCESS
|
|
|
|
free_error = manager.free(mem_objs)
|
|
|
|
assert free_error == L1Error.SUCCESS
|
|
|
|
manager.close()
|
|
|
|
def test_free_invalidates_memory_objects(self, basic_config, basic_layout):
|
|
"""Test that free invalidates the memory objects."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
error, mem_objs = manager.allocate(basic_layout, count=1)
|
|
assert error == L1Error.SUCCESS
|
|
assert mem_objs[0].is_valid()
|
|
|
|
free_error = manager.free(mem_objs)
|
|
assert free_error == L1Error.SUCCESS
|
|
|
|
# After freeing, the memory objects should be invalidated
|
|
for obj in mem_objs:
|
|
assert not obj.is_valid()
|
|
|
|
manager.close()
|
|
|
|
def test_free_empty_list_returns_success(self, basic_config):
|
|
"""Test that freeing an empty list returns SUCCESS."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
free_error = manager.free([])
|
|
|
|
assert free_error == L1Error.SUCCESS
|
|
|
|
manager.close()
|
|
|
|
def test_free_allows_reallocation(self, basic_config, basic_layout):
|
|
"""Test that freed memory can be reallocated."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
# Allocate
|
|
error1, mem_objs1 = manager.allocate(basic_layout, count=5)
|
|
assert error1 == L1Error.SUCCESS
|
|
|
|
# Free
|
|
free_error = manager.free(mem_objs1)
|
|
assert free_error == L1Error.SUCCESS
|
|
|
|
# Reallocate
|
|
error2, mem_objs2 = manager.allocate(basic_layout, count=5)
|
|
assert error2 == L1Error.SUCCESS
|
|
assert len(mem_objs2) == 5
|
|
|
|
manager.close()
|
|
|
|
def test_free_is_thread_safe(self, basic_config, basic_layout):
|
|
"""Test that free is thread-safe with concurrent operations."""
|
|
manager = L1MemoryManager(basic_config)
|
|
num_threads = 10
|
|
errors_list = []
|
|
lock = threading.Lock()
|
|
|
|
def allocate_and_free_task():
|
|
try:
|
|
error, mem_objs = manager.allocate(basic_layout, count=1)
|
|
if error == L1Error.SUCCESS:
|
|
free_error = manager.free(mem_objs)
|
|
with lock:
|
|
errors_list.append(free_error)
|
|
except Exception as e:
|
|
with lock:
|
|
errors_list.append(e)
|
|
|
|
threads = [
|
|
threading.Thread(target=allocate_and_free_task) for _ in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# All free operations should have succeeded
|
|
for err in errors_list:
|
|
if isinstance(err, Exception):
|
|
pytest.fail(f"Thread-safety error: {err}")
|
|
assert err == L1Error.SUCCESS
|
|
|
|
manager.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for L1MemoryManager.get_l1_memory_desc()
|
|
# =============================================================================
|
|
|
|
|
|
class TestGetL1MemoryDesc:
|
|
"""
|
|
Tests for L1MemoryManager.get_l1_memory_desc() method.
|
|
|
|
Per the docstring:
|
|
- Returns an L1MemoryDesc with ptr, size, and align_bytes
|
|
- Used by RDMA communication to register the underlying virtual memory space
|
|
"""
|
|
|
|
def test_get_l1_memory_desc_returns_desc(self, basic_config):
|
|
"""Test that get_l1_memory_desc returns an L1MemoryDesc."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
desc = manager.get_l1_memory_desc()
|
|
|
|
assert isinstance(desc, L1MemoryDesc)
|
|
|
|
manager.close()
|
|
|
|
def test_get_l1_memory_desc_returns_consistent_ptr(self, basic_config):
|
|
"""Test that get_l1_memory_desc returns the same pointer on multiple calls."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
desc1 = manager.get_l1_memory_desc()
|
|
desc2 = manager.get_l1_memory_desc()
|
|
|
|
assert desc1.ptr == desc2.ptr
|
|
|
|
manager.close()
|
|
|
|
def test_get_l1_memory_desc_size_matches_config(self, basic_config):
|
|
"""Test that the desc size matches the configured size."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
desc = manager.get_l1_memory_desc()
|
|
|
|
assert desc.size == basic_config.size_in_bytes
|
|
assert desc.align_bytes == basic_config.align_bytes
|
|
|
|
manager.close()
|
|
|
|
def test_get_l1_memory_desc_with_non_lazy_config(self, non_lazy_config):
|
|
"""Test get_l1_memory_desc with non-lazy (MixedMemoryAllocator)
|
|
configuration."""
|
|
manager = L1MemoryManager(non_lazy_config)
|
|
|
|
desc = manager.get_l1_memory_desc()
|
|
|
|
assert isinstance(desc, L1MemoryDesc)
|
|
assert desc.ptr != 0
|
|
assert desc.size == non_lazy_config.size_in_bytes
|
|
|
|
manager.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for L1MemoryManager integration
|
|
# =============================================================================
|
|
|
|
|
|
class TestL1MemoryManagerIntegration:
|
|
"""Integration tests for L1MemoryManager."""
|
|
|
|
def test_allocate_free_cycle(self, basic_config, basic_layout):
|
|
"""Test a complete allocate-free cycle."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
# Allocate
|
|
error, mem_objs = manager.allocate(basic_layout, count=3)
|
|
assert error == L1Error.SUCCESS
|
|
assert len(mem_objs) == 3
|
|
|
|
# Verify objects are valid
|
|
for obj in mem_objs:
|
|
assert obj.is_valid()
|
|
|
|
# Free
|
|
free_error = manager.free(mem_objs)
|
|
assert free_error == L1Error.SUCCESS
|
|
|
|
# Verify objects are invalidated
|
|
for obj in mem_objs:
|
|
assert not obj.is_valid()
|
|
|
|
manager.close()
|
|
|
|
def test_multiple_allocate_free_cycles(self, basic_config, basic_layout):
|
|
"""Test multiple allocate-free cycles."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
for i in range(5):
|
|
error, mem_objs = manager.allocate(basic_layout, count=2)
|
|
assert error == L1Error.SUCCESS, f"Cycle {i} failed"
|
|
assert len(mem_objs) == 2
|
|
|
|
free_error = manager.free(mem_objs)
|
|
assert free_error == L1Error.SUCCESS
|
|
|
|
manager.close()
|
|
|
|
def test_interleaved_allocate_and_free(self, basic_config, basic_layout):
|
|
"""Test interleaved allocation and free operations."""
|
|
manager = L1MemoryManager(basic_config)
|
|
|
|
# Allocate batch 1
|
|
err1, objs1 = manager.allocate(basic_layout, count=2)
|
|
assert err1 == L1Error.SUCCESS
|
|
|
|
# Allocate batch 2
|
|
err2, objs2 = manager.allocate(basic_layout, count=2)
|
|
assert err2 == L1Error.SUCCESS
|
|
|
|
# Free batch 1
|
|
free_err1 = manager.free(objs1)
|
|
assert free_err1 == L1Error.SUCCESS
|
|
|
|
# Allocate batch 3 (should reuse freed memory)
|
|
err3, objs3 = manager.allocate(basic_layout, count=2)
|
|
assert err3 == L1Error.SUCCESS
|
|
|
|
# Free remaining
|
|
manager.free(objs2)
|
|
manager.free(objs3)
|
|
|
|
manager.close()
|
|
|
|
def test_concurrent_allocate_and_free(self, basic_config, basic_layout):
|
|
"""
|
|
Test concurrent allocate and free operations from multiple threads.
|
|
This verifies the thread-safety guarantees of both methods.
|
|
"""
|
|
manager = L1MemoryManager(basic_config)
|
|
num_threads = 8
|
|
iterations = 20
|
|
exceptions = []
|
|
|
|
def worker():
|
|
try:
|
|
for _ in range(iterations):
|
|
# Allocate
|
|
error, mem_objs = manager.allocate(basic_layout, count=1)
|
|
if error == L1Error.SUCCESS:
|
|
# Briefly hold the memory
|
|
assert len(mem_objs) == 1
|
|
# Free
|
|
free_error = manager.free(mem_objs)
|
|
assert free_error == L1Error.SUCCESS
|
|
except Exception as e:
|
|
exceptions.append(e)
|
|
|
|
with ThreadPoolExecutor(max_workers=num_threads) as executor:
|
|
futures = [executor.submit(worker) for _ in range(num_threads)]
|
|
for future in as_completed(futures):
|
|
future.result() # Raises exception if worker failed
|
|
|
|
assert len(exceptions) == 0, f"Concurrency errors: {exceptions}"
|
|
|
|
manager.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for error code semantics
|
|
# =============================================================================
|
|
|
|
|
|
class TestErrorCodeSemantics:
|
|
"""Tests verifying error code semantics as documented."""
|
|
|
|
def test_success_error_code_value(self):
|
|
"""Test that SUCCESS is a valid enum member."""
|
|
assert L1Error.SUCCESS is not None
|
|
assert L1Error.SUCCESS.name == "SUCCESS"
|
|
|
|
def test_out_of_memory_error_code_value(self):
|
|
"""Test that OUT_OF_MEMORY is a valid enum member."""
|
|
assert L1Error.OUT_OF_MEMORY is not None
|
|
assert L1Error.OUT_OF_MEMORY.name == "OUT_OF_MEMORY"
|
|
|
|
def test_error_codes_are_distinct(self):
|
|
"""Test that error codes are distinct values."""
|
|
assert L1Error.SUCCESS != L1Error.OUT_OF_MEMORY
|