Implement Remote Fetch Batching and Python Bindings

MODULE.bazel

@@ -80,6 +80,9 @@ use_repo(
 
 bazel_dep(name = "abseil-cpp", version = "20260107.1", repo_name = "com_google_absl")
 
+# OSS crypto for global_registry_client.cc SHA256.
+bazel_dep(name = "boringssl", version = "0.20240913.0")
+
 llvm = use_extension("@xla//third_party/extensions:llvm.bzl", "llvm_extension")
 use_repo(llvm, "llvm-project")
 

tpu_raiden/api/jax/BUILD

@@ -36,10 +36,18 @@ py_library(
 py_test(
     name = "kv_cache_store_test",
     srcs = ["kv_cache_store_test.py"],
-    data = ["//tpu_raiden/kv_cache/global_registry:global_registry_server"],
+    data = [
+        "//tpu_raiden/kv_cache:raiden_orchestrator_main",
+        "//tpu_raiden/kv_cache/global_registry:global_registry_server",
+    ],
     deps = [
         ":kv_cache_store",
+        ":kv_cache_manager_jax_py",
+        "//pyglib:resources",
         "@com_google_absl_py//absl/testing:absltest",
+        "//third_party/py/portpicker",
+        "@jax//jax",
+        "@pypi//numpy",
         # buildcleaner: keep
         "//tpu_raiden/frameworks/jax:_raw_transfer",
     ],

tpu_raiden/api/jax/kv_cache_manager.py

@@ -50,6 +50,8 @@ def __init__(
       host_blocks_to_allocate: Optional[int] = None,
       parallelism: int = 4,
       node_id: int = 0,
+      listener_port: Optional[int] = None,
+      listener_controller_port: Optional[int] = None,
   ):
     """Instantiates the TransferEngine-based KVCacheManager.
 
@@ -65,8 +67,11 @@ def __init__(
         pool.
       parallelism: Number of parallel network copies per layer.
       node_id: Unique identifier for this host/node in the distributed mesh.
+      listener_port: Optional port for the internal KVCacheListener.
     """
     if host_blocks_to_allocate is not None:
+      # Legacy constructor doesn't support listener_port in this wrapper currently,
+      # but we can pass it if _impl supports it. Assuming it doesn't for now based on previous impl.
       self._impl = _impl.KVCacheManager(
           kv_caches,
           local_control_port if local_control_port > 0 else None,
@@ -89,9 +94,27 @@ def __init__(
           timeout_s=timeout_s,
           unsafe_skip_buffer_lock=unsafe_skip_buffer_lock,
           parallelism=parallelism,
+          listener_port=listener_port,
+          listener_controller_port=listener_controller_port,
       )
 
+  @property
+  def is_listener_active(self) -> bool:
+    """Returns True if any internal listener is active."""
+    return self._impl.is_listener_active
+
+  @property
+  def transfer_address(self) -> str:
+    """Returns the data plane transfer address."""
+    return self._impl.transfer_address
+
+  @property
+  def listener_address(self) -> str:
+    """Returns the listener address."""
+    return self._impl.listener_address
+
   def get_local_endpoints(self) -> List[Dict[str, Any]]:
+
     """Returns the active Raiden endpoint descriptors."""
     return self._impl.get_local_endpoints()