[tpu_raiden]Add PeregrineControlService TPU Raiden transport

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"],

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<PeregrineControlServiceImpl>(&raw_transport_)) {
   socket_workers_.reserve(parallelism_);
   for (int i = 0; i < parallelism_; ++i) {
     socket_workers_.push_back(

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<std::string>& peers,
@@ -223,6 +229,7 @@ class BlockTransport final {
   std::atomic<bool> scheduler_stopping_;
 
   lib::RawBufferTransport raw_transport_;
+  std::unique_ptr<PeregrineControlServiceImpl> peregrine_control_;
   std::vector<std::thread> socket_workers_;
 };
 

tpu_sync/transport/lib/raw_buffer_transport.cc

@@ -419,6 +419,13 @@ void RawBufferTransport::ConnectionWorker(int client_fd) {
   close(client_fd);
 }
 
+absl::StatusOr<RawBufferTransport::PspPeerKey>
+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_));

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<PspPeerKey> 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

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_

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 <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#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<std::vector<int>> AllocateBlocks(size_t num_blocks,
+                                                  uint64_t uuid = 0) override {
+    return std::vector<int>(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<grpc::Server> server = builder.BuildAndStart();
+  ASSERT_THAT(server, NotNull());
+
+  std::shared_ptr<grpc::Channel> 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

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"],
+)

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
+}