comms rdma reg cache

Author: amirafzali

tests/test_rdma_memory_cache.py

@@ -0,0 +1,109 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+import pytest
+import torch
+
+
+class MockRdmaMemory:
+    """Lightweight mock that records the tensor it was created with."""
+
+    def __init__(self, tensor: torch.Tensor) -> None:
+        self.data_ptr = tensor.data_ptr()
+        self.nbytes = tensor.nbytes
+
+
+@pytest.fixture(autouse=True)
+def _patch_rdma_memory(monkeypatch):
+    import torchstore.transport.torchcomms.cache as cache_mod
+
+    monkeypatch.setattr(cache_mod, "RdmaMemory", MockRdmaMemory)
+
+
+from torchstore.transport.torchcomms.cache import RdmaMemoryCache  # noqa: E402
+
+
+class TestRdmaMemoryCacheHitMiss:
+    def test_cache_hit_same_tensor(self):
+        cache = RdmaMemoryCache()
+        tensor = torch.randn(4, 4)
+
+        mem1 = cache.get_or_register(tensor)
+        mem2 = cache.get_or_register(tensor)
+
+        assert mem1 is mem2
+
+    def test_cache_miss_different_tensors(self):
+        cache = RdmaMemoryCache()
+        t1 = torch.randn(4, 4)
+        t2 = torch.randn(4, 4)
+
+        mem1 = cache.get_or_register(t1)
+        mem2 = cache.get_or_register(t2)
+
+        assert mem1 is not mem2
+        assert len(cache._cache) == 2
+
+
+class TestRdmaMemoryCacheClear:
+    def test_clear(self):
+        cache = RdmaMemoryCache()
+        t1 = torch.randn(2, 2)
+        t2 = torch.randn(3, 3)
+        cache.get_or_register(t1)
+        cache.get_or_register(t2)
+
+        assert len(cache._cache) == 2
+        cache.clear()
+        assert len(cache._cache) == 0
+        assert len(cache._storage_refs) == 0
+
+
+class TestRdmaMemoryCacheWeakrefEviction:
+    def test_evicts_when_tensor_is_deleted(self):
+        cache = RdmaMemoryCache()
+        t = torch.randn(4, 4)
+        key = (t.data_ptr(), t.nbytes)
+
+        cache.get_or_register(t)
+        assert key in cache._cache
+
+        del t
+        # refcount GC fires the weakref callback synchronously
+        assert key not in cache._cache
+        assert key not in cache._storage_refs
+
+    def test_survives_when_view_still_alive(self):
+        cache = RdmaMemoryCache()
+        t = torch.randn(4, 4)
+        view = t[0:2]  # shares the same storage
+        key = (t.data_ptr(), t.nbytes)
+
+        cache.get_or_register(t)
+        del t
+        # storage is still alive via 'view', so entry should persist
+        assert key in cache._cache
+
+        del view
+        # storage is now freed
+        assert key not in cache._cache
+
+    def test_new_tensor_after_eviction(self):
+        cache = RdmaMemoryCache()
+        t1 = torch.randn(4, 4)
+        mem1 = cache.get_or_register(t1)
+
+        del t1  # evicts
+
+        t2 = torch.randn(4, 4)
+        mem2 = cache.get_or_register(t2)
+
+        # should be a fresh registration, not the old one
+        assert mem2 is not mem1
+        assert len(cache._cache) == 1
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

torchstore/storage_volume.py

@@ -390,3 +390,4 @@ async def delete(self, key: str) -> None:
 
     def reset(self) -> None:
         self.kv = {}
+        self.transport_context.clear()

torchstore/transport/torchcomms/buffer.py

@@ -4,13 +4,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
 from dataclasses import dataclass
 from typing import Any, TYPE_CHECKING
 
 import torch
 
 from torchstore.transport.buffers import TransportBuffer
-from torchstore.transport.torchcomms.cache import RdmaTransportCache
+from torchstore.transport.torchcomms.cache import RdmaMemoryCache, RdmaTransportCache
 from torchstore.transport.types import Request
 
 try:
@@ -23,6 +24,13 @@
     from torchstore.transport.buffers import TransportContext
 
 
+# I'd like caching of user owned pinned memory to be toggleable.
+# yes, I understand we're introduce many "magic" vars. A follow up will
+# be to provide a clear TorchStrategy level config for these switches
+def _client_rdma_cache_enabled() -> bool:
+    return os.environ.get("TORCHSTORE_CLIENT_RDMA_CACHE", "1") == "1"
+
+
 @dataclass
 class RdmaContext:
     """Per-entry state for TorchComms RDMA batch operations."""
@@ -128,7 +136,11 @@ def __getstate__(self) -> dict[str, Any]:
 
     def _allocate_ctx(self, tensor: torch.Tensor) -> RdmaContext:
         self._assert_valid_tensor(tensor, tensor.dtype, tensor.shape)
