From 7be92d33b670c99b2e7cb9283220c7856f20b40c Mon Sep 17 00:00:00 2001 From: jcgu Date: Sat, 4 Jul 2026 23:50:45 -0700 Subject: [PATCH] Implement Remote Fetch Batching and Python Bindings PiperOrigin-RevId: 942711303 --- MODULE.bazel | 3 + tpu_raiden/api/jax/BUILD | 10 +- tpu_raiden/api/jax/kv_cache_manager.py | 23 + tpu_raiden/api/jax/kv_cache_store.py | 300 +++- tpu_raiden/api/jax/kv_cache_store.pyi | 94 +- tpu_raiden/api/jax/kv_cache_store_test.py | 1331 +++++++++++++++ tpu_raiden/api/torch/BUILD | 11 + tpu_raiden/api/torch/kv_cache_store.py | 239 ++- tpu_raiden/api/torch/kv_cache_store_test.py | 276 ++++ .../core/kv_cache_manager_with_transfer.cc | 8 +- tpu_raiden/frameworks/jax/BUILD | 3 + tpu_raiden/frameworks/jax/kv_cache_manager.cc | 63 +- tpu_raiden/frameworks/jax/kv_cache_manager.h | 22 +- tpu_raiden/frameworks/jax/kv_cache_store.pyi | 86 +- .../frameworks/jax/tpu_raiden_jax_module.cc | 192 ++- .../torch/tpu_raiden_torch_module.cc | 113 +- tpu_raiden/kv_cache/BUILD | 90 +- tpu_raiden/kv_cache/global_registry/BUILD | 5 +- .../global_registry/global_registry_client.cc | 8 +- .../global_registry/global_registry_client.h | 2 +- .../global_registry/global_registry_server.cc | 4 +- .../global_registry/global_registry_server.h | 4 +- .../global_registry_server_main.cc | 12 +- .../global_registry/global_registry_test.cc | 12 +- tpu_raiden/kv_cache/kv_cache_listener.cc | 67 +- tpu_raiden/kv_cache/kv_cache_listener.h | 4 +- tpu_raiden/kv_cache/kv_cache_manager_base.cc | 62 +- tpu_raiden/kv_cache/kv_cache_manager_base.h | 11 + tpu_raiden/kv_cache/kv_cache_store.cc | 372 ++++- tpu_raiden/kv_cache/kv_cache_store.h | 224 ++- .../kv_cache/kv_cache_store_internal.cc | 2 +- tpu_raiden/kv_cache/kv_cache_store_test.cc | 1465 ++++++++++++++++- tpu_raiden/kv_cache/lru_cache.h | 26 + tpu_raiden/kv_cache/lru_cache_test.cc | 23 + .../kv_cache/raiden_controller_embedded.cc | 660 ++++++++ .../kv_cache/raiden_controller_embedded.h | 111 ++ tpu_raiden/kv_cache/raiden_orchestrator.cc | 268 +++ tpu_raiden/kv_cache/raiden_orchestrator.h | 72 + .../kv_cache/raiden_orchestrator_main.cc | 42 + .../kv_cache/raiden_orchestrator_test.cc | 172 ++ tpu_raiden/rpc/BUILD | 20 + tpu_raiden/rpc/coordination_client.cc | 4 +- tpu_raiden/rpc/coordination_client.h | 2 +- tpu_raiden/rpc/coordination_client_test.cc | 14 +- tpu_raiden/rpc/raiden_service.proto | 27 + tpu_raiden/rpc/rpc_utils.cc | 147 ++ tpu_raiden/rpc/rpc_utils.h | 37 + tpu_raiden/transport/raw_buffer_transport.cc | 4 +- 48 files changed, 6546 insertions(+), 201 deletions(-) create mode 100644 tpu_raiden/api/torch/kv_cache_store_test.py create mode 100644 tpu_raiden/kv_cache/raiden_controller_embedded.cc create mode 100644 tpu_raiden/kv_cache/raiden_controller_embedded.h create mode 100644 tpu_raiden/kv_cache/raiden_orchestrator.cc create mode 100644 tpu_raiden/kv_cache/raiden_orchestrator.h create mode 100644 tpu_raiden/kv_cache/raiden_orchestrator_main.cc create mode 100644 tpu_raiden/kv_cache/raiden_orchestrator_test.cc create mode 100644 tpu_raiden/rpc/rpc_utils.cc create mode 100644 tpu_raiden/rpc/rpc_utils.h diff --git a/MODULE.bazel b/MODULE.bazel index d255c721..dca8611c 100644 --- a/MODULE.bazel +++ b/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") diff --git a/tpu_raiden/api/jax/BUILD b/tpu_raiden/api/jax/BUILD index 70f0db16..73e33b7a 100644 --- a/tpu_raiden/api/jax/BUILD +++ b/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", ], diff --git a/tpu_raiden/api/jax/kv_cache_manager.py b/tpu_raiden/api/jax/kv_cache_manager.py index d0972836..51b01989 100644 --- a/tpu_raiden/api/jax/kv_cache_manager.py +++ b/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() diff --git a/tpu_raiden/api/jax/kv_cache_store.py b/tpu_raiden/api/jax/kv_cache_store.py index 7b463a45..4bbbedc4 100644 --- a/tpu_raiden/api/jax/kv_cache_store.py +++ b/tpu_raiden/api/jax/kv_cache_store.py @@ -14,10 +14,18 @@ """Python wrapper for the compiled C++ KVCacheStore.""" +import enum from typing import Any from tpu_raiden.frameworks.jax import _tpu_raiden_jax as _impl +class BlockStatus(enum.Enum): + INIT = 0 + REMOTE = 1 + HOST = 2 + HBM = 3 + + class RaidenId: """Wrapper around compiled C++ RaidenId.""" @@ -59,62 +67,306 @@ def __repr__(self) -> str: ) +class RaidenBlockID: + """Wrapper around compiled C++ RaidenBlockID.""" + + def __init__( + self, + raiden_id: RaidenId | None = None, + host_block_id: int = -1, + status: BlockStatus = BlockStatus.INIT, + impl: Any = None, + ): + if impl is not None: + self._impl = impl + else: + if raiden_id is None: + raiden_id = RaidenId() + # Map Python enum to C++ enum + status_val = getattr(_impl.BlockStatus, status.name) + self._impl = _impl.RaidenBlockID( + raiden_id._impl, host_block_id, status_val # pylint: disable=protected-access + ) + + @property + def raiden_id(self) -> RaidenId: + return RaidenId(impl=self._impl.raiden_id) + + @property + def host_block_id(self) -> int: + return self._impl.host_block_id + + @property + def status(self) -> BlockStatus: + return BlockStatus[self._impl.status.name] + + @property + def job_name(self) -> str: + return self.raiden_id.job_name + + @property + def job_replica_id(self) -> str: + return self.raiden_id.job_replica_id + + @property + def data_name(self) -> str: + return self.raiden_id.data_name + + @property + def data_replica_idx(self) -> int: + return self.raiden_id.data_replica_idx + + def __repr__(self) -> str: + return ( + f"RaidenBlockID(raiden_id={self.raiden_id}," + f" host_block_id={self.host_block_id}, status={self.status})" + ) + + +class RemoteFetchConfig: + """Wrapper around compiled C++ RemoteFetchConfig.""" + + def __init__(self, impl: Any = None): + if impl is not None: + self._impl = impl + else: + self._impl = _impl.RemoteFetchConfig() + + @property + def orchestrator_address(self) -> str: + return self._impl.orchestrator_address + + @orchestrator_address.setter + def orchestrator_address(self, val: str): + self._impl.orchestrator_address = val + + @property + def controller_port(self) -> int: + return self._impl.controller_port + + @controller_port.setter + def controller_port(self, val: int): + self._impl.controller_port = val + + @property + def local_worker_port(self) -> int: + return self._impl.local_worker_port + + @local_worker_port.setter + def local_worker_port(self, val: int): + self._impl.local_worker_port = val + + @property + def bytes_per_block(self) -> int: + return self._impl.bytes_per_block + + @bytes_per_block.setter + def bytes_per_block(self, val: int): + self._impl.bytes_per_block = val + + @property + def num_shards(self) -> int: + return self._impl.num_shards + + @num_shards.setter + def num_shards(self, val: int): + self._impl.num_shards = val + + @property + def num_listeners(self) -> int: + return self._impl.num_listeners + + @num_listeners.setter + def num_listeners(self, val: int): + self._impl.num_listeners = val + + 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 = "", + raiden_id: RaidenId | None = None, + remote_config: Any | None = None, + ): + raw_raiden_id = _impl.RaidenId() + if raiden_id is not None: + raw_raiden_id = raiden_id._impl # pylint: disable=protected-access + + raw_remote_config = None + if remote_config is not None: + # Assuming remote_config is RemoteFetchConfig wrapper or raw impl + if hasattr(remote_config, "_impl"): + raw_remote_config = remote_config._impl + else: + raw_remote_config = remote_config # If it is directly the _impl type + + self._impl = _impl.KVCacheStore( + capacity=capacity, + global_registry_address=global_registry_address, + raiden_id=raw_raiden_id, + remote_config=raw_remote_config, + ) + + @property + def raiden_id(self) -> RaidenId: + """Returns the RaidenId associated with this store.""" + return RaidenId(impl=self._impl.raiden_id) def lookup( self, block_hashes: list[bytes], - ) -> list[tuple[bytes, list[RaidenId]]]: + enable_global: bool = False, + ) -> list[tuple[bytes, list[RaidenBlockID]]]: """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. + A list of tuples containing the block hash and a list of matching + RaidenBlockID 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] + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] final_res.append((hash_val, wrapped_slices)) return final_res def insert( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], on_host: bool, - ) -> tuple[bool, list[tuple[bytes, list[RaidenId]]]]: + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: + """Caches sharded buffers into host-RAM/HBM backing store. + + Args: + block_hashes: Incoming block hashes to insert. + slices: List of sharded buffer metadata corresponding to each block hash. + on_host: Whether the slices are located in host memory. + + Returns: + A tuple containing: + - bool: whether all blocks were successfully inserted (i.e. none already + existed). + - list: list of entries evicted from the LRU cache during insertion. + """ raw_slices = [] for slice_list in slices: - raw_slices.append( - [s._impl for s in slice_list] # pylint: disable=protected-access - ) + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_slices.append(converted_list) all_inserted, raw_evicted = self._impl.insert( block_hashes, raw_slices, on_host ) wrapped_evicted = [] for hash_val, raw_slices in raw_evicted: - wrapped_slices = [RaidenId(impl=rs) for rs in raw_slices] + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] wrapped_evicted.append((hash_val, wrapped_slices)) return all_inserted, wrapped_evicted + def insert_and_pin( + self, + block_hashes: list[bytes], + slices: list[list[RaidenBlockID]], + on_host: bool, + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: + """Pins existing block hashes, and inserts/pins new block hashes. + + Pins all existing block hashes, and inserts and pins new block hashes if + there is sufficient available space in the LRU cache. + + Args: + block_hashes: Incoming block hashes to insert and pin. + slices: List of sharded buffer metadata corresponding to each block hash. + on_host: Whether the slices are located in host memory. + + Returns: + A tuple containing: + - bool: whether the entire insert_and_pin operation succeeded (i.e. all + existing keys were pinned, all new keys inserted and pinned). + - list: list of entries evicted during insertion. + """ + raw_slices = [] + for slice_list in slices: + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_slices.append(converted_list) + all_inserted, raw_evicted = self._impl.insert_and_pin( + block_hashes, raw_slices, on_host + ) + wrapped_evicted = [] + for hash_val, raw_slices in raw_evicted: + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] + wrapped_evicted.append((hash_val, wrapped_slices)) + return all_inserted, wrapped_evicted + + def release_and_delete( + self, + block_hashes: list[bytes], + pending_evict_entries: ( + list[tuple[bytes, list[RaidenBlockID]]] | None + ) = None, + ) -> tuple[int, list[tuple[bytes, list[RaidenBlockID]]]]: + """Reverts an insert_and_pin operation. + + Unpins all block_hashes in the LRU cache, deletes any block_hash in REMOTE + status whose pin count is 0, and puts back evicted entries in reverse order + for each deleted remote block. + + Args: + block_hashes: Incoming block hashes to unpin and check for deletion. + pending_evict_entries: List of entries previously evicted during + insert_and_pin. + + Returns: + A tuple containing: + - int: number of remote blocks deleted. + - list: remaining evicted entries that were not restored. + """ + if pending_evict_entries is None: + pending_evict_entries = [] + raw_evicted_in = [] + for hash_val, slice_list in pending_evict_entries: + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_evicted_in.append((hash_val, converted_list)) + del_count, raw_evicted_out = self._impl.release_and_delete( + block_hashes, raw_evicted_in + ) + wrapped_evicted_out = [] + for hash_val, raw_slices in raw_evicted_out: + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] + wrapped_evicted_out.append((hash_val, wrapped_slices)) + return del_count, wrapped_evicted_out + def delete( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], ) -> None: raw_slices = [] for slice_list in slices: - raw_slices.append( - [s._impl for s in slice_list] # pylint: disable=protected-access - ) + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_slices.append(converted_list) self._impl.delete(block_hashes, raw_slices) def capacity(self) -> int: @@ -127,3 +379,17 @@ def pin(self, block_hashes: list[bytes]) -> bool: def release(self, block_hashes: list[bytes]) -> None: """Releases previously pinned block hashes, making them eligible for LRU eviction when capacity is exceeded.""" self._impl.release(block_hashes) + + def fetch_remote(self, block_hashes: list[bytes]) -> dict[bytes, Any]: + """Initiates remote fetch for the given block hashes.""" + # We return the raw _impl futures for now, or we could wrap them if needed. + return self._impl.fetch_remote(block_hashes) + + def poll_fetch_remote_status( + self, + ) -> tuple[list[bytes], list[bytes], list[bytes]]: + """Polls the status of remote fetches. + + Returns (done, failed, pending) block hashes. + """ + return self._impl.poll_fetch_remote_status() diff --git a/tpu_raiden/api/jax/kv_cache_store.pyi b/tpu_raiden/api/jax/kv_cache_store.pyi index 495fbc36..1981932e 100644 --- a/tpu_raiden/api/jax/kv_cache_store.pyi +++ b/tpu_raiden/api/jax/kv_cache_store.pyi @@ -1,5 +1,12 @@ +import enum from typing import Any +class BlockStatus(enum.Enum): + INIT = ... + REMOTE = ... + HOST = ... + HBM = ... + class RaidenId: job_name: str job_replica_id: str @@ -14,34 +21,103 @@ class RaidenId: impl: Any = ..., ) -> None: ... +class RaidenBlockID: + raiden_id: RaidenId + host_block_id: int + status: BlockStatus + def __init__( + self, + raiden_id: RaidenId | None = ..., + host_block_id: int = ..., + status: BlockStatus = ..., + impl: Any = ..., + ) -> None: ... + +class RemoteFetchConfig: + orchestrator_address: str + controller_port: int + local_worker_port: int + bytes_per_block: int + num_shards: int + def __init__(self) -> None: ... + +class FetchFuture: + def Await(self) -> None: ... + def IsDone(self) -> bool: ... + class KVCacheStore: - def __init__(self, capacity: int) -> None: ... + def __init__( + self, + capacity: int, + global_registry_address: str = ..., + raiden_id: RaidenId | None = ..., + remote_config: RemoteFetchConfig | None = ..., + ) -> None: ... + + @property + def raiden_id(self) -> RaidenId: ... + def lookup( self, block_hashes: list[bytes], - ) -> list[tuple[bytes, list[RaidenId]]]: + enable_global: bool = ..., + ) -> list[tuple[bytes, list[RaidenBlockID]]]: """Checks the LRU directory for cached block hashes. Args: block_hashes: Incoming block hashes to check. Returns: - A list of tuples containing the block hash and a list of matching RaidenId + A list of tuples containing the block hash and a list of matching RaidenBlockID replicas, halting immediately upon the first cache miss. """ ... def insert( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], on_host: bool, - ) -> tuple[bool, list[tuple[bytes, list[RaidenId]]]]: + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: """Inserts sharded buffers into the cache.""" ... + def insert_and_pin( + self, + block_hashes: list[bytes], + slices: list[list[RaidenBlockID]], + on_host: bool, + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: + """Pins existing block hashes, and inserts/pins new block hashes. + + Pins all existing block hashes, and inserts and pins new block hashes if + there is sufficient available space in the LRU cache. + + Returns: + - bool: whether the entire insert_and_pin operation succeeded (i.e. all + existing keys were pinned, all new keys inserted and pinned). + - list: list of entries evicted during insertion. + """ + ... + def release_and_delete( + self, + block_hashes: list[bytes], + pending_evict_entries: list[tuple[bytes, list[RaidenBlockID]]] + | None = ..., + ) -> tuple[int, list[tuple[bytes, list[RaidenBlockID]]]]: + """Reverts an insert_and_pin operation. + + Unpins all block_hashes in the LRU cache, deletes any block_hash in REMOTE + status whose pin count is 0, and puts back evicted entries in reverse order + for each deleted remote block. + + Returns: + - int: number of remote blocks deleted. + - list: remaining evicted entries that were not restored. + """ + ... def delete( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], ) -> None: """Deletes cached sharded buffers from host-RAM/HBM backing store entirely.""" ... @@ -54,3 +130,9 @@ class KVCacheStore: def release(self, block_hashes: list[bytes]) -> None: """Releases previously pinned block hashes, making them eligible for LRU eviction when capacity is exceeded.""" ... + def fetch_remote(self, block_hashes: list[bytes]) -> dict[bytes, FetchFuture]: + """Initiates remote fetch for the given block hashes.""" + ... + def poll_fetch_remote_status(self) -> tuple[list[bytes], list[bytes], list[bytes]]: + """Polls the status of remote fetches. Returns (done, failed, pending) block hashes.""" + ... diff --git a/tpu_raiden/api/jax/kv_cache_store_test.py b/tpu_raiden/api/jax/kv_cache_store_test.py index ef7601d4..858d68f1 100644 --- a/tpu_raiden/api/jax/kv_cache_store_test.py +++ b/tpu_raiden/api/jax/kv_cache_store_test.py @@ -12,14 +12,113 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import subprocess +import time +import unittest + +if not os.path.exists("/dev/accel0"): + os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4" + from absl.testing import absltest +import portpicker +from google3.pyglib import resources from tpu_raiden.api.jax import kv_cache_store +from tpu_raiden.api.jax.kv_cache_manager import KVCacheManager +import jax +import jax.numpy as jnp +import numpy as np + +# Global variables for subprocesses +_orchestrator_process = None +_registry_process = None +_orchestrator_port = None +_registry_port = None + + +def setUpModule(): + global _orchestrator_process, _registry_process + global _orchestrator_port, _registry_port + + _orchestrator_port = portpicker.pick_unused_port() + _registry_port = portpicker.pick_unused_port() + + orchestrator_binary = resources.GetResourceFilename( + "google3/third_party/tpu_raiden/tpu_raiden/kv_cache/raiden_orchestrator_main" + ) + registry_binary = resources.GetResourceFilename( + "google3/third_party/tpu_raiden/tpu_raiden/kv_cache/global_registry/global_registry_server" + ) + + print(f"Starting Orchestrator on port {_orchestrator_port}") + _orchestrator_process = subprocess.Popen([ + orchestrator_binary, + f"--port={_orchestrator_port}", + "--bind_ip=0.0.0.0", + "--alsologtostderr", + ]) + + print(f"Starting Registry on port {_registry_port}") + _registry_process = subprocess.Popen([ + registry_binary, + f"--port={_registry_port}", + "--bind_ip=0.0.0.0", + "--alsologtostderr", + ]) + + # Give them some time to start + time.sleep(2) + + +def tearDownModule(): + global _orchestrator_process, _registry_process + if _orchestrator_process: + _orchestrator_process.terminate() + _orchestrator_process.wait() + if _registry_process: + _registry_process.terminate() + _registry_process.wait() class KVCacheStoreTest(absltest.TestCase): + def _make_sharding(self, devices): + import jax.experimental.mesh_utils as mesh_utils + + n = len(devices) + if n >= 4: + device_mesh = mesh_utils.create_device_mesh((2, n // 2), devices) + else: + device_mesh = mesh_utils.create_device_mesh((1, n), devices) + mesh = jax.sharding.Mesh(device_mesh, ("x", "y")) + return jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec(None, None, "x", "y", None) + ) + + def setUp(self): + super().setUp() + try: + self.devices = jax.devices("tpu") + except Exception: + self.devices = jax.devices() + + self.num_devices = len(self.devices) + self.sharding = self._make_sharding(self.devices) + self.mesh = self.sharding.mesh + print(f"DEBUG_DEVICES: {self.devices}") + print(f"DEBUG_SHARDING: {self.sharding}") + + def _make_device_array(self, np_array, sharding=None): + if sharding is None: + sharding = self.sharding + return jax.device_put(jnp.array(np_array), sharding).block_until_ready() + def test_basic_tests(self): + print("TEST_BASIC_DEVICES:", jax.devices()) + print("TEST_BASIC_DEVICE_COUNT:", jax.device_count()) + print("TEST_BASIC_MESH_DEVICES:", self.mesh.devices) + print("TEST_BASIC_MESH_SHAPE:", self.mesh.shape) controller = kv_cache_store.KVCacheStore(capacity=20) self.assertEqual(controller.capacity(), 20) @@ -147,6 +246,1238 @@ 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.RaidenBlockID( + 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.RaidenBlockID( + kv_cache_store._impl.RaidenId("10.0.0.1:1234", "0", "kv_cache", 42) + ) + remote_id2 = kv_cache_store._impl.RaidenBlockID( + 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) + + def test_insert_and_pin_release_and_delete(self): + controller = kv_cache_store.KVCacheStore(capacity=2) + + local_hashes = [b"local_1", b"local_2"] + local_slices = [ + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("local_job", "0", "kv_cache", 0), + -1, + kv_cache_store.BlockStatus.HOST, + ) + ], + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("local_job", "0", "kv_cache", 1), + -1, + kv_cache_store.BlockStatus.HOST, + ) + ], + ] + self.assertTrue(controller.insert(local_hashes, local_slices, True)[0]) + + remote_hashes = [b"remote_1", b"remote_2"] + remote_slices = [ + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("remote_job", "0", "kv_cache", 0), + -1, + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("remote_job", "0", "kv_cache", 1), + -1, + kv_cache_store.BlockStatus.REMOTE, + ) + ], + ] + success, evicted = controller.insert_and_pin( + remote_hashes, remote_slices, True + ) + self.assertTrue(success) + self.assertLen(evicted, 2) + self.assertEmpty(controller.lookup([b"local_1"])) + + del_count, rem_evicted = controller.release_and_delete( + remote_hashes, evicted + ) + self.assertEqual(del_count, 2) + self.assertEmpty(rem_evicted) + self.assertLen(controller.lookup([b"local_1", b"local_2"]), 2) + + def test_fetch_remote_basic(self): + config = kv_cache_store.RemoteFetchConfig() + config.orchestrator_address = f"localhost:{_orchestrator_port}" + config.controller_port = 0 # Ephemeral + config.local_worker_port = 0 # Ephemeral + config.bytes_per_block = 1024 + config.num_shards = 1 + + controller = kv_cache_store.KVCacheStore(capacity=2, remote_config=config) + + remote_hashes = [b"remote_1"] + remote_slices = [[ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("remote_job", "0", "kv_cache", 0), + 5, # host_block_id + kv_cache_store.BlockStatus.REMOTE, + ) + ]] + self.assertTrue(controller.insert(remote_hashes, remote_slices, True)[0]) + + # Call fetch_remote + futures = controller.fetch_remote(remote_hashes) + self.assertIn(b"remote_1", futures) + future = futures[b"remote_1"] + self.assertFalse(future.IsDone()) + + # Poll status + done, failed, pending = controller.poll_fetch_remote_status() + self.assertEmpty(done) + self.assertEmpty(failed) + self.assertEqual(pending, [b"remote_1"]) + + def test_e2e_remote_fetch_cpu(self): + """Tests end-to-end remote fetch on CPU.""" + # Enable VLOG for debugging + # We can't easily set absl flags in python if they are C++ flags, but we can try environment variables + # or just rely on the fact that they might be logged to stderr anyway. + # Let's try setting environment variables for logging if possible. + + # 1. Setup Ports + sender_listener_port = portpicker.pick_unused_port() + receiver_listener_port = portpicker.pick_unused_port() + sender_controller_port = portpicker.pick_unused_port() + receiver_controller_port = portpicker.pick_unused_port() + + print(f"Sender Listener Port: {sender_listener_port}") + print(f"Receiver Listener Port: {receiver_listener_port}") + print(f"Sender Controller Port: {sender_controller_port}") + print(f"Receiver Controller Port: {receiver_controller_port}") + + # 2. Setup JAX arrays (CPU) + # Using small shape for testing + shape = (1, 128, 8, 8, 128) + sender_kv_caches = [ + self._make_device_array(np.ones(shape, dtype=np.float32)) + ] + receiver_kv_caches = [ + self._make_device_array(np.zeros(shape, dtype=np.float32)) + ] + + sender_mgr = KVCacheManager( + kv_caches=sender_kv_caches, + node_id=0, + local_control_port=0, # Ephemeral for Manager's own internal server + max_blocks=4, + num_slots=2, + listener_port=sender_listener_port, + listener_controller_port=sender_controller_port, + ) + self.assertTrue(sender_mgr.is_listener_active) + + receiver_mgr = KVCacheManager( + kv_caches=receiver_kv_caches, + node_id=0, + local_control_port=0, # Ephemeral for Manager's own internal server + max_blocks=4, + num_slots=2, + listener_port=receiver_listener_port, + listener_controller_port=receiver_controller_port, + ) + self.assertTrue(receiver_mgr.is_listener_active) + + bytes_per_block = 128 * 8 * 8 * 128 * 4 + local_shards_per_listener = self.num_devices + local_bytes_per_block = bytes_per_block // local_shards_per_listener + + # 4. Instantiate KVCacheStores (Control Plane) + sender_config = kv_cache_store.RemoteFetchConfig() + sender_config.orchestrator_address = f"localhost:{_orchestrator_port}" + sender_config.controller_port = sender_controller_port + sender_config.local_worker_port = sender_listener_port + sender_config.bytes_per_block = local_bytes_per_block + sender_config.num_shards = self.num_devices + sender_config.num_listeners = 1 + + sender_id = kv_cache_store.RaidenId("sender_job", "0", "kv_cache", 0) + sender_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address=f"localhost:{_registry_port}", + raiden_id=sender_id, + remote_config=sender_config, + ) + + receiver_config = kv_cache_store.RemoteFetchConfig() + receiver_config.orchestrator_address = f"localhost:{_orchestrator_port}" + receiver_config.controller_port = receiver_controller_port + receiver_config.local_worker_port = receiver_listener_port + receiver_config.bytes_per_block = local_bytes_per_block + receiver_config.num_shards = self.num_devices + receiver_config.num_listeners = 1 + + receiver_id = kv_cache_store.RaidenId("receiver_job", "0", "kv_cache", 0) + receiver_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address=f"localhost:{_registry_port}", + raiden_id=receiver_id, + remote_config=receiver_config, + ) + + # Give them some time to register with orchestrator + time.sleep(1) + + # 5. Prepare Data + hash_1 = b"hash_1" + # Sender has block 0 locally + sender_slices = [[ + kv_cache_store.RaidenBlockID( + sender_id, + 0, # host_block_id + kv_cache_store.BlockStatus.HOST, + ) + ]] + self.assertTrue(sender_store.insert([hash_1], sender_slices, True)[0]) + + # NOTE: We simulate global registry lookup here by inserting REMOTE status directly. + # When write / delete through in KVCacheStore is implemented, we can test with a real global registry. + # Receiver knows block is remote at sender + receiver_slices = [[ + kv_cache_store.RaidenBlockID( + sender_id, + 0, # host_block_id at sender + kv_cache_store.BlockStatus.REMOTE, + ) + ]] + # We need to insert it into receiver store so fetch_remote can look it up? + # Wait, fetch_remote takes hashes, and looks up in local store to find where it is. + self.assertTrue(receiver_store.insert([hash_1], receiver_slices, True)[0]) + + # 6. Trigger Fetch + futures = receiver_store.fetch_remote([hash_1]) + self.assertIn(hash_1, futures) + future = futures[hash_1] + + # 7. Wait for completion + # Await might block indefinitely if something is wrong, so let's use a timeout loop with poll + max_wait = 10 + start_time = time.time() + completed = False + while time.time() - start_time < max_wait: + done, failed, pending = receiver_store.poll_fetch_remote_status() + if hash_1 in done: + completed = True + break + if hash_1 in failed: + self.fail(f"Fetch failed for {hash_1}") + time.sleep(0.5) + + self.assertTrue(completed, "Fetch timed out") + self.assertTrue(future.IsDone()) + + # TODO: Verify data content if possible, but for now just success is good. + + def test_queue_flow_embedded(self): + print("test_queue_flow_embedded starting...") + """Verifies high-level Fetch API with embedded controller, including multiple blocks, layers, and shards.""" + sender_listener_port = portpicker.pick_unused_port() + receiver_listener_port = portpicker.pick_unused_port() + sender_controller_port = portpicker.pick_unused_port() + receiver_controller_port = portpicker.pick_unused_port() + + num_blocks = 4 + num_layers = 2 + shape = (num_blocks, 128, 8, 8, 128) + + sender_kv_caches = [] + receiver_kv_caches = [] + + for l in range(num_layers): + # Fill with non-zero values + sender_data = np.full(shape, fill_value=float(l + 1), dtype=np.float32) + sender_kv_caches.append(self._make_device_array(sender_data)) + print(f"Layer {l} sender_data sample: {sender_data[2, 0, 0, 0, :5]}") + + # Use jnp.zeros to create device buffers + buf = np.zeros(shape, dtype=np.float32) + receiver_kv_caches.append(self._make_device_array(buf)) + print(f"Layer {l} receiver buffer id before: {id(buf)}") + + # shape is (num_blocks, 128, 8, 8, 128) + # slice_byte_size in C++ is GetMajorSliceByteSize, which is size of one block. + # For shape (4, 128, 8, 8, 128) of float32, block size is 128 * 8 * 8 * 128 * 4 = 4194304 bytes. + bytes_per_block = 128 * 8 * 8 * 128 * 4 + + sender_mgr = KVCacheManager( + kv_caches=sender_kv_caches, + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=sender_listener_port, + listener_controller_port=sender_controller_port, + ) + + receiver_mgr = KVCacheManager( + kv_caches=receiver_kv_caches, + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=receiver_listener_port, + listener_controller_port=receiver_controller_port, + ) + + # Keep sender_mgr alive + print(f"Sender manager alive: {sender_mgr is not None}") + + sender_raiden_id = kv_cache_store.RaidenId("job1", "0", "data1", 0) + receiver_raiden_id = kv_cache_store.RaidenId("job2", "0", "data2", 0) + + local_shards_per_listener = self.num_devices + local_bytes_per_block = bytes_per_block // local_shards_per_listener + + sender_remote_config = kv_cache_store.RemoteFetchConfig() + sender_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + sender_remote_config.controller_port = sender_controller_port + sender_remote_config.local_worker_port = sender_listener_port + sender_remote_config.bytes_per_block = local_bytes_per_block + sender_remote_config.num_shards = self.num_devices + + receiver_remote_config = kv_cache_store.RemoteFetchConfig() + receiver_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + receiver_remote_config.controller_port = receiver_controller_port + receiver_remote_config.local_worker_port = receiver_listener_port + receiver_remote_config.bytes_per_block = local_bytes_per_block + receiver_remote_config.num_shards = self.num_devices + + sender_store = kv_cache_store.KVCacheStore( + capacity=100, + global_registry_address="", + raiden_id=sender_raiden_id, + remote_config=sender_remote_config, + ) + + receiver_store = kv_cache_store.KVCacheStore( + capacity=100, + global_registry_address="", + raiden_id=receiver_raiden_id, + remote_config=receiver_remote_config, + ) + + # Insert blocks into sender (HOST status) + print("Inserting into sender...") + for l in range(num_layers): + print( + f"Layer {l} sender buffer id before insert: {id(sender_kv_caches[l])}" + ) + + sender_store.insert( + [b"hash1", b"hash2"], + [ + [ + kv_cache_store.RaidenBlockID( + sender_raiden_id, 2, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_raiden_id, 3, kv_cache_store.BlockStatus.HOST + ) + ], + ], + on_host=True, + ) + print("Inserted into sender.") + for l in range(num_layers): + print( + f"Layer {l} sender buffer id after insert: {id(sender_kv_caches[l])}" + ) + + # NOTE: We simulate global registry lookup here by inserting REMOTE status directly. + # When write / delete through in KVCacheStore is implemented, we can test with a real global registry. + # Insert placeholders into receiver (REMOTE status) + receiver_store.insert( + [b"hash1", b"hash2"], + [ + [ + kv_cache_store.RaidenBlockID( + sender_raiden_id, 2, kv_cache_store.BlockStatus.REMOTE + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_raiden_id, 3, kv_cache_store.BlockStatus.REMOTE + ) + ], + ], + on_host=True, + ) + + # Manual D2H on sender to populate staging buffer + print("Running manual D2H on sender...") + sender_mgr.d2h([2, 3], [2, 3]).wait() + print("Manual D2H on sender completed.") + + # Trigger Remote Fetch + print("Triggering Remote Fetch...") + receiver_store.fetch_remote([b"hash1", b"hash2"]) + print("Remote Fetch Triggered.") + + # Wait for completion + done = False + for i in range(50): + done_recving, failed_recving, pending_recving = ( + receiver_store.poll_fetch_remote_status() + ) + print( + f"Poll {i}: done_recving={done_recving}," + f" failed_recving={failed_recving}, pending_recving={pending_recving}" + ) + if failed_recving: + self.fail(f"Transfer failed: {failed_recving}") + if done_recving: + done = True + break + time.sleep(0.1) + self.assertTrue(done, "Receiver did not finish transfer in time") + + # Manual H2D on receiver to populate JAX buffers + print("Running manual H2D on receiver...") + print("Calling receiver_mgr.h2d") + receiver_mgr.h2d([2, 3], [2, 3]).wait() + print("receiver_mgr.h2d finished") + print("Manual H2D on receiver completed.") + + # Give time for H2D copies to complete in background if needed + time.sleep(1) + + for l in range(num_layers): + print(f"Layer {l} receiver buffer id after: {id(receiver_kv_caches[l])}") + + # Verify data + for l in range(num_layers): + receiver_data = np.asarray(receiver_kv_caches[l]) + expected_data = np.zeros(shape, dtype=np.float32) + expected_data[2] = l + 1 + expected_data[3] = l + 1 + + np.testing.assert_array_equal(receiver_data[2], expected_data[2]) + np.testing.assert_array_equal(receiver_data[3], expected_data[3]) + + # [0] and [1] should remain zero + np.testing.assert_array_equal(receiver_data[0], expected_data[0]) + np.testing.assert_array_equal(receiver_data[1], expected_data[1]) + print(f"Layer {l} verification passed! Bytes are exact.") + + # Verify Status Upgrade in Receiver Store + # Verify Status Upgrade in Receiver Store + results = receiver_store.lookup([b"hash1", b"hash2"]) + self.assertEqual(len(results), 2) + for _, blocks in results: + self.assertEqual(len(blocks), 1) + self.assertEqual(blocks[0].status, kv_cache_store.BlockStatus.HOST) + + def test_queue_flow_multi_listener_end_to_end(self): + """Verifies remote fetch flow where both src and dst have multiple listeners.""" + # Robust port picking to avoid derived port collisions + all_picked_ports = set() + + def pick_non_colliding_base(size): + import socket + + while True: + base = portpicker.pick_unused_port() + derived_ports = [base + i for i in range(size)] + if not any(p in all_picked_ports for p in derived_ports): + bound_sockets = [] + success = True + for p in derived_ports: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("localhost", p)) + bound_sockets.append(s) + except socket.error: + success = False + break + for s in bound_sockets: + s.close() + if success: + for p in derived_ports: + all_picked_ports.add(p) + return base + + def pick_non_colliding_single(): + while True: + port = portpicker.pick_unused_port() + if port not in all_picked_ports: + all_picked_ports.add(port) + return port + + sender_listener_port_base = pick_non_colliding_base(2) + receiver_listener_port_base = pick_non_colliding_base(2) + + port1 = sender_listener_port_base + port2 = port1 + 1 # Consecutive + + receiver_port_base = receiver_listener_port_base + + sender_controller_port = pick_non_colliding_single() + receiver_controller_port = pick_non_colliding_single() + + num_blocks = 10 + num_layers = 1 + shape = (num_blocks, 128, 8, 8, 128) + bytes_per_block = 128 * 8 * 8 * 128 * 4 + + # Sender Data + sender_data_a = np.full(shape, fill_value=1.0, dtype=np.float32) + sender_data_b = np.full(shape, fill_value=2.0, dtype=np.float32) + + if self.num_devices > 1: + mid = self.num_devices // 2 + devices_a = self.devices[:mid] + devices_b = self.devices[mid:] + sharding_a = self._make_sharding(devices_a) + sharding_b = self._make_sharding(devices_b) + + sender_buf_a = self._make_device_array(sender_data_a, sharding_a) + sender_buf_b = self._make_device_array(sender_data_b, sharding_b) + receiver_buf_a = self._make_device_array( + np.zeros(shape, dtype=np.float32), sharding_a + ) + receiver_buf_b = self._make_device_array( + np.zeros(shape, dtype=np.float32), sharding_b + ) + + local_shards_per_listener = mid + num_shards_config = self.num_devices + else: + sender_buf_a = self._make_device_array(sender_data_a) + sender_buf_b = self._make_device_array(sender_data_b) + receiver_buf_a = self._make_device_array( + np.zeros(shape, dtype=np.float32) + ) + receiver_buf_b = self._make_device_array( + np.zeros(shape, dtype=np.float32) + ) + local_shards_per_listener = 1 + num_shards_config = 2 + + local_bytes_per_block = bytes_per_block // local_shards_per_listener + + # Sender Manager A (Shard 0) + sender_mgr_a = KVCacheManager( + kv_caches=[sender_buf_a], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=port1, + listener_controller_port=sender_controller_port, + ) + + # Sender Manager B (Shard 1) + sender_mgr_b = KVCacheManager( + kv_caches=[sender_buf_b], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=port2, + listener_controller_port=sender_controller_port, + ) + + # Receiver Manager A (Shard 0) + receiver_mgr_a = KVCacheManager( + kv_caches=[receiver_buf_a], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=receiver_port_base, + listener_controller_port=receiver_controller_port, + ) + + # Receiver Manager B (Shard 1) + receiver_mgr_b = KVCacheManager( + kv_caches=[receiver_buf_b], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=receiver_port_base + 1, + listener_controller_port=receiver_controller_port, + ) + + sender_id = kv_cache_store.RaidenId("sender_multi_job", "0", "kv_cache", 0) + receiver_id = kv_cache_store.RaidenId("receiver_multi_job", "0", "data2", 0) + + sender_remote_config = kv_cache_store.RemoteFetchConfig() + sender_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + sender_remote_config.controller_port = sender_controller_port + sender_remote_config.local_worker_port = port1 # Base + sender_remote_config.bytes_per_block = local_bytes_per_block + sender_remote_config.num_shards = num_shards_config # Total shards + sender_remote_config.num_listeners = 2 + + receiver_remote_config = kv_cache_store.RemoteFetchConfig() + receiver_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + receiver_remote_config.controller_port = receiver_controller_port + receiver_remote_config.local_worker_port = receiver_port_base # Base + receiver_remote_config.bytes_per_block = local_bytes_per_block + receiver_remote_config.num_shards = num_shards_config + receiver_remote_config.num_listeners = 2 + + sender_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=sender_id, + remote_config=sender_remote_config, + ) + + receiver_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=receiver_id, + remote_config=receiver_remote_config, + ) + + # Insert block 1 into sender + sender_store.insert( + [b"remote_hash1"], + [[ + kv_cache_store.RaidenBlockID( + sender_id, 1, kv_cache_store.BlockStatus.HOST + ) + ]], + on_host=True, + ) + + # Manual D2H on sender managers + sender_mgr_a.d2h([1], [1]).wait() + sender_mgr_b.d2h([1], [1]).wait() + + # Simulate Global Registry Lookup Result and Prepare InsertAndPin + allocated_blocks = [5] + + pin_hashes = [b"remote_hash1"] + pin_slices = [ + [ + kv_cache_store.RaidenBlockID( + sender_id, + allocated_blocks[0], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + ] + + success, evicted = receiver_store.insert_and_pin( + pin_hashes, pin_slices, True + ) + self.assertTrue(success) + self.assertEmpty(evicted) + + # Fetch Remote + futures = receiver_store.fetch_remote(pin_hashes) + self.assertEqual(len(futures), 1) + + # Wait for completion + done = False + for i in range(50): + done_recving, failed_recving, pending_recving = ( + receiver_store.poll_fetch_remote_status() + ) + print( + f"Poll {i}: done_recving={done_recving}," + f" failed_recving={failed_recving}, pending_recving={pending_recving}" + ) + if failed_recving: + self.fail(f"Transfer failed: {failed_recving}") + if len(done_recving) == 1: + done = True + break + time.sleep(0.1) + self.assertTrue(done, "Receiver did not finish transfer in time") + + # Manual H2D on receiver managers + receiver_mgr_a.h2d(allocated_blocks, allocated_blocks).wait() + receiver_mgr_b.h2d(allocated_blocks, allocated_blocks).wait() + + # Verify Data + rec_data_a = np.asarray(receiver_buf_a) + np.testing.assert_array_equal( + rec_data_a[allocated_blocks[0]], sender_data_a[1] + ) + + rec_data_b = np.asarray(receiver_buf_b) + np.testing.assert_array_equal( + rec_data_b[allocated_blocks[0]], sender_data_b[1] + ) + + def test_queue_flow_end_to_end(self): + """Verifies complete flow: simulated Lookup, InsertAndPin, and Fetch.""" + sender_listener_port = portpicker.pick_unused_port() + receiver_listener_port = portpicker.pick_unused_port() + sender_controller_port = portpicker.pick_unused_port() + receiver_controller_port = portpicker.pick_unused_port() + + num_blocks = 10 + num_layers = 1 # Match C++ test + shape = (num_blocks, 128, 8, 8, 128) + + # Setup JAX arrays + sender_kv_caches = [] + receiver_kv_caches = [] + + for l in range(num_layers): + # Fill sender with distinct data per block if needed, here just simple fill + sender_data = np.zeros(shape, dtype=np.float32) + sender_data[1] = 1.0 # Remote block 1 + sender_data[2] = 2.0 # Remote block 2 + sender_data[3] = 3.0 # Remote block 3 + sender_data[4] = 4.0 # Remote block 4 + sender_kv_caches.append(self._make_device_array(sender_data)) + + receiver_kv_caches.append( + self._make_device_array(np.zeros(shape, dtype=np.float32)) + ) + + bytes_per_block = 128 * 8 * 8 * 128 * 4 + local_shards_per_listener = self.num_devices + local_bytes_per_block = bytes_per_block // local_shards_per_listener + + sender_mgr = KVCacheManager( + kv_caches=sender_kv_caches, + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=sender_listener_port, + listener_controller_port=sender_controller_port, + ) + + receiver_mgr = KVCacheManager( + kv_caches=receiver_kv_caches, + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=receiver_listener_port, + listener_controller_port=receiver_controller_port, + ) + + sender_id = kv_cache_store.RaidenId("sender_job", "0", "kv_cache", 0) + receiver_id = kv_cache_store.RaidenId("receiver_job", "0", "data2", 0) + + sender_remote_config = kv_cache_store.RemoteFetchConfig() + sender_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + sender_remote_config.controller_port = sender_controller_port + sender_remote_config.local_worker_port = sender_listener_port + sender_remote_config.bytes_per_block = local_bytes_per_block + sender_remote_config.num_shards = self.num_devices + + receiver_remote_config = kv_cache_store.RemoteFetchConfig() + receiver_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + receiver_remote_config.controller_port = receiver_controller_port + receiver_remote_config.local_worker_port = receiver_listener_port + receiver_remote_config.bytes_per_block = local_bytes_per_block + receiver_remote_config.num_shards = self.num_devices + + sender_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=sender_id, + remote_config=sender_remote_config, + ) + + receiver_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=receiver_id, + remote_config=receiver_remote_config, + ) + + # Insert blocks into sender + sender_store.insert( + [b"remote_hash1", b"remote_hash2", b"remote_hash3", b"remote_hash4"], + [ + [ + kv_cache_store.RaidenBlockID( + sender_id, 1, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_id, 2, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_id, 3, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_id, 4, kv_cache_store.BlockStatus.HOST + ) + ], + ], + on_host=True, + ) + + # Insert local blocks into receiver + receiver_store.insert( + [b"local_hash1", b"local_hash2"], + [ + [ + kv_cache_store.RaidenBlockID( + receiver_id, 7, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + receiver_id, 8, kv_cache_store.BlockStatus.HOST + ) + ], + ], + on_host=True, + ) + + # Manual D2H on sender to populate staging buffer + sender_mgr.d2h([1, 2, 3, 4], [1, 2, 3, 4]).wait() + + # Simulate Global Registry Lookup Result and Prepare InsertAndPin + # Allocated local blocks in receiver for remote data + allocated_blocks = [5, 6, 7, 8] + + pin_hashes = [ + b"local_hash1", + b"local_hash2", + b"remote_hash1", + b"remote_hash2", + b"remote_hash3", + b"remote_hash4", + ] + + pin_slices = [ + [ + kv_cache_store.RaidenBlockID( + receiver_id, 7, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + receiver_id, 8, kv_cache_store.BlockStatus.HOST + ) + ], + # Simulated global hits: REMOTE status with local allocated block ID + [ + kv_cache_store.RaidenBlockID( + sender_id, + allocated_blocks[0], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_id, + allocated_blocks[1], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_id, + allocated_blocks[2], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender_id, + allocated_blocks[3], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + ] + + # Call InsertAndPin + success, evicted = receiver_store.insert_and_pin( + pin_hashes, pin_slices, True + ) + self.assertTrue(success) + self.assertEmpty(evicted) + + # Fetch Remote + remote_hashes = [ + b"remote_hash1", + b"remote_hash2", + b"remote_hash3", + b"remote_hash4", + ] + futures = receiver_store.fetch_remote(remote_hashes) + self.assertEqual(len(futures), 4) + + # Wait for completion + done = False + for _ in range(50): + done_recving, failed_recving, pending_recving = ( + receiver_store.poll_fetch_remote_status() + ) + if failed_recving: + self.fail(f"Transfer failed: {failed_recving}") + if len(done_recving) == 4: + done = True + break + time.sleep(0.1) + self.assertTrue(done, "Receiver did not finish transfer in time") + + # Manual H2D on receiver + receiver_mgr.h2d(allocated_blocks, allocated_blocks).wait() + + # Verify Data + for l in range(num_layers): + receiver_data = np.asarray(receiver_kv_caches[l]) + sender_data_np = np.asarray(sender_kv_caches[l]) + np.testing.assert_array_equal( + receiver_data[allocated_blocks[0]], sender_data_np[1] + ) + np.testing.assert_array_equal( + receiver_data[allocated_blocks[1]], sender_data_np[2] + ) + np.testing.assert_array_equal( + receiver_data[allocated_blocks[2]], sender_data_np[3] + ) + np.testing.assert_array_equal( + receiver_data[allocated_blocks[3]], sender_data_np[4] + ) + + # Verify Status Upgrade + for i, hash in enumerate(remote_hashes): + results = receiver_store.lookup([hash]) + self.assertEqual(len(results), 1) + blocks = results[0][1] + self.assertEqual(len(blocks), 1) + self.assertEqual(blocks[0].status, kv_cache_store.BlockStatus.HOST) + self.assertEqual(blocks[0].host_block_id, allocated_blocks[i]) + + def test_queue_flow_multi_remote_end_to_end(self): + """Verifies flow where blocks belong to multiple different remote RaidenIDs.""" + sender1_listener_port = portpicker.pick_unused_port() + sender2_listener_port = portpicker.pick_unused_port() + receiver_listener_port = portpicker.pick_unused_port() + + sender1_controller_port = portpicker.pick_unused_port() + sender2_controller_port = portpicker.pick_unused_port() + receiver_controller_port = portpicker.pick_unused_port() + + num_blocks = 10 + num_layers = 1 + shape = (num_blocks, 128, 8, 8, 128) + + # Sender 1 Data (ones) + sender1_data = np.ones(shape, dtype=np.float32) + sender1_buf = self._make_device_array(sender1_data) + + # Sender 2 Data (twos) + sender2_data = np.full(shape, fill_value=2.0, dtype=np.float32) + sender2_buf = self._make_device_array(sender2_data) + + # Receiver Buffers (zeros) + receiver_buf = self._make_device_array(np.zeros(shape, dtype=np.float32)) + + bytes_per_block = 128 * 8 * 8 * 128 * 4 + + sender1_mgr = KVCacheManager( + kv_caches=[sender1_buf], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=sender1_listener_port, + listener_controller_port=sender1_controller_port, + ) + + sender2_mgr = KVCacheManager( + kv_caches=[sender2_buf], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=sender2_listener_port, + listener_controller_port=sender2_controller_port, + ) + + receiver_mgr = KVCacheManager( + kv_caches=[receiver_buf], + node_id=0, + local_control_port=0, + max_blocks=num_blocks, + num_slots=2, + listener_port=receiver_listener_port, + listener_controller_port=receiver_controller_port, + ) + + bytes_per_block = 128 * 8 * 8 * 128 * 4 + local_shards_per_listener = self.num_devices + local_bytes_per_block = bytes_per_block // local_shards_per_listener + + sender1_id = kv_cache_store.RaidenId("sender1_job", "0", "kv_cache", 0) + sender2_id = kv_cache_store.RaidenId("sender2_job", "0", "kv_cache", 0) + receiver_id = kv_cache_store.RaidenId("receiver_job", "0", "data2", 0) + + sender1_remote_config = kv_cache_store.RemoteFetchConfig() + sender1_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + sender1_remote_config.controller_port = sender1_controller_port + sender1_remote_config.local_worker_port = sender1_listener_port + sender1_remote_config.bytes_per_block = local_bytes_per_block + sender1_remote_config.num_shards = self.num_devices + + sender2_remote_config = kv_cache_store.RemoteFetchConfig() + sender2_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + sender2_remote_config.controller_port = sender2_controller_port + sender2_remote_config.local_worker_port = sender2_listener_port + sender2_remote_config.bytes_per_block = local_bytes_per_block + sender2_remote_config.num_shards = self.num_devices + + receiver_remote_config = kv_cache_store.RemoteFetchConfig() + receiver_remote_config.orchestrator_address = ( + f"localhost:{_orchestrator_port}" + ) + receiver_remote_config.controller_port = receiver_controller_port + receiver_remote_config.local_worker_port = receiver_listener_port + receiver_remote_config.bytes_per_block = local_bytes_per_block + receiver_remote_config.num_shards = self.num_devices + + sender1_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=sender1_id, + remote_config=sender1_remote_config, + ) + + sender2_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=sender2_id, + remote_config=sender2_remote_config, + ) + + receiver_store = kv_cache_store.KVCacheStore( + capacity=10, + global_registry_address="", + raiden_id=receiver_id, + remote_config=receiver_remote_config, + ) + + # Insert blocks into senders + sender1_store.insert( + [b"remote_hash1", b"remote_hash2"], + [ + [ + kv_cache_store.RaidenBlockID( + sender1_id, 1, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender1_id, 2, kv_cache_store.BlockStatus.HOST + ) + ], + ], + on_host=True, + ) + + sender2_store.insert( + [b"remote_hash3", b"remote_hash4"], + [ + [ + kv_cache_store.RaidenBlockID( + sender2_id, 3, kv_cache_store.BlockStatus.HOST + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender2_id, 4, kv_cache_store.BlockStatus.HOST + ) + ], + ], + on_host=True, + ) + + # Manual D2H on senders + sender1_mgr.d2h([1, 2], [1, 2]).wait() + sender2_mgr.d2h([3, 4], [3, 4]).wait() + + # Simulate Global Registry Lookup Result and Prepare InsertAndPin + allocated_blocks = [5, 6, 7, 8] + + pin_hashes = [ + b"remote_hash1", + b"remote_hash2", + b"remote_hash3", + b"remote_hash4", + ] + + pin_slices = [ + [ + kv_cache_store.RaidenBlockID( + sender1_id, + allocated_blocks[0], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender1_id, + allocated_blocks[1], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender2_id, + allocated_blocks[2], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + sender2_id, + allocated_blocks[3], + kv_cache_store.BlockStatus.REMOTE, + ) + ], + ] + + # Call InsertAndPin + success, evicted = receiver_store.insert_and_pin( + pin_hashes, pin_slices, True + ) + self.assertTrue(success) + self.assertEmpty(evicted) + + # Fetch Remote + futures = receiver_store.fetch_remote(pin_hashes) + self.assertEqual(len(futures), 4) + + # Wait for completion + done = False + for _ in range(50): + done_recving, failed_recving, pending_recving = ( + receiver_store.poll_fetch_remote_status() + ) + if failed_recving: + self.fail(f"Transfer failed: {failed_recving}") + if len(done_recving) == 4: + done = True + break + time.sleep(0.1) + self.assertTrue(done, "Receiver did not finish transfer in time") + + # Manual H2D on receiver + receiver_mgr.h2d(allocated_blocks, allocated_blocks).wait() + + # Verify Data + rec_data = np.asarray(receiver_buf) + + # Blocks 5, 6 should come from Sender 1 (ones) + np.testing.assert_array_equal( + rec_data[allocated_blocks[0]], sender1_data[1] + ) + np.testing.assert_array_equal( + rec_data[allocated_blocks[1]], sender1_data[2] + ) + + # Blocks 7, 8 should come from Sender 2 (twos) + np.testing.assert_array_equal( + rec_data[allocated_blocks[2]], sender2_data[3] + ) + np.testing.assert_array_equal( + rec_data[allocated_blocks[3]], sender2_data[4] + ) + if __name__ == "__main__": absltest.main() diff --git a/tpu_raiden/api/torch/BUILD b/tpu_raiden/api/torch/BUILD index 4d81ed24..f78f734c 100644 --- a/tpu_raiden/api/torch/BUILD +++ b/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", ], ) diff --git a/tpu_raiden/api/torch/kv_cache_store.py b/tpu_raiden/api/torch/kv_cache_store.py index 95d31d1f..0f6b7629 100644 --- a/tpu_raiden/api/torch/kv_cache_store.py +++ b/tpu_raiden/api/torch/kv_cache_store.py @@ -12,12 +12,38 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Python wrapper for the compiled C++ KVCacheStore.""" - +import ctypes +import enum +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 +class BlockStatus(enum.Enum): + INIT = 0 + REMOTE = 1 + HOST = 2 + HBM = 3 + + class RaidenId: """Wrapper around compiled C++ RaidenId.""" @@ -59,62 +85,237 @@ def __repr__(self) -> str: ) +class RaidenBlockID: + """Wrapper around compiled C++ RaidenBlockID.""" + + def __init__( + self, + raiden_id: RaidenId | None = None, + host_block_id: int = -1, + status: BlockStatus = BlockStatus.INIT, + impl: Any = None, + ): + if impl is not None: + self._impl = impl + else: + if raiden_id is None: + raiden_id = RaidenId() + # Map Python enum to C++ enum + status_val = getattr(_impl.BlockStatus, status.name) + self._impl = _impl.RaidenBlockID( + raiden_id._impl, host_block_id, status_val # pylint: disable=protected-access + ) + + @property + def raiden_id(self) -> RaidenId: + return RaidenId(impl=self._impl.raiden_id) + + @property + def host_block_id(self) -> int: + return self._impl.host_block_id + + @property + def status(self) -> BlockStatus: + return BlockStatus[self._impl.status.name] + + @property + def job_name(self) -> str: + return self.raiden_id.job_name + + @property + def job_replica_id(self) -> str: + return self.raiden_id.job_replica_id + + @property + def data_name(self) -> str: + return self.raiden_id.data_name + + @property + def data_replica_idx(self) -> int: + return self.raiden_id.data_replica_idx + + def __repr__(self) -> str: + return ( + f"RaidenBlockID(raiden_id={self.raiden_id}," + f" host_block_id={self.host_block_id}, status={self.status})" + ) + + 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 = "", + raiden_id: RaidenId | None = None, + ): + raw_raiden_id = _impl.RaidenId() + if raiden_id is not None: + raw_raiden_id = raiden_id._impl # pylint: disable=protected-access + self._impl = _impl.KVCacheStore( + capacity=capacity, + global_registry_address=global_registry_address, + raiden_id=raw_raiden_id, + ) + + @property + def raiden_id(self) -> RaidenId: + """Returns the RaidenId associated with this store.""" + return RaidenId(impl=self._impl.raiden_id) def lookup( self, block_hashes: list[bytes], - ) -> list[tuple[bytes, list[RaidenId]]]: + enable_global: bool = False, + ) -> list[tuple[bytes, list[RaidenBlockID]]]: """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. + A list of tuples containing the block hash and a list of matching + RaidenBlockID 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] + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] final_res.append((hash_val, wrapped_slices)) return final_res def insert( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], on_host: bool, - ) -> tuple[bool, list[tuple[bytes, list[RaidenId]]]]: + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: + """Caches sharded buffers into host-RAM/HBM backing store. + + Args: + block_hashes: Incoming block hashes to insert. + slices: List of sharded buffer metadata corresponding to each block hash. + on_host: Whether the slices are located in host memory. + + Returns: + A tuple containing: + - bool: whether all blocks were successfully inserted (i.e. none already + existed). + - list: list of entries evicted from the LRU cache during insertion. + """ raw_slices = [] for slice_list in slices: - raw_slices.append( - [s._impl for s in slice_list] # pylint: disable=protected-access - ) + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_slices.append(converted_list) all_inserted, raw_evicted = self._impl.insert( block_hashes, raw_slices, on_host ) wrapped_evicted = [] for hash_val, raw_slices in raw_evicted: - wrapped_slices = [RaidenId(impl=rs) for rs in raw_slices] + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] + wrapped_evicted.append((hash_val, wrapped_slices)) + return all_inserted, wrapped_evicted + + def insert_and_pin( + self, + block_hashes: list[bytes], + slices: list[list[RaidenBlockID]], + on_host: bool, + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: + """Pins existing block hashes, and inserts/pins new block hashes. + + Pins all existing block hashes, and inserts and pins new block hashes if + there is sufficient available space in the LRU cache. + + Args: + block_hashes: Incoming block hashes to insert and pin. + slices: List of sharded buffer metadata corresponding to each block hash. + on_host: Whether the slices are located in host memory. + + Returns: + A tuple containing: + - bool: whether the entire insert_and_pin operation succeeded (i.e. all + existing keys were pinned, all new keys inserted and pinned). + - list: list of entries evicted during insertion. + """ + raw_slices = [] + for slice_list in slices: + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_slices.append(converted_list) + all_inserted, raw_evicted = self._impl.insert_and_pin( + block_hashes, raw_slices, on_host + ) + wrapped_evicted = [] + for hash_val, raw_slices in raw_evicted: + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] wrapped_evicted.append((hash_val, wrapped_slices)) return all_inserted, wrapped_evicted + def release_and_delete( + self, + block_hashes: list[bytes], + pending_evict_entries: ( + list[tuple[bytes, list[RaidenBlockID]]] | None + ) = None, + ) -> tuple[int, list[tuple[bytes, list[RaidenBlockID]]]]: + """Reverts an insert_and_pin operation. + + Unpins all block_hashes in the LRU cache, deletes any block_hash in REMOTE + status whose pin count is 0, and puts back evicted entries in reverse order + for each deleted remote block. + + Args: + block_hashes: Incoming block hashes to unpin and check for deletion. + pending_evict_entries: List of entries previously evicted during + insert_and_pin. + + Returns: + A tuple containing: + - int: number of remote blocks deleted. + - list: remaining evicted entries that were not restored. + """ + if pending_evict_entries is None: + pending_evict_entries = [] + raw_evicted_in = [] + for hash_val, slice_list in pending_evict_entries: + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_evicted_in.append((hash_val, converted_list)) + del_count, raw_evicted_out = self._impl.release_and_delete( + block_hashes, raw_evicted_in + ) + wrapped_evicted_out = [] + for hash_val, raw_slices in raw_evicted_out: + wrapped_slices = [RaidenBlockID(impl=rs) for rs in raw_slices] + wrapped_evicted_out.append((hash_val, wrapped_slices)) + return del_count, wrapped_evicted_out + def delete( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], ) -> None: raw_slices = [] for slice_list in slices: - raw_slices.append( - [s._impl for s in slice_list] # pylint: disable=protected-access - ) + converted_list = [] + for s in slice_list: + if isinstance(s, RaidenId): + s = RaidenBlockID(raiden_id=s) + converted_list.append(s._impl) # pylint: disable=protected-access + raw_slices.append(converted_list) self._impl.delete(block_hashes, raw_slices) def capacity(self) -> int: diff --git a/tpu_raiden/api/torch/kv_cache_store_test.py b/tpu_raiden/api/torch/kv_cache_store_test.py new file mode 100644 index 00000000..8aea5a09 --- /dev/null +++ b/tpu_raiden/api/torch/kv_cache_store_test.py @@ -0,0 +1,276 @@ +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# 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.torch import kv_cache_store + + +class KVCacheStoreTest(absltest.TestCase): + + def test_basic_tests(self): + controller = kv_cache_store.KVCacheStore(capacity=20) + self.assertEqual(controller.capacity(), 20) + + hashes = [b"6001", b"6002"] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + + # 1. Insert + self.assertTrue(controller.insert(hashes, slices, True)[0]) + self.assertFalse( + controller.insert(hashes, slices, True)[0] + ) # Already exists + + # 2. Lookup with a partial miss at the end + hashes_with_miss = [b"6001", b"6002", b"6003"] + lookup_res = controller.lookup(hashes_with_miss) + self.assertLen(lookup_res, 2) + self.assertEqual(lookup_res[0][0], b"6001") + self.assertLen(lookup_res[0][1], 1) + self.assertEqual(lookup_res[0][1][0].job_name, "inference_server") + self.assertEqual(lookup_res[0][1][0].job_replica_id, "0") + + # Lookup with an early miss + hashes_early_miss = [b"6001", b"6003", b"6002"] + lookup_res_early = controller.lookup(hashes_early_miss) + self.assertLen(lookup_res_early, 1) + self.assertEqual(lookup_res_early[0][0], b"6001") + + # 3. Delete + controller.delete(hashes, slices) + self.assertTrue( + controller.insert(hashes, slices, True)[0] + ) # Successful again + + def test_pin_and_release(self): + controller = kv_cache_store.KVCacheStore(capacity=2) + + hashes = [b"7001", b"7002"] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + # Pin both + self.assertTrue(controller.pin(hashes)) + + # Inserting a third element should fail to evict because both items are + # pinned. + hash_3 = [b"7003"] + slice_3 = [ + [kv_cache_store.RaidenId("inference_server", "2", "kv_cache", 0)] + ] + controller.insert(hash_3, slice_3, True) + + # Release 7001 + controller.release([b"7001"]) + + # Now inserting a fourth element (7004) should successfully evict 7001 + hash_4 = [b"7004"] + slice_4 = [ + [kv_cache_store.RaidenId("inference_server", "3", "kv_cache", 0)] + ] + controller.insert(hash_4, slice_4, True) + + self.assertEmpty(controller.lookup([b"7001", b"7002"])) + self.assertLen(controller.lookup([b"7002"]), 1) + + def test_partial_pin_rollback(self): + controller = kv_cache_store.KVCacheStore(capacity=2) + + hashes = [b"8001", b"8002"] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + # Attempt to pin a sequence with a missing hash (8003). + self.assertFalse(controller.pin([b"8001", b"8002", b"8003"])) + + # Now inserting two new items (8004, 8005) should successfully evict 8001 + # and 8002 because their pins were completely rolled back! + self.assertTrue( + controller.insert( + [b"8004", b"8005"], + [ + [ + kv_cache_store.RaidenId( + "inference_server", "2", "kv_cache", 0 + ) + ], + [ + kv_cache_store.RaidenId( + "inference_server", "3", "kv_cache", 0 + ) + ], + ], + True, + )[0] + ) + + self.assertEmpty(controller.lookup([b"8001", b"8002"])) + self.assertLen(controller.lookup([b"8004", b"8005"]), 2) + + def test_large_and_arbitrary_length_hashes(self): + controller = kv_cache_store.KVCacheStore(capacity=5) + + # Test both high-bit 8-byte hash and a very long arbitrary length hash + large_hash = b"\xff" * 8 + long_hash = b"a" * 100 + hashes = [large_hash, long_hash] + slices = [ + [kv_cache_store.RaidenId("inference_server", "0", "kv_cache", 0)], + [kv_cache_store.RaidenId("inference_server", "1", "kv_cache", 0)], + ] + + self.assertTrue(controller.insert(hashes, slices, True)[0]) + + lookup_res = controller.lookup(hashes) + self.assertLen(lookup_res, 2) + 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.RaidenBlockID( + 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.RaidenBlockID( + kv_cache_store._impl.RaidenId("10.0.0.1:1234", "0", "kv_cache", 42) + ) + remote_id2 = kv_cache_store._impl.RaidenBlockID( + 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) + + def test_insert_and_pin_release_and_delete(self): + controller = kv_cache_store.KVCacheStore(capacity=2) + + local_hashes = [b"local_1", b"local_2"] + local_slices = [ + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("local_job", "0", "kv_cache", 0), + -1, + kv_cache_store.BlockStatus.HOST, + ) + ], + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("local_job", "0", "kv_cache", 1), + -1, + kv_cache_store.BlockStatus.HOST, + ) + ], + ] + self.assertTrue(controller.insert(local_hashes, local_slices, True)[0]) + + remote_hashes = [b"remote_1", b"remote_2"] + remote_slices = [ + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("remote_job", "0", "kv_cache", 0), + -1, + kv_cache_store.BlockStatus.REMOTE, + ) + ], + [ + kv_cache_store.RaidenBlockID( + kv_cache_store.RaidenId("remote_job", "0", "kv_cache", 1), + -1, + kv_cache_store.BlockStatus.REMOTE, + ) + ], + ] + success, evicted = controller.insert_and_pin( + remote_hashes, remote_slices, True + ) + self.assertTrue(success) + self.assertLen(evicted, 2) + self.assertEmpty(controller.lookup([b"local_1"])) + + del_count, rem_evicted = controller.release_and_delete( + remote_hashes, evicted + ) + self.assertEqual(del_count, 2) + self.assertEmpty(rem_evicted) + self.assertLen(controller.lookup([b"local_1", b"local_2"]), 2) + + +if __name__ == "__main__": + absltest.main() diff --git a/tpu_raiden/core/kv_cache_manager_with_transfer.cc b/tpu_raiden/core/kv_cache_manager_with_transfer.cc index fbc8663b..bcc7aad3 100644 --- a/tpu_raiden/core/kv_cache_manager_with_transfer.cc +++ b/tpu_raiden/core/kv_cache_manager_with_transfer.cc @@ -1746,6 +1746,10 @@ absl::Status KVCacheManagerWithTransfer::OnBlocksReceived( bool found = false; std::vector accumulated_host_blocks; + // Forward to base class to notify listeners (e.g. RaidenController) + TF_RETURN_IF_ERROR( + kv_cache::KVCacheManagerBase::OnBlocksReceived(block_ids, uuid)); + { std::lock_guard lock(mu_); auto it = active_recv_entries_.find(uuid); @@ -1784,10 +1788,6 @@ absl::Status KVCacheManagerWithTransfer::OnBlocksReceived( } } - if (!found) { - // Forward to base class for direct pull operations - return RaidenManagerBase::OnBlocksReceived(block_ids, uuid); - } { std::lock_guard lock(mu_); diff --git a/tpu_raiden/frameworks/jax/BUILD b/tpu_raiden/frameworks/jax/BUILD index 0f917af2..d61cc661 100644 --- a/tpu_raiden/frameworks/jax/BUILD +++ b/tpu_raiden/frameworks/jax/BUILD @@ -44,6 +44,7 @@ cc_library( "//tpu_raiden/core:status_macros", "//tpu_raiden/core:tpu_utils", "//tpu_raiden/core:utils", + "//tpu_raiden/kv_cache:kv_cache_listener", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", @@ -69,6 +70,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":mock_nanobind", + "//tpu_raiden/kv_cache:kv_cache_listener", "@jax//jaxlib:py_client", "@nanobind", "@rules_python//python/cc:current_py_cc_headers", @@ -136,6 +138,7 @@ nanobind_extension( ":raw_transfer_jax", ":weight_synchronizer_jax", "//tpu_raiden/core:raw_transfer_core", + "//tpu_raiden/kv_cache:kv_cache_listener", "//tpu_raiden/kv_cache:kv_cache_store", "//tpu_raiden/weight_sync:weight_synchronizer_base", "@com_google_absl//absl/status", diff --git a/tpu_raiden/frameworks/jax/kv_cache_manager.cc b/tpu_raiden/frameworks/jax/kv_cache_manager.cc index 66c1483c..70188f13 100644 --- a/tpu_raiden/frameworks/jax/kv_cache_manager.cc +++ b/tpu_raiden/frameworks/jax/kv_cache_manager.cc @@ -45,7 +45,9 @@ #include "tpu_raiden/core/status_macros.h" #include "tpu_raiden/core/tpu_utils.h" #include "tpu_raiden/core/utils.h" +#include "tpu_raiden/kv_cache/kv_cache_listener.h" #ifndef WITHOUT_PYTHON + #include "tpu_raiden/frameworks/jax/utils.h" namespace nb = nanobind; @@ -166,15 +168,20 @@ KVCacheManager::KVCacheManager(UnpackedCache&& cache, KVCacheManager::KVCacheManager(nanobind::list kv_caches, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, - bool unsafe_skip_buffer_lock, int parallelism) + bool unsafe_skip_buffer_lock, int parallelism, + std::optional listener_port, + std::optional listener_controller_port) : KVCacheManager(UnpackAndMove(std::move(kv_caches)), node_id, local_control_port, max_blocks, num_slots, timeout_s, - unsafe_skip_buffer_lock, parallelism) {} + unsafe_skip_buffer_lock, parallelism, listener_port, + listener_controller_port) {} KVCacheManager::KVCacheManager(UnpackedCache&& cache, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, - bool unsafe_skip_buffer_lock, int parallelism) + bool unsafe_skip_buffer_lock, int parallelism, + std::optional listener_port, + std::optional listener_controller_port) : device_arrays_(std::move(cache.device_arrays)) { const char* enable_metrics_env = std::getenv("ENABLE_RAIDEN_METRICS"); if (enable_metrics_env != nullptr && @@ -183,8 +190,10 @@ KVCacheManager::KVCacheManager(UnpackedCache&& cache, int64_t node_id, } InitSubManagers(cache.layer_buffers, std::nullopt, std::nullopt, unsafe_skip_buffer_lock, parallelism, node_id, - local_control_port, max_blocks, num_slots, timeout_s); + local_control_port, max_blocks, num_slots, timeout_s, + listener_port, listener_controller_port); } + #endif KVCacheManager::KVCacheManager(size_t num_layers, size_t num_shards, @@ -220,7 +229,8 @@ void KVCacheManager::InitSubManagers( std::optional local_port, std::optional host_blocks_to_allocate, bool unsafe_skip_buffer_lock, int parallelism, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, - double timeout_s) { + double timeout_s, std::optional listener_port, + std::optional listener_controller_port) { if (layer_buffers.empty()) return; size_t num_layers = layer_buffers.size(); total_num_shards_ = layer_buffers[0].size(); @@ -349,6 +359,18 @@ void KVCacheManager::InitSubManagers( if (metrics_collector_) { sub_mgr->SetMetricsCollector(metrics_collector_); } + + if (listener_port.has_value()) { + int sub_listener_port = *listener_port + sub_idx; + int ctrl_port = listener_controller_port.value_or(local_control_port); + listeners_.push_back(std::make_unique( + sub_mgr.get(), sub_listener_port, ctrl_port)); + + LOG(INFO) << "Started KVCacheListener for sub-manager " << sub_idx + << " on port " << sub_listener_port + << ", reporting to controller port " << ctrl_port; + } + } catch (const std::exception& e) { // A consecutive (or requested) port was unavailable. Retry the whole // allocation from a fresh ephemeral base; for fixed ports, propagate. @@ -390,6 +412,37 @@ int64_t KVCacheManager::node_id() const { return sub_managers_.empty() ? 0 : sub_managers_[0]->node_id(); } +std::optional KVCacheManager::listener_port() const { + if (!listeners_.empty()) { + return listeners_[0]->listener_port(); + } + return std::nullopt; +} + +bool KVCacheManager::is_listener_active() const { + for (const auto& l : listeners_) { + if (l->is_active()) return true; + } + return false; +} + +std::string KVCacheManager::transfer_address() const { + auto port = local_port(); + if (!port.has_value()) return ""; + // Assuming local_ip() is available or deduce from sub_managers_ + return sub_managers_.empty() + ? "" + : sub_managers_[0]->local_ip() + ":" + std::to_string(*port); +} + +std::string KVCacheManager::listener_address() const { + auto port = listener_port(); + if (!port.has_value()) return ""; + return sub_managers_.empty() + ? "" + : sub_managers_[0]->local_ip() + ":" + std::to_string(*port); +} + uint8_t* KVCacheManager::GetHostPointer(size_t layer_idx, size_t shard_idx) { if (shard_idx >= global_shard_to_submanager_.size()) return nullptr; auto [sub_idx, local_shard] = global_shard_to_submanager_[shard_idx]; diff --git a/tpu_raiden/frameworks/jax/kv_cache_manager.h b/tpu_raiden/frameworks/jax/kv_cache_manager.h index 86be28bc..60809554 100644 --- a/tpu_raiden/frameworks/jax/kv_cache_manager.h +++ b/tpu_raiden/frameworks/jax/kv_cache_manager.h @@ -40,6 +40,8 @@ namespace tpu_raiden { class MetricsCollector; namespace kv_cache { +class KVCacheListener; + namespace jax { struct UnpackedCache { @@ -68,7 +70,10 @@ class KVCacheManager { KVCacheManager(nanobind::list kv_caches, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, - bool unsafe_skip_buffer_lock, int parallelism); + bool unsafe_skip_buffer_lock, int parallelism, + std::optional listener_port = std::nullopt, + std::optional listener_controller_port = std::nullopt); + #endif // FFI metadata constructor (cache-only by default) @@ -114,7 +119,13 @@ class KVCacheManager { std::vector get_local_endpoints() const; + std::optional listener_port() const; + bool is_listener_active() const; + std::string transfer_address() const; + std::string listener_address() const; + void SetSubmanagerShardsForTesting( + const std::vector>& assignment) { submanager_to_global_shards_ = assignment; } @@ -187,7 +198,9 @@ class KVCacheManager { KVCacheManager(UnpackedCache&& cache, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, double timeout_s, - bool unsafe_skip_buffer_lock, int parallelism); + bool unsafe_skip_buffer_lock, int parallelism, + std::optional listener_port = std::nullopt, + std::optional listener_controller_port = std::nullopt); std::optional device_arrays_; #endif @@ -197,13 +210,16 @@ class KVCacheManager { std::optional local_port, std::optional host_blocks_to_allocate, bool unsafe_skip_buffer_lock, int parallelism, int64_t node_id, int64_t local_control_port, int64_t max_blocks, int64_t num_slots, - double timeout_s); + double timeout_s, std::optional listener_port = std::nullopt, + std::optional listener_controller_port = std::nullopt); static constexpr uint64_t k48BitMask = 0xFFFFFFFFFFFFULL; std::atomic global_seq_counter_{1}; std::vector> sub_managers_; + std::vector> listeners_; std::vector> global_shard_to_submanager_; + std::vector> submanager_to_global_shards_; size_t total_num_shards_ = 0; diff --git a/tpu_raiden/frameworks/jax/kv_cache_store.pyi b/tpu_raiden/frameworks/jax/kv_cache_store.pyi index 074799c5..6cf949d0 100644 --- a/tpu_raiden/frameworks/jax/kv_cache_store.pyi +++ b/tpu_raiden/frameworks/jax/kv_cache_store.pyi @@ -1,5 +1,12 @@ +import enum from typing import Any +class BlockStatus(enum.Enum): + INIT = ... + REMOTE = ... + HOST = ... + HBM = ... + class RaidenId: job_name: str job_replica_id: str @@ -13,29 +20,94 @@ class RaidenId: data_replica_idx: int = 0, ) -> None: ... +class RaidenBlockID: + raiden_id: RaidenId + host_block_id: int + status: BlockStatus + def __init__( + self, + raiden_id: RaidenId = ..., + host_block_id: int = ..., + status: BlockStatus = ..., + ) -> None: ... + +class RemoteFetchConfig: + orchestrator_address: str + controller_port: int + local_worker_port: int + bytes_per_block: int + num_shards: int + def __init__(self) -> None: ... + +class FetchFuture: + def Await(self) -> None: ... + def IsDone(self) -> bool: ... + class KVCacheStore: def __init__( self, capacity: int, + global_registry_address: str = '', + raiden_id: RaidenId = ..., + remote_config: RemoteFetchConfig | None = None, ) -> None: ... + + @property + def raiden_id(self) -> RaidenId: ... + def lookup( self, block_hashes: list[bytes], - ) -> list[tuple[bytes, list[RaidenId]]]: + enable_global: bool = False, + ) -> list[tuple[bytes, list[RaidenBlockID]]]: """Checks the LRU directory for cached block hashes. Returns a list of all matched replica pairs prior to the first miss.""" ... def insert( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], on_host: bool, - ) -> tuple[bool, list[tuple[bytes, list[RaidenId]]]]: + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: """Caches sharded buffers into host-RAM/HBM backing store.""" ... + def insert_and_pin( + self, + block_hashes: list[bytes], + slices: list[list[RaidenBlockID]], + on_host: bool, + ) -> tuple[bool, list[tuple[bytes, list[RaidenBlockID]]]]: + """Pins existing block hashes, and inserts/pins new block hashes. + + Pins all existing block hashes, and inserts and pins new block hashes if + there is sufficient available space in the LRU cache. + + Returns: + - bool: whether the entire insert_and_pin operation succeeded (i.e. all + existing keys were pinned, all new keys inserted and pinned). + - list: list of entries evicted during insertion. + """ + ... + def release_and_delete( + self, + block_hashes: list[bytes], + pending_evict_entries: list[tuple[bytes, list[RaidenBlockID]]] + | None = ..., + ) -> tuple[int, list[tuple[bytes, list[RaidenBlockID]]]]: + """Reverts an insert_and_pin operation. + + Unpins all block_hashes in the LRU cache, deletes any block_hash in REMOTE + status whose pin count is 0, and puts back evicted entries in reverse order + for each deleted remote block. + + Returns: + - int: number of remote blocks deleted. + - list: remaining evicted entries that were not restored. + """ + ... def delete( self, block_hashes: list[bytes], - slices: list[list[RaidenId]], + slices: list[list[RaidenBlockID]], ) -> None: """Deletes cached sharded buffers from host-RAM/HBM backing store entirely.""" ... @@ -48,3 +120,9 @@ class KVCacheStore: def release(self, block_hashes: list[bytes]) -> None: """Releases previously pinned block hashes, making them eligible for LRU eviction when capacity is exceeded.""" ... + def fetch_remote(self, block_hashes: list[bytes]) -> dict[bytes, FetchFuture]: + """Initiates remote fetch for the given block hashes.""" + ... + def poll_fetch_remote_status(self) -> tuple[list[bytes], list[bytes], list[bytes]]: + """Polls the status of remote fetches. Returns (done, failed, pending) block hashes.""" + ... diff --git a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc index a5738e11..f262e37f 100644 --- a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc +++ b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc @@ -29,14 +29,15 @@ #include // IWYU pragma: keep #include // IWYU pragma: keep #include // IWYU pragma: keep +#include "xla/pjrt/status_casters.h" #include "tpu_raiden/core/raiden_future.h" #include "tpu_raiden/core/raw_transfer_core.h" #include "tpu_raiden/frameworks/jax/kv_cache_manager.h" #include "tpu_raiden/frameworks/jax/nb_statusor.h" // IWYU pragma: keep #include "tpu_raiden/frameworks/jax/raw_transfer_internal.h" #include "tpu_raiden/frameworks/jax/weight_synchronizer.h" +#include "tpu_raiden/kv_cache/kv_cache_listener.h" #include "tpu_raiden/kv_cache/kv_cache_store.h" -#include "xla/pjrt/status_casters.h" namespace nb = nanobind; @@ -58,8 +59,13 @@ std::vector ToStdStringVector( class KVCacheStoreWrapper { public: - explicit KVCacheStoreWrapper(size_t lru_capacity) { - controller_ = std::make_unique(lru_capacity); + explicit KVCacheStoreWrapper( + size_t lru_capacity, std::string global_registry_address = "", + RaidenId raiden_id = {}, + std::optional remote_config = std::nullopt) { + controller_ = std::make_unique( + lru_capacity, global_registry_address, std::move(raiden_id), + std::move(remote_config)); } KVCacheStore* operator->() { return controller_.get(); } KVCacheStore& operator*() { return *controller_; } @@ -110,12 +116,13 @@ NB_MODULE(_tpu_raiden_jax, m) { nb::arg("unsafe_skip_buffer_lock") = false, nb::arg("parallelism") = 1) .def(nb::init(), + bool, int, std::optional, std::optional>(), nb::arg("kv_caches"), nb::arg("node_id") = 0, nb::arg("local_control_port"), nb::arg("max_blocks"), nb::arg("num_slots"), nb::arg("timeout_s") = 120.0, nb::arg("unsafe_skip_buffer_lock") = true, - nb::arg("parallelism") = 4) + nb::arg("parallelism") = 4, nb::arg("listener_port") = nb::none(), + nb::arg("listener_controller_port") = nb::none()) // Use lambdas to wrap the returned raiden::PjRtCopyFuture into // KVCacheManagerFuture @@ -210,7 +217,15 @@ NB_MODULE(_tpu_raiden_jax, m) { .def_prop_ro( "local_control_port", &tpu_raiden::kv_cache::jax::KVCacheManager::local_control_port) + .def_prop_ro( + "is_listener_active", + &tpu_raiden::kv_cache::jax::KVCacheManager::is_listener_active) + .def_prop_ro("transfer_address", + &tpu_raiden::kv_cache::jax::KVCacheManager::transfer_address) + .def_prop_ro("listener_address", + &tpu_raiden::kv_cache::jax::KVCacheManager::listener_address) .def("get_local_endpoints", + [](const tpu_raiden::kv_cache::jax::KVCacheManager& self) { auto eps = self.get_local_endpoints(); nb::list py_eps; @@ -396,28 +411,82 @@ NB_MODULE(_tpu_raiden_jax, m) { // ========================================================================= nb::class_(m, "RaidenId") .def(nb::init(), - nb::arg("job_name"), nb::arg("job_replica_id") = "", - nb::arg("data_name"), nb::arg("data_replica_idx") = 0) + nb::arg("job_name") = "", nb::arg("job_replica_id") = "", + nb::arg("data_name") = "", nb::arg("data_replica_idx") = 0) .def_rw("job_name", &tpu_raiden::kv_cache::RaidenId::job_name) .def_rw("job_replica_id", &tpu_raiden::kv_cache::RaidenId::job_replica_id) .def_rw("data_name", &tpu_raiden::kv_cache::RaidenId::data_name) .def_rw("data_replica_idx", &tpu_raiden::kv_cache::RaidenId::data_replica_idx); + nb::enum_(m, "BlockStatus") + .value("INIT", tpu_raiden::kv_cache::BlockStatus::INIT) + .value("REMOTE", tpu_raiden::kv_cache::BlockStatus::REMOTE) + .value("HOST", tpu_raiden::kv_cache::BlockStatus::HOST) + .value("HBM", tpu_raiden::kv_cache::BlockStatus::HBM); + + nb::class_(m, "RaidenBlockID") + .def(nb::init(), + nb::arg("raiden_id") = tpu_raiden::kv_cache::RaidenId(), + nb::arg("host_block_id") = -1, + nb::arg("status") = tpu_raiden::kv_cache::BlockStatus::INIT) + .def_rw("raiden_id", &tpu_raiden::kv_cache::RaidenBlockID::raiden_id) + .def_rw("host_block_id", + &tpu_raiden::kv_cache::RaidenBlockID::host_block_id) + .def_rw("status", &tpu_raiden::kv_cache::RaidenBlockID::status); + + nb::class_(m, "RemoteFetchConfig") + .def(nb::init<>()) + .def_rw("orchestrator_address", + &tpu_raiden::kv_cache::RemoteFetchConfig::orchestrator_address) + .def_rw("controller_port", + &tpu_raiden::kv_cache::RemoteFetchConfig::controller_port) + .def_rw("local_worker_port", + &tpu_raiden::kv_cache::RemoteFetchConfig::local_worker_port) + .def_rw("bytes_per_block", + &tpu_raiden::kv_cache::RemoteFetchConfig::bytes_per_block) + .def_rw("num_shards", + &tpu_raiden::kv_cache::RemoteFetchConfig::num_shards) + .def_rw("num_listeners", + &tpu_raiden::kv_cache::RemoteFetchConfig::num_listeners); + + nb::class_(m, "FetchFuture") + .def("Await", + [](tpu_raiden::kv_cache::KVCacheStore::FetchFuture& self) { + nb::gil_scoped_release release; + absl::Status status = self.Await(); + if (!status.ok()) { + throw std::runtime_error("Fetch failed: " + + std::string(status.message())); + } + }) + .def("IsDone", &tpu_raiden::kv_cache::KVCacheStore::FetchFuture::IsDone); + nb::class_(m, "KVCacheStore") - .def(nb::init(), nb::arg("capacity")) + .def(nb::init>(), + nb::arg("capacity"), nb::arg("global_registry_address") = "", + nb::arg("raiden_id") = tpu_raiden::kv_cache::RaidenId(), + nb::arg("remote_config") = nb::none()) + .def_prop_ro( + "raiden_id", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self) { + return (*self).raiden_id(); + }, + "Returns the RaidenId associated with this store.") .def( "lookup", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, - const std::vector& block_hashes) { + const std::vector& block_hashes, bool enable_global) { auto hashes = ToStdStringVector(block_hashes); - auto res = self->Lookup(hashes); + auto res = self->Lookup(hashes, enable_global); if (!res.ok()) { throw std::runtime_error("KVCacheStore lookup failed: " + std::string(res.status().message())); } - std::vector>> + std::vector>> py_res; py_res.reserve(res.value().size()); for (const auto& pair : res.value()) { @@ -427,20 +496,41 @@ NB_MODULE(_tpu_raiden_jax, m) { } return py_res; }, - nb::arg("block_hashes"), + nb::arg("block_hashes"), nb::arg("enable_global") = false, "Checks the LRU directory for cached block hashes. Returns a list of " "all matched replica pairs prior to the first miss.") .def( "insert", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, const std::vector& block_hashes, - const std::vector>& - slices, + const std::vector< + std::vector>& slices, bool on_host) { auto hashes = ToStdStringVector(block_hashes); auto res = self->Insert(hashes, slices, on_host); - std::vector>> + std::vector>> + py_evicted; + py_evicted.reserve(res.second.size()); + for (const auto& pair : res.second) { + py_evicted.push_back(std::make_pair( + nb::bytes(pair.first.data(), pair.first.size()), + pair.second)); + } + return std::make_pair(res.first, py_evicted); + }, + nb::arg("block_hashes"), nb::arg("slices"), nb::arg("on_host")) + .def( + "insert_and_pin", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, + const std::vector& block_hashes, + const std::vector< + std::vector>& slices, + bool on_host) { + auto hashes = ToStdStringVector(block_hashes); + auto res = self->InsertAndPin(hashes, slices, on_host); + std::vector>> py_evicted; py_evicted.reserve(res.second.size()); for (const auto& pair : res.second) { @@ -451,12 +541,44 @@ NB_MODULE(_tpu_raiden_jax, m) { return std::make_pair(res.first, py_evicted); }, nb::arg("block_hashes"), nb::arg("slices"), nb::arg("on_host")) + .def( + "release_and_delete", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, + const std::vector& block_hashes, + const std::vector>>& + pending_evict_entries) { + auto hashes = ToStdStringVector(block_hashes); + std::vector>> + evicted; + evicted.reserve(pending_evict_entries.size()); + for (const auto& pair : pending_evict_entries) { + evicted.push_back(std::make_pair( + std::string(pair.first.c_str(), pair.first.size()), + pair.second)); + } + auto res = self->ReleaseAndDelete(hashes, evicted); + std::vector>> + py_rem_evicted; + py_rem_evicted.reserve(res.second.size()); + for (const auto& pair : res.second) { + py_rem_evicted.push_back(std::make_pair( + nb::bytes(pair.first.data(), pair.first.size()), + pair.second)); + } + return std::make_pair(res.first, py_rem_evicted); + }, + nb::arg("block_hashes"), + nb::arg("pending_evict_entries") = std::vector>>()) .def( "delete", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, const std::vector& block_hashes, - const std::vector>& - slices) { + const std::vector< + std::vector>& slices) { auto hashes = ToStdStringVector(block_hashes); self->Delete(hashes, slices); }, @@ -480,5 +602,35 @@ NB_MODULE(_tpu_raiden_jax, m) { auto hashes = ToStdStringVector(block_hashes); self->Release(hashes); }, - nb::arg("block_hashes")); + nb::arg("block_hashes")) + .def( + "fetch_remote", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, + const std::vector& block_hashes) { + auto hashes = ToStdStringVector(block_hashes); + auto res = self->FetchRemote(hashes); + nb::dict py_res; + for (const auto& pair : res) { + py_res[nb::bytes(pair.first.data(), pair.first.size())] = + pair.second; + } + return py_res; + }, + nb::arg("block_hashes")) + .def("poll_fetch_remote_status", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self) { + auto [done, failed, pending] = self->PollFetchRemoteStatus(); + + auto to_py_bytes_list = [](const std::vector& vec) { + nb::list py_list; + for (const auto& s : vec) { + py_list.append(nb::bytes(s.data(), s.size())); + } + return py_list; + }; + + return nb::make_tuple(to_py_bytes_list(done), + to_py_bytes_list(failed), + to_py_bytes_list(pending)); + }); } diff --git a/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc b/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc index 175208ec..1e11ee01 100644 --- a/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc +++ b/tpu_raiden/frameworks/torch/tpu_raiden_torch_module.cc @@ -58,8 +58,11 @@ std::vector ToStdStringVector( class KVCacheStoreWrapper { public: - explicit KVCacheStoreWrapper(size_t lru_capacity) { - controller_ = std::make_unique(lru_capacity); + explicit KVCacheStoreWrapper(size_t lru_capacity, + std::string global_registry_address = "", + RaidenId raiden_id = {}) { + controller_ = std::make_unique( + lru_capacity, global_registry_address, std::move(raiden_id)); } KVCacheStore* operator->() { return controller_.get(); } KVCacheStore& operator*() { return *controller_; } @@ -384,28 +387,53 @@ NB_MODULE(_tpu_raiden_torch, m) { // ========================================================================= nb::class_(m, "RaidenId") .def(nb::init(), - nb::arg("job_name"), nb::arg("job_replica_id") = "", - nb::arg("data_name"), nb::arg("data_replica_idx") = 0) + nb::arg("job_name") = "", nb::arg("job_replica_id") = "", + nb::arg("data_name") = "", nb::arg("data_replica_idx") = 0) .def_rw("job_name", &tpu_raiden::kv_cache::RaidenId::job_name) .def_rw("job_replica_id", &tpu_raiden::kv_cache::RaidenId::job_replica_id) .def_rw("data_name", &tpu_raiden::kv_cache::RaidenId::data_name) .def_rw("data_replica_idx", &tpu_raiden::kv_cache::RaidenId::data_replica_idx); + nb::enum_(m, "BlockStatus") + .value("INIT", tpu_raiden::kv_cache::BlockStatus::INIT) + .value("REMOTE", tpu_raiden::kv_cache::BlockStatus::REMOTE) + .value("HOST", tpu_raiden::kv_cache::BlockStatus::HOST) + .value("HBM", tpu_raiden::kv_cache::BlockStatus::HBM); + + nb::class_(m, "RaidenBlockID") + .def(nb::init(), + nb::arg("raiden_id") = tpu_raiden::kv_cache::RaidenId(), + nb::arg("host_block_id") = -1, + nb::arg("status") = tpu_raiden::kv_cache::BlockStatus::INIT) + .def_rw("raiden_id", &tpu_raiden::kv_cache::RaidenBlockID::raiden_id) + .def_rw("host_block_id", + &tpu_raiden::kv_cache::RaidenBlockID::host_block_id) + .def_rw("status", &tpu_raiden::kv_cache::RaidenBlockID::status); + nb::class_(m, "KVCacheStore") - .def(nb::init(), nb::arg("capacity")) + .def(nb::init(), + nb::arg("capacity"), nb::arg("global_registry_address") = "", + nb::arg("raiden_id") = tpu_raiden::kv_cache::RaidenId()) + .def_prop_ro( + "raiden_id", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self) { + return (*self).raiden_id(); + }, + "Returns the RaidenId associated with this store.") .def( "lookup", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, - const std::vector& block_hashes) { + const std::vector& block_hashes, bool enable_global) { auto hashes = ToStdStringVector(block_hashes); - auto res = self->Lookup(hashes); + auto res = self->Lookup(hashes, enable_global); if (!res.ok()) { throw std::runtime_error("KVCacheStore lookup failed: " + std::string(res.status().message())); } - std::vector>> + std::vector>> py_res; py_res.reserve(res.value().size()); for (const auto& pair : res.value()) { @@ -415,20 +443,41 @@ NB_MODULE(_tpu_raiden_torch, m) { } return py_res; }, - nb::arg("block_hashes"), + nb::arg("block_hashes"), nb::arg("enable_global") = false, "Checks the LRU directory for cached block hashes. Returns a list of " "all matched replica pairs prior to the first miss.") .def( "insert", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, const std::vector& block_hashes, - const std::vector>& - slices, + const std::vector< + std::vector>& slices, bool on_host) { auto hashes = ToStdStringVector(block_hashes); auto res = self->Insert(hashes, slices, on_host); - std::vector>> + std::vector>> + py_evicted; + py_evicted.reserve(res.second.size()); + for (const auto& pair : res.second) { + py_evicted.push_back(std::make_pair( + nb::bytes(pair.first.data(), pair.first.size()), + pair.second)); + } + return std::make_pair(res.first, py_evicted); + }, + nb::arg("block_hashes"), nb::arg("slices"), nb::arg("on_host")) + .def( + "insert_and_pin", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, + const std::vector& block_hashes, + const std::vector< + std::vector>& slices, + bool on_host) { + auto hashes = ToStdStringVector(block_hashes); + auto res = self->InsertAndPin(hashes, slices, on_host); + std::vector>> py_evicted; py_evicted.reserve(res.second.size()); for (const auto& pair : res.second) { @@ -439,12 +488,44 @@ NB_MODULE(_tpu_raiden_torch, m) { return std::make_pair(res.first, py_evicted); }, nb::arg("block_hashes"), nb::arg("slices"), nb::arg("on_host")) + .def( + "release_and_delete", + [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, + const std::vector& block_hashes, + const std::vector>>& + pending_evict_entries) { + auto hashes = ToStdStringVector(block_hashes); + std::vector>> + evicted; + evicted.reserve(pending_evict_entries.size()); + for (const auto& pair : pending_evict_entries) { + evicted.push_back(std::make_pair( + std::string(pair.first.c_str(), pair.first.size()), + pair.second)); + } + auto res = self->ReleaseAndDelete(hashes, evicted); + std::vector>> + py_rem_evicted; + py_rem_evicted.reserve(res.second.size()); + for (const auto& pair : res.second) { + py_rem_evicted.push_back(std::make_pair( + nb::bytes(pair.first.data(), pair.first.size()), + pair.second)); + } + return std::make_pair(res.first, py_rem_evicted); + }, + nb::arg("block_hashes"), + nb::arg("pending_evict_entries") = std::vector>>()) .def( "delete", [](tpu_raiden::kv_cache::KVCacheStoreWrapper& self, const std::vector& block_hashes, - const std::vector>& - slices) { + const std::vector< + std::vector>& slices) { auto hashes = ToStdStringVector(block_hashes); self->Delete(hashes, slices); }, diff --git a/tpu_raiden/kv_cache/BUILD b/tpu_raiden/kv_cache/BUILD index d66da5e8..70f89e6f 100644 --- a/tpu_raiden/kv_cache/BUILD +++ b/tpu_raiden/kv_cache/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") +load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test") load("@rules_python//python:defs.bzl", "py_library") package(default_visibility = ["//visibility:public"]) @@ -160,10 +160,73 @@ cc_test( ], ) +cc_test( + name = "raiden_orchestrator_test", + srcs = ["raiden_orchestrator_test.cc"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + deps = [ + ":raiden_orchestrator", + "//net/util:ports", + "//tpu_raiden/rpc:raiden_service_cc_proto", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "raiden_orchestrator", + srcs = ["raiden_orchestrator.cc"], + hdrs = ["raiden_orchestrator.h"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + visibility = ["//visibility:public"], + deps = [ + ":kv_cache_store", + "//tpu_raiden/rpc:raiden_service_cc_proto", + "//tpu_raiden/transport:raw_buffer_transport", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/hash", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + ], +) + +cc_binary( + name = "raiden_orchestrator_main", + srcs = ["raiden_orchestrator_main.cc"], + visibility = ["//visibility:public"], + deps = [ + ":raiden_orchestrator", + "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:initialize", + "@com_google_absl//absl/strings", + ], +) + cc_library( name = "kv_cache_store", - srcs = ["kv_cache_store.cc"], - hdrs = ["kv_cache_store.h"], + srcs = [ + "kv_cache_store.cc", + "raiden_controller_embedded.cc", + ], + hdrs = [ + "kv_cache_store.h", + "raiden_controller_embedded.h", + ], copts = [ "-fno-strict-aliasing", "-fexceptions", @@ -173,11 +236,22 @@ cc_library( deps = [ ":lru_cache", "//tpu_raiden/core:raw_transfer_core", + "//tpu_raiden/core:status_macros", + "//tpu_raiden/core:tpu_utils", + "//tpu_raiden/kv_cache/global_registry:global_registry_client_cc", + "//tpu_raiden/rpc:raiden_service_cc_proto", + "//tpu_raiden/rpc:rpc_utils", + "//tpu_raiden/transport:raw_buffer_transport", + "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", "@com_google_absl//absl/types:span", ], ) @@ -191,9 +265,17 @@ cc_test( ], features = ["-use_header_modules"], deps = [ + ":kv_cache_listener", + ":kv_cache_manager_base", ":kv_cache_store", + ":raiden_orchestrator", + "//net/util:ports", "//tpu_raiden/core:raw_transfer_core", + "//tpu_raiden/kv_cache/global_registry:global_registry_client_cc", + "//tpu_raiden/kv_cache/global_registry:global_registry_server_lib", + "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/time", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", ], @@ -222,7 +304,9 @@ cc_library( deps = [ ":kv_cache_manager_base", "//tpu_raiden/rpc:raiden_service_cc_proto", + "//tpu_raiden/rpc:rpc_utils", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", ], ) diff --git a/tpu_raiden/kv_cache/global_registry/BUILD b/tpu_raiden/kv_cache/global_registry/BUILD index 53594166..052d828c 100644 --- a/tpu_raiden/kv_cache/global_registry/BUILD +++ b/tpu_raiden/kv_cache/global_registry/BUILD @@ -56,10 +56,11 @@ cc_binary( srcs = ["global_registry_server_main.cc"], deps = [ ":global_registry_server_lib", - "//base", "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/flags:parse", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/time", ], ) @@ -71,7 +72,7 @@ cc_library( deps = [ ":global_registry_cc_grpc", ":global_registry_cc_proto", - "//third_party/openssl:crypto", + "@boringssl//:crypto", "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/tpu_raiden/kv_cache/global_registry/global_registry_client.cc b/tpu_raiden/kv_cache/global_registry/global_registry_client.cc index f24cbb85..52d6ee1a 100644 --- a/tpu_raiden/kv_cache/global_registry/global_registry_client.cc +++ b/tpu_raiden/kv_cache/global_registry/global_registry_client.cc @@ -26,10 +26,10 @@ #include "absl/strings/escaping.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" -#include "third_party/grpc/include/grpcpp/channel.h" -#include "third_party/grpc/include/grpcpp/client_context.h" -#include "third_party/grpc/include/grpcpp/support/status.h" -#include "third_party/openssl/sha.h" +#include "grpcpp/channel.h" +#include "grpcpp/client_context.h" +#include "grpcpp/support/status.h" +#include #include "tpu_raiden/kv_cache/global_registry/global_registry.grpc.pb.h" #include "tpu_raiden/kv_cache/global_registry/global_registry.pb.h" diff --git a/tpu_raiden/kv_cache/global_registry/global_registry_client.h b/tpu_raiden/kv_cache/global_registry/global_registry_client.h index bebfb370..c814bddf 100644 --- a/tpu_raiden/kv_cache/global_registry/global_registry_client.h +++ b/tpu_raiden/kv_cache/global_registry/global_registry_client.h @@ -26,7 +26,7 @@ #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" -#include "third_party/grpc/include/grpcpp/channel.h" +#include "grpcpp/channel.h" #include "tpu_raiden/kv_cache/global_registry/global_registry.grpc.pb.h" #include "tpu_raiden/kv_cache/global_registry/global_registry.pb.h" diff --git a/tpu_raiden/kv_cache/global_registry/global_registry_server.cc b/tpu_raiden/kv_cache/global_registry/global_registry_server.cc index c7879747..961647e2 100644 --- a/tpu_raiden/kv_cache/global_registry/global_registry_server.cc +++ b/tpu_raiden/kv_cache/global_registry/global_registry_server.cc @@ -25,8 +25,8 @@ #include "absl/synchronization/mutex.h" #include "absl/time/clock.h" #include "absl/time/time.h" -#include "third_party/grpc/include/grpcpp/server_context.h" -#include "third_party/grpc/include/grpcpp/support/status.h" +#include "grpcpp/server_context.h" +#include "grpcpp/support/status.h" #include "tpu_raiden/kv_cache/global_registry/global_registry.pb.h" namespace tpu_raiden { diff --git a/tpu_raiden/kv_cache/global_registry/global_registry_server.h b/tpu_raiden/kv_cache/global_registry/global_registry_server.h index b0926fd3..36748f3b 100644 --- a/tpu_raiden/kv_cache/global_registry/global_registry_server.h +++ b/tpu_raiden/kv_cache/global_registry/global_registry_server.h @@ -26,8 +26,8 @@ #include "absl/container/flat_hash_map.h" #include "absl/synchronization/mutex.h" #include "absl/time/time.h" -#include "third_party/grpc/include/grpcpp/server_context.h" -#include "third_party/grpc/include/grpcpp/support/status.h" +#include "grpcpp/server_context.h" +#include "grpcpp/support/status.h" #include "tpu_raiden/kv_cache/global_registry/global_registry.grpc.pb.h" #include "tpu_raiden/kv_cache/global_registry/global_registry.pb.h" diff --git a/tpu_raiden/kv_cache/global_registry/global_registry_server_main.cc b/tpu_raiden/kv_cache/global_registry/global_registry_server_main.cc index 16400179..c9be872d 100644 --- a/tpu_raiden/kv_cache/global_registry/global_registry_server_main.cc +++ b/tpu_raiden/kv_cache/global_registry/global_registry_server_main.cc @@ -19,13 +19,14 @@ #include #include -#include "base/init_google.h" +#include "absl/flags/parse.h" +#include "absl/log/initialize.h" #include "absl/flags/flag.h" #include "absl/log/log.h" #include "absl/time/time.h" -#include "third_party/grpc/include/grpcpp/security/server_credentials.h" -#include "third_party/grpc/include/grpcpp/server.h" -#include "third_party/grpc/include/grpcpp/server_builder.h" +#include "grpcpp/security/server_credentials.h" +#include "grpcpp/server.h" +#include "grpcpp/server_builder.h" #include "tpu_raiden/kv_cache/global_registry/global_registry_server.h" ABSL_FLAG(int32_t, port, 50051, "Port to listen on"); @@ -61,7 +62,8 @@ void RunServer() { } int main(int argc, char** argv) { - InitGoogle(argv[0], &argc, &argv, true); + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); RunServer(); return 0; } diff --git a/tpu_raiden/kv_cache/global_registry/global_registry_test.cc b/tpu_raiden/kv_cache/global_registry/global_registry_test.cc index 2f4fdadd..1758c956 100644 --- a/tpu_raiden/kv_cache/global_registry/global_registry_test.cc +++ b/tpu_raiden/kv_cache/global_registry/global_registry_test.cc @@ -22,12 +22,12 @@ #include #include "absl/time/clock.h" #include "absl/time/time.h" -#include "third_party/grpc/include/grpcpp/channel.h" -#include "third_party/grpc/include/grpcpp/create_channel.h" -#include "third_party/grpc/include/grpcpp/security/credentials.h" -#include "third_party/grpc/include/grpcpp/security/server_credentials.h" -#include "third_party/grpc/include/grpcpp/server.h" -#include "third_party/grpc/include/grpcpp/server_builder.h" +#include "grpcpp/channel.h" +#include "grpcpp/create_channel.h" +#include "grpcpp/security/credentials.h" +#include "grpcpp/security/server_credentials.h" +#include "grpcpp/server.h" +#include "grpcpp/server_builder.h" #include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" #include "tpu_raiden/kv_cache/global_registry/global_registry_server.h" diff --git a/tpu_raiden/kv_cache/kv_cache_listener.cc b/tpu_raiden/kv_cache/kv_cache_listener.cc index 82ba2f98..bcc19790 100644 --- a/tpu_raiden/kv_cache/kv_cache_listener.cc +++ b/tpu_raiden/kv_cache/kv_cache_listener.cc @@ -29,8 +29,11 @@ #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" #include "tpu_raiden/kv_cache/kv_cache_manager_base.h" #include "tpu_raiden/rpc/raiden_service.pb.h" +#include "tpu_raiden/rpc/rpc_utils.h" namespace tpu_raiden { namespace kv_cache { @@ -38,10 +41,35 @@ namespace kv_cache { using ::tpu_raiden::rpc::ControlRequest; using ::tpu_raiden::rpc::ControlResponse; -KVCacheListener::KVCacheListener(KVCacheManagerBase* engine, - int listener_port) - : engine_(engine), listener_port_(listener_port) { - server_fd_ = socket(AF_INET6, SOCK_STREAM, 0); +KVCacheListener::KVCacheListener(KVCacheManagerBase* engine, int listener_port, + int controller_port) + : engine_(engine), + listener_port_(listener_port), + controller_port_(controller_port) { + if (engine_ && controller_port_ > 0) { + engine_->SetBlocksReceivedCallback( + [this](const std::vector& block_ids, uint64_t uuid) { + ControlRequest req; + req.set_command(ControlRequest::COMMAND_TRANSFER_COMPLETED); + for (int id : block_ids) { + req.add_completed_block_ids(id); + } + + std::string controller_addr = + absl::StrCat("localhost:", controller_port_); + ControlResponse resp; + absl::Status status = rpc::SendRpcSync(controller_addr, req, resp); + if (!status.ok()) { + LOG(WARNING) << "Failed to notify controller at " << controller_addr + << ": " << status; + } else if (!resp.success()) { + LOG(WARNING) << "Controller failed to process completion: " + << resp.message(); + } + }); + } + + server_fd_ = socket(AF_INET, SOCK_STREAM, 0); if (server_fd_ < 0) { LOG(FATAL) << "Failed to create C++ KVCacheListener socket: " << std::strerror(errno); @@ -52,11 +80,10 @@ KVCacheListener::KVCacheListener(KVCacheManagerBase* engine, LOG(WARNING) << "setsockopt SO_REUSEADDR failed"; } - sockaddr_in6 address{ - .sin6_family = AF_INET6, - .sin6_port = htons(listener_port_), - .sin6_addr = in6addr_any, - }; + sockaddr_in address; + address.sin_family = AF_INET; + address.sin_port = htons(listener_port_); + address.sin_addr.s_addr = INADDR_ANY; if (bind(server_fd_, reinterpret_cast(&address), sizeof(address)) < 0) { @@ -72,7 +99,7 @@ KVCacheListener::KVCacheListener(KVCacheManagerBase* engine, socklen_t addr_len = sizeof(address); if (getsockname(server_fd_, reinterpret_cast(&address), &addr_len) == 0) { - listener_port_ = ntohs(address.sin6_port); + listener_port_ = ntohs(address.sin_port); } LOG(INFO) << "Native C++ KVCacheListener actively listening on port: " @@ -84,12 +111,12 @@ KVCacheListener::KVCacheListener(KVCacheManagerBase* engine, KVCacheListener::~KVCacheListener() { stopping_ = true; if (server_fd_ >= 0) { - int sock = socket(AF_INET6, SOCK_STREAM, 0); + int sock = socket(AF_INET, SOCK_STREAM, 0); if (sock >= 0) { - sockaddr_in6 serv_addr{}; - serv_addr.sin6_family = AF_INET6; - serv_addr.sin6_port = htons(listener_port_); - inet_pton(AF_INET6, "::1", &serv_addr.sin6_addr); + sockaddr_in serv_addr{}; + serv_addr.sin_family = AF_INET; + serv_addr.sin_port = htons(listener_port_); + inet_pton(AF_INET, "localhost", &serv_addr.sin_addr); if (connect(sock, reinterpret_cast(&serv_addr), sizeof(serv_addr)) == 0) { ControlRequest req; @@ -201,6 +228,16 @@ void KVCacheListener::ConnectionWorker(int client_fd) { LOG(ERROR) << "WaitForPendingWork failed during shutdown: " << status; } stopping_ = true; + } else if (req.command() == ControlRequest::COMMAND_GET_ENDPOINTS) { + LOG(INFO) << "C++ KVCacheListener received GET_ENDPOINTS command."; + std::string ip = engine_->local_ip(); + int port = engine_->local_port().value_or(0); + std::string endpoint = absl::StrContains(ip, ':') + ? absl::StrCat("[", ip, "]:", port) + : absl::StrCat(ip, ":", port); + for (size_t i = 0; i < engine_->num_shards(); ++i) { + resp.add_endpoints(endpoint); + } } else { resp.set_success(false); resp.set_message("COMMAND_UNSPECIFIED"); diff --git a/tpu_raiden/kv_cache/kv_cache_listener.h b/tpu_raiden/kv_cache/kv_cache_listener.h index f15ea3e8..627b8a8a 100644 --- a/tpu_raiden/kv_cache/kv_cache_listener.h +++ b/tpu_raiden/kv_cache/kv_cache_listener.h @@ -27,7 +27,8 @@ class KVCacheManagerBase; class KVCacheListener final { public: - KVCacheListener(KVCacheManagerBase* engine, int listener_port); + KVCacheListener(KVCacheManagerBase* engine, int listener_port, + int controller_port = -1); ~KVCacheListener(); KVCacheListener(const KVCacheListener&) = delete; @@ -42,6 +43,7 @@ class KVCacheListener final { KVCacheManagerBase* engine_; int listener_port_; + int controller_port_; int server_fd_ = -1; std::atomic stopping_{false}; diff --git a/tpu_raiden/kv_cache/kv_cache_manager_base.cc b/tpu_raiden/kv_cache/kv_cache_manager_base.cc index e5941ada..dcfd5d59 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_base.cc +++ b/tpu_raiden/kv_cache/kv_cache_manager_base.cc @@ -1099,6 +1099,14 @@ void KVCacheManagerBase::RegisterBlockReadinessCallback( cb(absl::OkStatus()); } +absl::Status KVCacheManagerBase::OnBlocksReceived( + const std::vector& block_ids, uint64_t uuid) { + if (blocks_received_cb_) { + blocks_received_cb_(block_ids, uuid); + } + return absl::OkStatus(); +} + absl::Status KVCacheManagerBase::PushKVCacheResharded( const tpu_raiden::rpc::StartTransferRequest& request) { // 1. Register the active plan so GetBlockChunks can use it @@ -1111,10 +1119,7 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( << ": uuid=" << request.uuid() << ", numa=" << numa; } - // 2. D2H to copy from device to host. - ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2h()); - - // 3. Group entries by dst_peer and collect unique block IDs + // 2. Group entries by dst_peer and collect unique block IDs std::map>> peer_transfers; for (const auto& [shard_idx, schedule] : request.shard_push_schedules()) { for (const auto& entry : schedule.entries()) { @@ -1123,13 +1128,7 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( } } - d2h_future.OnReady([this, request, peer_transfers, numa](auto status_or) { - if (!status_or.ok()) { - LOG(ERROR) << "D2H copy failed for resharded push uuid " << request.uuid() - << ": " << status_or.status().ToString(); - return; - } - + auto do_push = [this, request, peer_transfers, numa]() { for (size_t l = 0; l < num_layers_; ++l) { VLOG(1) << "StartPushInternal (H2H start layer " << l << "): uuid=" << request.uuid() << ", numa=" << numa; @@ -1173,6 +1172,31 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( } }); } + }; + + if (request.src_mem_type() == tpu_raiden::rpc::MEMORY_TYPE_DRAM || + buffer_holds_.empty()) { + VLOG(1) << "Skipping D2H for resharded push uuid " << request.uuid() + << " (Src is DRAM or CPU-only manager)"; + do_push(); + return absl::OkStatus(); + } + + for (size_t l = 0; l < num_layers_; ++l) { + VLOG(1) << "StartPushInternal (D2H start) layer " << l + << ": uuid=" << request.uuid() << ", numa=" << numa; + } + + // 3. D2H to copy from device to host. + ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2h()); + + d2h_future.OnReady([request, do_push](auto status_or) { + if (!status_or.ok()) { + LOG(ERROR) << "D2H copy failed for resharded push uuid " << request.uuid() + << ": " << status_or.status().ToString(); + return; + } + do_push(); }); return absl::OkStatus(); @@ -1277,8 +1301,19 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, } else { int found_src_shard = -1; if (sender_node_id != -1) { - found_src_shard = static_cast(sender_node_id); - } else { + auto it = schedules.find(static_cast(sender_node_id)); + if (it != schedules.end()) { + for (const auto& entry : it->second.entries()) { + if (static_cast(entry.dst_shard_idx()) == shard_idx && + static_cast(entry.dst_block_id()) == block_id) { + found_src_shard = static_cast(sender_node_id); + break; + } + } + } + } + + if (found_src_shard == -1) { for (const auto& [src_shard, src_schedule] : schedules) { for (const auto& entry : src_schedule.entries()) { if (static_cast(entry.dst_shard_idx()) == shard_idx && @@ -1319,6 +1354,7 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, } } + for (const auto& chunk : block_resolved_chunks) { if (accumulated_bytes >= total_bytes) break; size_t size_to_add = diff --git a/tpu_raiden/kv_cache/kv_cache_manager_base.h b/tpu_raiden/kv_cache/kv_cache_manager_base.h index 82c01003..0d36f6e1 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_base.h +++ b/tpu_raiden/kv_cache/kv_cache_manager_base.h @@ -102,6 +102,15 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { size_t layer_idx, size_t shard_idx, int block_id, uint64_t uuid, transport::BlockTransportDelegate::HostBlockReadyCallback cb) override; + using BlocksReceivedCallback = + std::function& block_ids, uint64_t uuid)>; + void SetBlocksReceivedCallback(BlocksReceivedCallback cb) { + blocks_received_cb_ = std::move(cb); + } + + absl::Status OnBlocksReceived(const std::vector& block_ids, + uint64_t uuid = 0) override; + // Async on-chip H2D offloads returning PJRT copy future E2E virtual absl::StatusOr H2d( const std::vector& src_offsets_major_dim = {}, @@ -284,6 +293,8 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { absl::flat_hash_map recv_callbacks_ ABSL_GUARDED_BY(recv_mu_); + BlocksReceivedCallback blocks_received_cb_; + mutable absl::Mutex plans_mu_; struct RegisteredPlan { tpu_raiden::rpc::StartTransferRequest request; diff --git a/tpu_raiden/kv_cache/kv_cache_store.cc b/tpu_raiden/kv_cache/kv_cache_store.cc index a410ffd6..35203833 100644 --- a/tpu_raiden/kv_cache/kv_cache_store.cc +++ b/tpu_raiden/kv_cache/kv_cache_store.cc @@ -15,55 +15,126 @@ #include "tpu_raiden/kv_cache/kv_cache_store.h" #include +#include #include #include #include #include +#include "grpcpp/grpcpp.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" +#include "grpcpp/security/credentials.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" #include "tpu_raiden/kv_cache/lru_cache.h" +#include "tpu_raiden/kv_cache/raiden_controller_embedded.h" namespace tpu_raiden { namespace kv_cache { -KVCacheStore::KVCacheStore(size_t capacity) : lru_cache_(capacity) {} +KVCacheStore::KVCacheStore(size_t capacity, std::string global_registry_address, + RaidenId raiden_id, + std::optional remote_config) + : lru_cache_(capacity), raiden_id_(std::move(raiden_id)) { + if (!global_registry_address.empty()) { + auto channel = grpc::CreateChannel(global_registry_address, + grpc::InsecureChannelCredentials()); + registry_client_ = + std::make_unique(channel); + } -KVCacheStore::~KVCacheStore() = default; + if (remote_config.has_value()) { + controller_ = std::make_unique( + this, remote_config->controller_port, + remote_config->orchestrator_address, remote_config->local_worker_port, + std::vector{}, remote_config->bytes_per_block, + remote_config->num_shards, remote_config->num_listeners); -absl::StatusOr>>> -KVCacheStore::Lookup(const std::vector& block_hashes) { - absl::MutexLock lock(mutex_); + absl::Status status = controller_->Start(); + if (!status.ok()) { + LOG(ERROR) << "Failed to start RaidenControllerEmbedded: " << status; + } - std::vector>> results; - results.reserve(block_hashes.size()); + completion_poller_thread_ = + std::thread(&KVCacheStore::CompletionPollerLoop, this); + } +} - for (const std::string& hash : block_hashes) { - std::vector* existing = lru_cache_.Get(hash); - if (!existing || existing->empty()) { - break; +KVCacheStore::~KVCacheStore() { + if (controller_) { + controller_->Stop(); + } + completion_queue_.Stop(); + if (completion_poller_thread_.joinable()) { + completion_poller_thread_.join(); + } +} + +absl::StatusOr KVCacheStore::Lookup( + const std::vector& block_hashes, bool enable_global) { + BlockSliceList results; + + size_t local_hits = 0; + size_t limit = 0; + { + absl::MutexLock lock(mutex_); + limit = std::min(block_hashes.size(), lru_cache_.available_space()); + results.reserve(limit); + for (size_t i = 0; i < limit; ++i) { + const std::string& hash = block_hashes[i]; + std::vector* existing = lru_cache_.Get(hash); + if (!existing || existing->empty()) { + break; + } + results.push_back(std::make_pair(hash, *existing)); + local_hits++; + } + } + + if (enable_global && local_hits < limit && registry_client_) { + std::vector remaining_hashes(block_hashes.begin() + local_hits, + block_hashes.begin() + limit); + auto global_results_or = registry_client_->Lookup(remaining_hashes); + if (global_results_or.ok()) { + const auto& global_results = global_results_or.value(); + for (size_t i = 0; i < global_results.size(); ++i) { + const auto& metadata = global_results[i]; + RaidenId remote_id; + remote_id.job_name = metadata.host_address(); + remote_id.job_replica_id = "0"; + remote_id.data_name = "kv_cache"; + remote_id.data_replica_idx = 0; // Do not abuse for block ID + results.push_back(std::make_pair( + remaining_hashes[i], std::vector{RaidenBlockID( + remote_id, -1, BlockStatus::REMOTE)})); + } + } else { + LOG(WARNING) << "Global registry lookup failed: " + << global_results_or.status().message(); } - results.push_back(std::make_pair(hash, *existing)); } return results; } -std::pair>>> -KVCacheStore::Insert(const std::vector& block_hashes, - const std::vector>& slices, - bool /*on_host*/) { +std::pair KVCacheStore::Insert( + const std::vector& block_hashes, + const std::vector>& slices, bool /*on_host*/) { absl::MutexLock lock(mutex_); - std::vector>> evicted_entries; + BlockSliceList evicted_entries; bool all_inserted = true; for (size_t i = 0; i < block_hashes.size(); ++i) { const std::string& hash = block_hashes[i]; if (lru_cache_.Contains(hash)) { + // NOTE(jcgu): This is technically true, as the key already exists in the + // cache. However, for the purpose of this insert, we treat it as if it + // was inserted. all_inserted = false; continue; } - std::optional>> evicted; + std::optional>> evicted; if (i < slices.size()) { evicted = lru_cache_.Put(hash, slices[i]); } else { @@ -77,8 +148,126 @@ KVCacheStore::Insert(const std::vector& block_hashes, return std::make_pair(all_inserted, std::move(evicted_entries)); } -void KVCacheStore::Delete(const std::vector& block_hashes, - const std::vector>& slices) { +// Categorizes block_hashes into existing (local) and new (remote) hashes. +// 1. Pins all existing block hashes. +// 2. Inserts new block hashes into the LRU cache if space permits. +// 3. Pins the newly inserted block hashes, with full rollback on failure. +std::pair KVCacheStore::InsertAndPin( + const std::vector& block_hashes, + const std::vector>& slices, bool /*on_host*/) { + absl::MutexLock lock(mutex_); + BlockSliceList evicted_entries; + + std::vector existing_indices; + std::vector new_indices; + for (size_t i = 0; i < block_hashes.size(); ++i) { + if (lru_cache_.Contains(block_hashes[i])) { + existing_indices.push_back(i); + } else { + new_indices.push_back(i); + } + } + + // 1. Pin all existing block_hashes + for (size_t idx = 0; idx < existing_indices.size(); ++idx) { + size_t i = existing_indices[idx]; + if (!lru_cache_.Pin(block_hashes[i])) { + for (size_t j = 0; j < idx; ++j) { + lru_cache_.Unpin(block_hashes[existing_indices[j]]); + } + return std::make_pair(false, std::move(evicted_entries)); + } + } + + // 2. Check if free space in lru_cache can hold all new block_hashes + if (lru_cache_.available_space() < new_indices.size()) { + for (size_t i : existing_indices) { + lru_cache_.Unpin(block_hashes[i]); + } + return std::make_pair(false, std::move(evicted_entries)); + } + + // Insert all new block hashes into the lru cache list + for (size_t i : new_indices) { + const std::string& hash = block_hashes[i]; + std::optional>> evicted; + if (i < slices.size()) { + evicted = lru_cache_.Put(hash, slices[i]); + } else { + evicted = lru_cache_.Put(hash, {}); + } + if (evicted.has_value()) { + evicted_entries.push_back(std::move(*evicted)); + } + } + + // 3. Pin the newly inserted block_hashes in the lru cache list + for (size_t idx = 0; idx < new_indices.size(); ++idx) { + size_t i = new_indices[idx]; + if (!lru_cache_.Pin(block_hashes[i])) { + for (size_t j = 0; j < idx; ++j) { + lru_cache_.Unpin(block_hashes[new_indices[j]]); + } + for (size_t j : existing_indices) { + lru_cache_.Unpin(block_hashes[j]); + } + for (size_t j : new_indices) { + lru_cache_.Erase(block_hashes[j]); + } + for (auto it = evicted_entries.rbegin(); it != evicted_entries.rend(); + ++it) { + lru_cache_.PutBack(it->first, std::move(it->second)); + } + return std::make_pair(false, BlockSliceList{}); + } + } + + return std::make_pair(true, std::move(evicted_entries)); +} + +// Reverts an InsertAndPin operation. +// 1. Unpins all block hashes in the LRU cache. +// 2. Erases remote block hashes whose pin count reaches 0. +// 3. Restores evicted entries to the back of the LRU cache in reverse order for +// each deleted remote block. Returns the number of deleted remote blocks and +// the remaining unrestored evicted entries. +std::pair KVCacheStore::ReleaseAndDelete( + const std::vector& block_hashes, + BlockSliceList pending_evict_entries) { + absl::MutexLock lock(mutex_); + size_t deleted_remote_blocks = 0; + for (const std::string& hash : block_hashes) { + lru_cache_.Unpin(hash); + auto* val = lru_cache_.Peek(hash); + if (val != nullptr && !val->empty() && + (*val)[0].status == BlockStatus::REMOTE && + lru_cache_.GetPinCount(hash) == 0) { + lru_cache_.Erase(hash); + deleted_remote_blocks++; + } + } + + if (deleted_remote_blocks > pending_evict_entries.size()) { + LOG(WARNING) << "Number of deleted remote blocks (" << deleted_remote_blocks + << ") exceeds number of pending evict entries (" + << pending_evict_entries.size() << ")."; + } + + size_t to_restore = + std::min(deleted_remote_blocks, pending_evict_entries.size()); + for (size_t i = 0; i < to_restore; ++i) { + auto& entry = pending_evict_entries.back(); + lru_cache_.PutBack(entry.first, std::move(entry.second)); + pending_evict_entries.pop_back(); + } + + return std::make_pair(deleted_remote_blocks, + std::move(pending_evict_entries)); +} + +void KVCacheStore::Delete( + const std::vector& block_hashes, + const std::vector>& slices) { absl::MutexLock lock(mutex_); for (const std::string& hash : block_hashes) { lru_cache_.Erase(hash); @@ -115,5 +304,148 @@ size_t KVCacheStore::capacity() const { return lru_cache_.capacity(); } +absl::flat_hash_map +KVCacheStore::FetchRemote(const std::vector& block_hashes) { + absl::flat_hash_map results; + absl::flat_hash_map grouped_reqs; + + { + absl::MutexLock lock(mutex_); + for (const auto& hash : block_hashes) { + std::vector* blocks = lru_cache_.Get(hash); + if (blocks) { + for (const auto& block : *blocks) { + if (block.status == BlockStatus::REMOTE) { + auto& item = grouped_reqs[block.raiden_id]; + item.src_raiden_id = block.raiden_id; + item.block_hashes.push_back(hash); + item.dst_block_ids.push_back(block.host_block_id); + } + } + } + } + } + + FetchRequest full_request; + for (auto& [raiden_id, item] : grouped_reqs) { + for (size_t i = 0; i < item.block_hashes.size(); ++i) { + const auto& hash = item.block_hashes[i]; + int dst_block_id = item.dst_block_ids[i]; + + std::shared_ptr state; + { + absl::MutexLock lock(fetch_mu_); + uint64_t fetch_id = next_fetch_id_++; + state = std::make_shared(); + state->fetch_id = fetch_id; + state->block_hash = hash; + state->pending_blocks.insert(dst_block_id); + block_to_fetch_[dst_block_id] = fetch_id; + active_fetches_[fetch_id] = state; + } + results[hash] = FetchFuture(state); + } + full_request.push_back(std::move(item)); + } + + if (!full_request.empty()) { + PushFetchWork(std::move(full_request)); + } + + return results; +} + +void KVCacheStore::CompletionPollerLoop() { + FetchCompletion completion; + while (PopFetchCompletion(completion)) { + { + absl::MutexLock lock(mutex_); + for (const auto& item : completion) { + if (item.success) { + std::vector* blocks = lru_cache_.Get(item.block_hash); + if (blocks) { + for (auto& block : *blocks) { + if (block.status == BlockStatus::REMOTE && + block.host_block_id == item.host_block_id) { + block.status = BlockStatus::HOST; + block.raiden_id = raiden_id_; + LOG(INFO) << "Upgraded block " << item.block_hash << " to HOST"; + } + } + } + } else { + LOG(ERROR) << "Fetch failed for " << item.block_hash << ": " + << item.error_message; + } + } + } + + { + absl::MutexLock lock(fetch_mu_); + for (const auto& item : completion) { + auto it = block_to_fetch_.find(item.host_block_id); + if (it != block_to_fetch_.end()) { + uint64_t fetch_id = it->second; + auto state_it = active_fetches_.find(fetch_id); + if (state_it != active_fetches_.end()) { + auto& state = state_it->second; + state->pending_blocks.erase(item.host_block_id); + if (!item.success) { + state->failed = true; + state->error_message += item.error_message + ";"; + } + + if (state->pending_blocks.empty()) { + if (state->failed) { + failed_fetches_.push_back(state->block_hash); + } else { + done_fetches_.push_back(state->block_hash); + } + state->notification.Notify(); + active_fetches_.erase(state_it); + } + } + block_to_fetch_.erase(it); + } + } + } + } +} + +absl::Status KVCacheStore::FetchFuture::Await() { + if (!state_) { + return absl::InvalidArgumentError("FetchFuture has no valid state."); + } + state_->notification.WaitForNotification(); + if (state_->failed) { + return absl::InternalError(state_->error_message); + } + return absl::OkStatus(); +} + +bool KVCacheStore::FetchFuture::IsDone() const { + if (!state_) return true; + return state_->notification.HasBeenNotified(); +} + +std::tuple, std::vector, + std::vector> +KVCacheStore::PollFetchRemoteStatus() { + absl::MutexLock lock(fetch_mu_); + std::vector pending; + pending.reserve(active_fetches_.size()); + for (const auto& [id, state] : active_fetches_) { + pending.push_back(state->block_hash); + } + + std::vector done = std::move(done_fetches_); + std::vector failed = std::move(failed_fetches_); + + done_fetches_.clear(); + failed_fetches_.clear(); + + return {std::move(done), std::move(failed), std::move(pending)}; +} + } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/kv_cache_store.h b/tpu_raiden/kv_cache/kv_cache_store.h index b4051afc..e13c4d33 100644 --- a/tpu_raiden/kv_cache/kv_cache_store.h +++ b/tpu_raiden/kv_cache/kv_cache_store.h @@ -17,15 +17,18 @@ #include #include +#include #include #include #include #include "absl/base/thread_annotations.h" #include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" +#include "absl/synchronization/notification.h" #include "absl/types/span.h" #include "tpu_raiden/core/raw_transfer_core.h" #include "tpu_raiden/kv_cache/lru_cache.h" @@ -33,6 +36,10 @@ namespace tpu_raiden { namespace kv_cache { +namespace global_registry { +class GlobalRegistryClient; +} + // Represents a microservice slice identifier / entity address hosting a replica // of a Key-Value cache block. struct RaidenId { @@ -49,38 +56,206 @@ struct RaidenId { } }; +struct RaidenIdHash { + size_t operator()(const RaidenId& id) const { + return absl::HashOf(id.job_name, id.job_replica_id, id.data_name, + id.data_replica_idx); + } +}; + +enum class BlockStatus { + INIT, + REMOTE, + HOST, + HBM, +}; + +struct RaidenBlockID { + RaidenId raiden_id; + int host_block_id = -1; + BlockStatus status = BlockStatus::INIT; + + RaidenBlockID() = default; + /* implicit */ RaidenBlockID(RaidenId id, int host_id = -1, + BlockStatus stat = BlockStatus::INIT) + : raiden_id(std::move(id)), host_block_id(host_id), status(stat) {} + + bool operator==(const RaidenBlockID& other) const { + return raiden_id == other.raiden_id && + host_block_id == other.host_block_id && status == other.status; + } +}; + +using BlockSliceList = + std::vector>>; + +struct FetchRequestItem { + RaidenId src_raiden_id; + std::vector block_hashes; + std::vector dst_block_ids; +}; + +using FetchRequest = std::vector; + +struct FetchCompletionItem { + std::string block_hash; + int host_block_id; + bool success; + std::string error_message; +}; + +using FetchCompletion = std::vector; + +template +class ThreadSafeQueue { + public: + ThreadSafeQueue() = default; + ~ThreadSafeQueue() = default; + + ThreadSafeQueue(const ThreadSafeQueue&) = delete; + ThreadSafeQueue& operator=(const ThreadSafeQueue&) = delete; + + void Push(T item) { + absl::MutexLock lock(&mutex_); + queue_.push(std::move(item)); + cond_.Signal(); + } + + bool Pop(T& item) { + absl::MutexLock lock(&mutex_); + while (queue_.empty() && !stopping_) { + cond_.Wait(&mutex_); + } + if (stopping_ && queue_.empty()) { + return false; + } + item = std::move(queue_.front()); + queue_.pop(); + return true; + } + + bool Empty() const { + absl::MutexLock lock(&mutex_); + return queue_.empty(); + } + + void Stop() { + absl::MutexLock lock(&mutex_); + stopping_ = true; + cond_.SignalAll(); + } + + private: + mutable absl::Mutex mutex_; + std::queue queue_ ABSL_GUARDED_BY(mutex_); + absl::CondVar cond_; + bool stopping_ = false; +}; + +using FetchWorkQueue = ThreadSafeQueue; +using FetchCompletionQueue = ThreadSafeQueue; + +class RaidenControllerEmbedded; + +struct RemoteFetchConfig { + std::string orchestrator_address; + int controller_port = 0; + int local_worker_port = 0; + size_t bytes_per_block = 0; + size_t num_shards = 0; + size_t num_listeners = 1; +}; + // KV Store that manages the indices and routing of prefix cache across serving // nodes and microservice slices. class KVCacheStore { public: - explicit KVCacheStore(size_t capacity); + struct FetchState { + uint64_t fetch_id; + std::string block_hash; + absl::flat_hash_set pending_blocks; + bool failed = false; + std::string error_message; + absl::Notification notification; + }; + + class FetchFuture { + public: + FetchFuture() = default; + explicit FetchFuture(std::shared_ptr state) + : state_(std::move(state)) {} + + absl::Status Await(); + bool IsDone() const; + + private: + std::shared_ptr state_; + }; + + explicit KVCacheStore( + size_t capacity, std::string global_registry_address = "", + RaidenId raiden_id = {}, + std::optional remote_config = std::nullopt); ~KVCacheStore(); KVCacheStore(const KVCacheStore&) = delete; KVCacheStore& operator=(const KVCacheStore&) = delete; + // Dedicated Fetch function to trigger remote fetching. + absl::flat_hash_map FetchRemote( + const std::vector& block_hashes); + + // Polls the status of all active fetches. + // Returns {done_block_hashes, failed_block_hashes, pending_block_hashes} + std::tuple, std::vector, + std::vector> + PollFetchRemoteStatus(); + // Authoritative KVCacheStore API implementations // Checks the LRU directory for cached block hashes. Returns a list of all - // matched replica pairs (block hash and vector of RaidenIds) encountered + // matched replica pairs (block hash and vector of RaidenBlockIDs) encountered // in sequence prior to the first miss. - absl::StatusOr>>> - Lookup(const std::vector& block_hashes); + // If enable_global is true, it will query the global registry for any + // misses after the local lookup. + absl::StatusOr Lookup( + const std::vector& block_hashes, bool enable_global = false); // Caches sharded buffers into host-RAM/HBM backing store. // Returns: // - bool: whether all blocks were successfully inserted (i.e. none already // existed) - // - std::vector>>: list of - // entries evicted from the LRU cache during insertion - std::pair>>> - Insert(const std::vector& block_hashes, - const std::vector>& slices, bool on_host); + // - BlockSliceList: list of entries evicted from the LRU cache during + // insertion + std::pair Insert( + const std::vector& block_hashes, + const std::vector>& slices, bool on_host); + + // Pins all existing block hashes, and inserts and pins new block hashes if + // there is sufficient available space in the LRU cache. + // Returns: + // - bool: whether the entire InsertAndPin operation succeeded (i.e. all + // existing keys were pinned, all new keys inserted and pinned) + // - BlockSliceList: list of entries evicted during insertion + std::pair InsertAndPin( + const std::vector& block_hashes, + const std::vector>& slices, bool on_host); + + // Reverts an InsertAndPin operation by unpinning all block_hashes in the + // LRU cache, deleting any block_hash in REMOTE status whose pin count is 0, + // and putting back evicted entries in reverse order for each deleted remote + // block. + // Returns: + // - size_t: number of remote blocks deleted + // - BlockSliceList: remaining evicted entries that were not restored + std::pair ReleaseAndDelete( + const std::vector& block_hashes, + BlockSliceList pending_evict_entries = {}); // Deletes cached sharded buffers from host-RAM/HBM backing store entirely. void Delete(const std::vector& block_hashes, - const std::vector>& slices); + const std::vector>& slices); // Pins cached block hashes in memory, protecting them against LRU eviction // while in active use. Returns true if all keys exist and were successfully @@ -95,10 +270,37 @@ class KVCacheStore { size_t capacity() const; + const RaidenId& raiden_id() const { return raiden_id_; } + + // Fetch Queue Accessors + bool PopFetchWork(FetchRequest& req) { return work_queue_.Pop(req); } + void PushFetchCompletion(FetchCompletion completion) { + completion_queue_.Push(std::move(completion)); + } + void PushFetchWork(FetchRequest req) { work_queue_.Push(std::move(req)); } + bool PopFetchCompletion(FetchCompletion& completion) { + return completion_queue_.Pop(completion); + } + private: mutable absl::Mutex mutex_; - mutable LRUCache> lru_cache_ + mutable LRUCache> lru_cache_ ABSL_GUARDED_BY(mutex_); + std::unique_ptr registry_client_; + RaidenId raiden_id_; + FetchWorkQueue work_queue_; + FetchCompletionQueue completion_queue_; + std::unique_ptr controller_; + void CompletionPollerLoop(); + std::thread completion_poller_thread_; + + mutable absl::Mutex fetch_mu_; + uint64_t next_fetch_id_ ABSL_GUARDED_BY(fetch_mu_) = 1; + absl::flat_hash_map> active_fetches_ + ABSL_GUARDED_BY(fetch_mu_); + absl::flat_hash_map block_to_fetch_ ABSL_GUARDED_BY(fetch_mu_); + std::vector done_fetches_ ABSL_GUARDED_BY(fetch_mu_); + std::vector failed_fetches_ ABSL_GUARDED_BY(fetch_mu_); }; } // namespace kv_cache diff --git a/tpu_raiden/kv_cache/kv_cache_store_internal.cc b/tpu_raiden/kv_cache/kv_cache_store_internal.cc index 79ff1560..7d1571ae 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_internal.cc +++ b/tpu_raiden/kv_cache/kv_cache_store_internal.cc @@ -32,7 +32,7 @@ #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" -#include "third_party/grpc/include/grpcpp/security/credentials.h" +#include "grpcpp/security/credentials.h" #include "tpu_raiden/core/raw_transfer_core.h" #include "tpu_raiden/core/status_macros.h" #include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" diff --git a/tpu_raiden/kv_cache/kv_cache_store_test.cc b/tpu_raiden/kv_cache/kv_cache_store_test.cc index ea91d792..f209c6fc 100644 --- a/tpu_raiden/kv_cache/kv_cache_store_test.cc +++ b/tpu_raiden/kv_cache/kv_cache_store_test.cc @@ -14,22 +14,71 @@ #include "tpu_raiden/kv_cache/kv_cache_store.h" +#include +#include #include +#include // NOLINT #include +#include "grpcpp/grpcpp.h" +#include "net/util/ports.h" +#include #include #include "absl/status/statusor.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "grpcpp/security/server_credentials.h" +#include "grpcpp/server.h" +#include "grpcpp/server_builder.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_server.h" +#include "tpu_raiden/kv_cache/kv_cache_listener.h" +#include "tpu_raiden/kv_cache/kv_cache_manager_base.h" +#include "tpu_raiden/kv_cache/raiden_controller_embedded.h" +#include "tpu_raiden/kv_cache/raiden_orchestrator.h" namespace tpu_raiden { namespace kv_cache { namespace { +int PickConsecutiveUnusedPortsOrDie(int count) { + for (int attempt = 0; attempt < 100; ++attempt) { + int base = net_util::PickUnusedPortOrDie(); + bool all_free = true; + for (int i = 0; i < count; ++i) { + int port = base + i; + int fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) { + all_free = false; + break; + } + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + addr.sin_addr.s_addr = INADDR_ANY; + int opt = 1; + setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)); + if (bind(fd, reinterpret_cast(&addr), sizeof(addr)) < 0) { + close(fd); + all_free = false; + break; + } + close(fd); + } + if (all_free) return base; + } + LOG(FATAL) << "Could not find " << count << " consecutive unused ports"; + return 0; +} + +// Verifies basic insert, lookup (including partial and early misses), and +// delete functionality. TEST(KVCacheStoreTest, BasicTests) { KVCacheStore controller(50); EXPECT_EQ(controller.capacity(), 50); std::vector hashes = {"4001", "4002"}; - std::vector> slices = { + std::vector> slices = { {RaidenId{"inference_server", "0", "kv_cache", 0}}, {RaidenId{"inference_server", "1", "kv_cache", 0}}}; @@ -59,11 +108,13 @@ TEST(KVCacheStoreTest, BasicTests) { controller.Insert(hashes, slices, true).first); // Succesful again } +// Verifies pinning blocks to prevent eviction and releasing them to allow +// eviction. TEST(KVCacheStoreTest, PinAndRelease) { KVCacheStore controller(2); std::vector hashes = {"101", "102"}; - std::vector> slices = { + std::vector> slices = { {RaidenId{"inference_server", "0", "kv_cache", 0}}, {RaidenId{"inference_server", "1", "kv_cache", 0}}}; @@ -77,7 +128,7 @@ TEST(KVCacheStoreTest, PinAndRelease) { // Inserting a third element should fail to evict because both existing items // are pinned std::vector hash_3 = {"103"}; - std::vector> slice_3 = { + std::vector> slice_3 = { {RaidenId{"inference_server", "2", "kv_cache", 0}}}; controller.Insert(hash_3, slice_3, true); @@ -87,7 +138,7 @@ TEST(KVCacheStoreTest, PinAndRelease) { // Now inserting a fourth element (104) should successfully evict 101 std::vector hash_4 = {"104"}; - std::vector> slice_4 = { + std::vector> slice_4 = { {RaidenId{"inference_server", "3", "kv_cache", 0}}}; controller.Insert(hash_4, slice_4, true); @@ -96,11 +147,13 @@ TEST(KVCacheStoreTest, PinAndRelease) { EXPECT_EQ(controller.Lookup({"102"})->size(), 1); } +// Verifies that if pinning a set of blocks fails (e.g. due to a missing block), +// none of them remain pinned. TEST(KVCacheStoreTest, PartialPinRollback) { KVCacheStore controller(2); std::vector hashes = {"201", "202"}; - std::vector> slices = { + std::vector> slices = { {RaidenId{"inference_server", "0", "kv_cache", 0}}, {RaidenId{"inference_server", "1", "kv_cache", 0}}}; @@ -114,11 +167,13 @@ TEST(KVCacheStoreTest, PartialPinRollback) { EXPECT_EQ(controller.GetPinCount("202"), 0); } +// Verifies that inserting beyond capacity evicts the least recently used (LRU) +// block and returns it. TEST(KVCacheStoreTest, EvictionTracking) { KVCacheStore controller(2); std::vector hashes_1_2 = {"101", "102"}; - std::vector> slices_1_2 = { + std::vector> slices_1_2 = { {RaidenId{"inference_server", "0", "kv_cache", 0}}, {RaidenId{"inference_server", "1", "kv_cache", 1}}}; @@ -129,7 +184,7 @@ TEST(KVCacheStoreTest, EvictionTracking) { // 2. Insert 103. Since capacity is 2, this must evict the LRU block (101). std::vector hash_3 = {"103"}; - std::vector> slice_3 = { + std::vector> slice_3 = { {RaidenId{"inference_server", "2", "kv_cache", 2}}}; auto res_3 = controller.Insert(hash_3, slice_3, true); @@ -137,8 +192,8 @@ TEST(KVCacheStoreTest, EvictionTracking) { ASSERT_EQ(res_3.second.size(), 1); EXPECT_EQ(res_3.second[0].first, "101"); ASSERT_EQ(res_3.second[0].second.size(), 1); - EXPECT_EQ(res_3.second[0].second[0].job_name, "inference_server"); - EXPECT_EQ(res_3.second[0].second[0].data_replica_idx, 0); + EXPECT_EQ(res_3.second[0].second[0].raiden_id.job_name, "inference_server"); + EXPECT_EQ(res_3.second[0].second[0].raiden_id.data_replica_idx, 0); // 3. Verify that lookup for 101 now misses, but 102 and 103 are present. EXPECT_EQ(controller.Lookup({"101"})->size(), 0); @@ -146,6 +201,1398 @@ TEST(KVCacheStoreTest, EvictionTracking) { EXPECT_EQ(controller.Lookup({"103"})->size(), 1); } +// Verifies looking up blocks in the global registry when they are not found +// locally. +TEST(KVCacheStoreTest, GlobalLookupFallback) { + // 1. Start a local registry server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + auto server = builder.BuildAndStart(); + std::string server_address = "localhost:" + std::to_string(port); + + // 2. Register some blocks in the registry + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::string hash1 = "global_hash_1"; + std::string host1 = "10.0.0.1:1234"; + int32_t block1 = 42; + + std::string hash2 = "global_hash_2"; + std::string host2 = "10.0.0.2:1234"; + int32_t block2 = 43; + + // For Case 2: Register a hash that will also be present locally, but with + // a different remote address in the registry. + std::string hash_shared = "shared_hash"; + std::string host_shared_remote = "10.0.0.9:1234"; + int32_t block_shared_remote = 99; + + ASSERT_TRUE( + registry_client + .Register({{hash1, host1, block1}, + {hash2, host2, block2}, + {hash_shared, host_shared_remote, block_shared_remote}}) + .ok()); + + // 3. Create KVCacheStore with the registry address + KVCacheStore store(50, server_address); + + // Insert blocks locally + std::vector local_hashes = {"local_only_hash", "shared_hash"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}, + {RaidenId{"local_job", "0", "kv_cache", 1}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Case 1: Full local hit, no global hit + { + auto lookup_res = store.Lookup({"local_only_hash"}, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.job_name, "local_job"); + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.data_replica_idx, 0); + } + + // Case 2: Both local and global has the same hit, but we return local hit + // results + { + auto lookup_res = store.Lookup({"shared_hash"}, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "shared_hash"); + // Should return local info, not remote info from registry + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.job_name, "local_job"); + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.data_replica_idx, 1); + } + + // Case 3: No local hit, only global hits + { + auto lookup_res = store.Lookup({"global_hash_1", "global_hash_2"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 2); + + EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.job_name, host1); + + EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); + EXPECT_EQ((*lookup_res)[1].second[0].raiden_id.job_name, host2); + } + + // 4. Lookup with enable_global = false + // It should stop at the first miss (which is the global hash if we query it) + // If we query {"local_only_hash", "global_hash_1"}, it should return + // local_only_hash and stop. + { + auto lookup_res = store.Lookup({"local_only_hash", "global_hash_1"}, + /*enable_global=*/false); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + } + + // 5. Lookup with enable_global = true + // It should return both local and global + { + auto lookup_res = + store.Lookup({"local_only_hash", "global_hash_1", "global_hash_2"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 3); + + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.job_name, "local_job"); + + EXPECT_EQ((*lookup_res)[1].first, "global_hash_1"); + EXPECT_EQ((*lookup_res)[1].second[0].raiden_id.job_name, host1); + + EXPECT_EQ((*lookup_res)[2].first, "global_hash_2"); + EXPECT_EQ((*lookup_res)[2].second[0].raiden_id.job_name, host2); + } + + // 6. Lookup with enable_global = true, but registry has a miss + // It should stop at the first miss in registry + { + auto lookup_res = store.Lookup( + {"local_only_hash", "global_hash_1", "missing_hash", "global_hash_2"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + ASSERT_EQ(lookup_res->size(), 2); // local_only_hash, global_hash_1 + EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); + EXPECT_EQ((*lookup_res)[1].first, "global_hash_1"); + } + + server->Shutdown(); +} + +// Verifies that local lookup still works even if the global registry is +// unreachable. +TEST(KVCacheStoreTest, GlobalLookupRegistryDown) { + // Create KVCacheStore with an unreachable registry address + KVCacheStore store(50, "invalid.address:12345"); + + // Insert one block locally + std::vector local_hashes = {"local_hash"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Lookup with enable_global = true. + // It should NOT fail even though the registry is down. It should return the + // local hit. + auto lookup_res = store.Lookup({"local_hash", "missing_hash"}, + /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 1); + EXPECT_EQ((*lookup_res)[0].first, "local_hash"); + EXPECT_EQ((*lookup_res)[0].second[0].raiden_id.job_name, "local_job"); +} + +// Verifies that lookup returns at most the capacity of the store. +TEST(KVCacheStoreTest, LookupCapLimit) { + KVCacheStore store(2); + + std::vector hashes = {"101", "102"}; + std::vector> slices = { + {RaidenId{"inference_server", "0", "kv_cache", 0}}, + {RaidenId{"inference_server", "1", "kv_cache", 1}}}; + + ASSERT_TRUE(store.Insert(hashes, slices, true).first); + + // Lookup 3 hashes, but capacity is 2. It should only return 2. + std::vector lookup_hashes = {"101", "102", "103"}; + auto lookup_res = store.Lookup(lookup_hashes); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "101"); + EXPECT_EQ((*lookup_res)[1].first, "102"); +} + +// Verifies that lookup returns at most the capacity of the store, even with +// global registry enabled. +TEST(KVCacheStoreTest, LookupCapLimitWithGlobal) { + // 1. Start a local registry server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + auto server = builder.BuildAndStart(); + std::string server_address = "localhost:" + std::to_string(port); + + // 2. Register some blocks in the registry + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::string hash1 = "global_hash_1"; + std::string host1 = "10.0.0.1:1234"; + int32_t block1 = 42; + + std::string hash2 = "global_hash_2"; + std::string host2 = "10.0.0.2:1234"; + int32_t block2 = 43; + + std::string hash3 = "global_hash_3"; + std::string host3 = "10.0.0.3:1234"; + int32_t block3 = 44; + + ASSERT_TRUE(registry_client + .Register({{hash1, host1, block1}, + {hash2, host2, block2}, + {hash3, host3, block3}}) + .ok()); + + // 3. Create KVCacheStore with capacity 2 + KVCacheStore store(2, server_address); + + // Lookup 3 hashes, but capacity is 2. It should only return 2. + std::vector lookup_hashes = {"global_hash_1", "global_hash_2", + "global_hash_3"}; + auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); + EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); + + server->Shutdown(); +} + +// Verifies capacity limits with a mix of local and global hits. +TEST(KVCacheStoreTest, LookupCapLimitMixed) { + // 1. Start a local registry server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + auto server = builder.BuildAndStart(); + std::string server_address = "localhost:" + std::to_string(port); + + // 2. Register some blocks in the registry + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::string hash2 = "global_hash_2"; + std::string host2 = "10.0.0.2:1234"; + int32_t block2 = 43; + + std::string hash3 = "global_hash_3"; + std::string host3 = "10.0.0.3:1234"; + int32_t block3 = 44; + + ASSERT_TRUE( + registry_client.Register({{hash2, host2, block2}, {hash3, host3, block3}}) + .ok()); + + // 3. Create KVCacheStore with capacity 2 + KVCacheStore store(2, server_address); + + // Insert 1 block locally + std::vector local_hashes = {"local_hash_1"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Lookup 3 hashes, but capacity is 2. It should only return 2 (1 local, 1 + // global). + std::vector lookup_hashes = {"local_hash_1", "global_hash_2", + "global_hash_3"}; + auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "local_hash_1"); + EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); + + server->Shutdown(); +} + +// Verifies that lookup is limited by available space (capacity - pinned +// blocks). +TEST(KVCacheStoreTest, LookupAvailableSpaceLimit) { + KVCacheStore store(3); + + std::vector hashes = {"101", "102", "103"}; + std::vector> slices = { + {RaidenId{"inference_server", "0", "kv_cache", 0}}, + {RaidenId{"inference_server", "1", "kv_cache", 1}}, + {RaidenId{"inference_server", "2", "kv_cache", 2}}}; + + ASSERT_TRUE(store.Insert(hashes, slices, true).first); + + // Pin 101. Pinned count = 1. Available space = 3 - 1 = 2. + EXPECT_TRUE(store.Pin({"101"})); + + // Lookup 4 hashes. Since available space is 2, it should only return the + // first 2. + std::vector lookup_hashes = {"101", "102", "103", "104"}; + auto lookup_res = store.Lookup(lookup_hashes); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + EXPECT_EQ((*lookup_res)[0].first, "101"); + EXPECT_EQ((*lookup_res)[1].first, "102"); +} + +// Verifies inserting blocks and pinning them atomically, and space limits. +TEST(KVCacheStoreTest, InsertAndPin) { + KVCacheStore store(2); + + // Insert local block + std::vector local_hashes = {"local_1"}; + std::vector> local_slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Execute InsertAndPin + std::vector> slices = { + {RaidenId{"local_job", "0", "kv_cache", 0}}, + {RaidenId{"remote_job", "0", "kv_cache", 42}}}; + auto res = store.InsertAndPin({"local_1", "remote_1"}, slices, true); + EXPECT_TRUE(res.first); + EXPECT_TRUE(res.second.empty()); + EXPECT_EQ(store.GetPinCount("local_1"), 1); + EXPECT_EQ(store.GetPinCount("remote_1"), 1); + + // Since capacity is 2 and both local_1 and remote_1 are pinned, available + // space is 0. Attempting to InsertAndPin remote_2 should fail due to lack + // of space. + auto res_fail = store.InsertAndPin({"remote_2"}, {}, true); + EXPECT_FALSE(res_fail.first); +} + +// Verifies PutBack puts items in the LRU position in the cache. +TEST(KVCacheStoreTest, LruCachePutBack) { + LRUCache cache(3); + cache.Put("A", 1); + cache.Put("B", 2); + // MRU to LRU is: B, A. + + // PutBack("C", 3) should add C to the back (LRU position). + cache.PutBack("C", 3); + // Now MRU to LRU should be: B, A, C. + // Let's verify by checking Evict(): first evicted item should be C! + auto evicted1 = cache.Evict(); + ASSERT_TRUE(evicted1.has_value()); + EXPECT_EQ(evicted1->first, "C"); + + auto evicted2 = cache.Evict(); + ASSERT_TRUE(evicted2.has_value()); + EXPECT_EQ(evicted2->first, "A"); + + auto evicted3 = cache.Evict(); + ASSERT_TRUE(evicted3.has_value()); + EXPECT_EQ(evicted3->first, "B"); +} + +// Verifies releasing pinned blocks and deleting them if they are remote +// placeholders. +TEST(KVCacheStoreTest, ReleaseAndDelete) { + KVCacheStore store(2); + + // Insert two local blocks (not remote) + std::vector local_hashes = {"local_1", "local_2"}; + std::vector> local_slices = { + {RaidenBlockID(RaidenId{"local_job", "0", "kv_cache", 0}, -1, + BlockStatus::HOST)}, + {RaidenBlockID(RaidenId{"local_job", "0", "kv_cache", 1}, -1, + BlockStatus::HOST)}}; + ASSERT_TRUE(store.Insert(local_hashes, local_slices, true).first); + + // Now InsertAndPin two remote blocks, which will evict local_1 and local_2. + std::vector remote_hashes = {"remote_1", "remote_2"}; + std::vector> remote_slices = { + {RaidenBlockID(RaidenId{"remote_job", "0", "kv_cache", 0}, -1, + BlockStatus::REMOTE)}, + {RaidenBlockID(RaidenId{"remote_job", "0", "kv_cache", 1}, -1, + BlockStatus::REMOTE)}}; + auto res = store.InsertAndPin(remote_hashes, remote_slices, true); + ASSERT_TRUE(res.first); + ASSERT_EQ(res.second.size(), 2); + EXPECT_EQ(store.GetPinCount("remote_1"), 1); + EXPECT_EQ(store.GetPinCount("remote_2"), 1); + EXPECT_EQ(store.Lookup({"local_1"})->size(), 0); + EXPECT_EQ(store.Lookup({"local_2"})->size(), 0); + + // Now call ReleaseAndDelete to revert InsertAndPin! + auto release_res = store.ReleaseAndDelete(remote_hashes, res.second); + EXPECT_EQ(release_res.first, 2); + EXPECT_TRUE(release_res.second.empty()); + + // remote_1 and remote_2 should be unpinned and deleted (since REMOTE) + EXPECT_EQ(store.GetPinCount("remote_1"), 0); + EXPECT_EQ(store.GetPinCount("remote_2"), 0); + EXPECT_EQ(store.Lookup({"remote_1"})->size(), 0); + EXPECT_EQ(store.Lookup({"remote_2"})->size(), 0); + + // local_1 and local_2 should be restored to the cache! + auto lookup_res = store.Lookup({"local_1", "local_2"}); + ASSERT_TRUE(lookup_res.ok()); + EXPECT_EQ(lookup_res->size(), 2); + + // Test non-remote block in ReleaseAndDelete: should unpin without deleting + store.InsertAndPin({"local_1"}, {local_slices[0]}, true); + EXPECT_EQ(store.GetPinCount("local_1"), 1); + auto res_non_remote = store.ReleaseAndDelete({"local_1"}); + EXPECT_EQ(res_non_remote.first, 0); + EXPECT_TRUE(res_non_remote.second.empty()); + EXPECT_EQ(store.GetPinCount("local_1"), 0); + EXPECT_EQ(store.Lookup({"local_1"})->size(), 1); + + // Test remote block pinned twice: after one ReleaseAndDelete, pin count is 1 + // so it should NOT be deleted! + store.InsertAndPin({"remote_1"}, {remote_slices[0]}, true); + store.Pin({"remote_1"}); // pin count is now 2 + EXPECT_EQ(store.GetPinCount("remote_1"), 2); + auto res_pinned = store.ReleaseAndDelete({"remote_1"}); + EXPECT_EQ(res_pinned.first, 0); // 0 deleted because pin count was 2 -> 1 + EXPECT_EQ(store.GetPinCount("remote_1"), 1); + EXPECT_EQ(store.Lookup({"remote_1"})->size(), 1); + store.Release({"remote_1"}); + store.Delete({"remote_1"}, {remote_slices[0]}); + + // Test partial restore: 1 deleted remote block with 2 evicted entries + BlockSliceList mock_evicted = {{"evict_1", {local_slices[0]}}, + {"evict_2", {local_slices[1]}}}; + store.InsertAndPin({"remote_2"}, {remote_slices[1]}, true); + auto res_partial = store.ReleaseAndDelete({"remote_2"}, mock_evicted); + EXPECT_EQ(res_partial.first, 1); + ASSERT_EQ(res_partial.second.size(), 1); + EXPECT_EQ(res_partial.second[0].first, "evict_1"); // evict_2 was restored! +} + +// --- ThreadSafeQueue Tests --- + +// Verifies basic queue operations (Push and Pop). +TEST(ThreadSafeQueueTest, BasicPushPop) { + ThreadSafeQueue q; + q.Push(1); + q.Push(2); + + int val = 0; + EXPECT_TRUE(q.Pop(val)); + EXPECT_EQ(val, 1); + EXPECT_TRUE(q.Pop(val)); + EXPECT_EQ(val, 2); + EXPECT_TRUE(q.Empty()); +} + +// Verifies blocking pop and stopping the queue unblocks it. +TEST(ThreadSafeQueueTest, BlockingPopAndStop) { + ThreadSafeQueue q; + std::atomic popped{false}; + int popped_val = 0; + + std::thread pop_thread([&]() { + if (q.Pop(popped_val)) { + popped = true; + } + }); + + // Give thread time to start and block + absl::SleepFor(absl::Milliseconds(100)); + EXPECT_FALSE(popped); + + q.Push(42); + pop_thread.join(); + + EXPECT_TRUE(popped); + EXPECT_EQ(popped_val, 42); +} + +// Verifies stopping the queue unblocks waiting threads. +TEST(ThreadSafeQueueTest, StopUnblocks) { + ThreadSafeQueue q; + std::atomic pop_result{true}; + + std::thread pop_thread([&]() { + int val = 0; + pop_result = q.Pop(val); + }); + + absl::SleepFor(absl::Milliseconds(100)); + q.Stop(); + pop_thread.join(); + + EXPECT_FALSE(pop_result); +} + +// --- RemoteFetchInteractionTest --- + +// Verifies controller registration with orchestrator. +TEST(RemoteFetchInteractionTest, ControllerRegistration) { + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + KVCacheStore store1(10, "", RaidenId{"job1", "0", "data1", 0}); + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_port1 = net_util::PickUnusedPortOrDie(); + + RaidenControllerEmbedded controller1( + &store1, ctrl_port1, orch_addr, worker_port1, {"localhost:10000"}, + /*bytes_per_block=*/1024, /*num_shards=*/1); + + // Start controller (this should trigger registration) + EXPECT_OK(controller1.Start()); + + // Give time for async registration + absl::SleepFor(absl::Milliseconds(500)); + + // Verify registration in orchestrator by trying to resolve it. + // We can use controller2 to resolve controller1. + KVCacheStore store2(10, "", RaidenId{"job2", "0", "data2", 0}); + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_port2 = net_util::PickUnusedPortOrDie(); + RaidenControllerEmbedded controller2( + &store2, ctrl_port2, orch_addr, worker_port2, {"localhost:20000"}, + /*bytes_per_block=*/1024, /*num_shards=*/1); + + EXPECT_OK(controller2.Start()); + absl::SleepFor(absl::Milliseconds(500)); + + // Now try to resolve store1's controller from controller2 + // We need to expose ResolveRemoteController or test it via the queue flow. + // Since ResolveRemoteController is private/internal, testing via Queue is + // better integration test. +} + +// Verifies end-to-end fetch flow for a single block using PushFetchWork. +TEST(RemoteFetchInteractionTest, QueueFlow) { + // Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // Setup Store 1 (Sender) + RaidenId id1{"job1", "0", "data1", 0}; + KVCacheStore store1(10, "", id1); + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager1(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port1, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener1(&manager1, worker_ctrl_port1); + + std::string worker_data_addr1 = absl::StrCat("localhost:", worker_data_port1); + RaidenControllerEmbedded controller1( + &store1, ctrl_port1, orch_addr, worker_ctrl_port1, {}, + /*bytes_per_block=*/1024, /*num_shards=*/1); + EXPECT_OK(controller1.Start()); + + // Setup Store 2 (Receiver/Fetcher) + RaidenId id2{"job2", "0", "data2", 0}; + KVCacheStore store2(10, "", id2); + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager2(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port2, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener2(&manager2, worker_ctrl_port2); + + std::string worker_data_addr2 = absl::StrCat("localhost:", worker_data_port2); + RaidenControllerEmbedded controller2( + &store2, ctrl_port2, orch_addr, worker_ctrl_port2, {}, + /*bytes_per_block=*/1024, /*num_shards=*/1); + EXPECT_OK(controller2.Start()); + + // Insert hash1 into store1 so negotiation finds it + std::vector> slices = { + {RaidenBlockID(id1, /*host_id=*/3, BlockStatus::HOST)}}; + store1.Insert({"hash1"}, slices, /*on_host=*/true); + + // Fill Source Buffer with pattern + uint8_t* src_ptr = manager1.GetHostPointer(0, 0); + size_t block_size = manager1.bytes_per_block(); + std::memset(src_ptr + 3 * block_size, 'A', block_size); + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // Simulate KVCacheStore2 pushing work + FetchRequest req; + FetchRequestItem item; + item.src_raiden_id = id1; // Fetch from store1 + item.block_hashes.push_back("hash1"); + item.dst_block_ids.push_back(5); + req.push_back(item); + + store2.PushFetchWork(req); + + // Wait for processing + absl::SleepFor( + absl::Seconds(2)); // Increased sleep to ensure transfer completes + + // Verify Destination Buffer + uint8_t* dst_ptr = manager2.GetHostPointer(0, 0); + for (size_t i = 0; i < block_size; ++i) { + EXPECT_EQ(dst_ptr[5 * block_size + i], 'A') << "Mismatch at index " << i; + } +} + +// Verifies fetch flow with multiple layers and shards. +TEST(RemoteFetchInteractionTest, QueueFlowMultiLayerMultiShard) { + // Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // Setup Store 1 (Sender) + RaidenId id1{"job1", "0", "data1", 0}; + KVCacheStore store1(10, "", id1); + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager1(/*num_layers=*/2, /*num_shards=*/2, + /*slice_byte_size=*/1024, worker_data_port1, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener1(&manager1, worker_ctrl_port1); + + std::string worker_data_addr1 = absl::StrCat("localhost:", worker_data_port1); + // Pass same address for both shards + RaidenControllerEmbedded controller1( + &store1, ctrl_port1, orch_addr, worker_ctrl_port1, + {worker_data_addr1, worker_data_addr1}, /*bytes_per_block=*/1024, + /*num_shards=*/2); + EXPECT_OK(controller1.Start()); + + // Setup Store 2 (Receiver/Fetcher) + RaidenId id2{"job2", "0", "data2", 0}; + KVCacheStore store2(10, "", id2); + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager2(/*num_layers=*/2, /*num_shards=*/2, + /*slice_byte_size=*/1024, worker_data_port2, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener2(&manager2, worker_ctrl_port2); + + std::string worker_data_addr2 = absl::StrCat("localhost:", worker_data_port2); + // Pass same address for both shards + RaidenControllerEmbedded controller2( + &store2, ctrl_port2, orch_addr, worker_ctrl_port2, + {worker_data_addr2, worker_data_addr2}, /*bytes_per_block=*/1024, + /*num_shards=*/2); + EXPECT_OK(controller2.Start()); + + // Insert hash1 into store1 so negotiation finds it + std::vector> slices = { + {RaidenBlockID(id1, /*host_id=*/3, BlockStatus::HOST)}}; + store1.Insert({"hash1"}, slices, /*on_host=*/true); + + // Fill Source Buffer with pattern for all layers and shards + size_t block_size = manager1.bytes_per_block(); + for (size_t l = 0; l < 2; ++l) { + for (size_t sh = 0; sh < 2; ++sh) { + uint8_t* src_ptr = manager1.GetHostPointer(l, sh); + std::memset(src_ptr + 3 * block_size, 'A' + l * 2 + sh, block_size); + } + } + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // Simulate KVCacheStore2 pushing work + FetchRequest req; + FetchRequestItem item; + item.src_raiden_id = id1; // Fetch from store1 + item.block_hashes.push_back("hash1"); + item.dst_block_ids.push_back(5); + req.push_back(item); + + store2.PushFetchWork(req); + + // Wait for processing + absl::SleepFor( + absl::Seconds(2)); // Increased sleep to ensure transfer completes + + // Verify Destination Buffer for all layers and shards + for (size_t l = 0; l < 2; ++l) { + for (size_t sh = 0; sh < 2; ++sh) { + uint8_t* dst_ptr = manager2.GetHostPointer(l, sh); + for (size_t i = 0; i < block_size; ++i) { + EXPECT_EQ(dst_ptr[5 * block_size + i], 'A' + l * 2 + sh) + << "Mismatch at layer " << l << " shard " << sh << " index " << i; + } + } + } +} + +// Verifies multiple concurrent fetch requests. +TEST(RemoteFetchInteractionTest, QueueFlowConcurrent) { + // Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // Setup Store 1 (Sender) + RaidenId id1{"job1", "0", "data1", 0}; + KVCacheStore store1(10, "", id1); + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager1(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port1, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener1(&manager1, worker_ctrl_port1); + + std::string worker_data_addr1 = absl::StrCat("localhost:", worker_data_port1); + RaidenControllerEmbedded controller1( + &store1, ctrl_port1, orch_addr, worker_ctrl_port1, {worker_data_addr1}, + /*bytes_per_block=*/1024, /*num_shards=*/1); + EXPECT_OK(controller1.Start()); + + // Setup Store 2 (Receiver/Fetcher) + RaidenId id2{"job2", "0", "data2", 0}; + KVCacheStore store2(10, "", id2); + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager2(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port2, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener2(&manager2, worker_ctrl_port2); + + std::string worker_data_addr2 = absl::StrCat("localhost:", worker_data_port2); + RaidenControllerEmbedded controller2( + &store2, ctrl_port2, orch_addr, worker_ctrl_port2, {worker_data_addr2}, + /*bytes_per_block=*/1024, /*num_shards=*/1); + EXPECT_OK(controller2.Start()); + + // Insert hashes into store1 + store1.Insert({"hash1"}, {{RaidenBlockID(id1, 3, BlockStatus::HOST)}}, true); + store1.Insert({"hash2"}, {{RaidenBlockID(id1, 4, BlockStatus::HOST)}}, true); + + // Fill Source Buffer with patterns + uint8_t* src_ptr = manager1.GetHostPointer(0, 0); + size_t block_size = manager1.bytes_per_block(); + std::memset(src_ptr + 3 * block_size, 'A', block_size); + std::memset(src_ptr + 4 * block_size, 'B', block_size); + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // Simulate KVCacheStore2 pushing multiple works concurrently (batched) + FetchRequest req; + FetchRequestItem item; + item.src_raiden_id = id1; + item.block_hashes.push_back("hash1"); + item.dst_block_ids.push_back(5); + item.block_hashes.push_back("hash2"); + item.dst_block_ids.push_back(6); + req.push_back(item); + + store2.PushFetchWork(req); + + // Wait for processing + absl::SleepFor( + absl::Seconds(3)); // Give it more time for concurrent requests + + // Verify Destination Buffer + uint8_t* dst_ptr = manager2.GetHostPointer(0, 0); + for (size_t i = 0; i < block_size; ++i) { + EXPECT_EQ(dst_ptr[5 * block_size + i], 'A') + << "Mismatch at index " << i << " for hash1"; + EXPECT_EQ(dst_ptr[6 * block_size + i], 'B') + << "Mismatch at index " << i << " for hash2"; + } +} + +// Verifies high-level Fetch API with embedded controller, including multiple +// blocks, layers, and shards. +TEST(RemoteFetchInteractionTest, QueueFlowEmbedded) { + // Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // Setup Store 1 (Sender) + RaidenId id1{"job1", "0", "data1", 0}; + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager1(/*num_layers=*/2, /*num_shards=*/2, + /*slice_byte_size=*/1024, worker_data_port1, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener1(&manager1, worker_ctrl_port1, ctrl_port1); + + // Enable embedded controller for Store 1 + RemoteFetchConfig config1{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port1, + .local_worker_port = worker_ctrl_port1, + .bytes_per_block = 1024, + .num_shards = 2, + }; + KVCacheStore store1(10, "", id1, config1); + + // Setup Store 2 (Receiver/Fetcher) + RaidenId id2{"job2", "0", "data2", 0}; + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager2(/*num_layers=*/2, /*num_shards=*/2, + /*slice_byte_size=*/1024, worker_data_port2, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener2(&manager2, worker_ctrl_port2, ctrl_port2); + + // Enable embedded controller for Store 2 + RemoteFetchConfig config2{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port2, + .local_worker_port = worker_ctrl_port2, + .bytes_per_block = 1024, + .num_shards = 2, + }; + KVCacheStore store2(10, "", id2, config2); + + // Insert hash1 and hash2 into store1 so negotiation finds them + std::vector> slices1_1 = { + {RaidenBlockID(id1, /*host_id=*/3, BlockStatus::HOST)}}; + store1.Insert({"hash1"}, slices1_1, /*on_host=*/true); + std::vector> slices1_2 = { + {RaidenBlockID(id1, /*host_id=*/4, BlockStatus::HOST)}}; + store1.Insert({"hash2"}, slices1_2, /*on_host=*/true); + + // Insert REMOTE placeholders into store2 + std::vector> slices2_1 = {{RaidenBlockID( + id1, /*host_id=*/5, BlockStatus::REMOTE)}}; // Destination block 5 + store2.Insert({"hash1"}, slices2_1, /*on_host=*/true); + std::vector> slices2_2 = {{RaidenBlockID( + id1, /*host_id=*/6, BlockStatus::REMOTE)}}; // Destination block 6 + store2.Insert({"hash2"}, slices2_2, /*on_host=*/true); + + // Fill Source Buffer with pattern for all layers and shards + size_t block_size = manager1.bytes_per_block(); + for (size_t l = 0; l < 2; ++l) { + for (size_t sh = 0; sh < 2; ++sh) { + uint8_t* src_ptr = manager1.GetHostPointer(l, sh); + std::memset(src_ptr + 3 * block_size, 'A' + l * 2 + sh, block_size); + std::memset(src_ptr + 4 * block_size, 'B' + l * 2 + sh, block_size); + } + } + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // Trigger Fetch + auto futures = store2.FetchRemote({"hash1", "hash2"}); + ASSERT_EQ(futures.size(), 2); + + ASSERT_TRUE(futures.contains("hash1")); + ASSERT_TRUE(futures.contains("hash2")); + + // Wait for completion using the futures + EXPECT_OK(futures["hash1"].Await()); + EXPECT_OK(futures["hash2"].Await()); + + // Verify Polling Status + auto [done, failed, pending] = store2.PollFetchRemoteStatus(); + EXPECT_EQ(done.size(), 2); + EXPECT_THAT(done, testing::UnorderedElementsAre("hash1", "hash2")); + EXPECT_EQ(failed.size(), 0); + EXPECT_EQ(pending.size(), 0); + + // Verify Destination Buffer for all layers and shards + for (size_t l = 0; l < 2; ++l) { + for (size_t sh = 0; sh < 2; ++sh) { + uint8_t* dst_ptr = manager2.GetHostPointer(l, sh); + for (size_t i = 0; i < block_size; ++i) { + EXPECT_EQ(dst_ptr[5 * block_size + i], 'A' + l * 2 + sh) + << "Mismatch at layer " << l << " shard " << sh << " index " << i + << " for hash1"; + EXPECT_EQ(dst_ptr[6 * block_size + i], 'B' + l * 2 + sh) + << "Mismatch at layer " << l << " shard " << sh << " index " << i + << " for hash2"; + } + } + } + + // Verify LRU Status Upgrade in Store 2 for hash1 + auto lookup_res1 = store2.Lookup({"hash1"}); + ASSERT_TRUE(lookup_res1.ok()); + ASSERT_EQ(lookup_res1->size(), 1); + const auto& block_list1 = (*lookup_res1)[0].second; + ASSERT_EQ(block_list1.size(), 1); + EXPECT_EQ(block_list1[0].status, BlockStatus::HOST); + EXPECT_EQ(block_list1[0].host_block_id, 5); + EXPECT_EQ(block_list1[0].raiden_id.job_name, "job2"); + + // Verify LRU Status Upgrade in Store 2 for hash2 + auto lookup_res2 = store2.Lookup({"hash2"}); + ASSERT_TRUE(lookup_res2.ok()); + ASSERT_EQ(lookup_res2->size(), 1); + const auto& block_list2 = (*lookup_res2)[0].second; + ASSERT_EQ(block_list2.size(), 1); + EXPECT_EQ(block_list2[0].status, BlockStatus::HOST); + EXPECT_EQ(block_list2[0].host_block_id, 6); + EXPECT_EQ(block_list2[0].raiden_id.job_name, "job2"); +} +// Verifies the complete flow: Lookup (local + global), InsertAndPin (allocating +// local blocks for remote placeholders), and Fetch. +TEST(RemoteFetchInteractionTest, QueueFlowEndToEnd) { + // 1. Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // 2. Setup Global Registry Server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int registry_port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + ®istry_port); + builder.RegisterService(service.get()); + auto registry_server = builder.BuildAndStart(); + std::string registry_addr = "localhost:" + std::to_string(registry_port); + + // 3. Setup Store 1 (Remote Sender) + // MUST use "kv_cache" as data_name and 0 as data_replica_idx to match Lookup + // hardcoding + RaidenId id1{"job1", "0", "kv_cache", 0}; + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager1(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port1, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener1(&manager1, worker_ctrl_port1, ctrl_port1); + + RemoteFetchConfig config1{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port1, + .local_worker_port = worker_ctrl_port1, + .bytes_per_block = 1024, + .num_shards = 1, + }; + KVCacheStore store1(10, registry_addr, id1, config1); + + // 4. Setup Store 2 (Local Receiver/Fetcher) + RaidenId id2{"job2", "0", "data2", 0}; + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager2(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port2, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener2(&manager2, worker_ctrl_port2, ctrl_port2); + + RemoteFetchConfig config2{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port2, + .local_worker_port = worker_ctrl_port2, + .bytes_per_block = 1024, + .num_shards = 1, + }; + KVCacheStore store2(10, registry_addr, id2, config2); + + // 5. Insert blocks into store1 and fill data + store1.Insert({"remote_hash1"}, {{RaidenBlockID(id1, 1, BlockStatus::HOST)}}, + true); + store1.Insert({"remote_hash2"}, {{RaidenBlockID(id1, 2, BlockStatus::HOST)}}, + true); + store1.Insert({"remote_hash3"}, {{RaidenBlockID(id1, 3, BlockStatus::HOST)}}, + true); + store1.Insert({"remote_hash4"}, {{RaidenBlockID(id1, 4, BlockStatus::HOST)}}, + true); + + size_t block_size = manager1.bytes_per_block(); + uint8_t* src_ptr = manager1.GetHostPointer(0, 0); + std::memset(src_ptr + 1 * block_size, 'A', block_size); + std::memset(src_ptr + 2 * block_size, 'B', block_size); + std::memset(src_ptr + 3 * block_size, 'C', block_size); + std::memset(src_ptr + 4 * block_size, 'D', block_size); + + // 6. Register Store 1 blocks in Global Registry + auto channel = + grpc::CreateChannel(registry_addr, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + // Use id1.job_name as host_address so resolution works + std::vector regs = { + {"remote_hash1", id1.job_name, 1}, + {"remote_hash2", id1.job_name, 2}, + {"remote_hash3", id1.job_name, 3}, + {"remote_hash4", id1.job_name, 4}, + }; + ASSERT_OK(registry_client.Register(regs)); + + // 7. Insert local blocks into store2 + store2.Insert({"local_hash1"}, {{RaidenBlockID(id2, 7, BlockStatus::HOST)}}, + true); + store2.Insert({"local_hash2"}, {{RaidenBlockID(id2, 8, BlockStatus::HOST)}}, + true); + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // 8. Lookup multiple hashes in store2 + std::vector lookup_hashes = {"local_hash1", "local_hash2", + "remote_hash1", "remote_hash2", + "remote_hash3", "remote_hash4"}; + + auto lookup_res_or = store2.Lookup(lookup_hashes, /*enable_global=*/true); + ASSERT_TRUE(lookup_res_or.ok()); + auto lookup_res = lookup_res_or.value(); + ASSERT_EQ(lookup_res.size(), 6); + + // 9. Prepare InsertAndPin arguments + std::vector pin_hashes; + std::vector> pin_slices; + + // Use hardcoded destination block IDs as in other tests + // Must NOT conflict with local blocks 7 and 8 in store2 + std::vector allocated_blocks = {1, 2, 3, 4}; + + size_t remote_count = 0; + for (auto& pair : lookup_res) { + pin_hashes.push_back(pair.first); + if (pair.second[0].status == BlockStatus::REMOTE) { + // Assign allocated local block ID + pair.second[0].host_block_id = allocated_blocks[remote_count++]; + } + pin_slices.push_back(pair.second); + } + + // 10. Call InsertAndPin + auto pin_res = store2.InsertAndPin(pin_hashes, pin_slices, true); + EXPECT_TRUE(pin_res.first); + EXPECT_TRUE(pin_res.second.empty()); // No evictions expected + + // Verify pins + for (const auto& hash : pin_hashes) { + EXPECT_EQ(store2.GetPinCount(hash), 1); + } + + // 11. Call Fetch for remote hashes + std::vector remote_hashes = {"remote_hash1", "remote_hash2", + "remote_hash3", "remote_hash4"}; + auto futures = store2.FetchRemote(remote_hashes); + ASSERT_EQ(futures.size(), 4); + + // 12. Wait for completion + for (const auto& hash : remote_hashes) { + ASSERT_TRUE(futures.contains(hash)); + EXPECT_OK(futures[hash].Await()); + } + + // 13. Verify Destination Buffer + uint8_t* dst_ptr = manager2.GetHostPointer(0, 0); + EXPECT_EQ(dst_ptr[allocated_blocks[0] * block_size], 'A'); + EXPECT_EQ(dst_ptr[allocated_blocks[1] * block_size], 'B'); + EXPECT_EQ(dst_ptr[allocated_blocks[2] * block_size], 'C'); + EXPECT_EQ(dst_ptr[allocated_blocks[3] * block_size], 'D'); + + // 14. Verify Status Upgrade in Store 2 + for (size_t i = 0; i < 4; ++i) { + auto lookup_res_after = store2.Lookup({remote_hashes[i]}); + ASSERT_TRUE(lookup_res_after.ok()); + ASSERT_EQ(lookup_res_after->size(), 1); + const auto& block_list = (*lookup_res_after)[0].second; + ASSERT_EQ(block_list.size(), 1); + EXPECT_EQ(block_list[0].status, BlockStatus::HOST); + EXPECT_EQ(block_list[0].host_block_id, allocated_blocks[i]); + } + + registry_server->Shutdown(); +} + +// Verifies the remote fetch flow where both src and dst have multiple +// listeners. +TEST(RemoteFetchInteractionTest, QueueFlowMultiListenerEndToEnd) { + // 1. Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // 2. Setup Global Registry Server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int registry_port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + ®istry_port); + builder.RegisterService(service.get()); + auto registry_server = builder.BuildAndStart(); + std::string registry_addr = "localhost:" + std::to_string(registry_port); + + // 3. Setup Store 1 (Remote Sender) with 2 listeners + RaidenId id1{"job1", "0", "kv_cache", 0}; + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1a = net_util::PickUnusedPortOrDie(); + int worker_data_port1b = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1_base = PickConsecutiveUnusedPortsOrDie(2); + int worker_ctrl_port1_b = worker_ctrl_port1_base + 1; + + // Each manager handles 1 shard + KVCacheManagerBase manager1a(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port1a, + /*host_blocks_to_allocate=*/10); + KVCacheManagerBase manager1b(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port1b, + /*host_blocks_to_allocate=*/10); + + KVCacheListener listener1a(&manager1a, worker_ctrl_port1_base, ctrl_port1); + KVCacheListener listener1b(&manager1b, worker_ctrl_port1_b, ctrl_port1); + + RemoteFetchConfig config1{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port1, + .local_worker_port = worker_ctrl_port1_base, + .bytes_per_block = 1024, + .num_shards = 2, // Total shards + .num_listeners = 2, + }; + KVCacheStore store1(10, registry_addr, id1, config1); + + // 4. Setup Store 2 (Local Receiver/Fetcher) with 2 listeners + RaidenId id2{"job2", "0", "data2", 0}; + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2a = net_util::PickUnusedPortOrDie(); + int worker_data_port2b = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2_base = PickConsecutiveUnusedPortsOrDie(2); + int worker_ctrl_port2_b = worker_ctrl_port2_base + 1; + + KVCacheManagerBase manager2a(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port2a, + /*host_blocks_to_allocate=*/10); + KVCacheManagerBase manager2b(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port2b, + /*host_blocks_to_allocate=*/10); + + KVCacheListener listener2a(&manager2a, worker_ctrl_port2_base, ctrl_port2); + KVCacheListener listener2b(&manager2b, worker_ctrl_port2_b, ctrl_port2); + + RemoteFetchConfig config2{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port2, + .local_worker_port = worker_ctrl_port2_base, + .bytes_per_block = 1024, + .num_shards = 2, // Total shards + .num_listeners = 2, + }; + KVCacheStore store2(10, registry_addr, id2, config2); + + // 5. Insert blocks into store1 and fill data + // For simplicity, let's insert into both sub-managers if needed, or just test + // routing. Wait, store1.Insert just records metadata in store1. The actual + // data is in managers. + store1.Insert({"remote_hash1"}, + {{RaidenBlockID(id1, 1, BlockStatus::HOST), + RaidenBlockID(id1, 1, BlockStatus::HOST)}}, + true); + + size_t block_size = 1024; + // Fill data in manager1a (Shard 0) + uint8_t* src_ptr_a = manager1a.GetHostPointer(0, 0); + std::memset(src_ptr_a + 1 * block_size, 'A', block_size); + + // Fill data in manager1b (Shard 1) + uint8_t* src_ptr_b = manager1b.GetHostPointer(0, 0); + std::memset(src_ptr_b + 1 * block_size, 'B', block_size); + + // 6. Register Store 1 blocks in Global Registry + auto channel = + grpc::CreateChannel(registry_addr, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::vector regs = { + {"remote_hash1", id1.job_name, 1}, + }; + ASSERT_OK(registry_client.Register(regs)); + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // 7. Lookup in store2 + auto lookup_res_or = store2.Lookup({"remote_hash1"}, /*enable_global=*/true); + ASSERT_TRUE(lookup_res_or.ok()); + auto lookup_res = lookup_res_or.value(); + ASSERT_EQ(lookup_res.size(), 1); + + // 8. InsertAndPin in store2 + std::vector pin_hashes = {"remote_hash1"}; + std::vector> pin_slices; + std::vector allocated_blocks = {1}; // Local ID 1 in store2 + + auto& pair = lookup_res[0]; + if (pair.second[0].status == BlockStatus::REMOTE) { + pair.second[0].host_block_id = allocated_blocks[0]; + } + pin_slices.push_back(pair.second); + + auto pin_res = store2.InsertAndPin(pin_hashes, pin_slices, true); + EXPECT_TRUE(pin_res.first); + + // 9. Fetch Remote + auto futures = store2.FetchRemote(pin_hashes); + ASSERT_EQ(futures.size(), 1); + + // 10. Wait for completion + // This might hang if there are bugs in multi-listener support! + EXPECT_OK(futures["remote_hash1"].Await()); + + // 11. Verify Destination Buffers + uint8_t* dst_ptr_a = manager2a.GetHostPointer(0, 0); + EXPECT_EQ(dst_ptr_a[allocated_blocks[0] * block_size], 'A'); + + uint8_t* dst_ptr_b = manager2b.GetHostPointer(0, 0); + EXPECT_EQ(dst_ptr_b[allocated_blocks[0] * block_size], 'B'); + + registry_server->Shutdown(); +} + +// Verifies end-to-end fetch flow where blocks belong to multiple different +// remote RaidenIDs. This tests that requests are correctly grouped by +// destination and dispatched. +TEST(RemoteFetchInteractionTest, QueueFlowMultiRemoteEndToEnd) { + // 1. Setup Orchestrator + int orch_port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(orch_port); + std::string orch_addr = absl::StrCat("localhost:", orch_port); + + // 2. Setup Global Registry Server + auto service = std::make_unique(); + grpc::ServerBuilder builder; + int registry_port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + ®istry_port); + builder.RegisterService(service.get()); + auto registry_server = builder.BuildAndStart(); + std::string registry_addr = "localhost:" + std::to_string(registry_port); + + // 3. Setup Store 1 (Remote Sender 1) + RaidenId id1{"job1", "0", "kv_cache", 0}; + int ctrl_port1 = net_util::PickUnusedPortOrDie(); + int worker_data_port1 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port1 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager1(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port1, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener1(&manager1, worker_ctrl_port1, ctrl_port1); + + RemoteFetchConfig config1{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port1, + .local_worker_port = worker_ctrl_port1, + .bytes_per_block = 1024, + .num_shards = 1, + }; + KVCacheStore store1(10, registry_addr, id1, config1); + + // 4. Setup Store 3 (Remote Sender 2) + RaidenId id3{"job3", "0", "kv_cache", 0}; + int ctrl_port3 = net_util::PickUnusedPortOrDie(); + int worker_data_port3 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port3 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager3(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port3, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener3(&manager3, worker_ctrl_port3, ctrl_port3); + + RemoteFetchConfig config3{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port3, + .local_worker_port = worker_ctrl_port3, + .bytes_per_block = 1024, + .num_shards = 1, + }; + KVCacheStore store3(10, registry_addr, id3, config3); + + // 5. Setup Store 2 (Local Receiver/Fetcher) + RaidenId id2{"job2", "0", "data2", 0}; + int ctrl_port2 = net_util::PickUnusedPortOrDie(); + int worker_data_port2 = net_util::PickUnusedPortOrDie(); + int worker_ctrl_port2 = net_util::PickUnusedPortOrDie(); + + KVCacheManagerBase manager2(/*num_layers=*/1, /*num_shards=*/1, + /*slice_byte_size=*/1024, worker_data_port2, + /*host_blocks_to_allocate=*/10); + KVCacheListener listener2(&manager2, worker_ctrl_port2, ctrl_port2); + + RemoteFetchConfig config2{ + .orchestrator_address = orch_addr, + .controller_port = ctrl_port2, + .local_worker_port = worker_ctrl_port2, + .bytes_per_block = 1024, + .num_shards = 1, + }; + KVCacheStore store2(10, registry_addr, id2, config2); + + // 6. Insert blocks into store1 and store3 and fill data + store1.Insert({"remote_hash1"}, {{RaidenBlockID(id1, 1, BlockStatus::HOST)}}, + true); + store1.Insert({"remote_hash2"}, {{RaidenBlockID(id1, 2, BlockStatus::HOST)}}, + true); + store3.Insert({"remote_hash3"}, {{RaidenBlockID(id3, 3, BlockStatus::HOST)}}, + true); + store3.Insert({"remote_hash4"}, {{RaidenBlockID(id3, 4, BlockStatus::HOST)}}, + true); + + size_t block_size = manager1.bytes_per_block(); + + // Fill Store 1 data + uint8_t* src_ptr1 = manager1.GetHostPointer(0, 0); + std::memset(src_ptr1 + 1 * block_size, 'A', block_size); + std::memset(src_ptr1 + 2 * block_size, 'B', block_size); + + // Fill Store 3 data + uint8_t* src_ptr3 = manager3.GetHostPointer(0, 0); + std::memset(src_ptr3 + 3 * block_size, 'C', block_size); + std::memset(src_ptr3 + 4 * block_size, 'D', block_size); + + // 7. Register blocks in Global Registry + auto channel = + grpc::CreateChannel(registry_addr, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + + std::vector regs = { + {"remote_hash1", id1.job_name, 1}, + {"remote_hash2", id1.job_name, 2}, + {"remote_hash3", id3.job_name, 3}, + {"remote_hash4", id3.job_name, 4}, + }; + ASSERT_OK(registry_client.Register(regs)); + + absl::SleepFor(absl::Milliseconds(500)); // Wait for registrations + + // 8. Lookup multiple hashes in store2 + std::vector lookup_hashes = {"remote_hash1", "remote_hash2", + "remote_hash3", "remote_hash4"}; + + auto lookup_res_or = store2.Lookup(lookup_hashes, /*enable_global=*/true); + ASSERT_TRUE(lookup_res_or.ok()); + auto lookup_res = lookup_res_or.value(); + ASSERT_EQ(lookup_res.size(), 4); + + // 9. Prepare InsertAndPin arguments + std::vector pin_hashes; + std::vector> pin_slices; + + std::vector allocated_blocks = {1, 2, 3, 4}; + + size_t remote_count = 0; + for (auto& pair : lookup_res) { + pin_hashes.push_back(pair.first); + if (pair.second[0].status == BlockStatus::REMOTE) { + pair.second[0].host_block_id = allocated_blocks[remote_count++]; + } + pin_slices.push_back(pair.second); + } + + // 10. Call InsertAndPin + auto pin_res = store2.InsertAndPin(pin_hashes, pin_slices, true); + EXPECT_TRUE(pin_res.first); + EXPECT_TRUE(pin_res.second.empty()); + + // 11. Call Fetch for remote hashes + auto futures = store2.FetchRemote(lookup_hashes); + ASSERT_EQ(futures.size(), 4); + + // 12. Wait for completion + for (const auto& hash : lookup_hashes) { + ASSERT_TRUE(futures.contains(hash)); + EXPECT_OK(futures[hash].Await()); + } + + // 13. Verify Destination Buffer + uint8_t* dst_ptr = manager2.GetHostPointer(0, 0); + EXPECT_EQ(dst_ptr[allocated_blocks[0] * block_size], 'A'); + EXPECT_EQ(dst_ptr[allocated_blocks[1] * block_size], 'B'); + EXPECT_EQ(dst_ptr[allocated_blocks[2] * block_size], 'C'); + EXPECT_EQ(dst_ptr[allocated_blocks[3] * block_size], 'D'); + + // 14. Verify Status Upgrade in Store 2 + for (size_t i = 0; i < 4; ++i) { + auto lookup_res_after = store2.Lookup({lookup_hashes[i]}); + ASSERT_TRUE(lookup_res_after.ok()); + ASSERT_EQ(lookup_res_after->size(), 1); + const auto& block_list = (*lookup_res_after)[0].second; + ASSERT_EQ(block_list.size(), 1); + EXPECT_EQ(block_list[0].status, BlockStatus::HOST); + EXPECT_EQ(block_list[0].host_block_id, allocated_blocks[i]); + } + + registry_server->Shutdown(); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/lru_cache.h b/tpu_raiden/kv_cache/lru_cache.h index 35414286..615c593a 100644 --- a/tpu_raiden/kv_cache/lru_cache.h +++ b/tpu_raiden/kv_cache/lru_cache.h @@ -53,6 +53,9 @@ class LRUCache { // Returns the number of elements currently stored in the cache. size_t size() const { return map_.size(); } + // Returns the available space (free + evictable capacity). + size_t available_space() const { return capacity_ - pinned_list_.size(); } + // Returns true if the cache is completely empty. bool empty() const { return map_.empty(); } @@ -97,6 +100,29 @@ class LRUCache { return evicted; } + // Inserts a key-value pair at the back of the lru_list_ (least recently + // used position). If the key already exists, its value is updated and moved + // to the back of lru_list_. + void PutBack(Key key, Value value) { + auto it = map_.find(key); + if (it != map_.end()) { + it->second->value = std::move(value); + if (it->second->pin_count > 0) { + pinned_list_.splice(pinned_list_.end(), pinned_list_, it->second); + } else { + lru_list_.splice(lru_list_.end(), lru_list_, it->second); + } + return; + } + + if (map_.size() >= capacity_) { + Evict(); + } + + lru_list_.push_back(CacheNode{std::move(key), std::move(value), 0}); + map_[lru_list_.back().key] = std::prev(lru_list_.end()); + } + // Retrieves a pointer to the stored value for the given key and promotes // the item to Most Recently Used (MRU). Returns nullptr if absent. Value* Get(const Key& key) { diff --git a/tpu_raiden/kv_cache/lru_cache_test.cc b/tpu_raiden/kv_cache/lru_cache_test.cc index 2053357a..651b6f40 100644 --- a/tpu_raiden/kv_cache/lru_cache_test.cc +++ b/tpu_raiden/kv_cache/lru_cache_test.cc @@ -144,6 +144,29 @@ TEST(LRUCacheTest, EraseAndClear) { EXPECT_EQ(cache.size(), 0); } +TEST(LRUCacheTest, PutBack) { + LRUCache cache(3); + cache.Put(1, "one"); + cache.Put(2, "two"); + // MRU to LRU is: 2, 1. + + // PutBack(3) should add 3 to the back (LRU position). + cache.PutBack(3, "three"); + // Now MRU to LRU should be: 2, 1, 3. + // Verify by checking Evict(): first evicted item should be 3. + auto evicted1 = cache.Evict(); + ASSERT_TRUE(evicted1.has_value()); + EXPECT_EQ(evicted1->first, 3); + + auto evicted2 = cache.Evict(); + ASSERT_TRUE(evicted2.has_value()); + EXPECT_EQ(evicted2->first, 1); + + auto evicted3 = cache.Evict(); + ASSERT_TRUE(evicted3.has_value()); + EXPECT_EQ(evicted3->first, 2); +} + } // namespace } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/raiden_controller_embedded.cc b/tpu_raiden/kv_cache/raiden_controller_embedded.cc new file mode 100644 index 00000000..58f6010f --- /dev/null +++ b/tpu_raiden/kv_cache/raiden_controller_embedded.cc @@ -0,0 +1,660 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/kv_cache/raiden_controller_embedded.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include // NOLINT +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "tpu_raiden/core/status_macros.h" +#include "tpu_raiden/core/tpu_utils.h" +#include "tpu_raiden/kv_cache/kv_cache_store.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" +#include "tpu_raiden/rpc/rpc_utils.h" +#include "tpu_raiden/transport/raw_buffer_transport.h" + +namespace tpu_raiden { +namespace kv_cache { + +using ::tpu_raiden::rpc::ControlRequest; +using ::tpu_raiden::rpc::ControlResponse; +using ::tpu_raiden::rpc::SendRpcSync; + +RaidenControllerEmbedded::RaidenControllerEmbedded( + KVCacheStore* store, int port, const std::string& orchestrator_address, + int local_worker_port, + const std::vector& local_worker_data_addresses, + size_t bytes_per_block, size_t num_shards, size_t num_listeners) + : store_(store), + + port_(port), + orchestrator_address_(orchestrator_address), + local_worker_port_(local_worker_port), + local_worker_data_addresses_(local_worker_data_addresses), + bytes_per_block_(bytes_per_block), + num_shards_(num_shards), + num_listeners_(num_listeners) {} + +RaidenControllerEmbedded::~RaidenControllerEmbedded() { Stop(); } + +absl::Status RaidenControllerEmbedded::Start() { + server_fd_ = socket(AF_INET6, SOCK_STREAM, 0); + if (server_fd_ < 0) { + return absl::InternalError(std::string("Failed to create socket: ") + + std::strerror(errno)); + } + + int opt = 1; + setsockopt(server_fd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)); + + sockaddr_in6 address{ + .sin6_family = AF_INET6, + .sin6_port = htons(port_), + .sin6_addr = in6addr_any, + }; + + if (bind(server_fd_, reinterpret_cast(&address), sizeof(address)) < + 0) { + close(server_fd_); + server_fd_ = -1; + return absl::InternalError(std::string("Bind failed: ") + + std::strerror(errno)); + } + + if (listen(server_fd_, 128) < 0) { + close(server_fd_); + server_fd_ = -1; + return absl::InternalError(std::string("Listen failed: ") + + std::strerror(errno)); + } + + socklen_t addr_len = sizeof(address); + if (getsockname(server_fd_, reinterpret_cast(&address), + &addr_len) == 0) { + port_ = ntohs(address.sin6_port); + } + + // Determine self IP address properly. + std::vector host_nics = GetLocalHostNicAddresses(); + std::string ctrl_ip = "localhost"; + for (const auto& nic : host_nics) { + if (nic.classification == NicClassification::kControlPlane) { + ctrl_ip = nic.ip_address; + break; + } + } + if (ctrl_ip == "localhost" && !host_nics.empty()) { + ctrl_ip = host_nics[0].ip_address; + } + + self_address_ = absl::StrContains(ctrl_ip, ':') + ? absl::StrCat("[", ctrl_ip, "]:", port_) + : absl::StrCat(ctrl_ip, ":", port_); + + LOG(INFO) << "RaidenControllerEmbedded listening on port: " << port_ + << " with address: " << self_address_; + + // Query Local Worker for its data endpoints + local_worker_data_addresses_.clear(); + listener_shard_counts_.assign(num_listeners_, 0); + for (size_t i = 0; i < num_listeners_; ++i) { + int port = local_worker_port_ + i; + std::string worker_addr = absl::StrCat("localhost:", port); + ControlRequest req; + req.set_command(ControlRequest::COMMAND_GET_ENDPOINTS); + ControlResponse resp; + absl::Status status = SendRpcSync(worker_addr, req, resp); + if (status.ok() && resp.success()) { + listener_shard_counts_[i] = resp.endpoints_size(); + for (const auto& addr : resp.endpoints()) { + local_worker_data_addresses_.push_back(addr); + LOG(INFO) << "Auto-discovered worker data address from listener " << i + << ": " << addr; + } + } else { + LOG(WARNING) << "Failed to query worker endpoints from listener " << i + << ": " << status << " " << resp.message(); + } + } + + absl::Status status = RegisterWithOrchestrator(); + + if (!status.ok()) { + LOG(ERROR) << "Failed to register with orchestrator: " << status; + // Might want to fail start if registration fails. + } + + listener_thread_ = std::thread(&RaidenControllerEmbedded::ListenerLoop, this); + poller_thread_ = + std::thread(&RaidenControllerEmbedded::WorkQueuePollerLoop, this); + + return absl::OkStatus(); +} + +void RaidenControllerEmbedded::Stop() { + stopping_ = true; + + // Wake up poller thread if it is blocked on PopFetchWork + if (store_) { + FetchRequest dummy; + store_->PushFetchWork(dummy); + } + + if (server_fd_ >= 0) { + int sock = socket(AF_INET6, SOCK_STREAM, 0); + if (sock >= 0) { + sockaddr_in6 serv_addr{}; + serv_addr.sin6_family = AF_INET6; + serv_addr.sin6_port = htons(port_); + inet_pton(AF_INET6, "::1", &serv_addr.sin6_addr); + connect(sock, reinterpret_cast(&serv_addr), sizeof(serv_addr)); + close(sock); + } + close(server_fd_); + server_fd_ = -1; + } + + if (listener_thread_.joinable()) listener_thread_.join(); + if (poller_thread_.joinable()) poller_thread_.join(); + + for (auto& t : worker_threads_) { + if (t.joinable()) t.join(); + } +} + +absl::Status RaidenControllerEmbedded::RegisterWithOrchestrator() { + if (orchestrator_address_.empty()) { + return absl::InvalidArgumentError("Orchestrator address is empty"); + } + + ControlRequest req; + req.set_command(ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg_req = req.mutable_register_work_unit_request(); + + const auto& self_id = store_->raiden_id(); + auto* unit = reg_req->mutable_unit(); + unit->set_job_name(self_id.job_name); + unit->set_job_replica_id(self_id.job_replica_id); + unit->set_data_name(self_id.data_name); + unit->set_data_replica_idx(self_id.data_replica_idx); + + reg_req->set_control_plane_rpc_address(self_address_); + + ControlResponse resp; + RETURN_IF_ERROR(SendRpcSync(orchestrator_address_, req, resp)); + + if (!resp.success()) { + return absl::InternalError( + absl::StrCat("Registration failed: ", resp.message())); + } + + return absl::OkStatus(); +} + +absl::StatusOr RaidenControllerEmbedded::ResolveRemoteController( + const RaidenId& remote_id) { + if (orchestrator_address_.empty()) { + return absl::InvalidArgumentError("Orchestrator address is empty"); + } + + ControlRequest req; + req.set_command(ControlRequest::COMMAND_RESOLVE_CONTROLLER); + auto* target = req.mutable_target_unit(); + target->set_job_name(remote_id.job_name); + target->set_job_replica_id(remote_id.job_replica_id); + target->set_data_name(remote_id.data_name); + target->set_data_replica_idx(remote_id.data_replica_idx); + + ControlResponse resp; + RETURN_IF_ERROR(SendRpcSync(orchestrator_address_, req, resp)); + + if (!resp.success()) { + return absl::NotFoundError( + absl::StrCat("Resolution failed: ", resp.message())); + } + + return resp.response_data(); +} + +absl::StatusOr +RaidenControllerEmbedded::NegotiateFetch(const std::string& remote_addr, + const FetchRequestItem& request, + uint64_t uuid) { + if (request.block_hashes.empty()) { + tpu_raiden::rpc::ControlResponse empty_resp; + return empty_resp; + } + + ControlRequest req; + req.set_command(ControlRequest::COMMAND_NEGOTIATE_FETCH); + auto* neg_req = req.mutable_fetch_negotiation_request(); + + const auto& src_id = request.src_raiden_id; + auto* src_unit = neg_req->mutable_src_unit(); + src_unit->set_job_name(src_id.job_name); + src_unit->set_job_replica_id(src_id.job_replica_id); + src_unit->set_data_name(src_id.data_name); + src_unit->set_data_replica_idx(src_id.data_replica_idx); + + const auto& self_id = store_->raiden_id(); + auto* dst_unit = neg_req->mutable_dst_unit(); + dst_unit->set_job_name(self_id.job_name); + dst_unit->set_job_replica_id(self_id.job_replica_id); + dst_unit->set_data_name(self_id.data_name); + dst_unit->set_data_replica_idx(self_id.data_replica_idx); + + for (const auto& addr : local_worker_data_addresses_) { + neg_req->add_dst_worker_data_addresses(addr); + } + + for (size_t i = 0; i < request.block_hashes.size(); ++i) { + neg_req->add_block_hashes(request.block_hashes[i]); + neg_req->add_dst_block_ids(request.dst_block_ids[i]); + } + neg_req->set_uuid(uuid); + + ControlResponse resp; + RETURN_IF_ERROR(SendRpcSync(remote_addr, req, resp)); + + if (!resp.success()) { + return absl::InternalError( + absl::StrCat("Negotiation failed: ", resp.message())); + } + + return resp; +} + +void RaidenControllerEmbedded::WorkQueuePollerLoop() { + while (!stopping_) { + FetchRequest req; + if (!store_->PopFetchWork(req)) { + if (stopping_) break; + continue; + } + + if (req.empty()) continue; + + for (const auto& batch_item : req) { + auto addr_or = ResolveRemoteController(batch_item.src_raiden_id); + if (!addr_or.ok()) { + LOG(ERROR) << "Failed to resolve remote controller for " + << batch_item.src_raiden_id.job_name << ": " + << addr_or.status(); + // TODO: Push failure to CompletionQueue + continue; + } + + // Record pending fetches before negotiating + { + absl::MutexLock lock(pending_mu_); + for (size_t i = 0; i < batch_item.block_hashes.size(); ++i) { + pending_fetches_[batch_item.dst_block_ids[i]] = + batch_item.block_hashes[i]; + } + } + + uint64_t uuid = absl::ToUnixMicros(absl::Now()); + auto neg_resp_or = NegotiateFetch(addr_or.value(), batch_item, uuid); + + if (!neg_resp_or.ok()) { + LOG(ERROR) << "Negotiation failed with " << addr_or.value() << ": " + << neg_resp_or.status(); + // TODO: Push failure to CompletionQueue + continue; + } + + ControlRequest base_start_req; + base_start_req.set_command(ControlRequest::COMMAND_START_TRANSFER); + auto* base_start_transfer = + base_start_req.mutable_start_transfer_request(); + if (neg_resp_or.value().has_start_transfer_request()) { + *base_start_transfer = neg_resp_or.value().start_transfer_request(); + } + base_start_transfer->set_is_sender(false); + base_start_transfer->set_use_block_chunks(false); + base_start_transfer->set_uuid(uuid); + base_start_transfer->set_expected_block_count( + batch_item.block_hashes.size()); + base_start_transfer->set_dst_mem_type(tpu_raiden::rpc::MEMORY_TYPE_DRAM); + + // Broadcast START_TRANSFER (is_sender=False) to all local listeners + size_t global_shard_offset = 0; + for (size_t i = 0; i < num_listeners_; ++i) { + size_t num_local_shards = listener_shard_counts_[i]; + if (num_local_shards == 0) continue; + + ControlRequest start_req_i = base_start_req; + auto* start_transfer_i = start_req_i.mutable_start_transfer_request(); + + if (start_transfer_i->shard_push_schedules().size() > 0) { + start_transfer_i->clear_shard_push_schedules(); + auto& schedules_i = *start_transfer_i->mutable_shard_push_schedules(); + for (size_t local_sh = 0; local_sh < num_local_shards; ++local_sh) { + size_t global_sh = global_shard_offset + local_sh; + auto it = base_start_transfer->shard_push_schedules().find( + static_cast(global_sh)); + if (it != base_start_transfer->shard_push_schedules().end()) { + schedules_i[static_cast(local_sh)] = it->second; + } + } + } + + int port = local_worker_port_ + i; + std::string worker_addr = absl::StrCat("localhost:", port); + + ControlResponse start_resp; + absl::Status start_status = + SendRpcSync(worker_addr, start_req_i, start_resp); + if (!start_status.ok() || !start_resp.success()) { + LOG(WARNING) << "Failed to broadcast START_TRANSFER to listener " << i + << ": " << start_status << " " << start_resp.message(); + } else { + VLOG(1) << "Broadcasted START_TRANSFER to listener " << i; + } + + global_shard_offset += num_local_shards; + } + + if (neg_resp_or.value().has_start_transfer_request()) { + ControlRequest exec_req; + exec_req.set_command(ControlRequest::COMMAND_EXECUTE_FETCH); + *exec_req.mutable_start_transfer_request() = + neg_resp_or.value().start_transfer_request(); + exec_req.mutable_start_transfer_request()->set_is_sender(true); + ControlResponse exec_resp; + absl::Status exec_status = + SendRpcSync(addr_or.value(), exec_req, exec_resp); + if (!exec_status.ok() || !exec_resp.success()) { + LOG(ERROR) << "Execute fetch failed: " << exec_status << " " + << exec_resp.message(); + } + } + } + } +} + +void RaidenControllerEmbedded::ListenerLoop() { + while (!stopping_) { + sockaddr_in6 client_addr{}; + socklen_t client_len = sizeof(client_addr); + int client_fd = accept( + server_fd_, reinterpret_cast(&client_addr), &client_len); + if (client_fd < 0) { + if (stopping_) break; + continue; + } + + worker_threads_.push_back(std::thread( + &RaidenControllerEmbedded::ConnectionWorker, this, client_fd)); + } +} + +void RaidenControllerEmbedded::ConnectionWorker(int client_fd) { + uint32_t net_len = 0; + if (!transport::RawBufferTransport::ReadExact(client_fd, &net_len, + sizeof(net_len)) + .ok()) { + close(client_fd); + return; + } + uint32_t payload_len = ntohl(net_len); + + std::vector buffer(payload_len); + if (!transport::RawBufferTransport::ReadExact(client_fd, buffer.data(), + payload_len) + .ok()) { + close(client_fd); + return; + } + + ControlRequest req; + if (!req.ParseFromString(absl::string_view(buffer.data(), buffer.size()))) { + LOG(ERROR) << "Failed to parse ControlRequest Protobuf"; + close(client_fd); + return; + } + + ControlResponse resp; + resp.set_success(true); + resp.set_message("SUCCESS"); + + if (req.command() == ControlRequest::COMMAND_NEGOTIATE_FETCH) { + if (req.has_fetch_negotiation_request()) { + const auto& neg_req = req.fetch_negotiation_request(); + + std::vector hashes(neg_req.block_hashes().begin(), + neg_req.block_hashes().end()); + auto lookup_res_or = store_->Lookup(hashes, /*enable_global=*/false); + + if (!lookup_res_or.ok()) { + resp.set_success(false); + resp.set_message(std::string(lookup_res_or.status().message())); + } else { + const auto& lookup_res = lookup_res_or.value(); + auto* neg_resp = resp.mutable_fetch_negotiation_response(); + + std::set found_set; + for (const auto& pair : lookup_res) { + neg_resp->add_found_block_hashes(pair.first); + found_set.insert(pair.first); + } + for (const auto& hash : hashes) { + if (found_set.find(hash) == found_set.end()) { + neg_resp->add_missing_block_hashes(hash); + } + } + + // Trigger PUSH via local worker if all/any blocks found. + if (!lookup_res.empty()) { + absl::flat_hash_map hash_to_dst_id; + for (int i = 0; i < neg_req.block_hashes_size(); ++i) { + if (i < neg_req.dst_block_ids_size()) { + hash_to_dst_id[neg_req.block_hashes(i)] = + neg_req.dst_block_ids(i); + } + } + + ControlRequest start_req; + start_req.set_command(ControlRequest::COMMAND_START_TRANSFER); + auto* start_transfer = start_req.mutable_start_transfer_request(); + start_transfer->set_is_sender(true); + start_transfer->set_use_block_chunks(false); + start_transfer->set_src_mem_type(tpu_raiden::rpc::MEMORY_TYPE_DRAM); + uint64_t uuid = neg_req.uuid(); + start_transfer->set_uuid(uuid); + + auto& schedules = *start_transfer->mutable_shard_push_schedules(); + + for (const auto& pair : lookup_res) { + const std::string& hash = pair.first; + const auto& local_blocks = pair.second; + if (local_blocks.empty()) continue; + + int src_block_id = local_blocks[0].host_block_id; + auto it = hash_to_dst_id.find(hash); + if (it == hash_to_dst_id.end()) { + LOG(WARNING) + << "Found hash in lookup but not in negotiation request: " + << hash; + continue; + } + int dst_block_id = it->second; + + absl::flat_hash_map peer_to_local_shard; + for (size_t sh = 0; sh < num_shards_; ++sh) { + auto* entry = schedules[sh].add_entries(); + std::string peer_addr; + if (sh < neg_req.dst_worker_data_addresses_size()) { + peer_addr = neg_req.dst_worker_data_addresses(sh); + entry->set_dst_peer(peer_addr); + } else { + LOG(ERROR) << "Missing peer address for shard " << sh; + continue; + } + int local_dst_sh = peer_to_local_shard[peer_addr]++; + entry->set_src_block_id(src_block_id); + entry->set_dst_block_id(dst_block_id); + entry->set_dst_shard_idx(local_dst_sh); + entry->set_size_bytes(bytes_per_block_); + entry->set_count(1); + entry->set_src_offset_bytes(0); + entry->set_dst_offset_bytes(0); + entry->set_src_stride_bytes(0); + entry->set_dst_stride_bytes(0); + } + } + + if (resp.success()) { + *resp.mutable_start_transfer_request() = + start_req.start_transfer_request(); + } + } + } + } else { + resp.set_success(false); + resp.set_message("Missing fetch_negotiation_request"); + } + } else if (req.command() == ControlRequest::COMMAND_EXECUTE_FETCH) { + if (!req.has_start_transfer_request()) { + resp.set_success(false); + resp.set_message("Missing start_transfer_request"); + } else { + const auto& start_req = req; + size_t global_shard_offset = 0; + for (size_t i = 0; i < num_listeners_; ++i) { + size_t num_local_shards = listener_shard_counts_[i]; + if (num_local_shards == 0) continue; + + ControlRequest start_req_i = start_req; + start_req_i.set_command(ControlRequest::COMMAND_START_TRANSFER); + auto* start_transfer_i = start_req_i.mutable_start_transfer_request(); + start_transfer_i->clear_shard_push_schedules(); + auto& schedules_i = *start_transfer_i->mutable_shard_push_schedules(); + + bool has_schedule = false; + for (size_t local_sh = 0; local_sh < num_local_shards; ++local_sh) { + size_t global_sh = global_shard_offset + local_sh; + auto it = + start_req.start_transfer_request().shard_push_schedules().find( + static_cast(global_sh)); + if (it != + start_req.start_transfer_request().shard_push_schedules().end()) { + schedules_i[static_cast(local_sh)] = it->second; + has_schedule = true; + } + } + + if (!has_schedule) { + global_shard_offset += num_local_shards; + continue; // Skip if this listener has no work in this request + } + + int port = local_worker_port_ + i; + std::string worker_addr = absl::StrCat("localhost:", port); + ControlResponse worker_resp; + absl::Status worker_status = + SendRpcSync(worker_addr, start_req_i, worker_resp); + + if (!worker_status.ok()) { + LOG(ERROR) << "Failed to trigger local worker PUSH for listener " << i + << ": " << worker_status; + resp.set_success(false); + resp.set_message("Failed to trigger local worker PUSH: " + + std::string(worker_status.message())); + break; + } else if (!worker_resp.success()) { + LOG(ERROR) << "Local worker PUSH failed for listener " << i << ": " + << worker_resp.message(); + resp.set_success(false); + resp.set_message("Local worker PUSH failed: " + + worker_resp.message()); + break; + } + + global_shard_offset += num_local_shards; + } + } + } else if (req.command() == ControlRequest::COMMAND_TRANSFER_COMPLETED) { + VLOG(1) << "Received COMMAND_TRANSFER_COMPLETED"; + FetchCompletion completion; + { + absl::MutexLock lock(pending_mu_); + for (int32_t block_id : req.completed_block_ids()) { + auto it = pending_fetches_.find(block_id); + if (it != pending_fetches_.end()) { + block_completion_counts_[block_id]++; + if (block_completion_counts_[block_id] >= num_listeners_) { + FetchCompletionItem item; + item.block_hash = it->second; + item.host_block_id = block_id; + item.success = true; + completion.push_back(item); + pending_fetches_.erase(it); + block_completion_counts_.erase(block_id); + VLOG(1) << "Fetch completed for block_id " << block_id + << " across all listeners"; + } else { + VLOG(1) << "Fetch completed for block_id " << block_id << " on " + << block_completion_counts_[block_id] << "/" + << num_listeners_ << " listeners"; + } + } else { + LOG(WARNING) << "Received completion for unknown block_id: " + << block_id; + } + } + } + + if (!completion.empty()) { + store_->PushFetchCompletion(std::move(completion)); + } + resp.set_success(true); + } else { + resp.set_success(false); + resp.set_message("Unsupported command"); + } + + std::string resp_str; + if (resp.SerializeToString(&resp_str)) { + uint32_t resp_net_len = htonl(resp_str.size()); + transport::RawBufferTransport::WriteExact(client_fd, &resp_net_len, + sizeof(resp_net_len)) + .IgnoreError(); + transport::RawBufferTransport::WriteExact(client_fd, resp_str.data(), + resp_str.size()) + .IgnoreError(); + } + + close(client_fd); +} + +} // namespace kv_cache +} // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/raiden_controller_embedded.h b/tpu_raiden/kv_cache/raiden_controller_embedded.h new file mode 100644 index 00000000..e1a0104d --- /dev/null +++ b/tpu_raiden/kv_cache/raiden_controller_embedded.h @@ -0,0 +1,111 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_KV_CACHE_RAIDEN_CONTROLLER_EMBEDDED_H_ +#define THIRD_PARTY_TPU_RAIDEN_KV_CACHE_RAIDEN_CONTROLLER_EMBEDDED_H_ + +#include +#include +#include // NOLINT +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "tpu_raiden/kv_cache/kv_cache_store.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" + +namespace tpu_raiden { +namespace kv_cache { + +// Embedded C++ implementation of RaidenController running in the control plane. +// Manages asynchronous remote fetching via FetchWorkQueue and +// FetchCompletionQueue. +class RaidenControllerEmbedded { + public: + RaidenControllerEmbedded( + KVCacheStore* store, int port, const std::string& orchestrator_address, + int local_worker_port, + const std::vector& local_worker_data_addresses, + size_t bytes_per_block, size_t num_shards, size_t num_listeners = 1); + ~RaidenControllerEmbedded(); + + RaidenControllerEmbedded(const RaidenControllerEmbedded&) = delete; + RaidenControllerEmbedded& operator=(const RaidenControllerEmbedded&) = delete; + + // Starts the background threads (Listener and Queue Poller). + absl::Status Start(); + + // Stops the background threads. + void Stop(); + + int port() const { return port_; } + + private: + // Thread function to poll FetchWorkQueue and process requests. + void WorkQueuePollerLoop(); + + // Thread function to listen for incoming control connections from other + // controllers. + void ListenerLoop(); + + // Worker thread function to handle a specific control connection. + void ConnectionWorker(int client_fd); + + // Helper to register this controller with the central orchestrator. + absl::Status RegisterWithOrchestrator(); + + // Helper to resolve a remote controller address via the orchestrator. + absl::StatusOr ResolveRemoteController( + const RaidenId& remote_id); + + // Helper to negotiate fetch with a remote source controller. + absl::StatusOr NegotiateFetch( + const std::string& remote_addr, const FetchRequestItem& request, + uint64_t uuid); + + KVCacheStore* store_; // Not owned. + int port_; + std::string orchestrator_address_; + std::string self_address_; + + int local_worker_port_; + std::vector local_worker_data_addresses_; + size_t bytes_per_block_; + size_t num_shards_; + size_t num_listeners_; + std::vector listener_shard_counts_; + + int server_fd_ = -1; + + std::atomic stopping_{false}; + + absl::Mutex pending_mu_; + absl::flat_hash_map pending_fetches_ + ABSL_GUARDED_BY(pending_mu_); + // Maps dst_block_id to number of listeners that have completed + absl::flat_hash_map block_completion_counts_ + ABSL_GUARDED_BY(pending_mu_); + // Maps dst_block_id to number of active fetch listeners + absl::flat_hash_map fetch_listeners_count_ + ABSL_GUARDED_BY(pending_mu_); + + std::thread poller_thread_; + std::thread listener_thread_; + std::vector worker_threads_; +}; + +} // namespace kv_cache +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_KV_CACHE_RAIDEN_CONTROLLER_EMBEDDED_H_ diff --git a/tpu_raiden/kv_cache/raiden_orchestrator.cc b/tpu_raiden/kv_cache/raiden_orchestrator.cc new file mode 100644 index 00000000..4de1b02e --- /dev/null +++ b/tpu_raiden/kv_cache/raiden_orchestrator.cc @@ -0,0 +1,268 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/kv_cache/raiden_orchestrator.h" + +#include +#include +#include +#include + +#include +#include +#include +#include // NOLINT +#include +#include + +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/synchronization/mutex.h" +#include "tpu_raiden/kv_cache/kv_cache_store.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" +#include "tpu_raiden/transport/raw_buffer_transport.h" + +namespace tpu_raiden { +namespace kv_cache { + +using ::tpu_raiden::rpc::ControlRequest; +using ::tpu_raiden::rpc::ControlResponse; + +RaidenOrchestrator::RaidenOrchestrator(int port, const std::string& bind_ip) + : port_(port) { + bool is_ipv4 = !bind_ip.empty() && absl::StrContains(bind_ip, '.'); + int family = is_ipv4 ? AF_INET : AF_INET6; + + server_fd_ = socket(family, SOCK_STREAM, 0); + if (server_fd_ < 0) { + LOG(FATAL) << "Failed to create RaidenOrchestrator socket: " + << std::strerror(errno); + } + + int opt = 1; + if (setsockopt(server_fd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt))) { + LOG(WARNING) << "setsockopt SO_REUSEADDR failed"; + } + + if (family == AF_INET) { + struct sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_port = htons(port_); + if (!bind_ip.empty()) { + inet_pton(AF_INET, bind_ip.c_str(), &address.sin_addr); + } else { + address.sin_addr.s_addr = INADDR_ANY; + } + if (bind(server_fd_, reinterpret_cast(&address), + sizeof(address)) < 0) { + LOG(FATAL) << "RaidenOrchestrator bind failed on port " << port_ << ": " + << std::strerror(errno); + } + } else { + struct sockaddr_in6 address{}; + address.sin6_family = AF_INET6; + address.sin6_port = htons(port_); + if (!bind_ip.empty()) { + inet_pton(AF_INET6, bind_ip.c_str(), &address.sin6_addr); + } else { + address.sin6_addr = in6addr_any; + } + if (bind(server_fd_, reinterpret_cast(&address), + sizeof(address)) < 0) { + LOG(FATAL) << "RaidenOrchestrator bind failed on port " << port_ << ": " + << std::strerror(errno); + } + } + + if (listen(server_fd_, 128) < 0) { + LOG(FATAL) << "RaidenOrchestrator listen failed: " << std::strerror(errno); + } + + // Retrieve actual port if 0 was passed + struct sockaddr_storage ss; + socklen_t addr_len = sizeof(ss); + if (getsockname(server_fd_, reinterpret_cast(&ss), &addr_len) == + 0) { + if (ss.ss_family == AF_INET) { + struct sockaddr_in* sin = reinterpret_cast(&ss); + port_ = ntohs(sin->sin_port); + } else if (ss.ss_family == AF_INET6) { + struct sockaddr_in6* sin6 = reinterpret_cast(&ss); + port_ = ntohs(sin6->sin6_port); + } + } + + LOG(INFO) << "RaidenOrchestrator actively listening on port: " << port_; + + listener_thread_ = std::thread(&RaidenOrchestrator::ListenerLoop, this); +} + +RaidenOrchestrator::~RaidenOrchestrator() { + stopping_ = true; + if (server_fd_ >= 0) { + // Wake up listener thread by connecting to self + int sock = socket(AF_INET6, SOCK_STREAM, 0); + if (sock >= 0) { + sockaddr_in6 serv_addr{}; + serv_addr.sin6_family = AF_INET6; + serv_addr.sin6_port = htons(port_); + inet_pton(AF_INET6, "::1", &serv_addr.sin6_addr); + connect(sock, reinterpret_cast(&serv_addr), sizeof(serv_addr)); + close(sock); + } + close(server_fd_); + } + + if (listener_thread_.joinable()) { + listener_thread_.join(); + } + + for (auto& t : worker_threads_) { + if (t.joinable()) { + t.join(); + } + } +} + +absl::Status RaidenOrchestrator::RegisterController( + const RaidenId& raiden_id, const std::string& address) { + absl::MutexLock lock(&mutex_); + registry_[raiden_id] = address; + LOG(INFO) << "Registered controller for " << raiden_id.job_name << ":" + << raiden_id.job_replica_id << " at " << address; + return absl::OkStatus(); +} + +absl::StatusOr RaidenOrchestrator::ResolveController( + const RaidenId& raiden_id) { + absl::MutexLock lock(&mutex_); + auto it = registry_.find(raiden_id); + if (it == registry_.end()) { + return absl::NotFoundError("Controller not found for specified RaidenId"); + } + return it->second; +} + +void RaidenOrchestrator::ListenerLoop() { + while (!stopping_) { + sockaddr_in6 client_addr{}; + socklen_t client_len = sizeof(client_addr); + int client_fd = accept( + server_fd_, reinterpret_cast(&client_addr), &client_len); + if (client_fd < 0) { + if (stopping_) break; + continue; + } + + worker_threads_.push_back( + std::thread(&RaidenOrchestrator::ConnectionWorker, this, client_fd)); + } +} + +void RaidenOrchestrator::ConnectionWorker(int client_fd) { + uint32_t net_len = 0; + if (!transport::RawBufferTransport::ReadExact(client_fd, &net_len, + sizeof(net_len)) + .ok()) { + close(client_fd); + return; + } + uint32_t payload_len = ntohl(net_len); + + std::vector buffer(payload_len); + if (!transport::RawBufferTransport::ReadExact(client_fd, buffer.data(), + payload_len) + .ok()) { + close(client_fd); + return; + } + + ControlRequest req; + if (!req.ParseFromString(absl::string_view(buffer.data(), buffer.size()))) { + LOG(ERROR) << "Failed to parse ControlRequest Protobuf"; + close(client_fd); + return; + } + + ControlResponse resp; + resp.set_success(true); + resp.set_message("SUCCESS"); + + if (req.command() == ControlRequest::COMMAND_REGISTER_WORK_UNIT) { + if (req.has_register_work_unit_request()) { + const auto& reg_req = req.register_work_unit_request(); + RaidenId id{.job_name = reg_req.unit().job_name(), + .job_replica_id = reg_req.unit().job_replica_id(), + .data_name = reg_req.unit().data_name(), + .data_replica_idx = reg_req.unit().data_replica_idx()}; + + // In Controller->Orchestrator context, control_plane_rpc_address is the + // address of the registering controller. + std::string addr = reg_req.control_plane_rpc_address(); + if (addr.empty()) { + resp.set_success(false); + resp.set_message( + "Missing control_plane_rpc_address (Controller Address)"); + } else { + absl::Status status = RegisterController(id, addr); + if (!status.ok()) { + resp.set_success(false); + resp.set_message(std::string(status.message())); + } + } + } else { + resp.set_success(false); + resp.set_message("Missing register_work_unit_request"); + } + } else if (req.command() == ControlRequest::COMMAND_RESOLVE_CONTROLLER) { + if (req.has_target_unit()) { + const auto& target = req.target_unit(); + RaidenId id{.job_name = target.job_name(), + .job_replica_id = target.job_replica_id(), + .data_name = target.data_name(), + .data_replica_idx = target.data_replica_idx()}; + + auto addr_or = ResolveController(id); + if (addr_or.ok()) { + resp.set_response_data(addr_or.value()); + } else { + resp.set_success(false); + resp.set_message(std::string(addr_or.status().message())); + } + } else { + resp.set_success(false); + resp.set_message("Missing target_unit for resolution"); + } + } else { + resp.set_success(false); + resp.set_message("Unsupported command"); + } + + std::string resp_str; + if (resp.SerializeToString(&resp_str)) { + uint32_t resp_net_len = htonl(resp_str.size()); + transport::RawBufferTransport::WriteExact(client_fd, &resp_net_len, + sizeof(resp_net_len)) + .IgnoreError(); + transport::RawBufferTransport::WriteExact(client_fd, resp_str.data(), + resp_str.size()) + .IgnoreError(); + } + + close(client_fd); +} + +} // namespace kv_cache +} // namespace tpu_raiden diff --git a/tpu_raiden/kv_cache/raiden_orchestrator.h b/tpu_raiden/kv_cache/raiden_orchestrator.h new file mode 100644 index 00000000..2a6514ea --- /dev/null +++ b/tpu_raiden/kv_cache/raiden_orchestrator.h @@ -0,0 +1,72 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_KV_CACHE_RAIDEN_ORCHESTRATOR_H_ +#define THIRD_PARTY_TPU_RAIDEN_KV_CACHE_RAIDEN_ORCHESTRATOR_H_ + +#include +#include +#include // NOLINT +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_map.h" +#include "absl/hash/hash.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/synchronization/mutex.h" +#include "tpu_raiden/kv_cache/kv_cache_store.h" + +namespace tpu_raiden { +namespace kv_cache { + +// Central orchestrator that acts as a directory service for Raiden Controllers. +// It maintains a mapping from RaidenId to Controller physical addresses. +class RaidenOrchestrator { + public: + explicit RaidenOrchestrator(int port, const std::string& bind_ip = ""); + ~RaidenOrchestrator(); + + RaidenOrchestrator(const RaidenOrchestrator&) = delete; + RaidenOrchestrator& operator=(const RaidenOrchestrator&) = delete; + + // Registers a controller address for a given RaidenId. + absl::Status RegisterController(const RaidenId& raiden_id, + const std::string& address); + + // Resolves the controller address for a given RaidenId. + absl::StatusOr ResolveController(const RaidenId& raiden_id); + + int port() const { return port_; } + + private: + void ListenerLoop(); + void ConnectionWorker(int client_fd); + + int port_; + int server_fd_ = -1; + std::atomic stopping_{false}; + std::thread listener_thread_; + std::vector worker_threads_; + + mutable absl::Mutex mutex_; + // Maps RaidenId to Controller Address (e.g., "ip:port") + absl::flat_hash_map registry_ + ABSL_GUARDED_BY(mutex_); +}; + +} // namespace kv_cache +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_KV_CACHE_RAIDEN_ORCHESTRATOR_H_ diff --git a/tpu_raiden/kv_cache/raiden_orchestrator_main.cc b/tpu_raiden/kv_cache/raiden_orchestrator_main.cc new file mode 100644 index 00000000..b886b23f --- /dev/null +++ b/tpu_raiden/kv_cache/raiden_orchestrator_main.cc @@ -0,0 +1,42 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "absl/flags/parse.h" +#include "absl/log/initialize.h" +#include "absl/flags/flag.h" +#include "absl/log/log.h" +#include "tpu_raiden/kv_cache/raiden_orchestrator.h" + +ABSL_FLAG(int, port, 9999, "Port to listen on"); +ABSL_FLAG(std::string, bind_ip, "", "IP to bind to"); + +int main(int argc, char* argv[]) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + + int port = absl::GetFlag(FLAGS_port); + std::string bind_ip = absl::GetFlag(FLAGS_bind_ip); + + LOG(INFO) << "Starting RaidenOrchestrator on port " << port; + tpu_raiden::kv_cache::RaidenOrchestrator orchestrator(port, bind_ip); + + // Keep alive until terminated + while (true) { + pause(); + } + + return 0; +} diff --git a/tpu_raiden/kv_cache/raiden_orchestrator_test.cc b/tpu_raiden/kv_cache/raiden_orchestrator_test.cc new file mode 100644 index 00000000..3b479dd1 --- /dev/null +++ b/tpu_raiden/kv_cache/raiden_orchestrator_test.cc @@ -0,0 +1,172 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/kv_cache/raiden_orchestrator.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "net/util/ports.h" +#include +#include +#include "absl/strings/str_cat.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" + +namespace tpu_raiden { +namespace kv_cache { +namespace { + +using ::tpu_raiden::rpc::ControlRequest; +using ::tpu_raiden::rpc::ControlResponse; + +// Helper to send a simple RPC to a Custom TCP server (Length + Proto) +absl::StatusOr SendRawRpc(const std::string& address, + const ControlRequest& req) { + size_t colon = address.rfind(':'); + if (colon == std::string::npos) + return absl::InvalidArgumentError("Invalid address"); + std::string host = address.substr(0, colon); + std::string port_str = address.substr(colon + 1); + + struct addrinfo hints; + struct addrinfo* result = nullptr; + std::memset(&hints, 0, sizeof(hints)); + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + + int ret = getaddrinfo(host.c_str(), port_str.c_str(), &hints, &result); + if (ret != 0 || result == nullptr) { + return absl::InvalidArgumentError( + absl::StrCat("getaddrinfo failed: ", gai_strerror(ret))); + } + + int sock = -1; + struct addrinfo* rp; + for (rp = result; rp != nullptr; rp = rp->ai_next) { + sock = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol); + if (sock < 0) continue; + + if (connect(sock, rp->ai_addr, rp->ai_addrlen) != -1) { + break; // Success + } + close(sock); + sock = -1; + } + + freeaddrinfo(result); + + if (sock < 0) { + return absl::InternalError("Connection Failed"); + } + + std::string payload; + if (!req.SerializeToString(&payload)) { + close(sock); + return absl::InternalError("Failed to serialize request"); + } + + uint32_t net_len = htonl(payload.size()); + if (send(sock, &net_len, sizeof(net_len), 0) != sizeof(net_len)) { + close(sock); + return absl::InternalError("Failed to send length"); + } + + if (send(sock, payload.data(), payload.size(), 0) != payload.size()) { + close(sock); + return absl::InternalError("Failed to send payload"); + } + + uint32_t resp_net_len = 0; + if (recv(sock, &resp_net_len, sizeof(resp_net_len), 0) != + sizeof(resp_net_len)) { + close(sock); + return absl::InternalError("Failed to receive length"); + } + + uint32_t resp_len = ntohl(resp_net_len); + std::vector resp_buffer(resp_len); + size_t total_recv = 0; + while (total_recv < resp_len) { + ssize_t n = + recv(sock, resp_buffer.data() + total_recv, resp_len - total_recv, 0); + if (n <= 0) { + close(sock); + return absl::InternalError("Failed to receive payload"); + } + total_recv += n; + } + + ControlResponse resp; + if (!resp.ParseFromString(absl::string_view(resp_buffer.data(), resp_len))) { + close(sock); + return absl::InternalError("Failed to parse response"); + } + + close(sock); + return resp; +} + +TEST(RaidenOrchestratorTest, RegistrationAndResolutionIsolated) { + int port = net_util::PickUnusedPortOrDie(); + RaidenOrchestrator orchestrator(port); + // Orchestrator starts in constructor. + + std::string orch_addr = absl::StrCat("127.0.0.1:", port); + + // 1. Register a unit + ControlRequest reg_req; + reg_req.set_command(ControlRequest::COMMAND_REGISTER_WORK_UNIT); + auto* reg_data = reg_req.mutable_register_work_unit_request(); + auto* raiden_id = reg_data->mutable_unit(); + raiden_id->set_job_name("test_job"); + raiden_id->set_job_replica_id("0"); + raiden_id->set_data_name("test_data"); + raiden_id->set_data_replica_idx(0); + reg_data->set_control_plane_rpc_address("10.0.0.1:5000"); + + auto reg_resp_or = SendRawRpc(orch_addr, reg_req); + ASSERT_OK(reg_resp_or); + EXPECT_TRUE(reg_resp_or->success()); + + // 2. Resolve the unit + ControlRequest res_req; + res_req.set_command(ControlRequest::COMMAND_RESOLVE_CONTROLLER); + auto* target_unit = res_req.mutable_target_unit(); + target_unit->set_job_name("test_job"); + target_unit->set_job_replica_id("0"); + target_unit->set_data_name("test_data"); + target_unit->set_data_replica_idx(0); + + auto res_resp_or = SendRawRpc(orch_addr, res_req); + ASSERT_OK(res_resp_or); + EXPECT_TRUE(res_resp_or->success()); + EXPECT_EQ(res_resp_or->response_data(), "10.0.0.1:5000"); + + // 3. Resolve non-existent unit + target_unit->set_data_name("non_existent"); + auto res_resp_fail = SendRawRpc(orch_addr, res_req); + ASSERT_OK(res_resp_fail); + EXPECT_FALSE(res_resp_fail->success()); +} + +} // namespace +} // namespace kv_cache +} // namespace tpu_raiden diff --git a/tpu_raiden/rpc/BUILD b/tpu_raiden/rpc/BUILD index f3a8c6d6..9778f913 100644 --- a/tpu_raiden/rpc/BUILD +++ b/tpu_raiden/rpc/BUILD @@ -38,6 +38,26 @@ cc_proto_library( deps = [":raiden_service_proto"], ) +cc_library( + name = "rpc_utils", + srcs = ["rpc_utils.cc"], + hdrs = ["rpc_utils.h"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + visibility = ["//visibility:public"], + deps = [ + ":raiden_service_cc_proto", + "//tpu_raiden/core:status_macros", + "//tpu_raiden/transport:raw_buffer_transport", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + ], +) + py_proto_library( name = "raiden_service_py_pb2", visibility = ["//visibility:public"], diff --git a/tpu_raiden/rpc/coordination_client.cc b/tpu_raiden/rpc/coordination_client.cc index 8af3e1a0..1f73c21e 100644 --- a/tpu_raiden/rpc/coordination_client.cc +++ b/tpu_raiden/rpc/coordination_client.cc @@ -21,8 +21,8 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" -#include "third_party/grpc/include/grpcpp/client_context.h" -#include "third_party/grpc/include/grpcpp/support/status.h" +#include "grpcpp/client_context.h" +#include "grpcpp/support/status.h" #include "tpu_raiden/rpc/coordination.grpc.pb.h" namespace tpu_raiden { diff --git a/tpu_raiden/rpc/coordination_client.h b/tpu_raiden/rpc/coordination_client.h index 39c7fca2..9d8b22b4 100644 --- a/tpu_raiden/rpc/coordination_client.h +++ b/tpu_raiden/rpc/coordination_client.h @@ -21,7 +21,7 @@ #include #include "absl/status/statusor.h" -#include "third_party/grpc/include/grpcpp/channel.h" +#include "grpcpp/channel.h" #include "tpu_raiden/rpc/coordination.grpc.pb.h" #include "tpu_raiden/rpc/coordination.pb.h" diff --git a/tpu_raiden/rpc/coordination_client_test.cc b/tpu_raiden/rpc/coordination_client_test.cc index 9a0f03fb..50f8edbb 100644 --- a/tpu_raiden/rpc/coordination_client_test.cc +++ b/tpu_raiden/rpc/coordination_client_test.cc @@ -21,13 +21,13 @@ #include #include -#include "third_party/grpc/include/grpcpp/create_channel.h" -#include "third_party/grpc/include/grpcpp/security/credentials.h" -#include "third_party/grpc/include/grpcpp/security/server_credentials.h" -#include "third_party/grpc/include/grpcpp/server.h" -#include "third_party/grpc/include/grpcpp/server_builder.h" -#include "third_party/grpc/include/grpcpp/server_context.h" -#include "third_party/grpc/include/grpcpp/support/status.h" +#include "grpcpp/create_channel.h" +#include "grpcpp/security/credentials.h" +#include "grpcpp/security/server_credentials.h" +#include "grpcpp/server.h" +#include "grpcpp/server_builder.h" +#include "grpcpp/server_context.h" +#include "grpcpp/support/status.h" #include "tpu_raiden/rpc/coordination.grpc.pb.h" #include "tpu_raiden/rpc/coordination.pb.h" diff --git a/tpu_raiden/rpc/raiden_service.proto b/tpu_raiden/rpc/raiden_service.proto index d476200d..34b41e7c 100644 --- a/tpu_raiden/rpc/raiden_service.proto +++ b/tpu_raiden/rpc/raiden_service.proto @@ -23,6 +23,7 @@ message RaidenIdProto { string job_name = 1; string job_replica_id = 2; string data_name = 3; + int32 data_replica_idx = 4; } // Coordinate intervals representing a contiguous slice along a single tensor @@ -100,6 +101,21 @@ message StartTransferRequest { bool use_block_chunks = 7; int64 expected_block_count = 8; string req_id = 9; + MemoryType src_mem_type = 10; +} + +message FetchNegotiationRequest { + repeated string block_hashes = 1; + RaidenIdProto src_unit = 2; + RaidenIdProto dst_unit = 3; + repeated string dst_worker_data_addresses = 4; + repeated int32 dst_block_ids = 5; + int64 uuid = 6; +} + +message FetchNegotiationResponse { + repeated string found_block_hashes = 1; + repeated string missing_block_hashes = 2; } message GetMetadataResponse { @@ -114,12 +130,20 @@ message ControlRequest { COMMAND_REGISTER_WORK_UNIT = 3; COMMAND_GET_METADATA = 4; COMMAND_REGISTER_TRANSFER_SCHEDULE = 5; + COMMAND_RESOLVE_CONTROLLER = 6; + COMMAND_NEGOTIATE_FETCH = 7; + COMMAND_GET_ENDPOINTS = 8; + COMMAND_TRANSFER_COMPLETED = 9; + COMMAND_EXECUTE_FETCH = 10; } Command command = 1; repeated string peers = 2; StartTransferRequest start_transfer_request = 3; RegisterWorkUnitRequest register_work_unit_request = 4; + RaidenIdProto target_unit = 5; + FetchNegotiationRequest fetch_negotiation_request = 6; + repeated int32 completed_block_ids = 7; } message ControlResponse { @@ -127,4 +151,7 @@ message ControlResponse { string message = 2; string response_data = 3; GetMetadataResponse get_metadata_response = 4; + FetchNegotiationResponse fetch_negotiation_response = 5; + repeated string endpoints = 6; + StartTransferRequest start_transfer_request = 7; } diff --git a/tpu_raiden/rpc/rpc_utils.cc b/tpu_raiden/rpc/rpc_utils.cc new file mode 100644 index 00000000..d95e6a3c --- /dev/null +++ b/tpu_raiden/rpc/rpc_utils.cc @@ -0,0 +1,147 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/rpc/rpc_utils.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "tpu_raiden/core/status_macros.h" +#include "tpu_raiden/transport/raw_buffer_transport.h" + +namespace tpu_raiden { +namespace rpc { + +absl::StatusOr SimpleConnect(absl::string_view peer) { + std::string host; + std::string port_str; + + if (!peer.empty() && peer.front() == '[') { + size_t closing_bracket = peer.find(']'); + if (closing_bracket == absl::string_view::npos || + closing_bracket + 1 >= peer.size() || + peer[closing_bracket + 1] != ':') { + return absl::InvalidArgumentError( + "Invalid IPv6 peer bracket string format"); + } + host = std::string(peer.substr(1, closing_bracket - 1)); + port_str = std::string(peer.substr(closing_bracket + 2)); + } else { + size_t colon = peer.rfind(':'); + if (colon == absl::string_view::npos) { + return absl::InvalidArgumentError( + "Invalid peer string format, missing port"); + } + host = std::string(peer.substr(0, colon)); + port_str = std::string(peer.substr(colon + 1)); + } + + struct addrinfo hints; + struct addrinfo* result = nullptr; + std::memset(&hints, 0, sizeof(hints)); + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + + int ret = getaddrinfo(host.c_str(), port_str.c_str(), &hints, &result); + if (ret != 0 || result == nullptr) { + return absl::InvalidArgumentError(absl::StrCat( + "getaddrinfo failed for host ", host, ": ", gai_strerror(ret))); + } + + int sock_fd = -1; + struct addrinfo* rp; + for (rp = result; rp != nullptr; rp = rp->ai_next) { + sock_fd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol); + if (sock_fd < 0) continue; + + if (connect(sock_fd, rp->ai_addr, rp->ai_addrlen) != -1) { + break; // Success + } + close(sock_fd); + sock_fd = -1; + } + + freeaddrinfo(result); + + if (sock_fd < 0) { + return absl::InternalError(absl::StrCat("Failed to connect to ", peer)); + } + + return sock_fd; +} + +absl::Status SendRpcSync(absl::string_view address, + const rpc::ControlRequest& req, + rpc::ControlResponse& resp) { + ASSIGN_OR_RETURN(int fd, SimpleConnect(address)); + + std::string payload; + if (!req.SerializeToString(&payload)) { + close(fd); + return absl::InternalError("Failed to serialize ControlRequest"); + } + + uint32_t net_len = htonl(payload.size()); + absl::Status status = + transport::RawBufferTransport::WriteExact(fd, &net_len, sizeof(net_len)); + if (!status.ok()) { + close(fd); + return status; + } + + status = transport::RawBufferTransport::WriteExact(fd, payload.data(), + payload.size()); + if (!status.ok()) { + close(fd); + return status; + } + + uint32_t resp_net_len = 0; + status = transport::RawBufferTransport::ReadExact(fd, &resp_net_len, + sizeof(resp_net_len)); + if (!status.ok()) { + close(fd); + return status; + } + + uint32_t resp_len = ntohl(resp_net_len); + std::vector resp_buffer(resp_len); + status = transport::RawBufferTransport::ReadExact(fd, resp_buffer.data(), + resp_len); + if (!status.ok()) { + close(fd); + return status; + } + + close(fd); + + if (!resp.ParseFromArray(resp_buffer.data(), resp_len)) { + return absl::InternalError("Failed to parse ControlResponse"); + } + + return absl::OkStatus(); +} + +} // namespace rpc +} // namespace tpu_raiden diff --git a/tpu_raiden/rpc/rpc_utils.h b/tpu_raiden/rpc/rpc_utils.h new file mode 100644 index 00000000..bbad7974 --- /dev/null +++ b/tpu_raiden/rpc/rpc_utils.h @@ -0,0 +1,37 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_RPC_RPC_UTILS_H_ +#define THIRD_PARTY_TPU_RAIDEN_RPC_RPC_UTILS_H_ + +#include "absl/status/status.h" +#include "absl/strings/string_view.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" + +namespace tpu_raiden { +namespace rpc { + +// Helper to establish a simple TCP connection to a peer address (host:port). +// Supports both IPv4 and IPv6. +absl::StatusOr SimpleConnect(absl::string_view peer); + +// Helper to send a ControlRequest and receive a ControlResponse synchronously. +absl::Status SendRpcSync(absl::string_view address, + const rpc::ControlRequest& req, + rpc::ControlResponse& resp); + +} // namespace rpc +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_RPC_RPC_UTILS_H_ diff --git a/tpu_raiden/transport/raw_buffer_transport.cc b/tpu_raiden/transport/raw_buffer_transport.cc index 2e8ce2db..d578420c 100644 --- a/tpu_raiden/transport/raw_buffer_transport.cc +++ b/tpu_raiden/transport/raw_buffer_transport.cc @@ -498,7 +498,9 @@ void RawBufferTransport::ConnectionWorker(int client_fd) { } if (ret == 0) continue; - if (!ProcessSingleRequest(client_fd).ok()) { + absl::Status s = ProcessSingleRequest(client_fd); + if (!s.ok()) { + LOG(ERROR) << "ConnectionWorker: ProcessSingleRequest failed: " << s; break; } }