chore: import upstream snapshot with attribution

python/ray/tests/rdt/test_nixl_memory_pool.py

@@ -0,0 +1,345 @@
+"""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__]))