From f7119229173fa237f155f585dbeb66ce5a7220bd Mon Sep 17 00:00:00 2001 From: Googler Date: Fri, 21 Aug 2026 10:42:41 -0700 Subject: [PATCH] [tpu_raiden]Add PeregrineControlService TPU Raiden transport PiperOrigin-RevId: 968577745 --- tpu_sync/transport/BUILD | 1 + tpu_sync/transport/block_transport.cc | 5 +- tpu_sync/transport/block_transport.h | 8 ++ tpu_sync/transport/lib/BUILD | 33 ++++++++ .../transport/lib/peregrine_control_service.h | 74 ++++++++++++++++++ .../lib/peregrine_control_service_test.cc | 78 +++++++++++++++++++ .../transport/lib/raw_buffer_transport.cc | 7 ++ tpu_sync/transport/lib/raw_buffer_transport.h | 10 +++ .../peregrine/src/internal/control/BUILD | 36 +++++++++ .../src/internal/control/service.proto | 33 ++++++++ 10 files changed, 284 insertions(+), 1 deletion(-) create mode 100644 tpu_sync/transport/lib/peregrine_control_service.h create mode 100644 tpu_sync/transport/lib/peregrine_control_service_test.cc create mode 100644 tpu_sync/transport/peregrine/src/internal/control/BUILD create mode 100644 tpu_sync/transport/peregrine/src/internal/control/service.proto diff --git a/tpu_sync/transport/BUILD b/tpu_sync/transport/BUILD index e76aa84a..fef20017 100644 --- a/tpu_sync/transport/BUILD +++ b/tpu_sync/transport/BUILD @@ -56,6 +56,7 @@ cc_library( "//tpu_sync/telemetry:metrics_backend", "//tpu_sync/transport/lib:chunk", "//tpu_sync/transport/lib:chunk_serializer", + "//tpu_sync/transport/lib:peregrine_control_service", "//tpu_sync/transport/lib:raw_buffer_transport", "//tpu_sync/transport/lib:transport_adapter", "//tpu_sync/transport/peregrine/src/api:socket_util", diff --git a/tpu_sync/transport/block_transport.cc b/tpu_sync/transport/block_transport.cc index 06153e52..1fe4e547 100644 --- a/tpu_sync/transport/block_transport.cc +++ b/tpu_sync/transport/block_transport.cc @@ -54,6 +54,7 @@ #include "tpu_sync/transport/buffer_push_task.h" #include "tpu_sync/transport/lib/chunk.h" #include "tpu_sync/transport/lib/chunk_serializer.h" +#include "tpu_sync/transport/lib/peregrine_control_service.h" #include "tpu_sync/transport/lib/raw_buffer_transport.h" #include "tpu_sync/transport/lib/transport_adapter.h" #include "tpu_sync/transport/peregrine/src/api/socket_util.h" @@ -221,7 +222,9 @@ BlockTransport::BlockTransport(BlockTransportDelegate* delegate, int local_port, [this](int client_fd, const lib::ChunkHeader& header) { return HandleCustomRequest(client_fd, header); }, - absl::GetFlag(FLAGS_raiden_transport_coalesce_window_bytes)) { + absl::GetFlag(FLAGS_raiden_transport_coalesce_window_bytes)), + peregrine_control_( + std::make_unique(&raw_transport_)) { socket_workers_.reserve(parallelism_); for (int i = 0; i < parallelism_; ++i) { socket_workers_.push_back( diff --git a/tpu_sync/transport/block_transport.h b/tpu_sync/transport/block_transport.h index 9da8c03b..4412404f 100644 --- a/tpu_sync/transport/block_transport.h +++ b/tpu_sync/transport/block_transport.h @@ -36,6 +36,7 @@ #include "tpu_sync/transport/block_transport_delegate.h" #include "tpu_sync/transport/buffer_push_task.h" #include "tpu_sync/transport/lib/chunk.h" +#include "tpu_sync/transport/lib/peregrine_control_service.h" #include "tpu_sync/transport/lib/raw_buffer_transport.h" #include "tpu_sync/transport/lib/transport_adapter.h" @@ -69,6 +70,12 @@ class BlockTransport final { // It is the first IP in `local_ips` if provided, otherwise "127.0.0.1". const std::string& bound_ip() const { return raw_transport_.bound_ip(); } + // Returns PeregrineService to register onto the host gRPC server. + ::peregrine::internal::control::PeregrineService::Service* + peregrine_control_service() { + return peregrine_control_.get(); + } + // Asynchronous Scatter-Gather Push void AsyncPush( const std::vector& peers, @@ -248,6 +255,7 @@ class BlockTransport final { std::atomic scheduler_stopping_; lib::RawBufferTransport raw_transport_; + std::unique_ptr peregrine_control_; std::vector socket_workers_; }; diff --git a/tpu_sync/transport/lib/BUILD b/tpu_sync/transport/lib/BUILD index 8619059d..7ae9b2c6 100644 --- a/tpu_sync/transport/lib/BUILD +++ b/tpu_sync/transport/lib/BUILD @@ -136,3 +136,36 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "peregrine_control_service", + hdrs = ["peregrine_control_service.h"], + deps = [ + ":raw_buffer_transport", + "//tpu_sync/transport/peregrine/src/internal/control:service_cc_grpc", + "//tpu_sync/transport/peregrine/src/internal/control:service_cc_proto", + "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "peregrine_control_service_test", + srcs = ["peregrine_control_service_test.cc"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + deps = [ + ":peregrine_control_service", + ":raw_buffer_transport", + ":raw_buffer_transport_delegate", + "//tpu_sync/transport/peregrine/src/internal/control:service_cc_grpc", + "//tpu_sync/transport/peregrine/src/internal/control:service_cc_proto", + "@com_github_grpc_grpc//:grpc++", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_sync/transport/lib/peregrine_control_service.h b/tpu_sync/transport/lib/peregrine_control_service.h new file mode 100644 index 00000000..98f7e262 --- /dev/null +++ b/tpu_sync/transport/lib/peregrine_control_service.h @@ -0,0 +1,74 @@ +// 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 TPU_SYNC_TRANSPORT_LIB_PEREGRINE_CONTROL_SERVICE_H_ +#define TPU_SYNC_TRANSPORT_LIB_PEREGRINE_CONTROL_SERVICE_H_ + +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "grpcpp/server_context.h" +#include "grpcpp/support/status.h" +#include "tpu_sync/transport/lib/raw_buffer_transport.h" +#include "tpu_sync/transport/peregrine/src/internal/control/service.grpc.pb.h" +#include "tpu_sync/transport/peregrine/src/internal/control/service.pb.h" + +namespace tpu_raiden::transport::lib { + +// Server-side gRPC implementation for PeregrineService. +// Handles incoming RPCs from connecting peers +class PeregrineControlServiceImpl final + : public ::peregrine::internal::control::PeregrineService::Service { + public: + explicit PeregrineControlServiceImpl(RawBufferTransport* transport) + : transport_(transport) {} + + grpc::Status ExchangePspKey( + grpc::ServerContext* context, + const ::peregrine::internal::control::PspKeyExchangeRequest* request, + ::peregrine::internal::control::PspKeyExchangeResponse* response) + override { + if (transport_ == nullptr) { + return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, + "RawBufferTransport is not initialized"); + } + if (request->client_spi() == 0) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, + "client_spi must be non-zero"); + } + if (request->client_key().size() != 16) { + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, + "client_key must be exactly 16 bytes"); + } + + auto server_rx_key = transport_->RegisterPspPeer(request->client_spi(), + request->client_key()); + if (!server_rx_key.ok()) { + return grpc::Status( + grpc::StatusCode::INTERNAL, + absl::StrCat("Failed to register PSP peer: ", + server_rx_key.status().message())); + } + + response->set_server_spi(server_rx_key->spi); + response->set_server_key(server_rx_key->key); + return grpc::Status::OK; + } + + private: + RawBufferTransport* const transport_; +}; + +} // namespace tpu_raiden::transport::lib + +#endif // TPU_SYNC_TRANSPORT_LIB_PEREGRINE_CONTROL_SERVICE_H_ diff --git a/tpu_sync/transport/lib/peregrine_control_service_test.cc b/tpu_sync/transport/lib/peregrine_control_service_test.cc new file mode 100644 index 00000000..b3802a23 --- /dev/null +++ b/tpu_sync/transport/lib/peregrine_control_service_test.cc @@ -0,0 +1,78 @@ +// 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_sync/transport/lib/peregrine_control_service.h" + +#include +#include +#include +#include + +#include +#include +#include "grpcpp/channel.h" +#include "grpcpp/client_context.h" +#include "grpcpp/server.h" +#include "grpcpp/server_builder.h" +#include "grpcpp/server_context.h" +#include "grpcpp/support/channel_arguments.h" +#include "grpcpp/support/status.h" +#include "tpu_sync/transport/lib/raw_buffer_transport.h" +#include "tpu_sync/transport/lib/raw_buffer_transport_delegate.h" +#include "tpu_sync/transport/peregrine/src/internal/control/service.grpc.pb.h" +#include "tpu_sync/transport/peregrine/src/internal/control/service.pb.h" + +namespace tpu_raiden::transport::lib { +namespace { + +using ::testing::NotNull; + +class FakeRawDelegate : public RawBufferTransportDelegate { + public: + uint8_t* GetHostPointer(size_t buffer_id, size_t shard_idx) override { + return nullptr; + } + size_t GetHostSize(size_t buffer_id, size_t shard_idx) override { return 0; } +}; + +TEST(PeregrineControlServiceTest, InProcessGrpcExchangePspKey) { + FakeRawDelegate raw_delegate; + RawBufferTransport transport(&raw_delegate, /*local_port=*/0); + PeregrineControlServiceImpl service(&transport); + + grpc::ServerBuilder builder; + builder.RegisterService(&service); + std::unique_ptr server = builder.BuildAndStart(); + ASSERT_THAT(server, NotNull()); + + std::shared_ptr channel = + server->InProcessChannel(grpc::ChannelArguments()); + auto stub = + ::peregrine::internal::control::PeregrineService::NewStub(channel); + + ::peregrine::internal::control::PspKeyExchangeRequest req; + req.set_client_spi(0x12345678); + req.set_client_key(std::string(16, 'z')); + ::peregrine::internal::control::PspKeyExchangeResponse resp; + grpc::ClientContext ctx; + + grpc::Status status = stub->ExchangePspKey(&ctx, req, &resp); + // TODO(yyd): update test once the implementation is done. + EXPECT_EQ(status.error_code(), grpc::StatusCode::INTERNAL); + + server->Shutdown(); +} + +} // namespace +} // namespace tpu_raiden::transport::lib diff --git a/tpu_sync/transport/lib/raw_buffer_transport.cc b/tpu_sync/transport/lib/raw_buffer_transport.cc index 8f421e59..1e0ecbf5 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.cc +++ b/tpu_sync/transport/lib/raw_buffer_transport.cc @@ -420,6 +420,13 @@ void RawBufferTransport::ConnectionWorker(int client_fd) { close(client_fd); } +absl::StatusOr +RawBufferTransport::RegisterPspPeer(uint32_t client_spi, + absl::string_view client_key) { + return absl::UnimplementedError( + "PSP key registration is not implemented yet."); +} + void RawBufferTransport::ListenerLoop() { while (!stopping_) { DCHECK(IsValidSocket(server_fd_)); diff --git a/tpu_sync/transport/lib/raw_buffer_transport.h b/tpu_sync/transport/lib/raw_buffer_transport.h index 77a4f943..922bdaef 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.h +++ b/tpu_sync/transport/lib/raw_buffer_transport.h @@ -122,6 +122,16 @@ class RawBufferTransport final { // Drops receive-progress counters belonging to the give `uuid`. void ForgetPushProgress(uint64_t uuid); + // TODO(yyd): Move this struct to psp_tcp_helper.h. + struct PspPeerKey { + uint32_t spi = 0; + std::string key; + }; + + // Registers incoming client PSP key and returns server's allocated RX key. + absl::StatusOr RegisterPspPeer(uint32_t client_spi, + absl::string_view client_key); + private: // Pushes a batch of buffers to the remote `peer`, by sending out a // `kOpBufferPushBatched ChunkHeader` followed by a `batch_size` sequence diff --git a/tpu_sync/transport/peregrine/src/internal/control/BUILD b/tpu_sync/transport/peregrine/src/internal/control/BUILD new file mode 100644 index 00000000..9ced9ab8 --- /dev/null +++ b/tpu_sync/transport/peregrine/src/internal/control/BUILD @@ -0,0 +1,36 @@ +# 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. + +load("@com_github_grpc_grpc//bazel:cc_grpc_library.bzl", "cc_grpc_library") +load("@com_google_protobuf//bazel:cc_proto_library.bzl", "cc_proto_library") +load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library") + +package(default_visibility = ["//visibility:public"]) + +proto_library( + name = "service_proto", + srcs = ["service.proto"], +) + +cc_proto_library( + name = "service_cc_proto", + deps = [":service_proto"], +) + +cc_grpc_library( + name = "service_cc_grpc", + srcs = [":service_proto"], + grpc_only = True, + deps = [":service_cc_proto"], +) diff --git a/tpu_sync/transport/peregrine/src/internal/control/service.proto b/tpu_sync/transport/peregrine/src/internal/control/service.proto new file mode 100644 index 00000000..75b443fd --- /dev/null +++ b/tpu_sync/transport/peregrine/src/internal/control/service.proto @@ -0,0 +1,33 @@ +// 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. + +edition = "2024"; + +package peregrine.internal.control; + +// Service for Peregrine control plane communication. +service PeregrineService { + // Client sends its local RX SPI and key; server returns its RX SPI and key. + rpc ExchangePspKey(PspKeyExchangeRequest) returns (PspKeyExchangeResponse) { } +} + +message PspKeyExchangeRequest { + uint32 client_spi = 1; + bytes client_key = 2; // 16-byte raw AES session key +} + +message PspKeyExchangeResponse { + uint32 server_spi = 1; + bytes server_key = 2; // 16-byte raw AES session key +}