Implement global registry lookup fallback for KVCacheStore.

tpu_raiden/api/jax/kv_cache_store.py

@@ -62,23 +62,27 @@ def __repr__(self) -> str:
 class KVCacheStore:
   """Wrapper around compiled C++ KVCacheStore."""
 
-  def __init__(self, capacity: int):
-    self._impl = _impl.KVCacheStore(capacity=capacity)
+  def __init__(self, capacity: int, global_registry_address: str = ""):
+    self._impl = _impl.KVCacheStore(
+        capacity=capacity, global_registry_address=global_registry_address
+    )
 
   def lookup(
       self,
       block_hashes: list[bytes],
+      enable_global: bool = False,
   ) -> list[tuple[bytes, list[RaidenId]]]:
     """Checks the LRU directory for cached block hashes.
 
     Args:
       block_hashes: Incoming block hashes to check.
+      enable_global: Whether to fallback to global registry on miss.
 
     Returns:
       A list of tuples containing the block hash and a list of matching RaidenId
       replicas, halting immediately upon the first cache miss.
     """
-    raw_res = self._impl.lookup(block_hashes)
+    raw_res = self._impl.lookup(block_hashes, enable_global)
     final_res = []
     for hash_val, raw_slices in raw_res:
       wrapped_slices = [RaidenId(impl=rs) for rs in raw_slices]

tpu_raiden/api/jax/kv_cache_store.pyi

@@ -15,10 +15,11 @@ class RaidenId:
   ) -> None: ...
 
 class KVCacheStore:
-  def __init__(self, capacity: int) -> None: ...
+  def __init__(self, capacity: int, global_registry_address: str = ...) -> None: ...
   def lookup(
       self,
       block_hashes: list[bytes],
+      enable_global: bool = ...,
   ) -> list[tuple[bytes, list[RaidenId]]]:
     """Checks the LRU directory for cached block hashes.
 

tpu_raiden/api/jax/kv_cache_store_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
+
 from absl.testing import absltest
 
 from tpu_raiden.api.jax import kv_cache_store
@@ -147,6 +149,73 @@ def test_large_and_arbitrary_length_hashes(self):
     self.assertEqual(lookup_res[0][0], large_hash)
     self.assertEqual(lookup_res[1][0], long_hash)
 
+  def test_global_lookup_case1_local_hit(self):
+    # Case 1: Full local hit, no global hit.
+    # We don't need a registry server for this because it shouldn't be queried.
+    controller = kv_cache_store.KVCacheStore(capacity=20)
+    hashes = [b"local_only"]
+    slices = [
+        [kv_cache_store.RaidenId("local_job", "0", "kv_cache", 0)],
+    ]
+    self.assertTrue(controller.insert(hashes, slices, True)[0])
+
+    res = controller.lookup(hashes, enable_global=True)
+    self.assertLen(res, 1)
+    self.assertEqual(res[0][0], b"local_only")
+    self.assertEqual(res[0][1][0].job_name, "local_job")
+    self.assertEqual(res[0][1][0].data_replica_idx, 0)
+
+  def test_global_lookup_case2_and_3_mocked(self):
+    # We mock _impl to simulate Case 2 and Case 3 because we don't
+    # have a running registry server in Python tests.
+    controller = kv_cache_store.KVCacheStore(capacity=20)
+
+    # Create a mock for the C++ impl
+    mock_impl = unittest.mock.MagicMock()
+    controller._impl = mock_impl
+
+    # Case 2: Both local and global have the same hit, but we return local.
+    local_id = kv_cache_store._impl.RaidenId("local_job", "0", "kv_cache", 1)
+    mock_impl.lookup.return_value = [(b"shared_hash", [local_id])]
+
+    res = controller.lookup([b"shared_hash"], enable_global=True)
+    self.assertLen(res, 1)
+    self.assertEqual(res[0][0], b"shared_hash")
+    self.assertEqual(res[0][1][0].job_name, "local_job")
+    self.assertEqual(res[0][1][0].data_replica_idx, 1)
+    mock_impl.lookup.assert_called_with([b"shared_hash"], True)
+
+    # Case 3: No local hit, only global hits.
+    remote_id1 = kv_cache_store._impl.RaidenId(
+        "10.0.0.1:1234", "0", "kv_cache", 42
+    )
+    remote_id2 = kv_cache_store._impl.RaidenId(
+        "10.0.0.2:1234", "0", "kv_cache", 43
+    )
+    mock_impl.lookup.return_value = [
+        (b"global_1", [remote_id1]),
+        (b"global_2", [remote_id2]),
+    ]
+
+    res = controller.lookup([b"global_1", b"global_2"], enable_global=True)
+    self.assertLen(res, 2)
+    self.assertEqual(res[0][0], b"global_1")
+    self.assertEqual(res[0][1][0].job_name, "10.0.0.1:1234")
+    self.assertEqual(res[0][1][0].data_replica_idx, 42)
+    self.assertEqual(res[1][0], b"global_2")
+    self.assertEqual(res[1][1][0].job_name, "10.0.0.2:1234")
+    self.assertEqual(res[1][1][0].data_replica_idx, 43)
+    mock_impl.lookup.assert_called_with([b"global_1", b"global_2"], True)
+
+  def test_global_lookup_error_ignored(self):
+    controller = kv_cache_store.KVCacheStore(
+        capacity=20, global_registry_address="invalid.address:12345"
+    )
+    hashes = [b"9001"]
+    # Should not fail, just return empty because the registry is down
+    res = controller.lookup(hashes, enable_global=True)
+    self.assertEmpty(res)
+
 
 if __name__ == "__main__":
   absltest.main()

tpu_raiden/api/torch/BUILD

@@ -44,6 +44,17 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         "//tpu_raiden/frameworks/torch:_tpu_raiden_torch",
+        "@torch_tpu//torch_tpu",
+        "@torch_tpu//torch_tpu:_loader",
+    ],
+)
+
+py_test(
+    name = "kv_cache_store_test",
+    srcs = ["kv_cache_store_test.py"],
+    deps = [
+        ":kv_cache_store",
+        "@com_google_absl_py//absl/testing:absltest",
     ],
 )
 

tpu_raiden/api/torch/kv_cache_store.py

@@ -12,9 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Python wrapper for the compiled C++ KVCacheStore."""
-
+import ctypes
+import os
+import pathlib
 from typing import Any