-        rdma_memory = RdmaMemory(tensor)
+        if _client_rdma_cache_enabled():
+            cache = self.storage_volume_ref.transport_context.get(RdmaMemoryCache)
+            rdma_memory = cache.get_or_register(tensor)
+        else:
+            rdma_memory = RdmaMemory(tensor)
         return RdmaContext(
             rdma_memory=rdma_memory,
             rdma_remote_buffer=rdma_memory.to_remote_buffer(),
@@ -185,6 +197,8 @@ async def handle_put_request(
         entries: list[tuple[Request, Any]],
     ) -> list[Any]:
         """Called by storage volume. Read from client's source RdmaMemory (put)."""
+        rdma_mem_cache = ctx.get(RdmaMemoryCache)
+
         results = []
         for entry, rdma_ctx in zip(entries, self._contexts, strict=True):
             _, maybe_tensor = entry
@@ -206,7 +220,7 @@ async def handle_put_request(
             self._assert_valid_tensor(maybe_tensor, rdma_ctx.dtype, rdma_ctx.shape)
 
             # TODO: replace sequential reads with true batch RDMA operations (coming to torchcomms)
-            receiving_buffer = RdmaMemory(maybe_tensor)
+            receiving_buffer = rdma_mem_cache.get_or_register(maybe_tensor)
             res = transport.read(
                 receiving_buffer.to_mutable_view(), rdma_ctx.rdma_remote_buffer
             )
@@ -226,6 +240,8 @@ async def handle_get_request(
         Note: SV determination of is_object is authoritative and mutates _contexts
         sent back to the client.
         """
+        rdma_mem_cache = ctx.get(RdmaMemoryCache)
+
         for entry, rdma_ctx in zip(entries, self._contexts, strict=True):
             _, data = entry
             if not isinstance(data, torch.Tensor):
@@ -243,6 +259,11 @@ async def handle_get_request(
                 )
                 contiguous_buffer.copy_(tensor)
                 tensor = contiguous_buffer
+                # staging copy is lost, don't cache it
+                rdma_memory = RdmaMemory(tensor)
+            else:
+                # stable tensor from SV, cache the registration
+                rdma_memory = rdma_mem_cache.get_or_register(tensor)
 
             transport = self._get_sv_transport(ctx, rdma_ctx.device_index)
 
@@ -252,7 +273,6 @@ async def handle_get_request(
                 )
             self._assert_valid_tensor(tensor, rdma_ctx.dtype, rdma_ctx.shape)
             # TODO: replace sequential writes with true batch RDMA operations (coming to torchcomms)
-            rdma_memory = RdmaMemory(tensor)
 
             res = transport.write(rdma_memory.to_view(), rdma_ctx.rdma_remote_buffer)
             if res != 0:

torchstore/transport/torchcomms/cache.py

@@ -5,14 +5,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+import weakref
 from functools import cache
 
 import torch
 
 from torchstore.transport.buffers import TransportCache
 
 try:
-    from torchcomms._transport import RdmaTransport
+    from torchcomms._transport import RdmaMemory, RdmaTransport
 
     torchcomms_available = True
 except ImportError:
@@ -82,3 +83,41 @@ def contains(self, key: str, device: torch.device | int) -> bool:
 
     def clear(self) -> None:
         self.transports.clear()
+
+
+class RdmaMemoryCache(TransportCache):
+    """Cache for RDMA memory registrations.
+
+    Keyed on (data_ptr, nbytes). Avoids repeated ibv_reg_mr kernel calls for stable tensors.
+
+    A weakref on each tensor's ``untyped_storage()`` auto-evicts the entry when
+    the underlying memory is released (i.e. no more live tensors on the process that reference this
+    storage). This makes the cache safe for both server-side (stable kv tensors)
+    and client-side (user-owned, transient) use.
+    """
+
+    def __init__(self) -> None:
+        # {(data_ptr, nbytes): RdmaMemory}
+        self._cache: dict[tuple[int, int], "RdmaMemory"] = {}
+        # parallel dict keeping the weakref alive so the callback can fire
+        self._storage_refs: dict[tuple[int, int], weakref.ref] = {}
+
+    def get_or_register(self, tensor: torch.Tensor) -> "RdmaMemory":
+        key = (tensor.data_ptr(), tensor.nbytes)
+        if key in self._cache:
+            return self._cache[key]
+
+        mem = RdmaMemory(tensor)
+        self._cache[key] = mem
+        self._storage_refs[key] = weakref.ref(
+            tensor.untyped_storage(), lambda _ref, _k=key: self._evict(_k)
+        )
+        return mem
+
+    def _evict(self, key: tuple[int, int]) -> None:
+        self._cache.pop(key, None)
+        self._storage_refs.pop(key, None)
+
+    def clear(self) -> None:
+        self._cache.clear()
+        self._storage_refs.clear()