Add WorkerService gRPC proto and core C++ client/service implementations.

tpu_raiden/core/controller/BUILD

@@ -0,0 +1,66 @@
+# 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("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "worker_service_client",
+    srcs = ["worker_service_client.cc"],
+    hdrs = ["worker_service_client.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tpu_raiden/proto:worker_service_cc_grpc",
+        "//tpu_raiden/proto:worker_service_cc_proto",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "worker_service_impl",
+    srcs = ["worker_service_impl.cc"],
+    hdrs = ["worker_service_impl.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tpu_raiden/core:host_memory_allocator",
+        "//tpu_raiden/proto:worker_service_cc_grpc",
+        "//tpu_raiden/proto:worker_service_cc_proto",
+        "//tpu_raiden/rpc:raiden_service_cc_proto",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+cc_test(
+    name = "worker_service_test",
+    srcs = ["worker_service_test.cc"],
+    deps = [
+        ":worker_service_client",
+        ":worker_service_impl",
+        "//tpu_raiden/proto:worker_service_cc_proto",
+        "//tpu_raiden/rpc:raiden_service_cc_proto",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tpu_raiden/core/controller/worker_service_client.cc

@@ -0,0 +1,60 @@
+// 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_raiden/core/controller/worker_service_client.h"
+
+#include <memory>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "third_party/grpc/include/grpcpp/client_context.h"
+#include "third_party/grpc/include/grpcpp/support/status.h"
+#include "tpu_raiden/proto/worker_service.grpc.pb.h"
+#include "tpu_raiden/proto/worker_service.pb.h"
+
+namespace tpu_raiden {
+namespace controller {
+
+WorkerServiceClient::WorkerServiceClient(std::shared_ptr<grpc::Channel> channel)
+    : stub_(proto::WorkerService::NewStub(channel)) {}
+
+absl::StatusOr<proto::CreateBuffersResponse> WorkerServiceClient::CreateBuffers(
+    const proto::CreateBuffersRequest& request) {
+  proto::CreateBuffersResponse response;
+  grpc::ClientContext context;
+
+  grpc::Status status = stub_->CreateBuffers(&context, request, &response);
+  if (!status.ok()) {
+    return absl::InternalError(
+        absl::StrCat("CreateBuffers RPC failed: ", status.error_message()));
+  }
+  return response;
+}
+
+absl::StatusOr<proto::DeleteBuffersResponse> WorkerServiceClient::DeleteBuffers(
+    const proto::DeleteBuffersRequest& request) {
+  proto::DeleteBuffersResponse response;
+  grpc::ClientContext context;
+
+  grpc::Status status = stub_->DeleteBuffers(&context, request, &response);
+  if (!status.ok()) {
+    return absl::InternalError(
+        absl::StrCat("DeleteBuffers RPC failed: ", status.error_message()));
+  }
+  return response;
+}
+
+}  // namespace controller
+}  // namespace tpu_raiden

tpu_raiden/core/controller/worker_service_client.h

@@ -0,0 +1,49 @@
+// 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_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_CLIENT_H_
+#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_CLIENT_H_
+
+#include <memory>
+
+#include "absl/status/statusor.h"
+#include "third_party/grpc/include/grpcpp/channel.h"
+#include "tpu_raiden/proto/worker_service.grpc.pb.h"
+#include "tpu_raiden/proto/worker_service.pb.h"
+
+namespace tpu_raiden {
+namespace controller {
+
+// Client for interacting with the WorkerService gRPC endpoint on a transfer
+// worker.
+class WorkerServiceClient {
+ public:
+  explicit WorkerServiceClient(std::shared_ptr<grpc::Channel> channel);
+
+  // Allocates sharded buffers on the remote transfer worker.
+  absl::StatusOr<proto::CreateBuffersResponse> CreateBuffers(
+      const proto::CreateBuffersRequest& request);
+
+  // Deallocates sharded buffers on the remote transfer worker.
+  absl::StatusOr<proto::DeleteBuffersResponse> DeleteBuffers(
+      const proto::DeleteBuffersRequest& request);
+
+ private:
+  std::unique_ptr<proto::WorkerService::Stub> stub_;
+};
+
+}  // namespace controller
+}  // namespace tpu_raiden
+
+#endif  // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_CLIENT_H_

tpu_raiden/core/controller/worker_service_impl.cc

@@ -0,0 +1,116 @@
+// 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_raiden/core/controller/worker_service_impl.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
+#include "third_party/grpc/include/grpcpp/server_context.h"
+#include "third_party/grpc/include/grpcpp/support/status.h"
+#include "tpu_raiden/core/host_memory_allocator.h"
+#include "tpu_raiden/proto/worker_service.pb.h"
+#include "tpu_raiden/rpc/raiden_service.pb.h"
+
+namespace tpu_raiden {
+namespace controller {
+
+WorkerServiceImpl::WorkerServiceImpl(
+    std::shared_ptr<HostMemoryAllocator> allocator)
+    : allocator_(allocator ? std::move(allocator)
+                           : std::make_shared<MallocHostMemoryAllocator>()) {}
+
+grpc::Status WorkerServiceImpl::CreateBuffers(
+    grpc::ServerContext* context, const proto::CreateBuffersRequest* request,
+    proto::CreateBuffersResponse* response) {
+  absl::MutexLock lock(mutex_);
+  for (const auto& spec : request->buffers()) {
+    if (spec.num_shards() <= 0 || spec.size_bytes() <= 0) {
+      response->set_success(false);
+      response->set_message("num_shards and size_bytes must be positive");
+      return grpc::Status::OK;
+    }
+    proto::BufferProto sharded_buf;
+    for (int32_t s = 0; s < spec.num_shards(); ++s) {
+      auto alloc_or = allocator_->Allocate(spec.size_bytes());
+      if (!alloc_or.ok()) {
+        response->set_success(false);
+        response->set_message(absl::StrCat("Failed to allocate memory: ",
+                                           alloc_or.status().message()));
+        return grpc::Status::OK;
+      }
+      BufferHandle handle = next_buffer_handle_++;
+      buffers_[handle] = std::move(*alloc_or);
+      sharded_buf.add_buffer_handles()->set_handle(handle);
+    }
+    *response->add_buffers() = std::move(sharded_buf);
+  }
+  response->set_success(true);
+  response->set_message("Buffers created successfully");
+  return grpc::Status::OK;
+}
+
+grpc::Status WorkerServiceImpl::DeleteBuffers(
+    grpc::ServerContext* context, const proto::DeleteBuffersRequest* request,
+    proto::DeleteBuffersResponse* response) {
+  absl::MutexLock lock(mutex_);
+  for (const auto& sharded_buffer : request->sharded_buffers()) {
+    if (sharded_buffer.buffer_handles().empty()) {
+      response->set_success(false);
+      response->set_message("BufferProto has no buffer_handles");
+      return grpc::Status::OK;
+    }
+    for (const auto& handle_proto : sharded_buffer.buffer_handles()) {
+      BufferHandle handle = handle_proto.handle();
+      auto it = buffers_.find(handle);
+      if (it == buffers_.end()) {
+        response->set_success(false);
+        response->set_message(
+            absl::StrCat("Buffer not found for deletion: handle ", handle));
+        return grpc::Status::OK;
+      }
+      buffers_.erase(it);
+    }
+  }
+  response->set_success(true);
+  response->set_message("Buffers deleted successfully");
+  return grpc::Status::OK;
+}
+
+absl::StatusOr<HostBufferAllocation> WorkerServiceImpl::GetBuffer(
+    BufferHandle handle) const {
+  absl::MutexLock lock(mutex_);
+  auto it = buffers_.find(handle);
+  if (it == buffers_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Buffer not found: handle ", handle));
+  }
+  return it->second;
+}
+
+size_t WorkerServiceImpl::GetBufferCount() const {
+  absl::MutexLock lock(mutex_);
+  return buffers_.size();
+}
+
+}  // namespace controller
+}  // namespace tpu_raiden

tpu_raiden/core/controller/worker_service_impl.h

@@ -0,0 +1,81 @@
+// 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_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_IMPL_H_
+#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_IMPL_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
+#include "third_party/grpc/include/grpcpp/server_context.h"
+#include "third_party/grpc/include/grpcpp/support/status.h"
+#include "tpu_raiden/core/host_memory_allocator.h"
+#include "tpu_raiden/proto/worker_service.grpc.pb.h"
+#include "tpu_raiden/proto/worker_service.pb.h"
+#include "tpu_raiden/rpc/raiden_service.pb.h"
+
+namespace tpu_raiden {
+namespace controller {
+
+using BufferHandle = uint64_t;
+
+// Implementation of the WorkerService gRPC service running on transfer workers.
+// Manages allocation and deallocation of sharded host memory buffers.
+class WorkerServiceImpl final : public proto::WorkerService::Service {
+ public:
+  // Constructs a WorkerServiceImpl with the given host memory allocator. If
+  // allocator is nullptr, defaults to MallocHostMemoryAllocator.
+  explicit WorkerServiceImpl(
+      std::shared_ptr<HostMemoryAllocator> allocator = nullptr);
+
+  grpc::Status CreateBuffers(grpc::ServerContext* context,
+                             const proto::CreateBuffersRequest* request,
+                             proto::CreateBuffersResponse* response) override;
+
+  grpc::Status DeleteBuffers(grpc::ServerContext* context,
+                             const proto::DeleteBuffersRequest* request,
+                             proto::DeleteBuffersResponse* response) override;
+
+  // Retrieves an allocated buffer shard for inspection or transfer operations.
+  // Returns NotFoundError if the buffer handle is invalid.
+  absl::StatusOr<HostBufferAllocation> GetBuffer(BufferHandle handle) const;
+
+  // Overload accepting BufferHandleProto for convenience.
+  absl::StatusOr<HostBufferAllocation> GetBuffer(
+      const proto::BufferHandleProto& handle_proto) const {
+    return GetBuffer(BufferHandle(handle_proto.handle()));
+  }
+
+  // Returns the total number of currently allocated buffer shards across all
+  // buffers.
+  size_t GetBufferCount() const;
+
+ private:
+  std::shared_ptr<HostMemoryAllocator> allocator_;
+  mutable absl::Mutex mutex_;
+  uint64_t next_buffer_handle_ ABSL_GUARDED_BY(mutex_) = 1;
+  absl::flat_hash_map<BufferHandle, HostBufferAllocation> buffers_
+      ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace controller
+}  // namespace tpu_raiden
+
+#endif  // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_IMPL_H_