diff --git a/tpu_raiden/core/controller/BUILD b/tpu_raiden/core/controller/BUILD new file mode 100644 index 0000000..d0ab5c8 --- /dev/null +++ b/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", + ], +) diff --git a/tpu_raiden/core/controller/worker_service_client.cc b/tpu_raiden/core/controller/worker_service_client.cc new file mode 100644 index 0000000..51f8e5e --- /dev/null +++ b/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 + +#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 channel) + : stub_(proto::WorkerService::NewStub(channel)) {} + +absl::StatusOr 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 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 diff --git a/tpu_raiden/core/controller/worker_service_client.h b/tpu_raiden/core/controller/worker_service_client.h new file mode 100644 index 0000000..984eec4 --- /dev/null +++ b/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 + +#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 channel); + + // Allocates sharded buffers on the remote transfer worker. + absl::StatusOr CreateBuffers( + const proto::CreateBuffersRequest& request); + + // Deallocates sharded buffers on the remote transfer worker. + absl::StatusOr DeleteBuffers( + const proto::DeleteBuffersRequest& request); + + private: + std::unique_ptr stub_; +}; + +} // namespace controller +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_CLIENT_H_ diff --git a/tpu_raiden/core/controller/worker_service_impl.cc b/tpu_raiden/core/controller/worker_service_impl.cc new file mode 100644 index 0000000..7ac28c5 --- /dev/null +++ b/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 +#include +#include +#include +#include + +#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 allocator) + : allocator_(allocator ? std::move(allocator) + : std::make_shared()) {} + +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 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 diff --git a/tpu_raiden/core/controller/worker_service_impl.h b/tpu_raiden/core/controller/worker_service_impl.h new file mode 100644 index 0000000..c3eda15 --- /dev/null +++ b/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 +#include +#include +#include +#include + +#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 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 GetBuffer(BufferHandle handle) const; + + // Overload accepting BufferHandleProto for convenience. + absl::StatusOr 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 allocator_; + mutable absl::Mutex mutex_; + uint64_t next_buffer_handle_ ABSL_GUARDED_BY(mutex_) = 1; + absl::flat_hash_map buffers_ + ABSL_GUARDED_BY(mutex_); +}; + +} // namespace controller +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_CORE_CONTROLLER_WORKER_SERVICE_IMPL_H_ diff --git a/tpu_raiden/core/controller/worker_service_test.cc b/tpu_raiden/core/controller/worker_service_test.cc new file mode 100644 index 0000000..d35d505 --- /dev/null +++ b/tpu_raiden/core/controller/worker_service_test.cc @@ -0,0 +1,138 @@ +// 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 +#include + +#include +#include +#include "absl/status/statusor.h" +#include "third_party/grpc/include/grpcpp/create_channel.h" +#include "third_party/grpc/include/grpcpp/security/credentials.h" +#include "third_party/grpc/include/grpcpp/security/server_credentials.h" +#include "third_party/grpc/include/grpcpp/server.h" +#include "third_party/grpc/include/grpcpp/server_builder.h" +#include "tpu_raiden/core/controller/worker_service_client.h" +#include "tpu_raiden/core/controller/worker_service_impl.h" +#include "tpu_raiden/proto/worker_service.pb.h" +#include "tpu_raiden/rpc/raiden_service.pb.h" + +namespace tpu_raiden { +namespace controller { +namespace { + +using ::testing::HasSubstr; + +class WorkerServiceTest : public ::testing::Test { + protected: + void SetUp() override { + service_ = std::make_unique(); + grpc::ServerBuilder builder; + int selected_port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &selected_port); + builder.RegisterService(service_.get()); + server_ = builder.BuildAndStart(); + + std::string server_address = "localhost:" + std::to_string(selected_port); + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + client_ = std::make_unique(channel); + + unit_.set_job_name("test_job"); + unit_.set_job_replica_id("0"); + unit_.set_data_name("test_data"); + } + + void TearDown() override { + if (server_) { + server_->Shutdown(); + } + } + + rpc::RaidenIdProto unit_; + std::unique_ptr service_; + std::unique_ptr server_; + std::unique_ptr client_; +}; + +TEST_F(WorkerServiceTest, CreateAndDeleteBuffersSuccess) { + proto::CreateBuffersRequest create_req; + *create_req.mutable_unit() = unit_; + auto* spec1 = create_req.add_buffers(); + spec1->set_num_shards(2); + spec1->set_size_bytes(1024); + auto* spec2 = create_req.add_buffers(); + spec2->set_num_shards(2); + spec2->set_size_bytes(1024); + + auto create_resp_or = client_->CreateBuffers(create_req); + ASSERT_TRUE(create_resp_or.ok()); + EXPECT_TRUE(create_resp_or->success()); + ASSERT_EQ(create_resp_or->buffers_size(), 2); + EXPECT_EQ(service_->GetBufferCount(), 4); + + const auto& buf1 = create_resp_or->buffers(0); + const auto& buf2 = create_resp_or->buffers(1); + ASSERT_EQ(buf1.buffer_handles_size(), 2); + ASSERT_EQ(buf2.buffer_handles_size(), 2); + + uint64_t handle1 = buf1.buffer_handles(0).handle(); + uint64_t handle2 = buf1.buffer_handles(1).handle(); + EXPECT_NE(handle1, handle2); + + auto alloc_or = service_->GetBuffer(BufferHandle(handle1)); + ASSERT_TRUE(alloc_or.ok()); + EXPECT_EQ(alloc_or->size, 1024); + EXPECT_NE(alloc_or->ptr, nullptr); + + proto::DeleteBuffersRequest delete_req; + *delete_req.mutable_unit() = unit_; + *delete_req.add_sharded_buffers() = buf1; + *delete_req.add_sharded_buffers() = buf2; + + auto delete_resp_or = client_->DeleteBuffers(delete_req); + ASSERT_TRUE(delete_resp_or.ok()); + EXPECT_TRUE(delete_resp_or->success()); + EXPECT_EQ(service_->GetBufferCount(), 0); +} + +TEST_F(WorkerServiceTest, CreateBuffersWithInvalidSpecFails) { + proto::CreateBuffersRequest create_req; + *create_req.mutable_unit() = unit_; + auto* spec = create_req.add_buffers(); + spec->set_num_shards(0); + spec->set_size_bytes(512); + + auto create_resp_or = client_->CreateBuffers(create_req); + ASSERT_TRUE(create_resp_or.ok()); + EXPECT_FALSE(create_resp_or->success()); + EXPECT_THAT(create_resp_or->message(), HasSubstr("must be positive")); +} + +TEST_F(WorkerServiceTest, DeleteNonExistentBufferFails) { + proto::DeleteBuffersRequest delete_req; + *delete_req.mutable_unit() = unit_; + auto* sharded_buf = delete_req.add_sharded_buffers(); + sharded_buf->add_buffer_handles()->set_handle(9999); + + auto delete_resp_or = client_->DeleteBuffers(delete_req); + ASSERT_TRUE(delete_resp_or.ok()); + EXPECT_FALSE(delete_resp_or->success()); + EXPECT_THAT(delete_resp_or->message(), HasSubstr("not found")); +} + +} // namespace +} // namespace controller +} // namespace tpu_raiden diff --git a/tpu_raiden/proto/BUILD b/tpu_raiden/proto/BUILD new file mode 100644 index 0000000..99b74ad --- /dev/null +++ b/tpu_raiden/proto/BUILD @@ -0,0 +1,53 @@ +# 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_github_grpc_grpc//bazel:python_rules.bzl", "py_grpc_library") +load("@com_google_protobuf//bazel:cc_proto_library.bzl", "cc_proto_library") +load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library") +load("@rules_python//python:proto.bzl", "py_proto_library") + +package(default_visibility = ["//visibility:public"]) + +proto_library( + name = "worker_service_proto", + srcs = ["worker_service.proto"], + deps = ["//tpu_raiden/rpc:raiden_service_proto"], +) + +cc_proto_library( + name = "worker_service_cc_proto", + visibility = ["//visibility:public"], + deps = [":worker_service_proto"], +) + +cc_grpc_library( + name = "worker_service_cc_grpc", + srcs = [":worker_service_proto"], + visibility = ["//visibility:public"], + deps = [":worker_service_cc_proto"], +) + +py_proto_library( + name = "worker_service_py_pb2", + visibility = ["//visibility:public"], + deps = [":worker_service_proto"], +) + +py_grpc_library( + name = "worker_service_py_pb2_grpc", + srcs = [":worker_service_proto"], + visibility = ["//visibility:public"], + deps = [":worker_service_py_pb2"], +) diff --git a/tpu_raiden/proto/worker_service.proto b/tpu_raiden/proto/worker_service.proto new file mode 100644 index 0000000..d8facd6 --- /dev/null +++ b/tpu_raiden/proto/worker_service.proto @@ -0,0 +1,85 @@ +// 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.proto; + +import "third_party/tpu_raiden/tpu_raiden/rpc/raiden_service.proto"; + +// Service providing a gRPC communication channel between Raiden Controller and +// Transfer Workers for buffer allocation and memory management. +// +// Assumption: The WorkerService runs on every KV cache manager instance or +// weight sync instance and is managed exclusively by its local controller. +// There is no concurrent multi-caller access or contention on buffer indices +// across separate controllers. +// TODO(justinlu): Investigate enforcing this contract by restricting RPC access +// so that only the local controller is permitted to invoke these methods. +service WorkerService { + // Allocates a list of sharded buffers on the transfer worker. + rpc CreateBuffers(CreateBuffersRequest) returns (CreateBuffersResponse); + + // Deallocates previously allocated sharded buffers on the transfer worker. + rpc DeleteBuffers(DeleteBuffersRequest) returns (DeleteBuffersResponse); +} + +// Uniquely identifies a buffer handle on a worker. +message BufferHandleProto { + uint64 handle = 1; +} + +// Represents a sharded buffer composed of multiple shard buffer handles. +message BufferProto { + repeated BufferHandleProto buffer_handles = 1; +} + +// Specification for creating/allocating a sharded buffer in host RAM. +message CreateBufferSpec { + // Number of shard partitions for this sharded buffer. + int32 num_shards = 1; + // Size of each buffer shard allocation in bytes. + int64 size_bytes = 2; +} + +// Request to allocate one or more sharded buffers on a transfer worker. +message CreateBuffersRequest { + // Target Raiden work unit (job/replica identifier). + tpu_raiden.rpc.RaidenIdProto unit = 1; + // List of sharded buffer specifications to allocate. + repeated CreateBufferSpec buffers = 2; +} + +// Response from allocating sharded buffers. +message CreateBuffersResponse { + bool success = 1; + string message = 2; + // Created sharded buffers, ordered corresponding to the request specs. + repeated BufferProto buffers = 3; +} + +// Request to deallocate previously allocated sharded buffers on a transfer +// worker. +message DeleteBuffersRequest { + // Target Raiden work unit (job/replica identifier). + tpu_raiden.rpc.RaidenIdProto unit = 1; + // List of sharded buffers to deallocate. + repeated BufferProto sharded_buffers = 2; +} + +// Response from deallocating sharded buffers. +message DeleteBuffersResponse { + bool success = 1; + string message = 2; +}