Cleaned up the current code per our latest design.

build.sh

@@ -153,7 +153,7 @@ if [ "${BUILD_TORCH}" = true ]; then
 fi
 
 BAZEL_TARGETS=(
-  "//tpu_raiden/weight_sync:weight_synchronizer_service_py_pb2"
+  "//tpu_raiden/rpc:raiden_service_py_pb2"
   "//tpu_raiden/rpc:coordination_py_pb2"
   "//tpu_raiden/rpc:coordination_py_pb2_grpc"
 )
@@ -232,7 +232,7 @@ echo "=== Building targets with Bazel ==="
 
 
 echo "=== Copying generated protobuf Python modules ==="
-cp -f "${WORKSPACE_DIR}/bazel-bin/weight_sync/weight_synchronizer_service_pb2.py" "${WORKSPACE_DIR}/weight_sync/" 2>/dev/null || true
+cp -f "${WORKSPACE_DIR}/bazel-bin/rpc/raiden_service_pb2.py" "${WORKSPACE_DIR}/rpc/" 2>/dev/null || true
 cp -f "${WORKSPACE_DIR}/bazel-bin/rpc/coordination_pb2.py" "${WORKSPACE_DIR}/rpc/" 2>/dev/null || true
 cp -f "${WORKSPACE_DIR}/bazel-bin/rpc/coordination_pb2_grpc.py" "${WORKSPACE_DIR}/rpc/" 2>/dev/null || true
 

tpu_raiden/api/jax/BUILD