+
+import torch_tpu
+from torch_tpu import _loader as _torch_tpu_loader
+
+Path = pathlib.Path
+
+
+def _load_torch_tpu_common() -> None:
+  _torch_tpu_loader.load()
+  common = Path(torch_tpu.__file__).resolve().parent / "common"
+  lib = common / "libpywrap_torch_tpu_common.so"
+  if lib.exists():
+    ctypes.CDLL(str(lib), mode=os.RTLD_GLOBAL | os.RTLD_NOW)
+
+
+_load_torch_tpu_common()
+
 from tpu_raiden.frameworks.torch import _tpu_raiden_torch as _impl
 
 
@@ -62,23 +80,27 @@ def __repr__(self) -> str:
 class KVCacheStore:
   """Wrapper around compiled C++ KVCacheStore."""
 
-  def __init__(self, capacity: int):
-    self._impl = _impl.KVCacheStore(capacity=capacity)
+  def __init__(self, capacity: int, global_registry_address: str = ""):
+    self._impl = _impl.KVCacheStore(
+        capacity=capacity, global_registry_address=global_registry_address
+    )
 
   def lookup(
       self,
       block_hashes: list[bytes],
+      enable_global: bool = False,
   ) -> list[tuple[bytes, list[RaidenId]]]:
     """Checks the LRU directory for cached block hashes.
 
     Args:
       block_hashes: Incoming block hashes to check.
+      enable_global: Whether to fallback to global registry on miss.
 
     Returns:
       A list of tuples containing the block hash and a list of matching RaidenId
       replicas, halting immediately upon the first cache miss.
     """
-    raw_res = self._impl.lookup(block_hashes)
+    raw_res = self._impl.lookup(block_hashes, enable_global)
     final_res = []
     for hash_val, raw_slices in raw_res:
       wrapped_slices = [RaidenId(impl=rs) for rs in raw_slices]