diff --git a/tpu_sync/transport/BUILD b/tpu_sync/transport/BUILD index 9d1176a6..f10809a2 100644 --- a/tpu_sync/transport/BUILD +++ b/tpu_sync/transport/BUILD @@ -50,6 +50,7 @@ cc_library( deps = [ ":block_transport_delegate", ":buffer_push_task", + ":peregrine_control_service", "//tpu_sync/core:status_macros", "//tpu_sync/core:tsl_platform_headers", "//tpu_sync/telemetry:metrics_api", @@ -71,6 +72,24 @@ cc_library( ], ) +cc_library( + name = "peregrine_control_service", + hdrs = ["peregrine_control_service.h"], + visibility = [ + "//tpu_sync/transport:__subpackages__", + ], + deps = [ + "//tpu_sync/transport/lib:raw_buffer_transport", + "//tpu_sync/transport/proto:peregrine_control_service_cc_grpc", + "//tpu_sync/transport/proto:peregrine_control_service_cc_proto", + "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + ], +) + cc_test( name = "block_transport_test", srcs = ["block_transport_test.cc"], @@ -99,6 +118,31 @@ cc_test( ], ) +cc_test( + name = "peregrine_control_service_test", + srcs = ["peregrine_control_service_test.cc"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + deps = [ + ":block_transport", + ":block_transport_delegate", + ":peregrine_control_service", + "//tpu_sync/transport/lib:raw_buffer_transport", + "//tpu_sync/transport/lib:raw_buffer_transport_delegate", + "//tpu_sync/transport/proto:peregrine_control_service_cc_grpc", + "//tpu_sync/transport/proto:peregrine_control_service_cc_proto", + "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) + cc_binary( name = "h2h_strided_bench", srcs = ["h2h_strided_bench.cc"], diff --git a/tpu_sync/transport/block_transport.cc b/tpu_sync/transport/block_transport.cc index 1d8bbe2b..44edbde1 100644 --- a/tpu_sync/transport/block_transport.cc +++ b/tpu_sync/transport/block_transport.cc @@ -54,6 +54,7 @@ #include "tpu_sync/transport/lib/chunk_serializer.h" #include "tpu_sync/transport/lib/raw_buffer_transport.h" #include "tpu_sync/transport/peregrine/src/api/socket_util.h" +#include "tpu_sync/transport/peregrine_control_service.h" ABSL_FLAG(size_t, raiden_transport_coalesce_window_bytes, 0, "Maximum size in bytes of the host-side coalescing buffer used " @@ -199,7 +200,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 347a0fa9..6831687f 100644 --- a/tpu_sync/transport/block_transport.h +++ b/tpu_sync/transport/block_transport.h @@ -36,6 +36,7 @@ #include "tpu_sync/transport/buffer_push_task.h" #include "tpu_sync/transport/lib/chunk.h" #include "tpu_sync/transport/lib/raw_buffer_transport.h" +#include "tpu_sync/transport/peregrine_control_service.h" namespace tpu_raiden { namespace transport { @@ -67,6 +68,11 @@ 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 PeregrineControlService to register onto the host gRPC server. + proto::PeregrineControlService::Service* peregrine_control_service() { + return peregrine_control_.get(); + } + // Asynchronous Scatter-Gather Push void AsyncPush( const std::vector& peers, @@ -223,6 +229,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/raw_buffer_transport.cc b/tpu_sync/transport/lib/raw_buffer_transport.cc index a5e842de..01999487 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.cc +++ b/tpu_sync/transport/lib/raw_buffer_transport.cc @@ -419,6 +419,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 759bff48..24b4d581 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.h +++ b/tpu_sync/transport/lib/raw_buffer_transport.h @@ -112,6 +112,16 @@ class RawBufferTransport final { // Drops receive-progress counters belonging to the give `uuid`. void ForgetPushProgress(uint64_t uuid); + struct PspPeerKey { + uint32_t spi = 0; + std::string key; + }; + + // Registers incoming client PSP key and returns server's allocated RX key. + // In Phase 1: returns UnimplementedError until Phase 2 PSP socket activation. + 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_control_service.h b/tpu_sync/transport/peregrine_control_service.h new file mode 100644 index 00000000..a09facc0 --- /dev/null +++ b/tpu_sync/transport/peregrine_control_service.h @@ -0,0 +1,70 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TRANSPORT_PEREGRINE_CONTROL_SERVICE_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TRANSPORT_PEREGRINE_CONTROL_SERVICE_H_ + +#include "absl/status/statusor.h" +#include "grpcpp/server_context.h" +#include "grpcpp/support/status.h" +#include "tpu_sync/transport/lib/raw_buffer_transport.h" +#include "tpu_sync/transport/proto/peregrine_control_service.grpc.pb.h" +#include "tpu_sync/transport/proto/peregrine_control_service.pb.h" + +namespace tpu_raiden::transport { + +// Server-side gRPC implementation for PeregrineControlService. +// Handles incoming ExchangePspKey RPCs from connecting peers, delegates key +// registration to data plane RawBufferTransport, and returns the server's +// allocated RX SPI and key. +class PeregrineControlServiceImpl final + : public proto::PeregrineControlService::Service { + public: + explicit PeregrineControlServiceImpl(lib::RawBufferTransport* transport) + : transport_(transport) {} + + grpc::Status ExchangePspKey( + grpc::ServerContext* context, const proto::PspKeyExchangeRequest* request, + proto::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_or = transport_->RegisterPspPeer(request->client_spi(), + request->client_key()); + if (!server_rx_key_or.ok()) { + return grpc::Status(server_rx_key_or.status()); + } + + response->set_server_spi(server_rx_key_or->spi); + response->set_server_key(server_rx_key_or->key); + return grpc::Status::OK; + } + + private: + lib::RawBufferTransport* const transport_; +}; + +} // namespace tpu_raiden::transport + +#endif // TPU_SYNC_TRANSPORT_PEREGRINE_CONTROL_SERVICE_H_ diff --git a/tpu_sync/transport/peregrine_control_service_test.cc b/tpu_sync/transport/peregrine_control_service_test.cc new file mode 100644 index 00000000..a7daab6d --- /dev/null +++ b/tpu_sync/transport/peregrine_control_service_test.cc @@ -0,0 +1,108 @@ +// 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/peregrine_control_service.h" + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#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/block_transport.h" +#include "tpu_sync/transport/block_transport_delegate.h" +#include "tpu_sync/transport/lib/raw_buffer_transport.h" +#include "tpu_sync/transport/lib/raw_buffer_transport_delegate.h" +#include "tpu_sync/transport/proto/peregrine_control_service.grpc.pb.h" +#include "tpu_sync/transport/proto/peregrine_control_service.pb.h" + +namespace tpu_raiden::transport { +namespace { + +using ::testing::NotNull; + +class FakeRawDelegate : public lib::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; } +}; + +class FakeBlockDelegate : public BlockTransportDelegate { + public: + absl::StatusOr> AllocateBlocks(size_t num_blocks, + uint64_t uuid = 0) override { + return std::vector(num_blocks, 0); + } + 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; } + int GetRemoteReadBlockId(int base_remote_id, int chunk_k) override { + return base_remote_id + chunk_k; + } + size_t num_layers() const override { return 1; } + size_t num_shards() const override { return 1; } + size_t slice_byte_size() const override { return 1024; } + size_t shard_factor() const override { return 1; } +}; + +TEST(PeregrineControlServiceTest, BlockTransportOwnsPeregrineControlService) { + FakeBlockDelegate block_delegate; + BlockTransport transport(&block_delegate, /*local_port=*/0); + + EXPECT_THAT(transport.peregrine_control_service(), NotNull()); +} + +TEST(PeregrineControlServiceTest, InProcessGrpcExchangePspKey) { + FakeRawDelegate raw_delegate; + lib::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 = proto::PeregrineControlService::NewStub(channel); + + proto::PspKeyExchangeRequest req; + req.set_client_spi(0x12345678); + req.set_client_key(std::string(16, 'z')); + proto::PspKeyExchangeResponse resp; + grpc::ClientContext ctx; + + grpc::Status status = stub->ExchangePspKey(&ctx, req, &resp); + // In Phase 1, RawBufferTransport::RegisterPspPeer returns UnimplementedError + EXPECT_EQ(status.error_code(), grpc::StatusCode::UNIMPLEMENTED); + + server->Shutdown(); +} + +} // namespace +} // namespace tpu_raiden::transport diff --git a/tpu_sync/transport/proto/BUILD b/tpu_sync/transport/proto/BUILD new file mode 100644 index 00000000..56ac9bd7 --- /dev/null +++ b/tpu_sync/transport/proto/BUILD @@ -0,0 +1,35 @@ +# 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 = "peregrine_control_service_proto", + srcs = ["peregrine_control_service.proto"], +) + +cc_proto_library( + name = "peregrine_control_service_cc_proto", + deps = [":peregrine_control_service_proto"], +) + +cc_grpc_library( + name = "peregrine_control_service_cc_grpc", + srcs = [":peregrine_control_service_proto"], + deps = [":peregrine_control_service_cc_proto"], +) diff --git a/tpu_sync/transport/proto/peregrine_control_service.proto b/tpu_sync/transport/proto/peregrine_control_service.proto new file mode 100644 index 00000000..468ba292 --- /dev/null +++ b/tpu_sync/transport/proto/peregrine_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. + +syntax = "proto3"; + +package tpu_raiden.transport.proto; + +// Service used to exchange SPI and session keys for TCP-over-PSP encryption. +service PeregrineControlService { + // 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 +}