@@ -60,7 +60,7 @@ py_test(
     tags = ["notap"],
     deps = [
         ":weight_synchronizer_jax_py",
-        "//tpu_raiden/weight_sync:weight_synchronizer_service_py_pb2",
+        "//tpu_raiden/rpc:raiden_service_py_pb2",
         "@com_google_absl_py//absl/testing:absltest",
         "@jax//jax",
         "@pypi//numpy",

tpu_raiden/api/jax/weight_synchronizer.py

@@ -29,7 +29,7 @@ def __init__(
       local_port: Optional[int] = None,
       parallelism: int = 1,
       unsafe_skip_buffer_lock: bool = False,
-      control_port: Optional[int] = None,
+      listener_port: Optional[int] = None,
   ):
     """Instantiates the Weight Synchronizer on a JAX weights list.
 
@@ -38,15 +38,14 @@ def __init__(
       local_port: Sockets server port for incoming pulls (inference mode).
       parallelism: Number of parallel network stream TCP sockets workers.
       unsafe_skip_buffer_lock: Skip PJRT buffer locks during weights unpack.
-      control_port: Sockets server port for incoming C++ Control Service
-        commands.
+      listener_port: Sockets server port for incoming C++ Listener commands.
     """
     self._impl = _weight_synchronizer.WeightSynchronizer(
         jax_arrays,
         local_port,
         parallelism,
         unsafe_skip_buffer_lock,
-        control_port,
+        listener_port,
     )
 
   def pull_weights(self, source: str) -> None:
@@ -129,14 +128,14 @@ def local_port(self) -> Optional[int]:
     return self._impl.local_port
 
   @property
-  def control_port(self) -> Optional[int]:
-    """Returns the active local port assigned to the C++ Control Service."""
-    return self._impl.control_port
+  def listener_port(self) -> Optional[int]:
+    """Returns the active local port assigned to the C++ Listener."""
+    return self._impl.listener_port
 
   @property
-  def is_control_service_active(self) -> bool:
-    """Returns whether the native C++ Control Service is actively running."""
-    return self._impl.is_control_service_active
+  def is_listener_active(self) -> bool:
+    """Returns whether the native C++ Listener is actively running."""
+    return self._impl.is_listener_active
 
   @property
   def num_layers(self) -> int:

tpu_raiden/api/jax/weight_synchronizer_test.py

@@ -25,7 +25,7 @@
 import numpy as np
 
 from tpu_raiden.api.jax import weight_synchronizer
-from tpu_raiden.weight_sync import weight_synchronizer_service_pb2
+from tpu_raiden.rpc import raiden_service_pb2
 
 WeightSynchronizer = weight_synchronizer.WeightSynchronizer
 
@@ -69,7 +69,7 @@ def test_push_synchronization(self):
         jax_arrays=src_arrs,
         local_port=0,
         unsafe_skip_buffer_lock=True,
-        control_port=0,
+        listener_port=0,
     )
     ws_dest1 = WeightSynchronizer(
         jax_arrays=dst1_arrs, local_port=0, unsafe_skip_buffer_lock=True
@@ -78,8 +78,8 @@ def test_push_synchronization(self):
         jax_arrays=dst2_arrs, local_port=0, unsafe_skip_buffer_lock=True
     )
 
-    req = weight_synchronizer_service_pb2.ControlRequest(
-        command=weight_synchronizer_service_pb2.ControlRequest.COMMAND_START_TRANSFER,
+    req = raiden_service_pb2.ControlRequest(
+        command=raiden_service_pb2.ControlRequest.COMMAND_START_TRANSFER,
         peers=[
             f"127.0.0.1:{ws_dest1.local_port}",
             f"127.0.0.1:{ws_dest2.local_port}",
@@ -88,12 +88,12 @@ def test_push_synchronization(self):
     payload = req.SerializeToString()
 
     sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM, 0)
-    sock.connect(("::1", ws_source.control_port))
+    sock.connect(("::1", ws_source.listener_port))
     sock.sendall(len(payload).to_bytes(4, "big") + payload)
 
     resp_len = int.from_bytes(sock.recv(4), "big")
     resp_bytes = sock.recv(resp_len)
-    resp = weight_synchronizer_service_pb2.ControlResponse()
+    resp = raiden_service_pb2.ControlResponse()
     resp.ParseFromString(resp_bytes)
     assert resp.success
     sock.close()
@@ -127,7 +127,7 @@ def test_pull_synchronization(self):
         jax_arrays=src_arrs,
         local_port=0,
         unsafe_skip_buffer_lock=True,
-        control_port=0,
+        listener_port=0,
     )
     ws_dest1 = WeightSynchronizer(
         jax_arrays=dst1_arrs, local_port=0, unsafe_skip_buffer_lock=True
@@ -137,19 +137,19 @@ def test_pull_synchronization(self):
     )
 
     # Self-push to populate ws_source's host buffer with current device weights
-    req = weight_synchronizer_service_pb2.ControlRequest(
-        command=weight_synchronizer_service_pb2.ControlRequest.COMMAND_START_TRANSFER,
+    req = raiden_service_pb2.ControlRequest(
+        command=raiden_service_pb2.ControlRequest.COMMAND_START_TRANSFER,
         peers=[f"127.0.0.1:{ws_source.local_port}"],
     )
     payload = req.SerializeToString()
 
     sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM, 0)
-    sock.connect(("::1", ws_source.control_port))
+    sock.connect(("::1", ws_source.listener_port))
     sock.sendall(len(payload).to_bytes(4, "big") + payload)
 
     resp_len = int.from_bytes(sock.recv(4), "big")
     resp_bytes = sock.recv(resp_len)
-    resp = weight_synchronizer_service_pb2.ControlResponse()
+    resp = raiden_service_pb2.ControlResponse()
     resp.ParseFromString(resp_bytes)
     assert resp.success
     sock.close()

tpu_raiden/core/BUILD

@@ -249,6 +249,7 @@ cc_library(
         ":status_macros",
         ":tpu_utils",
         "//tpu_raiden/kv_cache:kv_cache_manager_base",
+        "//tpu_raiden/rpc:raiden_service_cc_proto",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

tpu_raiden/core/kv_cache_manager_with_transfer.cc

@@ -535,15 +535,14 @@ int64_t KVCacheManagerWithTransfer::NotifyForRead(
 }
 
 absl::Status KVCacheManagerWithTransfer::RegisterActivePlan(
-    uint64_t uuid, const kv_cache::StartTransferRequest& request,
-    bool is_sender) {
+    uint64_t uuid, const rpc::StartTransferRequest& request, bool is_sender) {
   // 1. Call base class implementation to register the plan in active_plans_
   TF_RETURN_IF_ERROR(kv_cache::KVCacheManagerBase::RegisterActivePlan(
       uuid, request, is_sender));
 
   // 2. If we are the receiver and the destination memory type is HBM,
   //    populate active_recv_entries_ to enable automatic H2D copy!
-  if (!is_sender && request.dst_mem_type() == kv_cache::MEMORY_TYPE_HBM) {
+  if (!is_sender && request.dst_mem_type() == rpc::MEMORY_TYPE_HBM) {
     std::lock_guard<std::mutex> lock(mu_);
     RecvEntry recv_entry;
     std::string req_id = absl::StrCat("resharded_transfer_", uuid);

tpu_raiden/core/kv_cache_manager_with_transfer.h

@@ -194,7 +194,7 @@ class KVCacheManagerWithTransfer : public kv_cache::KVCacheManagerBase {
   CompleteReadRaw();
 
   absl::Status RegisterActivePlan(uint64_t uuid,
-                                  const kv_cache::StartTransferRequest& request,
+                                  const rpc::StartTransferRequest& request,
                                   bool is_sender) override;
 
   absl::Status OnBlocksReceived(const std::vector<int>& block_ids,

tpu_raiden/frameworks/jax/BUILD

@@ -524,7 +524,7 @@ py_library(
     srcs = ["resharding_planner.py"],
     visibility = ["//visibility:public"],
     deps = [
-        "//tpu_raiden/weight_sync:weight_synchronizer_service_py_pb2",
+        "//tpu_raiden/rpc:raiden_service_py_pb2",
         "@jax//jax",
     ],
 )

tpu_raiden/frameworks/jax/resharding_planner.py

@@ -31,7 +31,7 @@
 import itertools
 from typing import List, Tuple
 import jax
-from tpu_raiden.weight_sync import weight_synchronizer_service_pb2
+from tpu_raiden.rpc import raiden_service_pb2
 
 
 @dataclass
@@ -246,7 +246,7 @@ def make_resharding_plan_from_metadata(
 def compute_nd_shard_slices(
     global_shape: Tuple[int, ...],
     mesh_shape: Tuple[int, ...],
-) -> List[weight_synchronizer_service_pb2.NDSliceProto]:
+) -> List[raiden_service_pb2.NDSliceProto]:
   """Computes N-dimensional logical tensor bounding boxes for a sharded grid.
 
   This function derives the exact coordinate intervals along every dimension
@@ -279,11 +279,12 @@ def compute_nd_shard_slices(
 
   shard_slices = []
   for device_coord in itertools.product(*coordinate_ranges):
-    slice_proto = weight_synchronizer_service_pb2.NDSliceProto()
+    slice_proto = raiden_service_pb2.NDSliceProto()
     for d in range(rank):
       c = device_coord[d]
       start = c * tile_sizes[d]
-      # Ensure any exact remainders land nicely in the last physical mesh shard boundary
+      # Ensure any exact remainders land nicely in the last physical mesh shard
+      # boundary
       end = (
           (c + 1) * tile_sizes[d]
           if c < mesh_shape[d] - 1

tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc

@@ -261,7 +261,7 @@ NB_MODULE(_tpu_raiden_jax, m) {
            nb::arg("jax_arrays"), nb::arg("local_port") = nb::none(),
            nb::arg("parallelism") = 1,
            nb::arg("unsafe_skip_buffer_lock") = false,
-           nb::arg("control_port") = nb::none())
+           nb::arg("listener_port") = nb::none())
       .def(
           "PullWeights",
           [](WeightSynchronizer& self, absl::string_view source) {
@@ -363,9 +363,9 @@ NB_MODULE(_tpu_raiden_jax, m) {
           },
           nb::arg("layer_idx") = 0, nb::arg("shard_idx") = 0)
       .def_prop_ro("local_port", &WeightSynchronizer::local_port)
-      .def_prop_ro("control_port", &WeightSynchronizer::control_port)
-      .def_prop_ro("is_control_service_active",
-                   &WeightSynchronizer::is_control_service_active)
+      .def_prop_ro("listener_port", &WeightSynchronizer::listener_port)
+      .def_prop_ro("is_listener_active",
+                   &WeightSynchronizer::is_listener_active)
       .def_prop_ro("num_layers", &WeightSynchronizer::num_layers)
       .def_prop_ro("num_shards", &WeightSynchronizer::num_shards)
       .def_prop_ro("slice_byte_size", &WeightSynchronizer::slice_byte_size);

tpu_raiden/frameworks/jax/weight_synchronizer.cc

@@ -38,20 +38,20 @@ WeightSynchronizer::WeightSynchronizer(nanobind::list jax_arrays,
                                        std::optional<int> local_port,
                                        int parallelism,
                                        bool unsafe_skip_buffer_lock,
-                                       std::optional<int> control_port)
+                                       std::optional<int> listener_port)
     : WeightSynchronizer(UnpackAndMove(std::move(jax_arrays)), local_port,
-                         parallelism, unsafe_skip_buffer_lock, control_port) {}
+                         parallelism, unsafe_skip_buffer_lock, listener_port) {}
 
 WeightSynchronizer::WeightSynchronizer(UnpackedWeights&& weights,
                                        std::optional<int> local_port,
                                        int parallelism,
                                        bool unsafe_skip_buffer_lock,
-                                       std::optional<int> control_port)
+                                       std::optional<int> listener_port)
     : jax_arrays_(std::move(weights.jax_arrays)) {
   impl_ = std::make_unique<weight_sync::WeightSynchronizerBase>(
       weights.layer_buffers, local_port,
       /*external_host_ptrs=*/std::nullopt, unsafe_skip_buffer_lock, parallelism,
-      control_port);
+      listener_port);
 }
 
 #endif  // WITHOUT_PYTHON
@@ -87,11 +87,11 @@ const uint8_t* WeightSynchronizer::GetHostBufferPtr(size_t layer_idx,
 std::optional<int> WeightSynchronizer::local_port() const {
   return impl_->local_port();
 }
-std::optional<int> WeightSynchronizer::control_port() const {
-  return impl_->control_port();
+std::optional<int> WeightSynchronizer::listener_port() const {
+  return impl_->listener_port();
 }
-bool WeightSynchronizer::is_control_service_active() const {
-  return impl_->is_control_service_active();
+bool WeightSynchronizer::is_listener_active() const {
+  return impl_->is_listener_active();
 }
 size_t WeightSynchronizer::num_layers() const { return impl_->num_layers(); }
 size_t WeightSynchronizer::num_shards() const { return impl_->num_shards(); }

tpu_raiden/frameworks/jax/weight_synchronizer.h

@@ -56,7 +56,7 @@ class WeightSynchronizer {
   WeightSynchronizer(nanobind::list jax_arrays,
                      std::optional<int> local_port = std::nullopt,
                      int parallelism = 1, bool unsafe_skip_buffer_lock = false,
-                     std::optional<int> control_port = std::nullopt);
+                     std::optional<int> listener_port = std::nullopt);
 #endif
 
   ~WeightSynchronizer();
@@ -74,8 +74,8 @@ class WeightSynchronizer {
                                 size_t dst_offset_bytes, size_t size_bytes);
   const uint8_t* GetHostBufferPtr(size_t layer_idx, size_t shard_idx) const;
   std::optional<int> local_port() const;
-  std::optional<int> control_port() const;
-  bool is_control_service_active() const;
+  std::optional<int> listener_port() const;
+  bool is_listener_active() const;
   size_t num_layers() const;
   size_t num_shards() const;
   size_t slice_byte_size() const;
@@ -84,7 +84,7 @@ class WeightSynchronizer {
 #ifndef WITHOUT_PYTHON
   WeightSynchronizer(UnpackedWeights&& weights, std::optional<int> local_port,
                      int parallelism, bool unsafe_skip_buffer_lock,
-                     std::optional<int> control_port);
+                     std::optional<int> listener_port);
   std::optional<nanobind::list> jax_arrays_;
 #endif
   std::unique_ptr<weight_sync::WeightSynchronizerBase> impl_;

tpu_raiden/kv_cache/BUILD

@@ -12,11 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("@com_google_protobuf//bazel:cc_proto_library.bzl", "cc_proto_library")
-load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
 load("@rules_python//python:defs.bzl", "py_library")
-load("@rules_python//python:proto.bzl", "py_proto_library")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -59,7 +56,6 @@ cc_library(
     features = ["-use_header_modules"],
     visibility = ["//visibility:public"],
     deps = [
-        ":kv_cache_service_cc_proto",
         ":logical_block_manager",
         "//tpu_raiden/core:host_memory_allocator",
         "//tpu_raiden/core:numa_thread_pool",
@@ -69,6 +65,7 @@ cc_library(
         "//tpu_raiden/core:status_macros",
         "//tpu_raiden/core:tpu_utils",
         "//tpu_raiden/core:xla_raw_transfer_headers",
+        "//tpu_raiden/rpc:raiden_service_cc_proto",
         "//tpu_raiden/transport:block_transport",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -206,30 +203,13 @@ cc_test(
     ],
 )
 
-proto_library(
-    name = "kv_cache_service_proto",
-    srcs = ["kv_cache_service.proto"],
-)
-
-cc_proto_library(
-    name = "kv_cache_service_cc_proto",
-    visibility = ["//visibility:public"],
-    deps = [":kv_cache_service_proto"],
-)
-
-py_proto_library(
-    name = "kv_cache_service_py_pb2",
-    visibility = ["//visibility:public"],
-    deps = [":kv_cache_service_proto"],
-)
-
 py_library(
     name = "kv_cache_controller",
     srcs = ["kv_cache_controller.py"],
     visibility = ["//visibility:public"],
     deps = [
-        ":kv_cache_service_py_pb2",
         "//tpu_raiden/rpc:raiden_controller",
+        "//tpu_raiden/rpc:raiden_service_py_pb2",
     ],
 )
 
@@ -245,7 +225,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":kv_cache_manager_base",
-        ":kv_cache_service_cc_proto",
+        "//tpu_raiden/rpc:raiden_service_cc_proto",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
     ],