"""Unit tests for MemoryPoolManager. """ import sys import pytest import torch from ray.experimental.rdt.nixl_memory_pool import ( MemoryPoolManager, NixlOutOfMemoryError, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_tensor(values, dtype=torch.float32): """Create a contiguous CPU tensor.""" return torch.tensor(values, dtype=dtype) # --------------------------------------------------------------------------- # allocate_for_tensors — basic allocation and data copy # --------------------------------------------------------------------------- class TestAllocateForTensors: def test_single_tensor(self): t = _make_tensor([1.0, 2.0, 3.0]) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([t]) assert len(views) == 1 assert torch.equal(views[0], t) assert pool.has_block(t) def test_multiple_independent_tensors(self): t1 = _make_tensor([1.0, 2.0]) t2 = _make_tensor([3.0, 4.0, 5.0]) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([t1, t2]) assert len(views) == 2 assert torch.equal(views[0], t1) assert torch.equal(views[1], t2) assert pool.has_block(t1) assert pool.has_block(t2) def test_pool_views_are_backed_by_pool_tensor(self): """Returned views should be backed by the pool's internal tensor, not the source tensor's storage.""" t = _make_tensor([10.0, 20.0]) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([t]) # The view's storage should be the pool tensor's storage. assert ( views[0].untyped_storage().data_ptr() == pool.get_pool_tensor().untyped_storage().data_ptr() ) def test_data_is_copied_not_aliased(self): """Mutating the source tensor after allocation should not affect the pool copy.""" t = _make_tensor([1.0, 2.0, 3.0]) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([t]) original = views[0].clone() t[0] = 999.0 assert torch.equal(views[0], original) # --------------------------------------------------------------------------- # allocate_for_tensors — storage deduplication # --------------------------------------------------------------------------- class TestStorageDeduplication: def test_views_of_same_storage_share_one_block(self): """Two views of the same underlying storage should produce only one pool allocation.""" base = _make_tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) view_a = base[0:2] view_b = base[1:3] storage_size = base.untyped_storage().nbytes() # Pool is exactly one storage — a second allocation would OOM. pool = MemoryPoolManager(pool_size=storage_size, device=torch.device("cpu")) views = pool.allocate_for_tensors([view_a, view_b]) assert len(views) == 2 assert torch.equal(views[0], view_a) assert torch.equal(views[1], view_b) def test_duplicate_tensor_in_list(self): """The exact same tensor object appearing twice should deduplicate.""" t = _make_tensor([1.0, 2.0]) storage_size = t.untyped_storage().nbytes() pool = MemoryPoolManager(pool_size=storage_size, device=torch.device("cpu")) views = pool.allocate_for_tensors([t, t]) assert len(views) == 2 assert torch.equal(views[0], t) assert torch.equal(views[1], t) def test_cross_call_reuse(self): """A second allocate_for_tensors call with the same tensor should reuse the existing pool block (cache hit), not allocate a new one.""" t = _make_tensor([1.0, 2.0, 3.0]) storage_size = t.untyped_storage().nbytes() # Pool fits exactly one storage. pool = MemoryPoolManager(pool_size=storage_size, device=torch.device("cpu")) views1 = pool.allocate_for_tensors([t]) # Second call should hit cache, not OOM. views2 = pool.allocate_for_tensors([t]) assert torch.equal(views1[0], t) assert torch.equal(views2[0], t) def test_mixed_cache_hit_and_new_allocation(self): """One call with a mix of already-allocated and new tensors should only allocate for the new ones.""" t1 = _make_tensor([1.0, 2.0]) t2 = _make_tensor([3.0, 4.0, 5.0]) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) # Pre-allocate t1. pool.allocate_for_tensors([t1]) # Now allocate both — t1 should cache-hit, t2 should get new block. views = pool.allocate_for_tensors([t1, t2]) assert len(views) == 2 assert torch.equal(views[0], t1) assert torch.equal(views[1], t2) assert pool.has_block(t2) # --------------------------------------------------------------------------- # allocate_for_tensors — OOM # --------------------------------------------------------------------------- class TestOOM: def test_oom_single_tensor(self): t = _make_tensor([1.0, 2.0, 3.0]) # 12 bytes pool = MemoryPoolManager(pool_size=4, device=torch.device("cpu")) with pytest.raises(NixlOutOfMemoryError, match="out of memory"): pool.allocate_for_tensors([t]) def test_oom_does_not_corrupt_pool_state(self): """After an OOM error, the pool state should be unchanged — previously allocated blocks remain valid and no partial allocation leaks.""" t1 = _make_tensor([1.0, 2.0]) # 8 bytes t2 = _make_tensor([3.0, 4.0, 5.0]) # 12 bytes pool = MemoryPoolManager(pool_size=12, device=torch.device("cpu")) views1 = pool.allocate_for_tensors([t1]) assert torch.equal(views1[0], t1) # t2 doesn't fit in the remaining 4 bytes. with pytest.raises(NixlOutOfMemoryError): pool.allocate_for_tensors([t2]) # Pool should still be intact — t1's block is still valid. assert pool.has_block(t1) def test_atomic_allocation_failure(self): """When allocating multiple tensors atomically, if one doesn't fit, none should be allocated.""" t1 = _make_tensor([1.0]) # 4 bytes t2 = _make_tensor([1.0] * 100) # 400 bytes — won't fit pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) with pytest.raises(NixlOutOfMemoryError): pool.allocate_for_tensors([t1, t2]) # Neither tensor should have been tracked. assert not pool.has_block(t1) assert not pool.has_block(t2) # --------------------------------------------------------------------------- # free_tensors # --------------------------------------------------------------------------- class TestFreeTensors: def test_free_and_reallocate(self): """After freeing, the space should be reusable.""" t1 = _make_tensor([1.0, 2.0]) # 8 bytes pool = MemoryPoolManager(pool_size=8, device=torch.device("cpu")) pool.allocate_for_tensors([t1]) assert pool.has_block(t1) pool.free_tensors([t1]) assert not pool.has_block(t1) # Now a new tensor of the same size should fit. t2 = _make_tensor([3.0, 4.0]) views = pool.allocate_for_tensors([t2]) assert torch.equal(views[0], t2) def test_free_unknown_tensor_is_noop(self): """Freeing a tensor that was never allocated should not raise.""" t = _make_tensor([1.0]) pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) # Should not raise. pool.free_tensors([t]) def test_free_multiple_tensors(self): t1 = _make_tensor([1.0, 2.0]) t2 = _make_tensor([3.0, 4.0]) pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) pool.allocate_for_tensors([t1]) pool.allocate_for_tensors([t2]) pool.free_tensors([t1, t2]) assert not pool.has_block(t1) assert not pool.has_block(t2) def test_free_then_cross_call_reuse_is_broken(self): """After freeing, the same tensor should NOT get a cache hit — it should allocate a fresh block.""" t = _make_tensor([1.0, 2.0]) pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) pool.allocate_for_tensors([t]) pool.free_tensors([t]) assert not pool.has_block(t) # Re-allocate — should work (fresh allocation, not cache hit). views = pool.allocate_for_tensors([t]) assert torch.equal(views[0], t) assert pool.has_block(t) def test_double_free_is_noop(self): """Freeing an already-freed tensor should not raise or corrupt state.""" t = _make_tensor([1.0, 2.0]) pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) pool.allocate_for_tensors([t]) pool.free_tensors([t]) # Second free — should be a no-op. pool.free_tensors([t]) assert not pool.has_block(t) # --------------------------------------------------------------------------- # Block merging — allocation succeeds only after freed blocks are coalesced # --------------------------------------------------------------------------- class TestBlockMerging: def test_allocation_requires_merged_free_space(self): """After freeing adjacent blocks, the merged space should be usable for a single large allocation that wouldn't fit in either fragment.""" # Pool: 24 bytes, allocate three 8-byte tensors to fill it. t1 = _make_tensor([1.0, 2.0]) # 8 bytes t2 = _make_tensor([3.0, 4.0]) # 8 bytes t3 = _make_tensor([5.0, 6.0]) # 8 bytes pool = MemoryPoolManager(pool_size=24, device=torch.device("cpu")) pool.allocate_for_tensors([t1, t2, t3]) t_big = _make_tensor([7.0, 8.0, 9.0, 10.0]) # 16 bytes # Free only t1 — 8 bytes free, not enough for t_big (16 bytes). pool.free_tensors([t1]) with pytest.raises(NixlOutOfMemoryError): pool.allocate_for_tensors([t_big]) # Free t2 — now t1+t2 are adjacent and merged into 16 bytes free. pool.free_tensors([t2]) views = pool.allocate_for_tensors([t_big]) assert torch.equal(views[0], t_big) # --------------------------------------------------------------------------- # Edge cases # --------------------------------------------------------------------------- class TestEdgeCases: def test_empty_tensor_list(self): """allocate_for_tensors with an empty list should return an empty list.""" pool = MemoryPoolManager(pool_size=64, device=torch.device("cpu")) views = pool.allocate_for_tensors([]) assert views == [] def test_different_dtypes(self): """Tensors of different dtypes should each get their own block.""" t_f32 = torch.tensor([1.0], dtype=torch.float32) t_f64 = torch.tensor([1.0], dtype=torch.float64) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([t_f32, t_f64]) assert views[0].dtype == torch.float32 assert views[1].dtype == torch.float64 assert torch.equal(views[0], t_f32) assert torch.equal(views[1], t_f64) def test_view_with_storage_offset(self): """A tensor view with non-zero storage offset should be correctly mapped to the pool.""" base = _make_tensor([1.0, 2.0, 3.0, 4.0, 5.0]) view = base[2:4] # [3.0, 4.0], storage_offset = 2 pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([view]) assert torch.equal(views[0], view) assert views[0].shape == (2,) def test_multidimensional_tensor_shape_preserved(self): """Multi-dimensional tensor shapes should be preserved in pool views.""" t = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) views = pool.allocate_for_tensors([t]) assert views[0].shape == (3, 2) assert torch.equal(views[0], t) def test_allocate_multiple_preserves_request_order(self): """_allocate_multiple should return blocks in the same order as the input sizes, even though it allocates largest-first internally.""" pool = MemoryPoolManager(pool_size=1024, device=torch.device("cpu")) # Sizes in non-sorted order. sizes = [10, 50, 20, 40] result = pool._allocate_multiple(sizes) assert result is not None # Each result block should match the requested size, in order. for i, size in enumerate(sizes): assert result[i].size == size if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__]))