diff --git a/build.sh b/build.sh index b877da87..9fafb901 100755 --- a/build.sh +++ b/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 diff --git a/tpu_raiden/api/jax/BUILD b/tpu_raiden/api/jax/BUILD index 9ac0f96e..70f0db16 100644 --- a/tpu_raiden/api/jax/BUILD +++ b/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", diff --git a/tpu_raiden/api/jax/weight_synchronizer.py b/tpu_raiden/api/jax/weight_synchronizer.py index a61e0c0e..cb7cca98 100644 --- a/tpu_raiden/api/jax/weight_synchronizer.py +++ b/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: diff --git a/tpu_raiden/api/jax/weight_synchronizer_test.py b/tpu_raiden/api/jax/weight_synchronizer_test.py index 79e967b6..57af61e0 100644 --- a/tpu_raiden/api/jax/weight_synchronizer_test.py +++ b/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() diff --git a/tpu_raiden/core/BUILD b/tpu_raiden/core/BUILD index 2761809b..59f80f22 100644 --- a/tpu_raiden/core/BUILD +++ b/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", diff --git a/tpu_raiden/core/kv_cache_manager_with_transfer.cc b/tpu_raiden/core/kv_cache_manager_with_transfer.cc index 77f91f91..e82eda81 100644 --- a/tpu_raiden/core/kv_cache_manager_with_transfer.cc +++ b/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 lock(mu_); RecvEntry recv_entry; std::string req_id = absl::StrCat("resharded_transfer_", uuid); diff --git a/tpu_raiden/core/kv_cache_manager_with_transfer.h b/tpu_raiden/core/kv_cache_manager_with_transfer.h index 595be193..764993e8 100644 --- a/tpu_raiden/core/kv_cache_manager_with_transfer.h +++ b/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& block_ids, diff --git a/tpu_raiden/frameworks/jax/BUILD b/tpu_raiden/frameworks/jax/BUILD index 586f0499..2b2d3496 100644 --- a/tpu_raiden/frameworks/jax/BUILD +++ b/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", ], ) diff --git a/tpu_raiden/frameworks/jax/resharding_planner.py b/tpu_raiden/frameworks/jax/resharding_planner.py index 2099a6a3..cf7314fc 100644 --- a/tpu_raiden/frameworks/jax/resharding_planner.py +++ b/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 diff --git a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc index 7c69cfbe..56f2cf49 100644 --- a/tpu_raiden/frameworks/jax/tpu_raiden_jax_module.cc +++ b/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); diff --git a/tpu_raiden/frameworks/jax/weight_synchronizer.cc b/tpu_raiden/frameworks/jax/weight_synchronizer.cc index 3160d36e..b685742c 100644 --- a/tpu_raiden/frameworks/jax/weight_synchronizer.cc +++ b/tpu_raiden/frameworks/jax/weight_synchronizer.cc @@ -38,20 +38,20 @@ WeightSynchronizer::WeightSynchronizer(nanobind::list jax_arrays, std::optional local_port, int parallelism, bool unsafe_skip_buffer_lock, - std::optional control_port) + std::optional 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 local_port, int parallelism, bool unsafe_skip_buffer_lock, - std::optional control_port) + std::optional listener_port) : jax_arrays_(std::move(weights.jax_arrays)) { impl_ = std::make_unique( 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 WeightSynchronizer::local_port() const { return impl_->local_port(); } -std::optional WeightSynchronizer::control_port() const { - return impl_->control_port(); +std::optional 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(); } diff --git a/tpu_raiden/frameworks/jax/weight_synchronizer.h b/tpu_raiden/frameworks/jax/weight_synchronizer.h index 37c4cec8..af397889 100644 --- a/tpu_raiden/frameworks/jax/weight_synchronizer.h +++ b/tpu_raiden/frameworks/jax/weight_synchronizer.h @@ -56,7 +56,7 @@ class WeightSynchronizer { WeightSynchronizer(nanobind::list jax_arrays, std::optional local_port = std::nullopt, int parallelism = 1, bool unsafe_skip_buffer_lock = false, - std::optional control_port = std::nullopt); + std::optional 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 local_port() const; - std::optional control_port() const; - bool is_control_service_active() const; + std::optional 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 local_port, int parallelism, bool unsafe_skip_buffer_lock, - std::optional control_port); + std::optional listener_port); std::optional jax_arrays_; #endif std::unique_ptr impl_; diff --git a/tpu_raiden/kv_cache/BUILD b/tpu_raiden/kv_cache/BUILD index 54730e59..84aec42c 100644 --- a/tpu_raiden/kv_cache/BUILD +++ b/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", ], diff --git a/tpu_raiden/kv_cache/kv_cache_controller.py b/tpu_raiden/kv_cache/kv_cache_controller.py index 4f8c3045..169c7d48 100644 --- a/tpu_raiden/kv_cache/kv_cache_controller.py +++ b/tpu_raiden/kv_cache/kv_cache_controller.py @@ -29,133 +29,24 @@ from typing import Optional -from tpu_raiden.kv_cache import kv_cache_service_pb2 from tpu_raiden.rpc import raiden_controller +from tpu_raiden.rpc import raiden_service_pb2 from tpu_raiden.rpc.raiden_controller import RaidenId from tpu_raiden.rpc.raiden_controller import TransferPlan -def _raiden_id_to_proto( - unit: RaidenId, -) -> kv_cache_service_pb2.RaidenIdProto: - return kv_cache_service_pb2.RaidenIdProto( - job_name=unit.job_name, - job_replica_id=unit.job_replica_id, - data_name=unit.data_name, - ) - - -# TODO(raiden-dev): We should be able to use the same RPC client for both KV -# cache transfer and weights sync. So this can go to Raiden controller. class KVCacheWorkerRpcClient(raiden_controller.WorkerRpcClient): """Worker RPC client using KV Cache service proto.""" - def _encode_start_transfer( - self, target_id: RaidenId, transfer_plan: TransferPlan - ) -> Optional[bytes]: - if ( - target_id not in transfer_plan.src_units - and target_id not in transfer_plan.dst_units - ): - return None - - peers = [] - for dst in transfer_plan.dst_units: - dst_coords = transfer_plan.worker_data_addresses.get( - dst, ["127.0.0.1:8000"] - ) - peers.extend(dst_coords) - - req = kv_cache_service_pb2.ControlRequest( - command=kv_cache_service_pb2.ControlRequest.COMMAND_START_TRANSFER, - peers=peers, - ) - - is_sender = target_id in transfer_plan.src_units - start_req = kv_cache_service_pb2.StartTransferRequest( - src_units=[_raiden_id_to_proto(u) for u in transfer_plan.src_units], - dst_units=[_raiden_id_to_proto(u) for u in transfer_plan.dst_units], - uuid=transfer_plan.uuid, - is_sender=is_sender, - dst_mem_type=int(transfer_plan.dst_mem_type), - use_block_chunks=transfer_plan.use_block_chunks, - ) - - if transfer_plan.shard_push_schedules: - if target_id in transfer_plan.dst_units: - # Receiver path: send FILTERED plan, only containing entries for this receiver - target_endpoints = transfer_plan.worker_data_addresses.get( - target_id, [] - ) - for ( - src_unit, - push_schedules, - ) in transfer_plan.shard_push_schedules.items(): - src_replica_idx = int(src_unit.job_replica_id) - # We assume each source worker has only 1 shard (shard_idx = 0) - schedule = push_schedules.get(0) - if schedule: - schedule_proto = kv_cache_service_pb2.ShardPushScheduleProto() - for ( - dst_peer, - dst_shard_idx, - dst_offset, - src_offset, - size, - src_block_id, - dst_block_id, - ) in schedule: - if dst_peer in target_endpoints: - entry_proto = schedule_proto.entries.add() - entry_proto.dst_peer = dst_peer - entry_proto.dst_shard_idx = dst_shard_idx - entry_proto.dst_offset_bytes = dst_offset - entry_proto.src_offset_bytes = src_offset - entry_proto.size_bytes = size - entry_proto.src_block_id = src_block_id - entry_proto.dst_block_id = dst_block_id - if len(schedule_proto.entries) > 0: - start_req.shard_push_schedules[src_replica_idx].CopyFrom( - schedule_proto - ) - else: - # Sender path: only send local schedule - push_schedules = transfer_plan.shard_push_schedules.get(target_id) - if push_schedules: - for shard_idx, entries in push_schedules.items(): - schedule_proto = kv_cache_service_pb2.ShardPushScheduleProto() - for ( - dst_peer, - dst_shard_idx, - dst_offset, - src_offset, - size, - src_block_id, - dst_block_id, - ) in entries: - entry_proto = schedule_proto.entries.add() - entry_proto.dst_peer = dst_peer - entry_proto.dst_shard_idx = dst_shard_idx - entry_proto.dst_offset_bytes = dst_offset - entry_proto.src_offset_bytes = src_offset - entry_proto.size_bytes = size - entry_proto.src_block_id = src_block_id - entry_proto.dst_block_id = dst_block_id - start_req.shard_push_schedules[shard_idx].CopyFrom(schedule_proto) - - req.start_transfer_request.CopyFrom(start_req) - return req.SerializeToString() - - def _verify_response(self, resp_bytes: bytes) -> None: - resp = kv_cache_service_pb2.ControlResponse() - resp.ParseFromString(resp_bytes) - if not resp.success: - raise RuntimeError( - f"KV Cache remote native execution failed: {resp.message}" - ) - - def _encode_shutdown(self) -> bytes: - req = kv_cache_service_pb2.ControlRequest( - command=kv_cache_service_pb2.ControlRequest.COMMAND_SHUTDOWN + def __init__( + self, + endpoint_addresses: Optional[dict[RaidenId, str]] = None, + resolve_timeout: float = 300.0, + name_resolver: Optional[raiden_controller.NameResolver] = None, + ): + super().__init__( + endpoint_addresses=endpoint_addresses, + resolve_timeout=resolve_timeout, + name_resolver=name_resolver, + proto_module=raiden_service_pb2, ) - return req.SerializeToString() diff --git a/tpu_raiden/kv_cache/kv_cache_listener.cc b/tpu_raiden/kv_cache/kv_cache_listener.cc index 2383c8a5..5d894cb7 100644 --- a/tpu_raiden/kv_cache/kv_cache_listener.cc +++ b/tpu_raiden/kv_cache/kv_cache_listener.cc @@ -30,11 +30,14 @@ #include "absl/log/log.h" #include "absl/status/status.h" #include "tpu_raiden/kv_cache/kv_cache_manager_base.h" -#include "tpu_raiden/kv_cache/kv_cache_service.pb.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" namespace tpu_raiden { 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) { diff --git a/tpu_raiden/kv_cache/kv_cache_manager_base.cc b/tpu_raiden/kv_cache/kv_cache_manager_base.cc index c426ae98..1d9c155c 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_base.cc +++ b/tpu_raiden/kv_cache/kv_cache_manager_base.cc @@ -52,8 +52,8 @@ #include "tpu_raiden/core/raw_transfer_impl.h" #include "tpu_raiden/core/status_macros.h" #include "tpu_raiden/core/tpu_utils.h" -#include "tpu_raiden/kv_cache/kv_cache_service.pb.h" #include "tpu_raiden/kv_cache/logical_block_manager.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" #include "tpu_raiden/transport/block_transport.h" namespace tpu_raiden { @@ -1291,7 +1291,7 @@ absl::Status KVCacheManagerBase::WaitForBlockRead(size_t layer_idx, } absl::Status KVCacheManagerBase::PushKVCacheResharded( - const StartTransferRequest& request) { + const tpu_raiden::rpc::StartTransferRequest& request) { // 1. Register the active plan so GetBlockChunks can use it TF_RETURN_IF_ERROR( RegisterActivePlan(request.uuid(), request, /*is_sender=*/true)); @@ -1339,7 +1339,8 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( } absl::Status KVCacheManagerBase::RegisterActivePlan( - uint64_t uuid, const StartTransferRequest& request, bool is_sender) { + uint64_t uuid, const tpu_raiden::rpc::StartTransferRequest& request, + bool is_sender) { absl::MutexLock l(plans_mu_); if (auto [it, inserted] = active_plans_.try_emplace(uuid, RegisteredPlan{request, is_sender}); @@ -1358,7 +1359,7 @@ bool KVCacheManagerBase::IsDramDestination(uint64_t uuid) const { auto it = active_plans_.find(uuid); if (it != active_plans_.end()) { return it->second.request.dst_mem_type() == - tpu_raiden::kv_cache::MEMORY_TYPE_DRAM; + tpu_raiden::rpc::MEMORY_TYPE_DRAM; } return false; } diff --git a/tpu_raiden/kv_cache/kv_cache_manager_base.h b/tpu_raiden/kv_cache/kv_cache_manager_base.h index 28bf542a..6cd839e5 100644 --- a/tpu_raiden/kv_cache/kv_cache_manager_base.h +++ b/tpu_raiden/kv_cache/kv_cache_manager_base.h @@ -38,8 +38,8 @@ #include "tpu_raiden/core/numa_thread_pool.h" #include "tpu_raiden/core/raiden_manager_base.h" #include "tpu_raiden/core/raw_transfer_core.h" -#include "tpu_raiden/kv_cache/kv_cache_service.pb.h" #include "tpu_raiden/kv_cache/logical_block_manager.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" #include "tpu_raiden/transport/block_transport.h" namespace tpu_raiden { @@ -143,7 +143,8 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { // Executes a distributed resharding push transfer based on precise // centralized Controller schedules. - absl::Status PushKVCacheResharded(const StartTransferRequest& request); + absl::Status PushKVCacheResharded( + const tpu_raiden::rpc::StartTransferRequest& request); // Blocks until all pending asynchronous transfers/copies are complete. virtual absl::Status WaitForPendingWork() { return absl::OkStatus(); } @@ -207,9 +208,9 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { bool use_block_chunks(uint64_t uuid) const override; - virtual absl::Status RegisterActivePlan(uint64_t uuid, - const StartTransferRequest& request, - bool is_sender); + virtual absl::Status RegisterActivePlan( + uint64_t uuid, const tpu_raiden::rpc::StartTransferRequest& request, + bool is_sender); std::vector GetBlockChunks( size_t layer_idx, size_t shard_idx, int block_id, uint64_t uuid, @@ -284,7 +285,7 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { mutable absl::Mutex plans_mu_; struct RegisteredPlan { - StartTransferRequest request; + tpu_raiden::rpc::StartTransferRequest request; bool is_sender = false; }; absl::flat_hash_map active_plans_ diff --git a/tpu_raiden/rpc/BUILD b/tpu_raiden/rpc/BUILD index af4a1a31..1cd7d1d4 100644 --- a/tpu_raiden/rpc/BUILD +++ b/tpu_raiden/rpc/BUILD @@ -27,6 +27,23 @@ proto_library( srcs = ["coordination.proto"], ) +proto_library( + name = "raiden_service_proto", + srcs = ["raiden_service.proto"], +) + +cc_proto_library( + name = "raiden_service_cc_proto", + visibility = ["//visibility:public"], + deps = [":raiden_service_proto"], +) + +py_proto_library( + name = "raiden_service_py_pb2", + visibility = ["//visibility:public"], + deps = [":raiden_service_proto"], +) + cc_proto_library( name = "coordination_cc_proto", deps = [":coordination_proto"], @@ -63,7 +80,7 @@ py_library( name = "raiden_controller", srcs = ["raiden_controller.py"], deps = [ - "//tpu_raiden/weight_sync:weight_synchronizer_service_py_pb2", + ":raiden_service_py_pb2", ], ) diff --git a/tpu_raiden/rpc/raiden_controller.py b/tpu_raiden/rpc/raiden_controller.py index 754018ea..cea3084b 100644 --- a/tpu_raiden/rpc/raiden_controller.py +++ b/tpu_raiden/rpc/raiden_controller.py @@ -38,7 +38,7 @@ import typing from typing import Any, Optional, Union -from tpu_raiden.weight_sync import weight_synchronizer_service_pb2 +from tpu_raiden.rpc import raiden_service_pb2 class NameResolver(typing.Protocol): @@ -191,6 +191,7 @@ def __init__( endpoint_addresses: Optional[dict[RaidenId, str]] = None, resolve_timeout: float = 300.0, name_resolver: Optional[NameResolver] = None, + proto_module: Optional[Any] = None, ): """Instantiates RPC Client with an optional initial endpoint mapping. @@ -199,11 +200,13 @@ def __init__( resolve_timeout: Maximum duration in seconds to wait for a pending worker task to self-register before raising a Timeout RuntimeError. name_resolver: Interface for resolving remote coordinates (e.g. BNS). + proto_module: Optional protobuf module to use for ControlRequest/Response. """ self._endpoints = endpoint_addresses or {} self._pending_endpoints: dict[RaidenId, asyncio.Future[str]] = {} self._resolve_timeout = resolve_timeout self._name_resolver = name_resolver + self._proto_module = proto_module or raiden_service_pb2 def register_worker_endpoint( self, worker_name: RaidenId, rpc_address: str @@ -297,6 +300,13 @@ async def start_transfer( finally: sock.close() + def _raiden_id_to_proto(self, unit: RaidenId) -> Any: + return self._proto_module.RaidenIdProto( + job_name=unit.job_name, + job_replica_id=unit.job_replica_id, + data_name=unit.data_name, + ) + def _encode_start_transfer( self, target_id: RaidenId, transfer_plan: TransferPlan ) -> Optional[bytes]: @@ -309,56 +319,10 @@ def _encode_start_transfer( Returns: Serialized binary bytes payload, or None for no-op execution. """ - raise NotImplementedError("Subclasses must override _encode_start_transfer") - - def _verify_response(self, resp_bytes: bytes) -> None: - """Validates demarshaled remote response bytes returned from C++ workers. - - Args: - resp_bytes: Raw binary bytes received over the TCP socket. - - Raises: - RuntimeError: If remote execution reports explicit failure status. - """ - raise NotImplementedError("Subclasses must override _verify_response") - - def get_worker_endpoints(self) -> dict[RaidenId, str]: - """Returns active read-only snapshot of known registered Worker RPC endpoints.""" - return dict(self._endpoints) - - async def shutdown_workers(self) -> None: - """Dispatches remote shutdown signaling payloads to all registered worker daemons.""" - for _, addr in list(self._endpoints.items()): - try: - sock = connect_socket(addr, timeout=10.0, resolver=self._name_resolver) - payload = self._encode_shutdown() - sock.sendall(len(payload).to_bytes(4, "big") + payload) - sock.close() - except Exception: # pylint: disable=broad-except - pass - - def _encode_shutdown(self) -> bytes: - """Serializes domain-specific binary command for remote shutdown signaling.""" - raise NotImplementedError("Subclasses must override _encode_shutdown") - - -class WeightSyncWorkerRpcClient(WorkerRpcClient): - """Concrete domain subclass for state-of-the-art Weight Synchronizer Protobuf serialization.""" - - def _encode_start_transfer( - self, target_id: RaidenId, transfer_plan: TransferPlan - ) -> Optional[bytes]: - """Serializes ControlRequest protobuf specifically for native WeightSynchronizer servicer endpoints. - - Args: - target_id: Target worker RaidenId coordinate. - transfer_plan: Top-level distributed Collective Transfer execution plan. - - Returns: - Serialized binary ControlRequest Protobuf bytes payload, or None for no-op - Destination execution. - """ - if target_id not in transfer_plan.src_units: + if ( + target_id not in transfer_plan.src_units + and target_id not in transfer_plan.dst_units + ): return None peers = [] @@ -368,65 +332,140 @@ def _encode_start_transfer( ) peers.extend(dst_coords) - req = weight_synchronizer_service_pb2.ControlRequest( - command=weight_synchronizer_service_pb2.ControlRequest.COMMAND_START_TRANSFER, + req = self._proto_module.ControlRequest( + command=self._proto_module.ControlRequest.COMMAND_START_TRANSFER, peers=peers, ) - start_req = weight_synchronizer_service_pb2.StartTransferRequest( - src_units=[_raiden_id_to_proto(u) for u in transfer_plan.src_units], - dst_units=[_raiden_id_to_proto(u) for u in transfer_plan.dst_units], + is_sender = target_id in transfer_plan.src_units + start_req = self._proto_module.StartTransferRequest( + src_units=[ + self._raiden_id_to_proto(u) for u in transfer_plan.src_units + ], + dst_units=[ + self._raiden_id_to_proto(u) for u in transfer_plan.dst_units + ], + uuid=transfer_plan.uuid, + is_sender=is_sender, + dst_mem_type=int(transfer_plan.dst_mem_type), + use_block_chunks=transfer_plan.use_block_chunks, ) if transfer_plan.shard_push_schedules: - push_schedules = transfer_plan.shard_push_schedules.get(target_id) - if push_schedules: - for shard_idx, entries in push_schedules.items(): - schedule_proto = ( - weight_synchronizer_service_pb2.ShardPushScheduleProto() - ) - for ( - dst_peer, - dst_shard_idx, - dst_offset, - src_offset, - size, - ) in entries: - entry_proto = schedule_proto.entries.add() - entry_proto.dst_peer = dst_peer - entry_proto.dst_shard_idx = dst_shard_idx - entry_proto.dst_offset_bytes = dst_offset - entry_proto.src_offset_bytes = src_offset - entry_proto.size_bytes = size - start_req.shard_push_schedules[shard_idx].CopyFrom(schedule_proto) + if target_id in transfer_plan.dst_units: + # Receiver path: send FILTERED plan, only containing entries for this + # receiver + target_endpoints = transfer_plan.worker_data_addresses.get( + target_id, [] + ) + for ( + src_unit, + push_schedules, + ) in transfer_plan.shard_push_schedules.items(): + src_replica_idx = int(src_unit.job_replica_id) + # We assume each source worker has only 1 shard (shard_idx = 0) + schedule = push_schedules.get(0) + if schedule: + schedule_proto = self._proto_module.ShardPushScheduleProto() + for ( + dst_peer, + dst_shard_idx, + dst_offset, + src_offset, + size, + src_block_id, + dst_block_id, + ) in schedule: + if dst_peer in target_endpoints: + entry_proto = schedule_proto.entries.add() + entry_proto.dst_peer = dst_peer + entry_proto.dst_shard_idx = dst_shard_idx + entry_proto.dst_offset_bytes = dst_offset + entry_proto.src_offset_bytes = src_offset + entry_proto.size_bytes = size + entry_proto.src_block_id = src_block_id + entry_proto.dst_block_id = dst_block_id + if len(schedule_proto.entries) > 0: + start_req.shard_push_schedules[src_replica_idx].CopyFrom( + schedule_proto + ) + else: + # Sender path: only send local schedule + push_schedules = transfer_plan.shard_push_schedules.get(target_id) + if push_schedules: + for shard_idx, entries in push_schedules.items(): + schedule_proto = self._proto_module.ShardPushScheduleProto() + for ( + dst_peer, + dst_shard_idx, + dst_offset, + src_offset, + size, + src_block_id, + dst_block_id, + ) in entries: + entry_proto = schedule_proto.entries.add() + entry_proto.dst_peer = dst_peer + entry_proto.dst_shard_idx = dst_shard_idx + entry_proto.dst_offset_bytes = dst_offset + entry_proto.src_offset_bytes = src_offset + entry_proto.size_bytes = size + entry_proto.src_block_id = src_block_id + entry_proto.dst_block_id = dst_block_id + start_req.shard_push_schedules[shard_idx].CopyFrom(schedule_proto) req.start_transfer_request.CopyFrom(start_req) return req.SerializeToString() def _verify_response(self, resp_bytes: bytes) -> None: - """Validates demarshaled ControlResponse protobuf status returned from C++ WeightSynchronizer servers. - - Args: - resp_bytes: Raw binary Protobuf bytes received over the TCP socket. - - Raises: - RuntimeError: If remote execution reports explicit failure status. - """ - resp = weight_synchronizer_service_pb2.ControlResponse() + """Validates demarshaled remote response bytes returned from C++ workers.""" + resp = self._proto_module.ControlResponse() resp.ParseFromString(resp_bytes) if not resp.success: raise RuntimeError( - f"Weight Synchronizer remote native execution failed: {resp.message}" + f"Raiden remote native execution failed: {resp.message}" ) + def get_worker_endpoints(self) -> dict[RaidenId, str]: + """Returns active read-only snapshot of known registered Worker RPC endpoints.""" + return dict(self._endpoints) + + async def shutdown_workers(self) -> None: + """Dispatches remote shutdown signaling payloads to all registered worker daemons.""" + for _, addr in list(self._endpoints.items()): + try: + sock = connect_socket(addr, timeout=10.0, resolver=self._name_resolver) + payload = self._encode_shutdown() + sock.sendall(len(payload).to_bytes(4, "big") + payload) + sock.close() + except Exception: # pylint: disable=broad-except + pass + def _encode_shutdown(self) -> bytes: - """Serializes ControlRequest protobuf for remote WeightSynchronizer shutdown signaling.""" - req = weight_synchronizer_service_pb2.ControlRequest( - command=weight_synchronizer_service_pb2.ControlRequest.COMMAND_SHUTDOWN + """Serializes domain-specific binary command for remote shutdown signaling.""" + req = self._proto_module.ControlRequest( + command=self._proto_module.ControlRequest.COMMAND_SHUTDOWN ) return req.SerializeToString() +class WeightSyncWorkerRpcClient(WorkerRpcClient): + """Concrete domain subclass for state-of-the-art Weight Synchronizer Protobuf serialization.""" + + def __init__( + self, + endpoint_addresses: Optional[dict[RaidenId, str]] = None, + resolve_timeout: float = 300.0, + name_resolver: Optional[NameResolver] = None, + ): + super().__init__( + endpoint_addresses=endpoint_addresses, + resolve_timeout=resolve_timeout, + name_resolver=name_resolver, + proto_module=raiden_service_pb2, + ) + + class RaidenFuture: """Future representing an asynchronous transfer execution.""" @@ -583,7 +622,7 @@ def register_work_unit( shard_nd_slices: Optional[ Union[ list[list[tuple[int, int]]], - list[weight_synchronizer_service_pb2.NDSliceProto], + list[raiden_service_pb2.NDSliceProto], ] ] = None, itemsize: Optional[int] = None, @@ -825,10 +864,10 @@ def __init__( Args: controller: High-level RaidenController instance managing transfer plans. proto_module: Optional protobuf module to use for ControlRequest/Response. - Defaults to weight_synchronizer_service_pb2. + Defaults to raiden_service_pb2. """ self._controller = controller - self._proto_module = proto_module or weight_synchronizer_service_pb2 + self._proto_module = proto_module or raiden_service_pb2 self._sock = create_server_socket(controller.port) self._stopped = False self._thread = None @@ -967,16 +1006,6 @@ def _handle_conn( conn.close() -def _raiden_id_to_proto( - unit: RaidenId, -) -> weight_synchronizer_service_pb2.RaidenIdProto: - return weight_synchronizer_service_pb2.RaidenIdProto( - job_name=unit.job_name, - job_replica_id=unit.job_replica_id, - data_name=unit.data_name, - ) - - class RaidenControllerClientFacade: """Client-side stub encapsulating real remote Network RPCs to a centralized RaidenControllerServer.""" @@ -989,7 +1018,7 @@ def __init__( """Accepts Controller server coordinate 'ip:port'.""" self._address = controller_address self._name_resolver = name_resolver - self._proto_module = proto_module or weight_synchronizer_service_pb2 + self._proto_module = proto_module or raiden_service_pb2 def _raiden_id_to_proto( self, diff --git a/tpu_raiden/rpc/raiden_controller_test.py b/tpu_raiden/rpc/raiden_controller_test.py index b71c1562..ae9d249f 100644 --- a/tpu_raiden/rpc/raiden_controller_test.py +++ b/tpu_raiden/rpc/raiden_controller_test.py @@ -27,14 +27,24 @@ # limitations under the License. """Tests for Raiden Controller high-level transfer API under rpc/.""" +import asyncio from absl.testing import absltest from tpu_raiden.rpc import raiden_controller +class DummyWorkerRpcClient(raiden_controller.WorkerRpcClient): + + async def start_transfer(self, target_id, transfer_plan) -> None: + pass + + class RaidenControllerTest(absltest.TestCase): def test_dynamic_balancing_and_overlap_planner(self): - controller = raiden_controller.RaidenController(port=10000) + dummy_client = DummyWorkerRpcClient() + controller = raiden_controller.RaidenController( + port=10000, worker_rpc_client=dummy_client + ) src_unit_0 = raiden_controller.RaidenId( job_name="sampler", @@ -67,11 +77,7 @@ def test_dynamic_balancing_and_overlap_planner(self): src_units=[src_unit_0, src_unit_1], dst_units=[target_unit], ) - wait_task = future_1.wait() - try: - wait_task.send(None) - except StopIteration: - pass + asyncio.run(future_1.wait()) self.assertTrue(future_1.done()) self.assertEqual(future_1.session_id, 0) @@ -153,17 +159,13 @@ async def start_transfer(self, orchestrator_id, plan) -> None: dst_units=[dst], ) - wait_task = future.wait() - try: - wait_task.send(None) - except StopIteration: - pass + asyncio.run(future.wait()) self.assertEqual( recorded_actions, [ - ("start", [src]), ("start", [dst]), + ("start", [src]), ], ) diff --git a/tpu_raiden/kv_cache/kv_cache_service.proto b/tpu_raiden/rpc/raiden_service.proto similarity index 97% rename from tpu_raiden/kv_cache/kv_cache_service.proto rename to tpu_raiden/rpc/raiden_service.proto index b9a9542f..208f9892 100644 --- a/tpu_raiden/kv_cache/kv_cache_service.proto +++ b/tpu_raiden/rpc/raiden_service.proto @@ -14,10 +14,10 @@ syntax = "proto3"; -package tpu_raiden.kv_cache; +package tpu_raiden.rpc; option java_multiple_files = true; -option java_outer_classname = "KVCacheServiceProto"; +option java_outer_classname = "RaidenServiceProto"; message RaidenIdProto { string job_name = 1; diff --git a/tpu_raiden/weight_sync/BUILD b/tpu_raiden/weight_sync/BUILD index 3a14cc2c..1ea85de3 100644 --- a/tpu_raiden/weight_sync/BUILD +++ b/tpu_raiden/weight_sync/BUILD @@ -12,37 +12,19 @@ # 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:proto.bzl", "py_proto_library") package(default_visibility = ["//visibility:public"]) -proto_library( - name = "weight_synchronizer_service_proto", - srcs = ["weight_synchronizer_service.proto"], -) - -cc_proto_library( - name = "weight_synchronizer_service_cc_proto", - deps = [":weight_synchronizer_service_proto"], -) - -py_proto_library( - name = "weight_synchronizer_service_py_pb2", - deps = [":weight_synchronizer_service_proto"], -) - cc_library( name = "weight_synchronizer_base", srcs = [ "weight_synchronizer_base.cc", - "weight_synchronizer_control_service.cc", + "weight_synchronizer_listener.cc", ], hdrs = [ "weight_synchronizer_base.h", - "weight_synchronizer_control_service.h", + "weight_synchronizer_listener.h", ], copts = [ "-fno-strict-aliasing", @@ -51,11 +33,11 @@ cc_library( features = ["-use_header_modules"], visibility = ["//visibility:public"], deps = [ - ":weight_synchronizer_service_cc_proto", "//tpu_raiden/core:raiden_manager_base", "//tpu_raiden/core:raw_transfer_core", "//tpu_raiden/core:status_macros", "//tpu_raiden/core:xla_raw_transfer_headers", + "//tpu_raiden/rpc:raiden_service_cc_proto", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/log", "@com_google_absl//absl/status", @@ -86,8 +68,8 @@ cc_test( ) cc_test( - name = "weight_synchronizer_control_service_test", - srcs = ["weight_synchronizer_control_service_test.cc"], + name = "weight_synchronizer_listener_test", + srcs = ["weight_synchronizer_listener_test.cc"], copts = [ "-fno-strict-aliasing", "-fexceptions", @@ -95,7 +77,7 @@ cc_test( features = ["-use_header_modules"], deps = [ ":weight_synchronizer_base", - ":weight_synchronizer_service_cc_proto", + "//tpu_raiden/rpc:raiden_service_cc_proto", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", ], diff --git a/tpu_raiden/weight_sync/weight_synchronizer_base.cc b/tpu_raiden/weight_sync/weight_synchronizer_base.cc index 88b05a7e..624aaec1 100644 --- a/tpu_raiden/weight_sync/weight_synchronizer_base.cc +++ b/tpu_raiden/weight_sync/weight_synchronizer_base.cc @@ -40,8 +40,8 @@ #include "xla/tsl/platform/statusor.h" #include "tpu_raiden/core/raiden_manager_base.h" #include "tpu_raiden/core/raw_transfer_core.h" -#include "tpu_raiden/weight_sync/weight_synchronizer_control_service.h" -#include "tpu_raiden/weight_sync/weight_synchronizer_service.pb.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" +#include "tpu_raiden/weight_sync/weight_synchronizer_listener.h" ABSL_FLAG(size_t, raiden_weight_sync_host_buffer_scratchpad_size, 256 * 1024, "Amount of scratchpad to allocate to host buffers for resharding " @@ -187,7 +187,7 @@ WeightSynchronizerBase::WeightSynchronizerBase( std::optional local_port, std::optional> external_host_ptrs, bool unsafe_skip_buffer_lock, int parallelism, - std::optional control_port) + std::optional listener_port) : tpu_raiden::RaidenManagerBase( layer_buffers.size(), layer_buffers.empty() ? 0 : layer_buffers[0].size(), @@ -278,16 +278,16 @@ WeightSynchronizerBase::WeightSynchronizerBase( buffer_holds_.push_back(std::move(hold_info)); } - if (control_port) { - control_service_ = - std::make_unique(this, *control_port); + if (listener_port) { + listener_ = + std::make_unique(this, *listener_port); } } WeightSynchronizerBase::WeightSynchronizerBase( size_t num_layers, size_t num_shards, size_t slice_byte_size, std::optional local_port, std::optional host_blocks_to_allocate, - int parallelism, std::optional control_port) + int parallelism, std::optional listener_port) : tpu_raiden::RaidenManagerBase(num_layers, num_shards, slice_byte_size, local_port, parallelism) { physical_size_ = slice_byte_size_; @@ -325,22 +325,22 @@ WeightSynchronizerBase::WeightSynchronizerBase( layers_.push_back(std::move(layer_info)); } - if (control_port) { - control_service_ = - std::make_unique(this, *control_port); + if (listener_port) { + listener_ = + std::make_unique(this, *listener_port); } } -std::optional WeightSynchronizerBase::control_port() const { - if (control_service_) { - return control_service_->control_port(); +std::optional WeightSynchronizerBase::listener_port() const { + if (listener_) { + return listener_->listener_port(); } return std::nullopt; } -bool WeightSynchronizerBase::is_control_service_active() const { - if (control_service_) { - return control_service_->is_active(); +bool WeightSynchronizerBase::is_listener_active() const { + if (listener_) { + return listener_->is_active(); } return false; } @@ -524,7 +524,7 @@ absl::Status WeightSynchronizerBase::PushWeights( } absl::Status WeightSynchronizerBase::PushWeightsResharded( - const tpu_raiden::weight_sync::StartTransferRequest& request) { + const tpu_raiden::rpc::StartTransferRequest& request) { const auto& schedules = request.shard_push_schedules(); TF_ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2h()); TF_RETURN_IF_ERROR(d2h_future.Await()); diff --git a/tpu_raiden/weight_sync/weight_synchronizer_base.h b/tpu_raiden/weight_sync/weight_synchronizer_base.h index 5612b6e7..62aa2ffa 100644 --- a/tpu_raiden/weight_sync/weight_synchronizer_base.h +++ b/tpu_raiden/weight_sync/weight_synchronizer_base.h @@ -32,10 +32,13 @@ #include "tpu_raiden/core/raw_transfer_core.h" namespace tpu_raiden { +namespace rpc { +class StartTransferRequest; +} // namespace rpc + namespace weight_sync { -class WeightSynchronizerControlService; -class StartTransferRequest; +class WeightSynchronizerListener; class WeightSynchronizerBase : public tpu_raiden::RaidenManagerBase { public: @@ -46,17 +49,17 @@ class WeightSynchronizerBase : public tpu_raiden::RaidenManagerBase { std::optional> external_host_ptrs = std::nullopt, bool unsafe_skip_buffer_lock = false, int parallelism = 1, - std::optional control_port = std::nullopt); + std::optional listener_port = std::nullopt); // CPU-only constructor for remote workers and mock E2E testing WeightSynchronizerBase( size_t num_layers, size_t num_shards, size_t slice_byte_size, std::optional local_port = std::nullopt, std::optional host_blocks_to_allocate = std::nullopt, - int parallelism = 1, std::optional control_port = std::nullopt); + int parallelism = 1, std::optional listener_port = std::nullopt); - std::optional control_port() const; - bool is_control_service_active() const; + std::optional listener_port() const; + bool is_listener_active() const; ~WeightSynchronizerBase() override; @@ -80,7 +83,7 @@ class WeightSynchronizerBase : public tpu_raiden::RaidenManagerBase { * delivery to all remote peers. */ absl::Status PushWeightsResharded( - const tpu_raiden::weight_sync::StartTransferRequest& request); + const tpu_raiden::rpc::StartTransferRequest& request); // Inference server pulls current weights from the source peer E2E (network // pull + H2D) @@ -103,7 +106,7 @@ class WeightSynchronizerBase : public tpu_raiden::RaidenManagerBase { size_t size_bytes); protected: - std::unique_ptr control_service_; + std::unique_ptr listener_; const PJRT_Api* c_api_ = nullptr; const PJRT_RawBuffer_Extension* extension_ = nullptr; size_t physical_size_ = 0; diff --git a/tpu_raiden/weight_sync/weight_synchronizer_control_service.cc b/tpu_raiden/weight_sync/weight_synchronizer_listener.cc similarity index 57% rename from tpu_raiden/weight_sync/weight_synchronizer_control_service.cc rename to tpu_raiden/weight_sync/weight_synchronizer_listener.cc index e9fc18ab..0ed4f972 100644 --- a/tpu_raiden/weight_sync/weight_synchronizer_control_service.cc +++ b/tpu_raiden/weight_sync/weight_synchronizer_listener.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "tpu_raiden/weight_sync/weight_synchronizer_control_service.h" +#include "tpu_raiden/weight_sync/weight_synchronizer_listener.h" #include #include @@ -28,18 +28,18 @@ #include #include "absl/log/log.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" #include "tpu_raiden/weight_sync/weight_synchronizer_base.h" -#include "tpu_raiden/weight_sync/weight_synchronizer_service.pb.h" namespace tpu_raiden { namespace weight_sync { -WeightSynchronizerControlService::WeightSynchronizerControlService( - WeightSynchronizerBase* engine, int control_port) - : engine_(engine), control_port_(control_port) { +WeightSynchronizerListener::WeightSynchronizerListener( + WeightSynchronizerBase* engine, int listener_port) + : engine_(engine), listener_port_(listener_port) { server_fd_ = socket(AF_INET6, SOCK_STREAM, 0); if (server_fd_ < 0) { - LOG(FATAL) << "Failed to create C++ Control Service socket: " + LOG(FATAL) << "Failed to create C++ Listener socket: " << std::strerror(errno); } @@ -51,46 +51,45 @@ WeightSynchronizerControlService::WeightSynchronizerControlService( sockaddr_in6 address{}; address.sin6_family = AF_INET6; address.sin6_addr = in6addr_any; - address.sin6_port = htons(control_port_); + address.sin6_port = htons(listener_port_); if (bind(server_fd_, reinterpret_cast(&address), sizeof(address)) < 0) { - LOG(FATAL) << "C++ Control Service bind failed on port " << control_port_ - << ": " << std::strerror(errno); + LOG(FATAL) << "C++ Listener bind failed on port " << listener_port_ << ": " + << std::strerror(errno); } if (listen(server_fd_, 128) < 0) { - LOG(FATAL) << "C++ Control Service listen failed: " << std::strerror(errno); + LOG(FATAL) << "C++ Listener listen failed: " << std::strerror(errno); } socklen_t addr_len = sizeof(address); if (getsockname(server_fd_, reinterpret_cast(&address), &addr_len) == 0) { - control_port_ = ntohs(address.sin6_port); + listener_port_ = ntohs(address.sin6_port); } - LOG(INFO) << "Native C++ WeightSynchronizerControlService actively listening " + LOG(INFO) << "Native C++ WeightSynchronizerListener actively listening " "on port: " - << control_port_; + << listener_port_; listener_thread_ = - std::thread(&WeightSynchronizerControlService::ListenerLoop, this); + std::thread(&WeightSynchronizerListener::ListenerLoop, this); } -WeightSynchronizerControlService::~WeightSynchronizerControlService() { +WeightSynchronizerListener::~WeightSynchronizerListener() { stopping_ = true; 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(control_port_); + serv_addr.sin6_port = htons(listener_port_); serv_addr.sin6_addr = in6addr_loopback; if (connect(sock, reinterpret_cast(&serv_addr), sizeof(serv_addr)) == 0) { - tpu_raiden::weight_sync::ControlRequest req; - req.set_command( - tpu_raiden::weight_sync::ControlRequest::COMMAND_SHUTDOWN); + tpu_raiden::rpc::ControlRequest req; + req.set_command(tpu_raiden::rpc::ControlRequest::COMMAND_SHUTDOWN); std::string payload; if (req.SerializeToString(&payload)) { uint32_t net_len = htonl(payload.size()); @@ -114,7 +113,7 @@ WeightSynchronizerControlService::~WeightSynchronizerControlService() { } } -void WeightSynchronizerControlService::ListenerLoop() { +void WeightSynchronizerListener::ListenerLoop() { while (!stopping_) { sockaddr_in6 client_addr{}; socklen_t client_len = sizeof(client_addr); @@ -126,11 +125,11 @@ void WeightSynchronizerControlService::ListenerLoop() { } worker_threads_.push_back(std::thread( - &WeightSynchronizerControlService::ConnectionWorker, this, client_fd)); + &WeightSynchronizerListener::ConnectionWorker, this, client_fd)); } } -void WeightSynchronizerControlService::ConnectionWorker(int client_fd) { +void WeightSynchronizerListener::ConnectionWorker(int client_fd) { uint32_t net_len = 0; if (read(client_fd, &net_len, sizeof(net_len)) != sizeof(net_len)) { close(client_fd); @@ -150,53 +149,64 @@ void WeightSynchronizerControlService::ConnectionWorker(int client_fd) { total_read += n; } - tpu_raiden::weight_sync::ControlRequest req; + tpu_raiden::rpc::ControlRequest req; if (!req.ParseFromString(absl::string_view(buffer.data(), buffer.size()))) { LOG(ERROR) << "Failed to parse ControlRequest Protobuf"; close(client_fd); return; } - tpu_raiden::weight_sync::ControlResponse resp; + tpu_raiden::rpc::ControlResponse resp; resp.set_success(true); resp.set_message("SUCCESS"); if (req.command() == - tpu_raiden::weight_sync::ControlRequest::COMMAND_START_TRANSFER) { - if (req.has_start_transfer_request() && - !req.start_transfer_request().shard_push_schedules().empty()) { - LOG(INFO) << "C++ Control Service received START_TRANSFER with " - "shard_push_schedules"; + tpu_raiden::rpc::ControlRequest::COMMAND_START_TRANSFER) { + if (req.has_start_transfer_request()) { const auto& start_req = req.start_transfer_request(); - absl::Status status = engine_->PushWeightsResharded(start_req); - if (!status.ok()) { - resp.set_success(false); - resp.set_message(std::string(status.message())); - LOG(ERROR) << "PushWeightsResharded native execution failed: " - << status; - } - } else { - std::vector peers(req.peers().begin(), req.peers().end()); - LOG(INFO) << "C++ Control Service received START_TRANSFER request to " - << peers.size() << " peers"; - if (!peers.empty()) { - absl::Status status = engine_->PushWeights(peers); - if (!status.ok()) { - resp.set_success(false); - resp.set_message(std::string(status.message())); - LOG(ERROR) << "PushWeights native execution failed: " << status; + if (start_req.is_sender()) { + LOG(INFO) << "C++ Listener received START_TRANSFER (Sender)"; + if (!start_req.shard_push_schedules().empty()) { + LOG(INFO) << "C++ Listener executing PushWeightsResharded"; + absl::Status status = engine_->PushWeightsResharded(start_req); + if (!status.ok()) { + resp.set_success(false); + resp.set_message(std::string(status.message())); + LOG(ERROR) << "PushWeightsResharded native execution failed: " + << status; + } + } else { + std::vector peers(req.peers().begin(), + req.peers().end()); + LOG(INFO) << "C++ Listener executing PushWeights to " << peers.size() + << " peers"; + if (!peers.empty()) { + absl::Status status = engine_->PushWeights(peers); + if (!status.ok()) { + resp.set_success(false); + resp.set_message(std::string(status.message())); + LOG(ERROR) << "PushWeights native execution failed: " << status; + } + } } + } else { + LOG(INFO) << "C++ Listener received START_TRANSFER (Receiver) - no-op " + "for weight sync"; } + } else { + resp.set_success(false); + resp.set_message("Missing start_transfer_request"); + LOG(ERROR) << "Missing start_transfer_request in START_TRANSFER command"; } } else if (req.command() == - tpu_raiden::weight_sync::ControlRequest::COMMAND_SHUTDOWN) { - LOG(INFO) << "C++ Control Service received SHUTDOWN command. Initiating " + tpu_raiden::rpc::ControlRequest::COMMAND_SHUTDOWN) { + LOG(INFO) << "C++ Listener received SHUTDOWN command. Initiating " "clean exit."; stopping_ = true; } else { resp.set_success(false); resp.set_message("COMMAND_UNSPECIFIED"); - LOG(WARNING) << "C++ Control Service received unknown or unspecified " + LOG(WARNING) << "C++ Listener received unknown or unspecified " "Protobuf command"; } diff --git a/tpu_raiden/weight_sync/weight_synchronizer_control_service.h b/tpu_raiden/weight_sync/weight_synchronizer_listener.h similarity index 68% rename from tpu_raiden/weight_sync/weight_synchronizer_control_service.h rename to tpu_raiden/weight_sync/weight_synchronizer_listener.h index d1c2ee4f..5769dc71 100644 --- a/tpu_raiden/weight_sync/weight_synchronizer_control_service.h +++ b/tpu_raiden/weight_sync/weight_synchronizer_listener.h @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -#ifndef THIRD_PARTY_TPU_RAIDEN_WEIGHT_SYNC_WEIGHT_SYNCHRONIZER_CONTROL_SERVICE_H_ -#define THIRD_PARTY_TPU_RAIDEN_WEIGHT_SYNC_WEIGHT_SYNCHRONIZER_CONTROL_SERVICE_H_ +#ifndef THIRD_PARTY_TPU_RAIDEN_WEIGHT_SYNC_WEIGHT_SYNCHRONIZER_LISTENER_H_ +#define THIRD_PARTY_TPU_RAIDEN_WEIGHT_SYNC_WEIGHT_SYNCHRONIZER_LISTENER_H_ #include #include @@ -28,18 +28,16 @@ class WeightSynchronizerBase; // TCP Socket Server Daemon that runs natively in C++ to accept Control-Plane // management RPC commands (like PushWeights and Shutdown) directly from the // RL Coordinator or Controller task, bypassing Python servicer overhead. -class WeightSynchronizerControlService final { +class WeightSynchronizerListener final { public: - WeightSynchronizerControlService(WeightSynchronizerBase* engine, - int control_port); - ~WeightSynchronizerControlService(); + WeightSynchronizerListener(WeightSynchronizerBase* engine, int listener_port); + ~WeightSynchronizerListener(); - WeightSynchronizerControlService(const WeightSynchronizerControlService&) = + WeightSynchronizerListener(const WeightSynchronizerListener&) = delete; + WeightSynchronizerListener& operator=(const WeightSynchronizerListener&) = delete; - WeightSynchronizerControlService& operator=( - const WeightSynchronizerControlService&) = delete; - int control_port() const { return control_port_; } + int listener_port() const { return listener_port_; } bool is_active() const { return !stopping_; } private: @@ -47,7 +45,7 @@ class WeightSynchronizerControlService final { void ConnectionWorker(int client_fd); WeightSynchronizerBase* engine_; - int control_port_; + int listener_port_; int server_fd_ = -1; std::atomic stopping_{false}; @@ -58,4 +56,4 @@ class WeightSynchronizerControlService final { } // namespace weight_sync } // namespace tpu_raiden -#endif // THIRD_PARTY_TPU_RAIDEN_WEIGHT_SYNC_WEIGHT_SYNCHRONIZER_CONTROL_SERVICE_H_ +#endif // THIRD_PARTY_TPU_RAIDEN_WEIGHT_SYNC_WEIGHT_SYNCHRONIZER_LISTENER_H_ diff --git a/tpu_raiden/weight_sync/weight_synchronizer_control_service_test.cc b/tpu_raiden/weight_sync/weight_synchronizer_listener_test.cc similarity index 83% rename from tpu_raiden/weight_sync/weight_synchronizer_control_service_test.cc rename to tpu_raiden/weight_sync/weight_synchronizer_listener_test.cc index bf9a65fb..57aa6554 100644 --- a/tpu_raiden/weight_sync/weight_synchronizer_control_service_test.cc +++ b/tpu_raiden/weight_sync/weight_synchronizer_listener_test.cc @@ -26,7 +26,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "tpu_raiden/weight_sync/weight_synchronizer_control_service.h" +#include "tpu_raiden/weight_sync/weight_synchronizer_listener.h" #include #include @@ -40,14 +40,21 @@ #include #include +#include "tpu_raiden/rpc/raiden_service.pb.h" #include "tpu_raiden/weight_sync/weight_synchronizer_base.h" -#include "tpu_raiden/weight_sync/weight_synchronizer_service.pb.h" namespace tpu_raiden { namespace weight_sync { + +using ::tpu_raiden::rpc::ControlRequest; +using ::tpu_raiden::rpc::ControlResponse; +using ::tpu_raiden::rpc::ShardPushEntryProto; +using ::tpu_raiden::rpc::ShardPushScheduleProto; +using ::tpu_raiden::rpc::StartTransferRequest; + namespace { -int ConnectToControlPort(int port) { +int ConnectToListenerPort(int port) { int sock = socket(AF_INET6, SOCK_STREAM, 0); if (sock < 0) { return -1; @@ -67,23 +74,23 @@ int ConnectToControlPort(int port) { return sock; } -TEST(WeightSynchronizerControlServiceTest, PushWeightsCommandSuccess) { +TEST(WeightSynchronizerListenerTest, PushWeightsCommandSuccess) { WeightSynchronizerBase engine( /*num_layers=*/1, /*num_shards=*/1, /*slice_byte_size=*/128, /*local_port=*/0, /*host_blocks_to_allocate=*/std::nullopt, - /*parallelism=*/1, /*control_port=*/std::nullopt); + /*parallelism=*/1, /*listener_port=*/std::nullopt); - WeightSynchronizerControlService control_service(&engine, /*control_port=*/0); - ASSERT_GT(control_service.control_port(), 0); - EXPECT_TRUE(control_service.is_active()); + WeightSynchronizerListener listener(&engine, /*listener_port=*/0); + ASSERT_GT(listener.listener_port(), 0); + EXPECT_TRUE(listener.is_active()); - int sock = ConnectToControlPort(control_service.control_port()); + int sock = ConnectToListenerPort(listener.listener_port()); ASSERT_GE(sock, 0); WeightSynchronizerBase dst_engine( /*num_layers=*/1, /*num_shards=*/1, /*slice_byte_size=*/128, /*local_port=*/0, /*host_blocks_to_allocate=*/1, - /*parallelism=*/1, /*control_port=*/std::nullopt); + /*parallelism=*/1, /*listener_port=*/std::nullopt); ASSERT_TRUE(dst_engine.local_port().has_value()); ControlRequest req; @@ -113,16 +120,16 @@ TEST(WeightSynchronizerControlServiceTest, PushWeightsCommandSuccess) { close(sock); } -TEST(WeightSynchronizerControlServiceTest, ShutdownCommandStopsService) { +TEST(WeightSynchronizerListenerTest, ShutdownCommandStopsService) { WeightSynchronizerBase engine( /*num_layers=*/1, /*num_shards=*/1, /*slice_byte_size=*/128, /*local_port=*/0, /*host_blocks_to_allocate=*/std::nullopt, - /*parallelism=*/1, /*control_port=*/std::nullopt); + /*parallelism=*/1, /*listener_port=*/std::nullopt); - WeightSynchronizerControlService control_service(&engine, /*control_port=*/0); - EXPECT_TRUE(control_service.is_active()); + WeightSynchronizerListener listener(&engine, /*listener_port=*/0); + EXPECT_TRUE(listener.is_active()); - int sock = ConnectToControlPort(control_service.control_port()); + int sock = ConnectToListenerPort(listener.listener_port()); ASSERT_GE(sock, 0); ControlRequest req; @@ -151,23 +158,23 @@ TEST(WeightSynchronizerControlServiceTest, ShutdownCommandStopsService) { close(sock); } -TEST(WeightSynchronizerControlServiceTest, PushWeightsReshardedSuccess) { +TEST(WeightSynchronizerListenerTest, PushWeightsReshardedSuccess) { WeightSynchronizerBase src_engine( /*num_layers=*/1, /*num_shards=*/4, /*slice_byte_size=*/16, /*local_port=*/0, /*host_blocks_to_allocate=*/std::nullopt, - /*parallelism=*/1, /*control_port=*/std::nullopt); + /*parallelism=*/1, /*listener_port=*/std::nullopt); - WeightSynchronizerControlService control_service(&src_engine, - /*control_port=*/0); - ASSERT_GT(control_service.control_port(), 0); + WeightSynchronizerListener listener(&src_engine, + /*listener_port=*/0); + ASSERT_GT(listener.listener_port(), 0); - int sock = ConnectToControlPort(control_service.control_port()); + int sock = ConnectToListenerPort(listener.listener_port()); ASSERT_GE(sock, 0); WeightSynchronizerBase dst_engine( /*num_layers=*/1, /*num_shards=*/4, /*slice_byte_size=*/16, /*local_port=*/0, /*host_blocks_to_allocate=*/1, - /*parallelism=*/1, /*control_port=*/std::nullopt); + /*parallelism=*/1, /*listener_port=*/std::nullopt); ASSERT_TRUE(dst_engine.local_port().has_value()); std::string dst_peer = "127.0.0.1:" + std::to_string(*dst_engine.local_port()); diff --git a/tpu_raiden/weight_sync/weight_synchronizer_service.proto b/tpu_raiden/weight_sync/weight_synchronizer_service.proto deleted file mode 100644 index 38a018ea..00000000 --- a/tpu_raiden/weight_sync/weight_synchronizer_service.proto +++ /dev/null @@ -1,107 +0,0 @@ -// 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. - -syntax = "proto3"; - -package tpu_raiden.weight_sync; - -option java_multiple_files = true; -option java_outer_classname = "WeightSynchronizerServiceProto"; - -message RaidenIdProto { - string job_name = 1; - string job_replica_id = 2; - string data_name = 3; -} - -// Coordinate intervals representing a contiguous slice along a single tensor -// dimension. -message NDSliceDimensionProto { - // Inclusive start coordinate. - int64 start = 1; - // Exclusive end coordinate. - int64 end = 2; -} - -// Slices representing a multi-dimensional bounding box inside a distributed -// tensor. -message NDSliceProto { - // Dimension intervals ordered from outer-most dimension to inner-most minor - // dimension. - repeated NDSliceDimensionProto dimensions = 1; -} - -message RegisterWorkUnitRequest { - RaidenIdProto unit = 1; - repeated string shards = 2; - string control_plane_rpc_address = 3; - // Multi-dimensional bounding boxes owned by each logical device shard in - // row-major order. - repeated NDSliceProto shard_nd_slices = 4; - // Byte size of a single array scalar element (e.g., 4 for float32). - int32 itemsize = 5; -} - -// Defines an explicit 1D linear memory push operation from a local source to a -// remote destination shard. -message ShardPushEntryProto { - // Network coordinate "ip:port" or BNS of the remote Destination Worker Data - // servicer. - string dst_peer = 1; - // Logical shard partition index managed by the target Destination Worker - // process. - int32 dst_shard_idx = 2; - // Exact linear memory byte offset to paste data into within the target - // destination buffer. - int64 dst_offset_bytes = 3; - // Exact linear memory byte offset to copy data out of within the local source - // buffer. - int64 src_offset_bytes = 4; - // Number of continuous bytes to transmit across the TCP stream. - int64 size_bytes = 5; -} - -// A collection of network push entries orchestrated by a specific source TPU -// shard. -message ShardPushScheduleProto { - repeated ShardPushEntryProto entries = 1; -} - -message StartTransferRequest { - repeated RaidenIdProto src_units = 1; - repeated RaidenIdProto dst_units = 2; - // Centralized transfer execution plans mapped by local Source Worker logical - // shard partition index. - map shard_push_schedules = 3; -} - -message ControlRequest { - enum Command { - COMMAND_UNSPECIFIED = 0; - COMMAND_START_TRANSFER = 1; - COMMAND_SHUTDOWN = 2; - COMMAND_REGISTER_WORK_UNIT = 3; - } - Command command = 1; - repeated string peers = 2; - - RegisterWorkUnitRequest register_work_unit_request = 3; - StartTransferRequest start_transfer_request = 4; -} - -message ControlResponse { - bool success = 1; - string message = 2; - string response_data = 3; -}