From ed6388e7deb656ddfa06fa6f8946d611fae1ce90 Mon Sep 17 00:00:00 2001 From: Danna Wang Date: Sun, 23 Aug 2026 00:07:07 -0700 Subject: [PATCH] Add support for persistent storage backend offloading and recalling in TPU Raiden PiperOrigin-RevId: 969268085 --- tpu_sync/api/common.py | 1 + tpu_sync/core/BUILD | 1 + tpu_sync/core/controller/BUILD | 5 + tpu_sync/core/controller/raiden_controller.cc | 70 +++- tpu_sync/core/controller/raiden_controller.h | 40 +- .../core/controller/raiden_controller_test.cc | 48 +++ .../core/controller/worker_service_client.cc | 23 ++ .../core/controller/worker_service_client.h | 4 + .../core/controller/worker_service_impl.cc | 170 +++++++- .../core/controller/worker_service_impl.h | 7 + .../controller/worker_service_server_test.cc | 22 ++ tpu_sync/core/host_memory_allocator.cc | 51 +++ tpu_sync/core/host_memory_allocator.h | 103 ++++- tpu_sync/core/host_memory_allocator_test.cc | 248 ++++++++++++ tpu_sync/core/kv_manager_holder.h | 187 +++++++++ tpu_sync/frameworks/jax/kv_cache_store.pyi | 1 + .../frameworks/jax/tpu_raiden_jax_module.cc | 3 +- .../torch/tpu_raiden_torch_module.cc | 3 +- tpu_sync/kv_cache/BUILD | 2 + tpu_sync/kv_cache/kv_cache_manager_base.cc | 259 +++++++++++- tpu_sync/kv_cache/kv_cache_manager_base.h | 84 +++- tpu_sync/kv_cache/kv_cache_store.cc | 50 ++- tpu_sync/kv_cache/kv_cache_store.h | 2 + tpu_sync/kv_cache/kv_cache_store_backend.h | 4 + tpu_sync/kv_cache/storage/BUILD | 85 ++++ tpu_sync/kv_cache/storage/k5_backend.cc | 159 ++++++++ tpu_sync/kv_cache/storage/k5_backend.h | 94 +++++ tpu_sync/kv_cache/storage/storage.cc | 132 +++++++ tpu_sync/kv_cache/storage/storage.h | 154 ++++++++ tpu_sync/kv_cache/storage/storage_test.cc | 374 ++++++++++++++++++ tpu_sync/proto/worker_service.proto | 77 ++++ 31 files changed, 2438 insertions(+), 25 deletions(-) create mode 100644 tpu_sync/kv_cache/storage/BUILD create mode 100644 tpu_sync/kv_cache/storage/k5_backend.cc create mode 100644 tpu_sync/kv_cache/storage/k5_backend.h create mode 100644 tpu_sync/kv_cache/storage/storage.cc create mode 100644 tpu_sync/kv_cache/storage/storage.h create mode 100644 tpu_sync/kv_cache/storage/storage_test.cc diff --git a/tpu_sync/api/common.py b/tpu_sync/api/common.py index 398899d7..ea8aafb8 100644 --- a/tpu_sync/api/common.py +++ b/tpu_sync/api/common.py @@ -29,3 +29,4 @@ class BlockStatus(enum.Enum): HOST_AND_HBM = ( 4 # Resident in both local Host DRAM and TPU HBM device memory. ) + STORAGE = 5 # Resident in persistent storage backend. diff --git a/tpu_sync/core/BUILD b/tpu_sync/core/BUILD index 60c6b3e0..ca7c3ee1 100644 --- a/tpu_sync/core/BUILD +++ b/tpu_sync/core/BUILD @@ -276,6 +276,7 @@ cc_library( ":raiden_transfer_endpoint", ":raw_transfer_core", ":status_macros", + "//tpu_sync/kv_cache/storage", "//tpu_sync/rpc:raiden_service_cc_proto", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/tpu_sync/core/controller/BUILD b/tpu_sync/core/controller/BUILD index 74c4918a..322e7f2d 100644 --- a/tpu_sync/core/controller/BUILD +++ b/tpu_sync/core/controller/BUILD @@ -39,11 +39,14 @@ cc_library( hdrs = ["worker_service_impl.h"], visibility = ["//visibility:public"], deps = [ + "//third_party/protobuf", "//tpu_sync/core:host_memory_allocator", "//tpu_sync/core:kv_manager_holder", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/core:raw_transfer_core", "//tpu_sync/core:transfer_program_reshard", + "//tpu_sync/kv_cache/storage", + "//tpu_sync/kv_cache/storage:k5_backend", "//tpu_sync/proto:transfer_program_cc_proto", "//tpu_sync/proto:worker_service_cc_grpc", "//tpu_sync/proto:worker_service_cc_proto", @@ -91,6 +94,7 @@ cc_library( "//tpu_sync/core:status_macros", "//tpu_sync/kv_cache:logical_block_manager", "//tpu_sync/kv_cache:raiden_id", + "//tpu_sync/kv_cache/storage", "//tpu_sync/proto:controller_service_cc_grpc", "//tpu_sync/proto:controller_service_cc_proto", "//tpu_sync/proto:worker_service_cc_proto", @@ -245,6 +249,7 @@ cc_test( "//tpu_sync/core:kv_manager_holder", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/kv_cache:raiden_id", + "//tpu_sync/kv_cache/storage", "//tpu_sync/proto:worker_service_cc_proto", "//tpu_sync/rpc:raiden_service_cc_proto", "@com_google_absl//absl/container:flat_hash_map", diff --git a/tpu_sync/core/controller/raiden_controller.cc b/tpu_sync/core/controller/raiden_controller.cc index 032df558..4b2d78b4 100644 --- a/tpu_sync/core/controller/raiden_controller.cc +++ b/tpu_sync/core/controller/raiden_controller.cc @@ -487,7 +487,8 @@ absl::StatusOr<::tpu_sync::proto::TransferBuffersRequest> RaidenController::BuildTransferBuffersRequest( absl::Span src_buffers, absl::Span dst_buffers, absl::Span staging_host_buffers, - absl::Span copy_sizes) { + absl::Span copy_sizes, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec) { if (src_buffers.empty() || src_buffers.size() != dst_buffers.size()) { return absl::InvalidArgumentError( "Source and destination buffers must have the same non-zero length"); @@ -536,6 +537,10 @@ RaidenController::BuildTransferBuffersRequest( added_buf->set_index(buf.index()); } + if (storage_spec.has_value()) { + *transfer->mutable_storage_spec() = *storage_spec; + } + return request; } @@ -543,9 +548,10 @@ tsl::Future<> RaidenController::TransferBuffers( absl::string_view worker_id, absl::Span src_buffers, absl::Span dst_buffers, absl::Span staging_host_buffers, - absl::Span copy_sizes) { + absl::Span copy_sizes, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec) { auto request_or = BuildTransferBuffersRequest( - src_buffers, dst_buffers, staging_host_buffers, copy_sizes); + src_buffers, dst_buffers, staging_host_buffers, copy_sizes, storage_spec); if (!request_or.ok()) { return tsl::Future<>(request_or.status()); } @@ -568,7 +574,8 @@ tsl::Future<> RaidenController::TransferBuffers( tsl::Future<> RaidenController::TransferBuffers( absl::Span src_buffers, absl::Span dst_buffers, absl::Span staging_host_buffers, - absl::Span copy_sizes) { + absl::Span copy_sizes, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec) { if (src_buffers.empty() || src_buffers.size() != dst_buffers.size()) { return tsl::Future<>(absl::InvalidArgumentError( "Source and destination buffers must have the same non-zero length")); @@ -751,10 +758,27 @@ tsl::Future<> RaidenController::TransferBuffers( if (worker_src.empty()) continue; + std::optional<::tpu_sync::proto::StorageTransferSpec> worker_storage_spec; + if (storage_spec.has_value()) { + if (workers.size() > 1) { + ::tpu_sync::proto::StorageTransferSpec rank_spec = *storage_spec; + rank_spec.clear_keys_by_buffer_index(); + auto it = + storage_spec->keys_by_buffer_index().find(static_cast(w)); + if (it != storage_spec->keys_by_buffer_index().end()) { + (*rank_spec.mutable_keys_by_buffer_index())[0] = it->second; + } + worker_storage_spec = std::move(rank_spec); + } else { + worker_storage_spec = storage_spec; + } + } + // Every worker owns a shard of every block, so the (host) staging offsets // are identical across workers. - auto req_or = BuildTransferBuffersRequest( - worker_src, worker_dst, request_staging, worker_copy_sizes); + auto req_or = + BuildTransferBuffersRequest(worker_src, worker_dst, request_staging, + worker_copy_sizes, worker_storage_spec); if (!req_or.ok()) { return tsl::Future<>(req_or.status()); } @@ -783,6 +807,16 @@ tsl::Future<> RaidenController::TransferBuffers( return aggregate_future; } +absl::Status RaidenController::ExecuteTransferBuffersgRPCSync( + absl::Span src_buffers, absl::Span dst_buffers, + absl::Span staging_host_buffers, + absl::Span copy_sizes, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec) { + return TransferBuffers(src_buffers, dst_buffers, staging_host_buffers, + copy_sizes, storage_spec) + .Await(); +} + void RaidenController::SetReadRemoteHooks( core::controller::RaidenControllerServiceImpl::ValidateAndPinCallback validate_and_pin, @@ -1132,5 +1166,29 @@ RaidenController::SubmitTransferProgram( return registration->worker_service_client->SubmitTransferProgram(request); } +absl::Status RaidenController::RegisterBackends( + const std::vector<::tpu_sync::proto::BackendConfig>& configs) { + if (!worker_registry_) { + return absl::FailedPreconditionError("Worker registry is not initialized"); + } + ::tpu_sync::proto::RegisterBackendsRequest request; + for (const auto& config : configs) { + *request.add_configs() = config; + } + for (const auto& reg : worker_registry_->GetRegisteredWorkers()) { + if (!reg.worker_service_client) continue; + auto resp_or = reg.worker_service_client->RegisterBackends(request).Await(); + if (!resp_or.ok()) { + return resp_or.status(); + } + if (!resp_or->success()) { + return absl::InternalError( + absl::StrCat("Worker registration failed on worker ", reg.worker_id, + ": ", resp_or->error_message())); + } + } + return absl::OkStatus(); +} + } // namespace controller } // namespace tpu_raiden diff --git a/tpu_sync/core/controller/raiden_controller.h b/tpu_sync/core/controller/raiden_controller.h index 6ba1f902..0b753c9d 100644 --- a/tpu_sync/core/controller/raiden_controller.h +++ b/tpu_sync/core/controller/raiden_controller.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -36,6 +37,7 @@ #include "tpu_sync/core/controller/worker_service_client.h" #include "tpu_sync/kv_cache/logical_block_manager.h" #include "tpu_sync/kv_cache/raiden_id.h" +#include "tpu_sync/kv_cache/storage/storage.h" #include "tpu_sync/proto/controller_service.grpc.pb.h" #include "tpu_sync/proto/worker_service.pb.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -143,7 +145,9 @@ class RaidenController { absl::string_view worker_id, absl::Span src_buffers, absl::Span dst_buffers, absl::Span staging_host_buffers = {}, - absl::Span copy_sizes = {}); + absl::Span copy_sizes = {}, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec = + std::nullopt); // Broadcast transfer to all registered workers (staging_host_buffers as // above). @@ -151,7 +155,33 @@ class RaidenController { absl::Span src_buffers, absl::Span dst_buffers, absl::Span staging_host_buffers = {}, - absl::Span copy_sizes = {}); + absl::Span copy_sizes = {}, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec = + std::nullopt); + + // Synchronous transfer worker that executes TransferBuffers and waits for + // completion. + absl::Status ExecuteTransferBuffersgRPCSync( + absl::Span src_buffers, + absl::Span dst_buffers, + absl::Span staging_host_buffers = {}, + absl::Span copy_sizes = {}, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec = + std::nullopt); + + // Registers storage backends dynamically on all registered workers. + absl::Status RegisterBackends( + const std::vector<::tpu_sync::proto::BackendConfig>& configs); + + void SetMapper(std::shared_ptr mapper) { + absl::MutexLock lock(&mutex_); + mapper_ = std::move(mapper); + } + + std::shared_ptr mapper() const { + absl::MutexLock lock(&mutex_); + return mapper_; + } // Reads blocks from a remote source, receiver-initiated: this controller's // own workers pull the bytes, so the write window belongs to the destination @@ -254,7 +284,9 @@ class RaidenController { absl::Span src_buffers, absl::Span dst_buffers, absl::Span staging_host_buffers = {}, - absl::Span copy_sizes = {}); + absl::Span copy_sizes = {}, + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec = + std::nullopt); private: RaidenController(const ::tpu_sync::rpc::RaidenIdProto& unit, int num_blocks, @@ -278,6 +310,8 @@ class RaidenController { std::vector<::tpu_sync::proto::BufferProto> all_sharded_buffers_; std::shared_ptr worker_registry_; mutable absl::Mutex mutex_; + std::shared_ptr mapper_ + ABSL_GUARDED_BY(mutex_); std::unique_ptr block_manager_ ABSL_GUARDED_BY(mutex_); std::string raiden_controller_address_; diff --git a/tpu_sync/core/controller/raiden_controller_test.cc b/tpu_sync/core/controller/raiden_controller_test.cc index 7192a0e7..4368dc0d 100644 --- a/tpu_sync/core/controller/raiden_controller_test.cc +++ b/tpu_sync/core/controller/raiden_controller_test.cc @@ -44,6 +44,7 @@ #include "tpu_sync/core/controller/test_util.h" #include "tpu_sync/core/kv_manager_holder.h" #include "tpu_sync/core/raiden_transfer_endpoint.h" +#include "tpu_sync/kv_cache/storage/storage.h" #include "tpu_sync/proto/worker_service.pb.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -1412,6 +1413,53 @@ TEST_F(RaidenControllerTest, TransferBuffersRemoteDramToLocalHbmSuccess) { EXPECT_THAT(mock_mgr.last_dst_offsets, ElementsAre(1)); } +TEST_F(RaidenControllerTest, MapperAndRegisterBackends) { + MockTransferManager mock_mgr; + test_server_->service->SetTransferManager(KVManagerHolder(&mock_mgr)); + + TF_ASSERT_OK_AND_ASSIGN( + auto controller, + RaidenController::Create(unit_, /*num_blocks=*/5, /*num_shards=*/1, + /*shard_size_bytes=*/512, "")); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + EXPECT_EQ(controller->mapper(), nullptr); + auto mapper = std::make_shared( + "/tmp/storage", "test_model", 1, 0); + controller->SetMapper(mapper); + EXPECT_EQ(controller->mapper(), mapper); + + std::vector<::tpu_sync::proto::BackendConfig> configs; + ::tpu_sync::proto::BackendConfig config; + config.set_name("PosixBackend"); + config.set_scheme("posix_test"); + configs.push_back(config); + + auto status = controller->RegisterBackends(configs); + EXPECT_TRUE(status.ok()); +} + +TEST_F(RaidenControllerTest, ExecuteTransferBuffersgRPCSyncSuccess) { + MockTransferManager mock_mgr; + test_server_->service->SetTransferManager(KVManagerHolder(&mock_mgr)); + + TF_ASSERT_OK_AND_ASSIGN( + auto controller, + RaidenController::Create(unit_, /*num_blocks=*/5, /*num_shards=*/1, + /*shard_size_bytes=*/512, "")); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + TF_ASSERT_OK_AND_ASSIGN(auto src_buffers, controller->AllocateBuffers(1)); + src_buffers[0].set_memory_type(::tpu_sync::rpc::MEMORY_TYPE_HBM); + + TF_ASSERT_OK_AND_ASSIGN(auto dst_buffers, controller->AllocateBuffers(1)); + dst_buffers[0].set_memory_type(::tpu_sync::rpc::MEMORY_TYPE_DRAM); + + auto status = + controller->ExecuteTransferBuffersgRPCSync(src_buffers, dst_buffers); + EXPECT_TRUE(status.ok()); +} + } // namespace } // namespace controller } // namespace tpu_raiden diff --git a/tpu_sync/core/controller/worker_service_client.cc b/tpu_sync/core/controller/worker_service_client.cc index 7a1f6997..655df7a8 100644 --- a/tpu_sync/core/controller/worker_service_client.cc +++ b/tpu_sync/core/controller/worker_service_client.cc @@ -122,5 +122,28 @@ tsl::Future<> WorkerServiceClient::TransferBuffers( return future; } +tsl::Future<::tpu_sync::proto::RegisterBackendsResponse> +WorkerServiceClient::RegisterBackends( + const ::tpu_sync::proto::RegisterBackendsRequest& request) { + auto [promise, future] = + tsl::MakePromise<::tpu_sync::proto::RegisterBackendsResponse>(); + auto context = std::make_shared(); + auto response = + std::make_shared<::tpu_sync::proto::RegisterBackendsResponse>(); + + stub_->async()->RegisterBackends( + context.get(), &request, response.get(), + [context, response, + promise = std::move(promise).ToShared()](grpc::Status status) { + if (!status.ok()) { + promise->Set(absl::InternalError(absl::StrCat( + "RegisterBackends RPC failed: ", status.error_message()))); + } else { + promise->Set(std::move(*response)); + } + }); + return future; +} + } // namespace controller } // namespace tpu_raiden diff --git a/tpu_sync/core/controller/worker_service_client.h b/tpu_sync/core/controller/worker_service_client.h index 05bb5934..9f69c1a8 100644 --- a/tpu_sync/core/controller/worker_service_client.h +++ b/tpu_sync/core/controller/worker_service_client.h @@ -47,6 +47,10 @@ class WorkerServiceClient { tsl::Future<> TransferBuffers( const ::tpu_sync::proto::TransferBuffersRequest& request); + // Registers storage backends on the remote transfer worker asynchronously. + tsl::Future<::tpu_sync::proto::RegisterBackendsResponse> RegisterBackends( + const ::tpu_sync::proto::RegisterBackendsRequest& request); + // Submits a transfer program and resolves with the full response. The // reshard coordinator needs success + message verbatim for its // abandon-claim contract, so admission verdicts are not collapsed into a diff --git a/tpu_sync/core/controller/worker_service_impl.cc b/tpu_sync/core/controller/worker_service_impl.cc index 06f3de7e..aa5e595e 100644 --- a/tpu_sync/core/controller/worker_service_impl.cc +++ b/tpu_sync/core/controller/worker_service_impl.cc @@ -28,12 +28,14 @@ #include "absl/strings/str_cat.h" #include "absl/synchronization/mutex.h" #include "grpcpp/server_context.h" -#include "grpcpp/support/status.h" +#include "google/protobuf/repeated_ptr_field.h" #include "tpu_sync/core/host_memory_allocator.h" #include "tpu_sync/core/kv_manager_holder.h" #include "tpu_sync/core/raiden_transfer_endpoint.h" #include "tpu_sync/core/raw_transfer_core.h" #include "tpu_sync/core/transfer_program_reshard.h" +#include "tpu_sync/kv_cache/storage/k5_backend.h" +#include "tpu_sync/kv_cache/storage/storage.h" #include "tpu_sync/proto/transfer_program.pb.h" #include "tpu_sync/proto/worker_service.pb.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -43,6 +45,14 @@ namespace controller { namespace { +bool HasStage(const google::protobuf::RepeatedPtrField& pipeline, + absl::string_view stage) { + for (const auto& s : pipeline) { + if (s == stage) return true; + } + return false; +} + bool IsD2H(::tpu_sync::rpc::MemoryType src_mem_type, ::tpu_sync::rpc::MemoryType dst_mem_type) { return src_mem_type == ::tpu_sync::rpc::MEMORY_TYPE_HBM && @@ -138,6 +148,109 @@ grpc::Status WorkerServiceImpl::TransferBuffers( absl::MutexLock lock(mutex_); const auto& transfer = request->transfer(); + if (transfer.has_storage_spec()) { + if (!transfer_manager_) { + response->set_success(false); + response->set_message( + "Transfer manager is not configured on WorkerService"); + return grpc::Status::OK; + } + const auto& storage_spec = transfer.storage_spec(); + std::string scheme = storage_spec.scheme(); + std::shared_ptr driver = + transfer_manager_.GetBackend(scheme); + if (!driver) { + LOG(ERROR) << "[Worker] No storage backend registered for scheme: " + << scheme; + response->set_success(false); + response->set_message( + absl::StrCat("No storage backend registered for scheme: ", scheme)); + return grpc::Status::OK; + } + + // Unpack per-block descriptors from keys_by_buffer_index map + std::vector block_keys; + std::vector src_offsets; + std::vector dst_offsets; + std::vector copy_sizes; + + size_t num_buffers = storage_spec.keys_by_buffer_index().size(); + block_keys.reserve(num_buffers); + src_offsets.reserve(num_buffers); + dst_offsets.reserve(num_buffers); + copy_sizes.reserve(num_buffers); + + for (size_t i = 0; i < num_buffers; ++i) { + auto it = + storage_spec.keys_by_buffer_index().find(static_cast(i)); + if (it == storage_spec.keys_by_buffer_index().end()) { + response->set_success(false); + response->set_message( + absl::StrCat("Missing StorageKeyDescriptor for buffer_index: ", i)); + return grpc::Status::OK; + } + kv_cache::storage::BlockKey key; + key.resolved_key = it->second.storage_key(); + block_keys.push_back(key); + + int64_t src_off = (i < transfer.src_offsets().size()) + ? transfer.src_offsets(i) + : (i < transfer.src_buffers_size() && + transfer.src_buffers(i).has_index()) + ? transfer.src_buffers(i).index() + : 0; + int64_t dst_off = (i < transfer.dst_offsets().size()) + ? transfer.dst_offsets(i) + : (i < transfer.dst_buffers_size() && + transfer.dst_buffers(i).has_index()) + ? transfer.dst_buffers(i).index() + : 0; + int64_t copy_sz = (i < transfer.copy_sizes().size()) + ? transfer.copy_sizes(i) + : transfer_manager_.bytes_per_block(); + + src_offsets.push_back(src_off); + dst_offsets.push_back(dst_off); + copy_sizes.push_back(copy_sz); + } + + absl::Status status; + switch (storage_spec.direction()) { + case ::tpu_sync::proto::TRANSFER_DIR_OFFLOAD: { + auto fut_or = transfer_manager_.D2hWriteToBackend( + driver, block_keys, src_offsets, dst_offsets, copy_sizes); + if (fut_or.ok()) { + status = fut_or.value().Await(); + } else { + status = fut_or.status(); + } + break; + } + + case ::tpu_sync::proto::TRANSFER_DIR_RECALL: { + auto fut_or = transfer_manager_.H2dReadFromBackend( + driver, block_keys, src_offsets, dst_offsets, copy_sizes); + if (fut_or.ok()) { + status = fut_or.value().Await(); + } else { + status = fut_or.status(); + } + break; + } + + default: + status = absl::InvalidArgumentError( + "Unrecognized storage transfer direction"); + break; + } + + response->set_success(status.ok()); + if (!status.ok()) { + response->set_message(std::string(status.message())); + } + return grpc::Status::OK; + } + ::tpu_sync::rpc::MemoryType src_mem_type = transfer.src_mem_type(); ::tpu_sync::rpc::MemoryType dst_mem_type = transfer.dst_mem_type(); @@ -427,5 +540,60 @@ grpc::Status WorkerServiceImpl::AbortTransfer( "AbortTransfer is not implemented"); } +grpc::Status WorkerServiceImpl::RegisterBackends( + grpc::ServerContext* context, + const ::tpu_sync::proto::RegisterBackendsRequest* request, + ::tpu_sync::proto::RegisterBackendsResponse* response) { + absl::MutexLock lock(mutex_); + if (!transfer_manager_) { + response->set_success(false); + response->set_error_message( + "Transfer manager is not configured on WorkerService"); + return grpc::Status::OK; + } + for (const auto& config : request->configs()) { + std::string name = config.name(); + std::string scheme = config.scheme(); + if (scheme.empty()) { + auto scheme_it = config.properties().find("scheme"); + if (scheme_it != config.properties().end()) { + scheme = scheme_it->second; + } + } + if (name == "PosixBackend") { + if (scheme.empty()) { + response->set_success(false); + response->set_error_message("Missing 'scheme' in backend config"); + return grpc::Status::OK; + } + auto backend = std::make_shared(); + transfer_manager_.RegisterBackend(scheme, backend); + LOG(INFO) + << "[Worker] Dynamically registered storage backend for scheme '" + << scheme << "' (name: " << name << ")"; + } else if (name == "K5Backend" || name == "K5BackendMock") { + auto uds_it = config.properties().find("uds_path"); + if (scheme.empty() || uds_it == config.properties().end()) { + response->set_success(false); + response->set_error_message( + "Missing 'scheme' or 'uds_path' in K5Backend config"); + return grpc::Status::OK; + } + std::string uds_path = uds_it->second; + auto backend = + std::make_shared(uds_path); + transfer_manager_.RegisterBackend(scheme, backend); + LOG(INFO) << "[Worker] Dynamically registered K5 backend for scheme '" + << scheme << "' and UDS: " << uds_path; + } else { + response->set_success(false); + response->set_error_message("Unsupported backend name: " + name); + return grpc::Status::OK; + } + } + response->set_success(true); + return grpc::Status::OK; +} + } // namespace controller } // namespace tpu_raiden diff --git a/tpu_sync/core/controller/worker_service_impl.h b/tpu_sync/core/controller/worker_service_impl.h index 100e6c14..608e4442 100644 --- a/tpu_sync/core/controller/worker_service_impl.h +++ b/tpu_sync/core/controller/worker_service_impl.h @@ -72,6 +72,13 @@ class WorkerServiceImpl final const ::tpu_sync::proto::TransferBuffersRequest* request, ::tpu_sync::proto::TransferBuffersResponse* response) override; + // Dynamically configures storage backends (e.g. POSIX disk, K5 thick client) + // on this worker node. + grpc::Status RegisterBackends( + grpc::ServerContext* context, + const ::tpu_sync::proto::RegisterBackendsRequest* request, + ::tpu_sync::proto::RegisterBackendsResponse* response) override; + // Normalizes a pool-reshard transfer program, lowers it to the // byte-identical StartTransferRequest the framed entry would deliver, and // drives the same pool executor operations. Unsupported completion diff --git a/tpu_sync/core/controller/worker_service_server_test.cc b/tpu_sync/core/controller/worker_service_server_test.cc index d68d32ba..1d716bed 100644 --- a/tpu_sync/core/controller/worker_service_server_test.cc +++ b/tpu_sync/core/controller/worker_service_server_test.cc @@ -78,6 +78,28 @@ TEST(WorkerServiceServerTest, StartServerWithInvalidPortFails) { EXPECT_THAT(status, StatusIs(absl::StatusCode::kInvalidArgument)); } +TEST(WorkerServiceServerTest, RegisterBackendsViaClient) { + WorkerServiceServer& server = WorkerServiceServer::GetInstance(); + ABSL_ASSERT_OK(server.StartServer(/*host_allocator=*/nullptr, /*port=*/0)); + int port = server.GetRaidenWorkerPort(); + EXPECT_GT(port, 0); + + std::string server_address = "localhost:" + std::to_string(port); + auto channel = + grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); + WorkerServiceClient client(channel); + + ::tpu_sync::proto::RegisterBackendsRequest req; + auto* config = req.add_configs(); + config->set_name("PosixBackend"); + config->set_scheme("posix_test"); + + auto resp_or = client.RegisterBackends(req).Await(); + ABSL_ASSERT_OK(resp_or); + // Expect false if transfer_manager is null, or true if configured. + // The RPC should succeed without transport errors. +} + } // namespace } // namespace controller } // namespace tpu_raiden diff --git a/tpu_sync/core/host_memory_allocator.cc b/tpu_sync/core/host_memory_allocator.cc index 22ff8d81..a0344e65 100644 --- a/tpu_sync/core/host_memory_allocator.cc +++ b/tpu_sync/core/host_memory_allocator.cc @@ -447,4 +447,55 @@ SharedMemoryHostMemoryAllocator::AllocateDmaMappedForDevice( return alloc_or; } +absl::StatusOr +SharedMemoryHostMemoryAllocator::GetSharedMemoryInfo(const void* ptr) const { + if (mapped_ptr_ == nullptr || mapped_ptr_ == MAP_FAILED) { + return absl::FailedPreconditionError("Shared memory is not mapped"); + } + const uint8_t* p = static_cast(ptr); + const uint8_t* start = static_cast(mapped_ptr_); + const uint8_t* end = start + mapped_size_; + if (p < start || p >= end) { + return absl::InvalidArgumentError( + "Pointer does not belong to this allocator"); + } + SharedMemoryInfo info; + info.shm_key = shm_key_; + info.size = mapped_size_; + info.offset = p - start; + info.base_ptr = mapped_ptr_; + info.fd = shm_fd_; + return info; +} + +HostBufferAllocator CreateHostMemoryAllocator(xla::PjRtClient* client, + size_t num_blocks, + size_t block_size) { + SharedMemoryHeader expected_schema; + expected_schema.magic = 0x52414944454E5348; // "RAIDENSH" + expected_schema.version = 1; + std::strncpy(expected_schema.model_uid, "mock_model_uid", + sizeof(expected_schema.model_uid)); + expected_schema.global_mesh_shape[0] = 1; // Single rank dummy mesh + expected_schema.num_blocks = num_blocks; + expected_schema.block_size = block_size; + expected_schema.total_payload_bytes = num_blocks * block_size; + expected_schema.reference_count = 0; + + const char* shm_key_env = std::getenv("RAIDEN_SHM_KEY"); + std::string shm_key = + (shm_key_env != nullptr) ? shm_key_env : "raiden_shm_key"; + + auto shm_alloc_or = + SharedMemoryHostMemoryAllocator::Create(client, shm_key, expected_schema); + if (!shm_alloc_or.ok()) { + LOG(ERROR) << "Failed to create SharedMemoryHostMemoryAllocator: " + << shm_alloc_or.status().ToString(); + return nullptr; + } + auto shm_alloc = std::shared_ptr( + std::move(shm_alloc_or).value()); + return HostBufferAllocator(std::move(shm_alloc)); +} + } // namespace tpu_raiden diff --git a/tpu_sync/core/host_memory_allocator.h b/tpu_sync/core/host_memory_allocator.h index 59ff508d..fd5ba25d 100644 --- a/tpu_sync/core/host_memory_allocator.h +++ b/tpu_sync/core/host_memory_allocator.h @@ -34,11 +34,78 @@ struct HostBufferAllocation { std::shared_ptr owner; }; -// Returns a HostBufferAllocation of at least the requested size for a given -// device. If device is nullptr, it allocates on the default/current NUMA node. -// On failure, returns a non-OK status. -using HostBufferAllocator = std::function( - size_t, const xla::PjRtDevice*)>; +struct SharedMemoryInfo { + std::string shm_key; + size_t size = 0; + size_t offset = 0; + void* base_ptr = nullptr; + int fd = -1; +}; + +class HostMemoryAllocator; + +// A wrapper class behaving like a callback while keeping allocator references. +class HostBufferAllocator { + public: + HostBufferAllocator() = default; + HostBufferAllocator(std::nullptr_t) {} // NOLINT + + // Construct from std::function + HostBufferAllocator( // NOLINT + std::function< + absl::StatusOr(size_t, const xla::PjRtDevice*)> + alloc_fn) + : alloc_fn_(std::move(alloc_fn)) {} + + // Construct from shared_ptr to HostMemoryAllocator + HostBufferAllocator( + std::shared_ptr allocator); // NOLINT + + // Construct from general lambdas + template < + typename F, + typename = std::enable_if_t< + !std::is_same_v, HostBufferAllocator> && + !std::is_same_v, + std::shared_ptr> && + std::is_invocable_r_v, F, size_t, + const xla::PjRtDevice*>>> + HostBufferAllocator(F&& f) // NOLINT + : alloc_fn_(std::forward(f)) {} + + absl::StatusOr operator()( + size_t size, const xla::PjRtDevice* device) const { + if (!alloc_fn_) { + return absl::FailedPreconditionError("Allocator is not initialized"); + } + return alloc_fn_(size, device); + } + + explicit operator bool() const { return static_cast(alloc_fn_); } + + std::shared_ptr host_memory_allocator() const { + return allocator_; + } + + friend bool operator==(const HostBufferAllocator& a, std::nullptr_t) { + return !a.alloc_fn_; + } + friend bool operator==(std::nullptr_t, const HostBufferAllocator& a) { + return !a.alloc_fn_; + } + friend bool operator!=(const HostBufferAllocator& a, std::nullptr_t) { + return static_cast(a.alloc_fn_); + } + friend bool operator!=(std::nullptr_t, const HostBufferAllocator& a) { + return static_cast(a.alloc_fn_); + } + + private: + std::function(size_t, + const xla::PjRtDevice*)> + alloc_fn_; + std::shared_ptr allocator_; +}; // High-performance host memory allocator that allocates DMA-capable pinned // memory using PJRT APIs or standard fallback allocations. @@ -66,8 +133,26 @@ class HostMemoryAllocator { size_t size_bytes, const xla::PjRtDevice* device) { return AllocateDmaMapped(size_bytes); } + + // Returns shared memory info (fd, offset, size, base pointer) for a given + // host pointer if the memory was allocated via shared memory. + virtual absl::StatusOr GetSharedMemoryInfo( + const void* ptr) const { + return absl::UnimplementedError( + "GetSharedMemoryInfo is not supported for this allocator."); + } }; +inline HostBufferAllocator::HostBufferAllocator( + std::shared_ptr allocator) + : allocator_(allocator) { + if (allocator) { + alloc_fn_ = [allocator](size_t size, const xla::PjRtDevice* device) { + return allocator->AllocateDmaMappedForDevice(size, device); + }; + } +} + class XlaHostMemoryAllocator : public HostMemoryAllocator { public: static absl::StatusOr> Create( @@ -121,6 +206,9 @@ class SharedMemoryHostMemoryAllocator : public HostMemoryAllocator { absl::StatusOr AllocateDmaMappedForDevice( size_t size_bytes, const xla::PjRtDevice* device) override; + absl::StatusOr GetSharedMemoryInfo( + const void* ptr) const override; + private: SharedMemoryHostMemoryAllocator(xla::PjRtClient* client, absl::string_view shm_key, @@ -135,6 +223,11 @@ class SharedMemoryHostMemoryAllocator : public HostMemoryAllocator { bool dma_mapped_ = false; }; +// Factory to create allocators wrapped in HostBufferAllocator +HostBufferAllocator CreateHostMemoryAllocator(xla::PjRtClient* client, + size_t num_blocks, + size_t block_size); + } // namespace tpu_raiden #endif // THIRD_PARTY_TPU_RAIDEN_CORE_HOST_MEMORY_ALLOCATOR_H_ diff --git a/tpu_sync/core/host_memory_allocator_test.cc b/tpu_sync/core/host_memory_allocator_test.cc index 4395a05e..cae6da94 100644 --- a/tpu_sync/core/host_memory_allocator_test.cc +++ b/tpu_sync/core/host_memory_allocator_test.cc @@ -14,11 +14,17 @@ #include "tpu_sync/core/host_memory_allocator.h" +#include #include +#include +#include +#include #include #include +#include #include +#include #include #include "absl/strings/str_format.h" @@ -163,5 +169,247 @@ TEST(HostMemoryAllocatorTest, SharedMemoryColdAndWarmBoot) { shm_unlink(shm_key.c_str()); } +TEST(HostMemoryAllocatorTest, SharedMemoryInfoAndHostBufferAllocator) { + std::string shm_key = "/test_raiden_shm_info_" + std::to_string(getpid()); + shm_unlink(shm_key.c_str()); + + SharedMemoryHeader schema = {}; + schema.magic = 0x52414944454E; + schema.version = 1; + schema.num_blocks = 64; + schema.block_size = 4096; + + auto allocator_or = + SharedMemoryHostMemoryAllocator::Create(nullptr, shm_key, schema); + ASSERT_TRUE(allocator_or.ok()); + std::shared_ptr shm_alloc = + std::move(allocator_or).value(); + + // Test HostBufferAllocator wrapping shared_ptr + HostBufferAllocator wrapper(shm_alloc); + EXPECT_TRUE(static_cast(wrapper)); + EXPECT_EQ(wrapper.host_memory_allocator(), shm_alloc); + + TF_ASSERT_OK_AND_ASSIGN(HostBufferAllocation alloc, wrapper(4096, nullptr)); + EXPECT_NE(alloc.ptr, nullptr); + EXPECT_EQ(alloc.size, 4096); + + // Test GetSharedMemoryInfo + auto info_or = shm_alloc->GetSharedMemoryInfo(alloc.ptr); + ASSERT_TRUE(info_or.ok()); + SharedMemoryInfo info = info_or.value(); + EXPECT_EQ(info.shm_key, shm_key); + EXPECT_GT(info.size, 0); + EXPECT_EQ(info.offset, sizeof(SharedMemoryHeader)); + EXPECT_EQ(info.base_ptr, static_cast(alloc.ptr) - + sizeof(SharedMemoryHeader)); + EXPECT_GE(info.fd, 0); + + // Invalid pointer test + uint8_t dummy_ptr[16]; + EXPECT_FALSE(shm_alloc->GetSharedMemoryInfo(dummy_ptr).ok()); + + shm_unlink(shm_key.c_str()); +} + +TEST(HostMemoryAllocatorTest, CreateHostMemoryAllocatorFactory) { + std::string shm_key = "/test_create_factory_" + std::to_string(getpid()); + shm_unlink(shm_key.c_str()); + setenv("RAIDEN_SHM_KEY", shm_key.c_str(), 1); + + auto host_alloc = CreateHostMemoryAllocator(nullptr, /*num_blocks=*/4, + /*block_size=*/4096); + ASSERT_TRUE(static_cast(host_alloc)); + ASSERT_NE(host_alloc.host_memory_allocator(), nullptr); + + TF_ASSERT_OK_AND_ASSIGN(HostBufferAllocation alloc, + host_alloc(4 * 4096, nullptr)); + EXPECT_NE(alloc.ptr, nullptr); + EXPECT_EQ(alloc.size, 4 * 4096); + + auto info_or = + host_alloc.host_memory_allocator()->GetSharedMemoryInfo(alloc.ptr); + ASSERT_TRUE(info_or.ok()); + EXPECT_EQ(info_or->shm_key, shm_key); + EXPECT_GE(info_or->fd, 0); + EXPECT_EQ(info_or->offset, sizeof(SharedMemoryHeader)); + + unsetenv("RAIDEN_SHM_KEY"); + shm_unlink(shm_key.c_str()); +} + +namespace { +bool SendFdOverUds(int sock, int fd, void* buf, size_t buflen) { + struct msghdr msg = {0}; + struct iovec iov[1]; + iov[0].iov_base = buf; + iov[0].iov_len = buflen; + msg.msg_iov = iov; + msg.msg_iovlen = 1; + + char cmsg_buf[CMSG_SPACE(sizeof(int))]; + msg.msg_control = cmsg_buf; + msg.msg_controllen = sizeof(cmsg_buf); + + struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg); + cmsg->cmsg_level = SOL_SOCKET; + cmsg->cmsg_type = SCM_RIGHTS; + cmsg->cmsg_len = CMSG_LEN(sizeof(int)); + *(reinterpret_cast(CMSG_DATA(cmsg))) = fd; + + return sendmsg(sock, &msg, 0) == static_cast(buflen); +} + +int RecvFdOverUds(int sock, void* buf, size_t buflen) { + struct msghdr msg = {0}; + struct iovec iov[1]; + iov[0].iov_base = buf; + iov[0].iov_len = buflen; + msg.msg_iov = iov; + msg.msg_iovlen = 1; + + char cmsg_buf[CMSG_SPACE(sizeof(int))]; + msg.msg_control = cmsg_buf; + msg.msg_controllen = sizeof(cmsg_buf); + + if (recvmsg(sock, &msg, 0) < 0) return -1; + + struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg); + if (cmsg && cmsg->cmsg_level == SOL_SOCKET && cmsg->cmsg_type == SCM_RIGHTS) { + return *(reinterpret_cast(CMSG_DATA(cmsg))); + } + return -1; +} +} // namespace + +TEST(HostMemoryAllocatorTest, CrossProcessSharedMemoryIpcWithFdTransfer) { + int sv[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0); + + std::string shm_key = "/test_raiden_shm_uds_" + std::to_string(getpid()); + shm_unlink(shm_key.c_str()); + + SharedMemoryHeader schema = {}; + schema.magic = 0x52414944454E; + schema.version = 1; + absl::SNPrintF(schema.model_uid, sizeof(schema.model_uid), "uds_ipc_model"); + schema.num_blocks = 4; + schema.block_size = 4096; + + auto allocator_or = + SharedMemoryHostMemoryAllocator::Create(nullptr, shm_key, schema); + ASSERT_TRUE(allocator_or.ok()); + auto allocator = std::move(allocator_or).value(); + + TF_ASSERT_OK_AND_ASSIGN(HostBufferAllocation alloc, + allocator->Allocate(4096)); + std::memset(alloc.ptr, 0x42, 4096); + + auto info_or = allocator->GetSharedMemoryInfo(alloc.ptr); + ASSERT_TRUE(info_or.ok()); + SharedMemoryInfo info = info_or.value(); + ASSERT_GE(info.fd, 0); + + pid_t pid = fork(); + ASSERT_NE(pid, -1); + + if (pid == 0) { // Child: Receiver daemon + close(sv[0]); + + struct { + size_t offset; + size_t size; + } meta; + + int recved_fd = RecvFdOverUds(sv[1], &meta, sizeof(meta)); + if (recved_fd < 0) _exit(1); + + void* mapped = mmap(nullptr, meta.offset + meta.size, + PROT_READ | PROT_WRITE, MAP_SHARED, recved_fd, 0); + if (mapped == MAP_FAILED) _exit(2); + + uint8_t* payload = static_cast(mapped) + meta.offset; + for (size_t i = 0; i < meta.size; ++i) { + if (payload[i] != 0x42) _exit(3); + } + + // Modify buffer from child + std::memset(payload, 0x77, meta.size); + + char ack = 'K'; + (void)send(sv[1], &ack, 1, 0); + + munmap(mapped, meta.offset + meta.size); + close(recved_fd); + close(sv[1]); + _exit(0); + } else { // Parent: Sender + close(sv[1]); + + struct { + size_t offset; + size_t size; + } meta = {info.offset, alloc.size}; + + ASSERT_TRUE(SendFdOverUds(sv[0], info.fd, &meta, sizeof(meta))); + + char ack = 0; + (void)recv(sv[0], &ack, 1, 0); + EXPECT_EQ(ack, 'K'); + + int status = 0; + waitpid(pid, &status, 0); + EXPECT_EQ(WEXITSTATUS(status), 0); + + // Verify parent sees child's zero-copy modification + for (size_t i = 0; i < 4096; ++i) { + EXPECT_EQ(alloc.ptr[i], 0x77); + } + + close(sv[0]); + shm_unlink(shm_key.c_str()); + } +} + +TEST(HostMemoryAllocatorTest, FileStorageOffloadAndRecall) { + TF_ASSERT_OK_AND_ASSIGN(auto allocator, HostMemoryAllocator::Create(nullptr)); + + const size_t kSize = 1024 * 1024; // 1MB + TF_ASSERT_OK_AND_ASSIGN(HostBufferAllocation alloc, + allocator->Allocate(kSize)); + ASSERT_NE(alloc.ptr, nullptr); + + for (size_t i = 0; i < kSize; ++i) { + alloc.ptr[i] = static_cast(i & 0xFF); + } + + std::string temp_file = + "/tmp/raiden_allocator_test_" + std::to_string(getpid()) + ".bin"; + { + std::ofstream ofs(temp_file, std::ios::binary); + ASSERT_TRUE(ofs.is_open()); + ofs.write(reinterpret_cast(alloc.ptr), kSize); + ASSERT_TRUE(ofs.good()); + } + + // Corrupt original buffer + std::memset(alloc.ptr, 0, kSize); + + // Recall back from disk + { + std::ifstream ifs(temp_file, std::ios::binary); + ASSERT_TRUE(ifs.is_open()); + ifs.read(reinterpret_cast(alloc.ptr), kSize); + ASSERT_TRUE(ifs.good()); + } + + // Verify contents + for (size_t i = 0; i < kSize; ++i) { + ASSERT_EQ(alloc.ptr[i], static_cast(i & 0xFF)); + } + + std::remove(temp_file.c_str()); +} + } // namespace } // namespace tpu_raiden diff --git a/tpu_sync/core/kv_manager_holder.h b/tpu_sync/core/kv_manager_holder.h index 547f815b..22e00e5b 100644 --- a/tpu_sync/core/kv_manager_holder.h +++ b/tpu_sync/core/kv_manager_holder.h @@ -32,6 +32,7 @@ #include "tpu_sync/core/raiden_transfer_endpoint.h" #include "tpu_sync/core/raw_transfer_core.h" #include "tpu_sync/core/status_macros.h" +#include "tpu_sync/kv_cache/storage/storage.h" #include "tpu_sync/rpc/raiden_service.pb.h" namespace tpu_raiden { @@ -194,6 +195,73 @@ template inline constexpr bool has_pool_reshard_register_recv_v = has_pool_reshard_register_recv::value; +template +struct has_d2h_write_to_backend : std::false_type {}; + +template +struct has_d2h_write_to_backend< + T, std::void_t().D2hWriteToBackend( + std::declval>(), + std::declval&>(), + std::declval&>(), + std::declval&>(), + std::declval&>()))>> : std::true_type {}; + +template +inline constexpr bool has_d2h_write_to_backend_v = + has_d2h_write_to_backend::value; + +template +struct has_h2d_read_from_backend : std::false_type {}; + +template +struct has_h2d_read_from_backend< + T, std::void_t().H2dReadFromBackend( + std::declval>(), + std::declval&>(), + std::declval&>(), + std::declval&>(), + std::declval&>()))>> : std::true_type {}; + +template +inline constexpr bool has_h2d_read_from_backend_v = + has_h2d_read_from_backend::value; + +template +struct has_register_backend : std::false_type {}; + +template +struct has_register_backend< + T, std::void_t().RegisterBackend( + std::declval(), + std::declval>()))>> + : std::true_type {}; + +template +inline constexpr bool has_register_backend_v = has_register_backend::value; + +template +struct has_get_backend : std::false_type {}; + +template +struct has_get_backend().GetBackend( + std::declval()))>> + : std::true_type {}; + +template +inline constexpr bool has_get_backend_v = has_get_backend::value; + +template +struct has_bytes_per_block : std::false_type {}; + +template +struct has_bytes_per_block< + T, std::void_t().bytes_per_block())>> + : std::true_type {}; + +template +inline constexpr bool has_bytes_per_block_v = has_bytes_per_block::value; + } // namespace internal // Type-erased wrapper for any KV Cache Manager or Transfer Manager @@ -259,6 +327,24 @@ class KVManagerHolder { virtual absl::Status PoolReshardRegisterRecv( const tpu_sync::rpc::StartTransferRequest& request, absl::Span chip_block_ids) = 0; + virtual absl::StatusOr D2hWriteToBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_device_offsets, + const std::vector& dst_host_offsets, + const std::vector& copy_sizes) = 0; + virtual absl::StatusOr H2dReadFromBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_host_offsets, + const std::vector& dst_device_offsets, + const std::vector& copy_sizes) = 0; + virtual void RegisterBackend( + const std::string& scheme, + std::shared_ptr backend) = 0; + virtual std::shared_ptr GetBackend( + const std::string& scheme) const = 0; + virtual int64_t bytes_per_block() const = 0; }; template @@ -450,6 +536,58 @@ class KVManagerHolder { "transfer manager."); } } + absl::StatusOr D2hWriteToBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_device_offsets, + const std::vector& dst_host_offsets, + const std::vector& copy_sizes) override { + if constexpr (internal::has_d2h_write_to_backend_v) { + return impl_->D2hWriteToBackend(backend, block_keys, src_device_offsets, + dst_host_offsets, copy_sizes); + } else { + return absl::UnimplementedError( + "D2hWriteToBackend is not implemented by the underlying transfer " + "manager."); + } + } + absl::StatusOr H2dReadFromBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_host_offsets, + const std::vector& dst_device_offsets, + const std::vector& copy_sizes) override { + if constexpr (internal::has_h2d_read_from_backend_v) { + return impl_->H2dReadFromBackend(backend, block_keys, src_host_offsets, + dst_device_offsets, copy_sizes); + } else { + return absl::UnimplementedError( + "H2dReadFromBackend is not implemented by the underlying transfer " + "manager."); + } + } + void RegisterBackend( + const std::string& scheme, + std::shared_ptr backend) override { + if constexpr (internal::has_register_backend_v) { + impl_->RegisterBackend(scheme, std::move(backend)); + } + } + std::shared_ptr GetBackend( + const std::string& scheme) const override { + if constexpr (internal::has_get_backend_v) { + return impl_->GetBackend(scheme); + } else { + return nullptr; + } + } + int64_t bytes_per_block() const override { + if constexpr (internal::has_bytes_per_block_v) { + return impl_->bytes_per_block(); + } else { + return 0; + } + } private: absl::StatusOr> SafeCastOffsets( @@ -612,6 +750,55 @@ class KVManagerHolder { return self_->PoolReshardRegisterRecv(request, chip_block_ids); } + absl::StatusOr D2hWriteToBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_device_offsets, + const std::vector& dst_host_offsets, + const std::vector& copy_sizes) const { + if (!self_) { + return absl::InternalError("KVManagerHolder is null"); + } + return self_->D2hWriteToBackend(backend, block_keys, src_device_offsets, + dst_host_offsets, copy_sizes); + } + + absl::StatusOr H2dReadFromBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_host_offsets, + const std::vector& dst_device_offsets, + const std::vector& copy_sizes) const { + if (!self_) { + return absl::InternalError("KVManagerHolder is null"); + } + return self_->H2dReadFromBackend(backend, block_keys, src_host_offsets, + dst_device_offsets, copy_sizes); + } + + void RegisterBackend( + const std::string& scheme, + std::shared_ptr backend) const { + if (self_) { + self_->RegisterBackend(scheme, std::move(backend)); + } + } + + std::shared_ptr GetBackend( + const std::string& scheme) const { + if (!self_) { + return nullptr; + } + return self_->GetBackend(scheme); + } + + int64_t bytes_per_block() const { + if (!self_) { + return 0; + } + return self_->bytes_per_block(); + } + explicit operator bool() const { return self_ != nullptr; } bool operator==(std::nullptr_t) const { return self_ == nullptr; } bool operator!=(std::nullptr_t) const { return self_ != nullptr; } diff --git a/tpu_sync/frameworks/jax/kv_cache_store.pyi b/tpu_sync/frameworks/jax/kv_cache_store.pyi index 1b468237..27218c39 100644 --- a/tpu_sync/frameworks/jax/kv_cache_store.pyi +++ b/tpu_sync/frameworks/jax/kv_cache_store.pyi @@ -10,6 +10,7 @@ class BlockStatus(enum.Enum): # peer leaves no local entry at all -- its landing blocks are freed and no # host copy is kept. HOST_AND_HBM = ... + STORAGE = ... class RaidenId: job_name: str diff --git a/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc b/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc index 5ae1f20e..0448b502 100644 --- a/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc +++ b/tpu_sync/frameworks/jax/tpu_raiden_jax_module.cc @@ -465,7 +465,8 @@ NB_MODULE(_tpu_raiden_jax, m) { .value("REMOTE", tpu_raiden::kv_cache::BlockStatus::REMOTE) .value("HBM", tpu_raiden::kv_cache::BlockStatus::HBM) .value("HOST", tpu_raiden::kv_cache::BlockStatus::HOST) - .value("HOST_AND_HBM", tpu_raiden::kv_cache::BlockStatus::HOST_AND_HBM); + .value("HOST_AND_HBM", tpu_raiden::kv_cache::BlockStatus::HOST_AND_HBM) + .value("STORAGE", tpu_raiden::kv_cache::BlockStatus::STORAGE); nb::class_(m, "RaidenBlockId") .def(nb::init(m, "RaidenBlockId") .def(nb::init p, - raiden::BufferHolders holds) + raiden::BufferHolders holds = {}) : total_chunks(total_chunks), promise(std::move(p)), combined_holds(std::move(holds)) {} @@ -2644,5 +2644,262 @@ void KVCacheManagerBase::UpdateAllocatedOccupancyMetric() const { static_cast(total_host_dram)); } +void KVCacheManagerBase::RegisterBackend( + const std::string& scheme, std::shared_ptr backend) { + absl::MutexLock lock(&storage_backends_mu_); + storage_backends_[scheme] = std::move(backend); +} + +std::shared_ptr KVCacheManagerBase::GetBackend( + const std::string& scheme) const { + absl::MutexLock lock(&storage_backends_mu_); + auto it = storage_backends_.find(scheme); + if (it == storage_backends_.end()) return nullptr; + return it->second; +} + +absl::StatusOr> KVCacheManagerBase::ToHostBlockIds( + const std::vector& offsets) { + std::vector block_ids; + block_ids.reserve(offsets.size()); + size_t block_sz = bytes_per_block(); + for (int64_t offset : offsets) { + if (offset % block_sz != 0) { + return absl::InvalidArgumentError( + "Offset is not aligned to bytes_per_block"); + } + block_ids.push_back(static_cast(offset / block_sz)); + } + return block_ids; +} + +// Save Path (implicit HBM -> DRAM -> Storage) +absl::StatusOr KVCacheManagerBase::D2hWriteToBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_device_offsets, + const std::vector& dst_host_offsets, + const std::vector& copy_sizes) { + TF_ASSIGN_OR_RETURN(std::vector staging_block_ids, + ToHostBlockIds(dst_host_offsets)); + + size_t num_chunks = copy_sizes.size(); + auto [promise, aggregate_future] = xla::MakePromise(); + + // Helper structure to hold the first-stage DMA future and target staging slot + // index + struct ChunkD2h { + raiden::PjRtCopyFuture d2h_fut; + int staging_block_id; + }; + std::vector chunks; + chunks.reserve(num_chunks); + raiden::BufferHolders all_holds; + + // Step 1: Dispatch local PCIe HBM-to-DRAM DMA copies for each chunk + for (size_t i = 0; i < num_chunks; ++i) { + size_t block_sz = bytes_per_block(); + int64_t src_block_idx = src_device_offsets[i] / block_sz; + int64_t dst_block_idx = dst_host_offsets[i] / block_sz; + int64_t size_blocks = copy_sizes[i] / block_sz; + + // Trigger local DMA copy chunk (non-blocking) + TF_ASSIGN_OR_RETURN( + auto chunk_futures, + DispatchD2hChunks({src_block_idx}, {dst_block_idx}, {size_blocks})); + raiden::PjRtCopyFuture d2h_fut = + raiden::JoinPjRtCopyFutures(absl::MakeSpan(chunk_futures)); + // Accumulate reference holds to prevent buffer reclaim + for (const auto& h : d2h_fut.holds) { + all_holds.push_back(h); + } + chunks.push_back({std::move(d2h_fut), staging_block_ids[i]}); + } + + // Aggregate completion state + auto state = std::make_shared( + num_chunks, std::move(promise), std::move(all_holds)); + + // Step 2: Attach completion callbacks. + // As each block's local D2H copy completes, we schedule its storage write + // task to run asynchronously on the background `push_pool_`. + for (size_t i = 0; i < num_chunks; ++i) { + int staging_block_id = chunks[i].staging_block_id; + chunks[i].d2h_fut.OnReady( + [this, backend, state, staging_block_id, + key = block_keys[i]](absl::StatusOr status) { + if (!status.ok()) { + state->SetError(status.status()); + state->MarkChunkComplete(); + return; + } + + // Dispatch synchronous block write to the push thread pool + push_pool_->Schedule([this, backend, state, staging_block_id, key]() { + absl::Status s = + WriteSingleBlockToBackendSync(backend, key, staging_block_id); + if (!s.ok()) state->SetError(s); + state->MarkChunkComplete(); + }); + }); + } + return raiden::PjRtCopyFuture(std::move(aggregate_future), + state->combined_holds); +} + +// Load Path (implicit Storage -> DRAM -> HBM) +absl::StatusOr KVCacheManagerBase::H2dReadFromBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_host_offsets, + const std::vector& dst_device_offsets, + const std::vector& copy_sizes) { + TF_ASSIGN_OR_RETURN(std::vector staging_block_ids, + ToHostBlockIds(src_host_offsets)); + + size_t num_chunks = src_host_offsets.size(); + auto [promise, aggregate_future] = xla::MakePromise(); + auto state = + std::make_shared(num_chunks, std::move(promise)); + + // Step 1: Schedule storage read tasks. + // Since reading files from backends (e.g. Lustre/GCS) involves blocking I/O + // syscalls, we dispatch them to the background `pull_pool_` so they do not + // block the worker thread. + for (size_t i = 0; i < num_chunks; ++i) { + int staging_block_id = staging_block_ids[i]; + size_t block_sz = bytes_per_block(); + int64_t staging_block_idx = src_host_offsets[i] / block_sz; + int64_t device_block_idx = dst_device_offsets[i] / block_sz; + int64_t size_blocks = copy_sizes[i] / block_sz; + storage::BlockKey key = block_keys[i]; + + pull_pool_->Schedule([this, backend, state, key, staging_block_id, + staging_block_idx, device_block_idx, size_blocks]() { + // 1. Synchronously read block from KVBackend storage into DRAM staging + // buffer + absl::Status read_status = + ReadSingleBlockFromBackendSync(backend, key, staging_block_id); + if (!read_status.ok()) { + state->SetError(read_status); + state->MarkChunkComplete(); + return; + } + // 2. Immediately trigger Host DRAM -> TPU HBM DMA copy (second-stage + // pipeline) + auto h2d_fut_or = + H2d({staging_block_idx}, {device_block_idx}, {size_blocks}); + if (!h2d_fut_or.ok()) { + state->SetError(h2d_fut_or.status()); + state->MarkChunkComplete(); + return; + } + h2d_fut_or->OnReady( + [state](absl::StatusOr status) { + if (!status.ok()) state->SetError(status.status()); + state->MarkChunkComplete(); + }); + }); + } + return raiden::PjRtCopyFuture(std::move(aggregate_future), /*holds=*/{}); +} + +// Explicit Write-back +absl::Status KVCacheManagerBase::WriteToBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& host_offsets) { + TF_ASSIGN_OR_RETURN(std::vector staging_block_ids, + ToHostBlockIds(host_offsets)); + + size_t num_blocks = host_offsets.size(); + std::atomic completed(0); + absl::Notification all_done; + std::vector statuses(num_blocks); + + // Dispatch all block writes to the push thread pool concurrently + for (size_t i = 0; i < num_blocks; ++i) { + int staging_block_id = staging_block_ids[i]; + push_pool_->Schedule([this, backend, key = block_keys[i], staging_block_id, + i, &completed, &statuses, &all_done, num_blocks]() { + statuses[i] = + WriteSingleBlockToBackendSync(backend, key, staging_block_id); + if (++completed == num_blocks) all_done.Notify(); + }); + } + // Block calling thread until all writes are completed + all_done.WaitForNotification(); + for (const auto& s : statuses) { + if (!s.ok()) return s; + } + return absl::OkStatus(); +} + +// Helper: Synchronously offloads a single block from local DRAM staging buffer +// to the storage backend. Blocks the calling background thread until the +// backend's async callback resolves. +absl::Status KVCacheManagerBase::WriteSingleBlockToBackendSync( + std::shared_ptr backend, const storage::BlockKey& key, + int staging_block_id) { + absl::Notification done; + absl::Status status; + // Resolve host DRAM raw memory pointer for the staging slot (layer=0, shard=0 + // in mock) + uint8_t* host_ptr = + GetBlockHostPointer(/*layer_idx=*/0, /*shard_idx=*/0, staging_block_id); + size_t size = bytes_per_block(); + + storage::StorageBufferDescriptor buffer; + buffer.ptr = host_ptr; + if (host_allocator_ && host_allocator_.host_memory_allocator() != nullptr) { + auto info_or = + host_allocator_.host_memory_allocator()->GetSharedMemoryInfo(host_ptr); + if (info_or.ok()) { + buffer.fd = info_or.value().fd; + buffer.offset = info_or.value().offset; + } + } + + backend->WriteAsync(key, buffer, size, [&](const absl::Status& s) { + status = s; + done.Notify(); + }); + done.WaitForNotification(); + return status; +} + +// Helper: Synchronously reads a single block from the storage backend into a +// local DRAM staging buffer. Blocks the calling background thread until the +// backend's async callback resolves. +absl::Status KVCacheManagerBase::ReadSingleBlockFromBackendSync( + std::shared_ptr backend, const storage::BlockKey& key, + int staging_block_id) { + absl::Notification done; + absl::Status status; + // Resolve host DRAM raw memory pointer for the staging slot (layer=0, shard=0 + // in mock) + uint8_t* host_ptr = + GetBlockHostPointer(/*layer_idx=*/0, /*shard_idx=*/0, staging_block_id); + size_t size = bytes_per_block(); + + storage::StorageBufferDescriptor buffer; + buffer.ptr = host_ptr; + if (host_allocator_ && host_allocator_.host_memory_allocator() != nullptr) { + auto info_or = + host_allocator_.host_memory_allocator()->GetSharedMemoryInfo(host_ptr); + if (info_or.ok()) { + buffer.fd = info_or.value().fd; + buffer.offset = info_or.value().offset; + } + } + + backend->ReadAsync(key, buffer, size, [&](const absl::Status& s) { + status = s; + done.Notify(); + }); + done.WaitForNotification(); + return status; +} + } // namespace kv_cache } // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.h b/tpu_sync/kv_cache/kv_cache_manager_base.h index b29d4f60..1117c0c2 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.h +++ b/tpu_sync/kv_cache/kv_cache_manager_base.h @@ -46,6 +46,7 @@ #include "tpu_sync/core/raw_transfer_core.h" #include "tpu_sync/kv_cache/logical_block_manager.h" #include "tpu_sync/kv_cache/pool_layout.h" +#include "tpu_sync/kv_cache/storage/storage.h" #include "tpu_sync/rpc/raiden_service.pb.h" #include "tpu_sync/transport/block_transport.h" #include "tpu_sync/transport/block_transport_delegate.h" @@ -253,6 +254,63 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { // Blocks until all pending asynchronous transfers/copies are complete. virtual absl::Status WaitForPendingWork() { return absl::OkStatus(); } + // Registers a storage driver plugin (e.g. Lustre POSIX mount or K5 driver). + void RegisterBackend(const std::string& scheme, + std::shared_ptr backend); + + // Returns the storage driver registered for the specified URI scheme, or + // nullptr if none. + std::shared_ptr GetBackend( + const std::string& scheme) const; + + // Save Path (implicit HBM -> Host DRAM staging -> Storage write) + // + // Coordinates the two-stage pipeline to offload blocks: + // 1. Launches async device-to-host (D2H) copy from TPU HBM into local Host + // DRAM staging slots. + // 2. Registers a callback via `.OnReady` on the D2H future. + // 3. When the DMA finishes for a block, schedules a task on `push_pool_` + // thread pool + // to synchronously write the staging DRAM buffer contents to the target + // `KVBackend`. + // + absl::StatusOr D2hWriteToBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_device_offsets, + const std::vector& dst_host_offsets, + const std::vector& copy_sizes); + + // Load Path (implicit Storage read -> Host DRAM staging -> HBM copy) + // + // Coordinates the two-stage pipeline to recall blocks: + // 1. For each block, schedules a task on `pull_pool_` thread pool to read + // the block + // data from `KVBackend` into the host DRAM staging buffer (blocking FS + // call). + // 2. Once the disk read completes on the pull thread, immediately + // triggers the Host-to-Device (H2D) + // DMA copy using block index arrays to promote the block to HBM. + // 3. Chain block callbacks to notify the aggregate future when H2D copies + // finish. + // + // Returns an aggregate PjRtCopyFuture that signals complete when the + // final stage of H2D is done for all blocks. + absl::StatusOr H2dReadFromBackend( + std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& src_host_offsets, + const std::vector& dst_device_offsets, + const std::vector& copy_sizes); + + // Explicit DRAM -> Storage write + // Synchronously offloads a batch of blocks from local DRAM staging + // buffers to the storage backend. Blocks the calling thread until all + // disk write tasks complete. + absl::Status WriteToBackend(std::shared_ptr backend, + const std::vector& block_keys, + const std::vector& host_offsets); + virtual absl::StatusOr H2hReadExplicit( std::string peer, const std::vector& src_block_ids, const std::vector& local_block_ids, @@ -285,12 +343,11 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { const std::vector& dst_offsets = {}, const std::vector& copy_sizes = {}, int64_t device_id = -1); - // Layer-wise copies using caller-owned host buffers. The raw copy operation - // captures the supplied address when issued, so independent calls can overlap - // across layers. - // NOTE: These functions are temporary. Long-term, KVCacheManager should own - // these host buffers to enable serving prefix cache lookups directly from - // RAM. + // Layer-wise copies using caller-owned host buffers. The raw copy + // operation captures the supplied address when issued, so independent + // calls can overlap across layers. NOTE: These functions are temporary. + // Long-term, KVCacheManager should own these host buffers to enable + // serving prefix cache lookups directly from RAM. absl::Status ConfigureHostStagingSlots(int64_t num_slots, int64_t max_major_per_slot); @@ -596,6 +653,21 @@ class KVCacheManagerBase : public tpu_raiden::RaidenManagerBase { // Initializes background worker thread if RAIDEN_ENABLE_ASYNC_DISPATCH is // enabled. void InitBackgroundWorker(); + + // Synchronous wrappers that wait on absl::Notification. + // Used as helper tasks scheduled on `push_pool_` or `pull_pool_`. + absl::Status WriteSingleBlockToBackendSync( + std::shared_ptr backend, const storage::BlockKey& key, + int staging_block_id); + absl::Status ReadSingleBlockFromBackendSync( + std::shared_ptr backend, const storage::BlockKey& key, + int staging_block_id); + absl::StatusOr> ToHostBlockIds( + const std::vector& offsets); + + mutable absl::Mutex storage_backends_mu_; + absl::flat_hash_map> + storage_backends_ ABSL_GUARDED_BY(storage_backends_mu_); }; } // namespace kv_cache diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index 909edafa..bf9bc5be 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -1041,9 +1041,43 @@ absl::Status KVCacheStore::SaveLocal( ::tpu_sync::rpc::MEMORY_TYPE_DRAM); } + std::string storage_scheme; + for (const auto& b : backends_) { + if (b && !b->scheme().empty()) { + storage_scheme = b->scheme(); + break; + } + } + + std::optional<::tpu_sync::proto::StorageTransferSpec> storage_spec; + if (!storage_scheme.empty()) { + ::tpu_sync::proto::StorageTransferSpec spec; + spec.set_direction(::tpu_sync::proto::TRANSFER_DIR_OFFLOAD); + spec.set_scheme(storage_scheme); + + auto ctrl_mapper = (raiden_controller_ != nullptr) + ? raiden_controller_->mapper() + : nullptr; + int tp_size = (ctrl_mapper != nullptr) ? ctrl_mapper->tp_size() : 1; + + for (size_t i = 0; i < block_hashes.size(); ++i) { + for (int r = 0; r < tp_size; ++r) { + storage::BlockKey resolved_key = + (ctrl_mapper != nullptr) + ? ctrl_mapper->MapKey(block_hashes[i], r) + : storage::BlockKey{block_hashes[i], block_hashes[i]}; + ::tpu_sync::proto::StorageKeyDescriptor key_desc; + key_desc.set_storage_key(resolved_key.resolved_key); + int32_t buf_idx = static_cast(i * tp_size + r); + (*spec.mutable_keys_by_buffer_index())[buf_idx] = key_desc; + } + } + storage_spec = std::move(spec); + } + tsl::Future<> future = raiden_controller_->TransferBuffers( src_buffers, dst_buffers, /*staging_host_buffers=*/{}, - /*copy_sizes=*/{}); + /*copy_sizes=*/{}, storage_spec); { absl::MutexLock lock(mutex_); @@ -1051,6 +1085,7 @@ absl::Status KVCacheStore::SaveLocal( .future = std::move(future), .block_hashes = block_hashes, .host_block_ids = host_block_ids, + .target_scheme = storage_scheme, }); } @@ -2139,6 +2174,19 @@ void KVCacheStore::PollSavesInternal(std::vector ready_saves) { .raiden_id = raiden_id_, .block_id = state.host_block_ids[i], }); + if (!state.target_scheme.empty()) { + RaidenId storage_id{ + .job_name = raiden_id_.job_name, + .job_replica_id = raiden_id_.job_replica_id, + .data_name = state.target_scheme, + .data_replica_idx = raiden_id_.data_replica_idx, + }; + write_through_regs.push_back({ + .prefix_hash = hash, + .raiden_id = storage_id, + .block_id = state.host_block_ids[i], + }); + } } } done_saves_.push_back(hash); diff --git a/tpu_sync/kv_cache/kv_cache_store.h b/tpu_sync/kv_cache/kv_cache_store.h index d407e473..652ee016 100644 --- a/tpu_sync/kv_cache/kv_cache_store.h +++ b/tpu_sync/kv_cache/kv_cache_store.h @@ -45,6 +45,7 @@ #include "tpu_sync/kv_cache/lru_cache.h" #include "tpu_sync/kv_cache/raiden_id.h" #include "tpu_sync/kv_cache/reshard/reshard_service.h" +#include "tpu_sync/kv_cache/storage/storage.h" namespace tpu_raiden { @@ -563,6 +564,7 @@ class KVCacheStore { tsl::Future<> future; std::vector block_hashes; std::vector host_block_ids; + std::string target_scheme; }; struct LoadState { diff --git a/tpu_sync/kv_cache/kv_cache_store_backend.h b/tpu_sync/kv_cache/kv_cache_store_backend.h index 0b552a40..f21380e3 100644 --- a/tpu_sync/kv_cache/kv_cache_store_backend.h +++ b/tpu_sync/kv_cache/kv_cache_store_backend.h @@ -42,6 +42,7 @@ enum class BlockStatus { HBM, HOST, HOST_AND_HBM, + STORAGE, }; struct RaidenBlockId { @@ -110,6 +111,9 @@ class KVCacheStoreBackend { // "GlobalMemoryPoolingBackend"). virtual std::string name() const = 0; + // URI scheme identifying the backend (e.g., "lustre", "k5", "posix"). + virtual std::string scheme() const { return ""; } + // Resolves cached block hashes in sequence. // Returns a list of matched (block_hash, RaidenBlockId) pairs up to the first // miss. diff --git a/tpu_sync/kv_cache/storage/BUILD b/tpu_sync/kv_cache/storage/BUILD new file mode 100644 index 00000000..0b18b65a --- /dev/null +++ b/tpu_sync/kv_cache/storage/BUILD @@ -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. + +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "storage", + srcs = ["storage.cc"], + hdrs = ["storage.h"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = [ + "-use_header_modules", + "-layering_check", + ], + deps = [ + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + ], +) + +cc_library( + name = "k5_backend", + srcs = ["k5_backend.cc"], + hdrs = ["k5_backend.h"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = [ + "-use_header_modules", + "-layering_check", + ], + deps = [ + ":storage", + "//tpu_sync/kv_cache:kv_cache_store_backend", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "storage_test", + srcs = ["storage_test.cc"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = [ + "-use_header_modules", + "-layering_check", + ], + tags = [ + "no_oss", + ], + deps = [ + ":k5_backend", + ":storage", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/synchronization", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_sync/kv_cache/storage/k5_backend.cc b/tpu_sync/kv_cache/storage/k5_backend.cc new file mode 100644 index 00000000..769f0af8 --- /dev/null +++ b/tpu_sync/kv_cache/storage/k5_backend.cc @@ -0,0 +1,159 @@ +// 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/kv_cache/storage/k5_backend.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/kv_cache/storage/storage.h" + +namespace tpu_raiden { +namespace kv_cache { +namespace storage { + +K5BackendMock::K5BackendMock(std::string uds_socket_path) + : uds_socket_path_(std::move(uds_socket_path)) {} + +absl::Status K5BackendMock::SendFdAndMetadata(int fd, size_t offset, + size_t size, + const std::string& resolved_key, + bool is_write) { + int sock = socket(AF_UNIX, SOCK_STREAM, 0); + if (sock < 0) { + return absl::InternalError( + absl::StrCat("socket failed: ", std::strerror(errno))); + } + + struct sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, uds_socket_path_.c_str(), + sizeof(addr.sun_path) - 1); + + if (connect(sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) { + close(sock); + return absl::InternalError( + absl::StrCat("connect failed on UDS: ", std::strerror(errno))); + } + + std::string metadata = + absl::StrCat("op=", is_write ? "write" : "read", " offset=", offset, + " size=", size, " key=", resolved_key); + + struct iovec iov[1]; + iov[0].iov_base = const_cast(metadata.data()); + iov[0].iov_len = metadata.size(); + + union { + struct cmsghdr cm; + char control[CMSG_SPACE(sizeof(int))]; + } control_un; + + struct msghdr msg; + std::memset(&msg, 0, sizeof(msg)); + msg.msg_iov = iov; + msg.msg_iovlen = 1; + msg.msg_control = control_un.control; + msg.msg_controllen = sizeof(control_un.control); + + struct cmsghdr* cmptr = CMSG_FIRSTHDR(&msg); + cmptr->cmsg_len = CMSG_LEN(sizeof(int)); + cmptr->cmsg_level = SOL_SOCKET; + cmptr->cmsg_type = SCM_RIGHTS; + *(reinterpret_cast(CMSG_DATA(cmptr))) = fd; + + ssize_t sent = sendmsg(sock, &msg, 0); + if (sent < 0) { + close(sock); + return absl::InternalError( + absl::StrCat("sendmsg failed: ", std::strerror(errno))); + } + + char reply[16]; + std::memset(reply, 0, sizeof(reply)); + ssize_t recved = recv(sock, reply, sizeof(reply) - 1, 0); + close(sock); + + if (recved <= 0 || absl::string_view(reply) != "OK") { + return absl::InternalError( + absl::StrCat("daemon verification failed: ", reply)); + } + + return absl::OkStatus(); +} + +void K5BackendMock::WriteAsync( + const BlockKey& key, StorageBufferDescriptor src_buffer, size_t size, + std::function callback) { + absl::Status status = + SendFdAndMetadata(src_buffer.fd, src_buffer.offset, size, + key.resolved_key, /*is_write=*/true); + callback(status); +} + +void K5BackendMock::ReadAsync( + const BlockKey& key, StorageBufferDescriptor dst_buffer, size_t size, + std::function callback) { + absl::Status status = + SendFdAndMetadata(dst_buffer.fd, dst_buffer.offset, size, + key.resolved_key, /*is_write=*/false); + callback(status); +} + +absl::StatusOr K5BackendMock::Exists(const BlockKey& key) { + // Check if file exists inside scratch directory by matching hash path + struct stat st; + if (stat(key.resolved_key.c_str(), &st) == 0) { + return true; + } + return false; +} + +// --- K5BlockNameMapper Implementation --- + +K5BlockNameMapper::K5BlockNameMapper(absl::string_view root_dir, + absl::string_view model_name, int tp_size, + int rank) + : root_dir_(root_dir), + model_name_(model_name), + tp_size_(tp_size), + rank_(rank) {} + +BlockKey K5BlockNameMapper::MapKey(const std::string& block_hash, + int rank) const { + int target_rank = (rank == -1) ? rank_ : rank; + std::string k5_block_name = absl::StrCat("k5_", model_name_, "_tp", tp_size_, + "_r", target_rank, "_", block_hash); + std::string resolved_path = + absl::StrCat(root_dir_, "/", k5_block_name, ".bin"); + return BlockKey{block_hash, resolved_path}; +} + +} // namespace storage +} // namespace kv_cache +} // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/storage/k5_backend.h b/tpu_sync/kv_cache/storage/k5_backend.h new file mode 100644 index 00000000..f879ec15 --- /dev/null +++ b/tpu_sync/kv_cache/storage/k5_backend.h @@ -0,0 +1,94 @@ +// 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_KV_CACHE_STORAGE_K5_BACKEND_H_ +#define THIRD_PARTY_TPU_RAIDEN_KV_CACHE_STORAGE_K5_BACKEND_H_ + +#include +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "tpu_sync/kv_cache/kv_cache_store_backend.h" +#include "tpu_sync/kv_cache/storage/storage.h" + +namespace tpu_raiden { +namespace kv_cache { +namespace storage { + +class K5BackendMock : public KVBackend { + public: + explicit K5BackendMock(std::string uds_socket_path); + ~K5BackendMock() override = default; + + std::string scheme() const override { return "k5"; } + + void ReadAsync(const BlockKey& key, StorageBufferDescriptor dst_buffer, + size_t size, + std::function callback) override; + + void WriteAsync(const BlockKey& key, StorageBufferDescriptor src_buffer, + size_t size, + std::function callback) override; + + absl::StatusOr Exists(const BlockKey& key) override; + + private: + absl::Status SendFdAndMetadata(int fd, size_t offset, size_t size, + const std::string& resolved_key, + bool is_write); + + std::string uds_socket_path_; +}; + +using K5Backend = K5BackendMock; + +// K5BlockNameMapper implements the K5 layout mapping policy. +// In production K5 uses logical block/chunk names (keys) instead of files. +// For this mock, we map the logical block name to a file path under the hood +// to simulate persistence in the local sandbox. +class K5BlockNameMapper : public BlockKeyMapper { + public: + K5BlockNameMapper(absl::string_view root_dir, absl::string_view model_name, + int tp_size, int rank); + + BlockKey MapKey(const std::string& block_hash, int rank) const override; + + private: + std::string root_dir_; + std::string model_name_; + int tp_size_; + int rank_; +}; + +} // namespace storage + +using storage::K5Backend; +using storage::K5BackendMock; +using storage::K5BlockNameMapper; + +} // namespace kv_cache + +using kv_cache::storage::K5Backend; +using kv_cache::storage::K5BackendMock; +using kv_cache::storage::K5BlockNameMapper; + +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_KV_CACHE_STORAGE_K5_BACKEND_H_ diff --git a/tpu_sync/kv_cache/storage/storage.cc b/tpu_sync/kv_cache/storage/storage.cc new file mode 100644 index 00000000..bb4b2f94 --- /dev/null +++ b/tpu_sync/kv_cache/storage/storage.cc @@ -0,0 +1,132 @@ +// 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/kv_cache/storage/storage.h" + +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" + +namespace tpu_raiden { +namespace kv_cache { +namespace storage { + +namespace fs = std::filesystem; + +// --- PosixBackend Implementation --- + +// Asynchronously reads block data from the local file system. +// Fires the callback immediately after completing the synchronous POSIX read. +void PosixBackend::ReadAsync( + const BlockKey& key, StorageBufferDescriptor dst_buffer, size_t size, + std::function callback) { + std::ifstream file(key.resolved_key, std::ios::binary); + if (!file) { + callback(absl::NotFoundError( + absl::StrCat("Failed to open file for reading: ", key.resolved_key))); + return; + } + + // Read exact byte size into the host DRAM staging buffer + file.read(reinterpret_cast(dst_buffer.ptr), size); + if (!file) { + callback(absl::DataLossError( + absl::StrCat("Failed to read expected ", size, + " bytes from file: ", key.resolved_key))); + return; + } + callback(absl::OkStatus()); +} + +// Asynchronously writes block data to the local file system. +// Creates parent directories if missing (simulating directory creation on +// Lustre). +void PosixBackend::WriteAsync( + const BlockKey& key, StorageBufferDescriptor src_buffer, size_t size, + std::function callback) { + std::error_code ec; + fs::path filepath(key.resolved_key); + + // Ensure the target directory structure exists + fs::create_directories(filepath.parent_path(), ec); + if (ec) { + callback(absl::InternalError( + absl::StrCat("Failed to create directories for: ", key.resolved_key, + ", error: ", ec.message()))); + return; + } + + std::ofstream file(key.resolved_key, std::ios::binary); + if (!file) { + callback(absl::InternalError( + absl::StrCat("Failed to open file for writing: ", key.resolved_key))); + return; + } + + // Write the byte buffer content to disk + file.write(reinterpret_cast(src_buffer.ptr), size); + if (!file) { + callback(absl::DataLossError(absl::StrCat( + "Failed to write ", size, " bytes to file: ", key.resolved_key))); + return; + } + callback(absl::OkStatus()); +} + +// Verifies if the file exists on the local filesystem. +absl::StatusOr PosixBackend::Exists(const BlockKey& key) { + std::error_code ec; + bool exists = fs::exists(key.resolved_key, ec); + if (ec) { + return absl::InternalError(absl::StrCat( + "fs::exists failed for: ", key.resolved_key, ", msg: ", ec.message())); + } + return exists; +} + +// --- PosixPathMapper Implementation --- + +PosixPathMapper::PosixPathMapper(absl::string_view root_dir, + absl::string_view model_name, int tp_size, + int rank) + : root_dir_(root_dir), + model_name_(model_name), + tp_size_(tp_size), + rank_(rank) {} + +// Resolves a global content hash to a structured local path. +// This matches the layout schema: +// /_tp_r/.bin +BlockKey PosixPathMapper::MapKey(const std::string& block_hash, + int rank) const { + // If rank is omitted (-1), resolve the path using the current worker's rank. + int target_rank = (rank == -1) ? rank_ : rank; + std::string resolved_path = + absl::StrCat(root_dir_, "/", model_name_, "_tp", tp_size_, "_r", + target_rank, "/", block_hash, ".bin"); + return BlockKey{block_hash, resolved_path}; +} + +} // namespace storage +} // namespace kv_cache +} // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/storage/storage.h b/tpu_sync/kv_cache/storage/storage.h new file mode 100644 index 00000000..2670f633 --- /dev/null +++ b/tpu_sync/kv_cache/storage/storage.h @@ -0,0 +1,154 @@ +// 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_KV_CACHE_STORAGE_STORAGE_H_ +#define THIRD_PARTY_TPU_RAIDEN_KV_CACHE_STORAGE_STORAGE_H_ + +#include +#include +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" + +namespace tpu_raiden { +namespace kv_cache { +namespace storage { + +// BlockKey represents a resolved key on persistent storage. +// It bundles: +// - `block_hash`: The content-addressable hash identifying the cache block +// data. +// - `resolved_key`: The target absolute storage path (e.g. Lustre mount +// filepath or GCS URI). +struct BlockKey { + std::string block_hash; + std::string resolved_key; // Absolute resolved file path +}; + +// StorageBufferDescriptor aggregates virtual memory pointer and physical IPC +// segment descriptors (e.g. shared memory file descriptor and offset) for I/O +// operations. +struct StorageBufferDescriptor { + uint8_t* ptr = nullptr; // Raw CPU virtual memory address + int fd = -1; // Underlying file descriptor (e.g. for IPC/mmap) + size_t offset = 0; // Byte offset within the FD segment +}; + +// Stateless backend I/O driver interface. +// Production drivers (like LustrePOSIXBackend or GcsObjectBackend) will +// implement this interface. +class KVBackend : public std::enable_shared_from_this { + public: + virtual ~KVBackend() = default; + + // Returns the storage scheme string (e.g. "local_disk", "lustre", "k5"). + virtual std::string scheme() const = 0; + + // Reads `size` bytes asynchronously from the backend key into `dst_buffer`. + // Invokes `callback` upon completion. + virtual void ReadAsync(const BlockKey& key, + StorageBufferDescriptor dst_buffer, size_t size, + std::function callback) = 0; + + // Writes `size` bytes asynchronously from `src_buffer` to the backend key. + // Invokes `callback` upon completion. + virtual void WriteAsync( + const BlockKey& key, StorageBufferDescriptor src_buffer, size_t size, + std::function callback) = 0; + + // Checks whether the block exists in persistent storage. + virtual absl::StatusOr Exists(const BlockKey& key) = 0; +}; + +// PosixBackend implements KVBackend using standard C++ filesystem APIs. +// Simulates a POSIX mount directory on the worker nodes. +class PosixBackend : public KVBackend { + public: + explicit PosixBackend(std::string scheme_name = "local_disk") + : scheme_(std::move(scheme_name)) {} + + std::string scheme() const override { return scheme_; } + + void ReadAsync(const BlockKey& key, StorageBufferDescriptor dst_buffer, + size_t size, + std::function callback) override; + + void WriteAsync(const BlockKey& key, StorageBufferDescriptor src_buffer, + size_t size, + std::function callback) override; + + absl::StatusOr Exists(const BlockKey& key) override; + + private: + std::string scheme_ = "local_disk"; +}; + +// BlockKeyMapper defines the coordinator-side path resolution policy. +// It maps block hash identifiers and rank distributions to absolute storage +// locations. +class BlockKeyMapper { + public: + virtual ~BlockKeyMapper() = default; + + // Maps the block hash and worker rank to a resolved storage path key. + virtual BlockKey MapKey(const std::string& block_hash, int rank) const = 0; + virtual int tp_size() const { return 1; } +}; + +// PosixPathMapper implements the V4 layout mapping policy. +// Resolves files into flat directories structured by model name, tensor +// parallel size, and rank: Layout: +// `/_tp_r/.bin` +class PosixPathMapper : public BlockKeyMapper { + public: + PosixPathMapper(absl::string_view root_dir, absl::string_view model_name, + int tp_size, int rank); + + BlockKey MapKey(const std::string& block_hash, int rank) const override; + int tp_size() const override { return tp_size_; } + + private: + std::string root_dir_; + std::string model_name_; + int tp_size_; + int rank_; +}; + +} // namespace storage + +// Namespace aliases for convenience and backward compatibility +using storage::BlockKey; +using storage::BlockKeyMapper; +using storage::KVBackend; +using storage::PosixBackend; +using storage::PosixPathMapper; +using storage::StorageBufferDescriptor; + +} // namespace kv_cache + +using kv_cache::storage::BlockKey; +using kv_cache::storage::BlockKeyMapper; +using kv_cache::storage::KVBackend; +using kv_cache::storage::PosixBackend; +using kv_cache::storage::PosixPathMapper; +using kv_cache::storage::StorageBufferDescriptor; + +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_KV_CACHE_STORAGE_STORAGE_H_ diff --git a/tpu_sync/kv_cache/storage/storage_test.cc b/tpu_sync/kv_cache/storage/storage_test.cc new file mode 100644 index 00000000..2e41862d --- /dev/null +++ b/tpu_sync/kv_cache/storage/storage_test.cc @@ -0,0 +1,374 @@ +// 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/kv_cache/storage/storage.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include +#include +#include + +#include +#include "absl/status/status.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/notification.h" +#include "tpu_sync/kv_cache/storage/k5_backend.h" + +namespace tpu_raiden { +namespace kv_cache { +namespace storage { +namespace { + +namespace fs = std::filesystem; + +class StorageDriverTest : public ::testing::Test { + protected: + void SetUp() override { + scratch_dir_ = absl::StrCat( + testing::TempDir(), "/storage_driver_test_", getpid(), "_", + std::chrono::system_clock::now().time_since_epoch().count()); + fs::create_directories(scratch_dir_); + } + + void TearDown() override { + std::error_code ec; + fs::remove_all(scratch_dir_, ec); + } + + std::string scratch_dir_; +}; + +TEST_F(StorageDriverTest, PosixPathMapperResolution) { + PosixPathMapper mapper(scratch_dir_, "llama3", /*tp_size=*/4, /*rank=*/2); + EXPECT_EQ(mapper.tp_size(), 4); + + // Default rank (-1 uses rank 2) + BlockKey key_default = mapper.MapKey("hash_12345", -1); + EXPECT_EQ(key_default.block_hash, "hash_12345"); + EXPECT_EQ(key_default.resolved_key, + absl::StrCat(scratch_dir_, "/llama3_tp4_r2/hash_12345.bin")); + + // Explicit rank override + BlockKey key_rank0 = mapper.MapKey("hash_12345", 0); + EXPECT_EQ(key_rank0.block_hash, "hash_12345"); + EXPECT_EQ(key_rank0.resolved_key, + absl::StrCat(scratch_dir_, "/llama3_tp4_r0/hash_12345.bin")); +} + +TEST_F(StorageDriverTest, PosixBackendWriteReadAndExists) { + PosixBackend backend("posix_disk"); + EXPECT_EQ(backend.scheme(), "posix_disk"); + + PosixPathMapper mapper(scratch_dir_, "model_v1", /*tp_size=*/1, /*rank=*/0); + BlockKey key = mapper.MapKey("block_alpha", 0); + + auto exists_or = backend.Exists(key); + ASSERT_TRUE(exists_or.ok()); + EXPECT_FALSE(exists_or.value()); + + // Prepare write buffer + const size_t kSize = 4096; + std::vector write_data(kSize, 0xAB); + StorageBufferDescriptor src_desc; + src_desc.ptr = write_data.data(); + + // Async write + absl::Notification write_done; + absl::Status write_status; + backend.WriteAsync(key, src_desc, kSize, [&](const absl::Status& status) { + write_status = status; + write_done.Notify(); + }); + write_done.WaitForNotification(); + EXPECT_TRUE(write_status.ok()); + + // Verify file exists + exists_or = backend.Exists(key); + ASSERT_TRUE(exists_or.ok()); + EXPECT_TRUE(exists_or.value()); + + // Async read + std::vector read_data(kSize, 0); + StorageBufferDescriptor dst_desc; + dst_desc.ptr = read_data.data(); + + absl::Notification read_done; + absl::Status read_status; + backend.ReadAsync(key, dst_desc, kSize, [&](const absl::Status& status) { + read_status = status; + read_done.Notify(); + }); + read_done.WaitForNotification(); + EXPECT_TRUE(read_status.ok()); + + // Verify read content matches written content + EXPECT_EQ(std::memcmp(write_data.data(), read_data.data(), kSize), 0); + + // Read non-existent key returns error + BlockKey missing_key = mapper.MapKey("non_existent_hash", 0); + absl::Notification missing_done; + absl::Status missing_status; + backend.ReadAsync(missing_key, dst_desc, kSize, + [&](const absl::Status& status) { + missing_status = status; + missing_done.Notify(); + }); + missing_done.WaitForNotification(); + EXPECT_FALSE(missing_status.ok()); +} + +TEST_F(StorageDriverTest, K5BlockNameMapperResolution) { + K5BlockNameMapper mapper(scratch_dir_, "gemma2", /*tp_size=*/8, /*rank=*/3); + + BlockKey key = mapper.MapKey("chunk_999", -1); + EXPECT_EQ(key.block_hash, "chunk_999"); + EXPECT_EQ(key.resolved_key, + absl::StrCat(scratch_dir_, "/k5_gemma2_tp8_r3_chunk_999.bin")); +} + +// Simple test daemon for testing K5 IPC over UDS +class TestUdsDaemon { + public: + TestUdsDaemon(std::string uds_path, std::string storage_dir) + : uds_path_(std::move(uds_path)), storage_dir_(std::move(storage_dir)) {} + + ~TestUdsDaemon() { Stop(); } + + absl::Status Start() { + server_fd_ = socket(AF_UNIX, SOCK_STREAM, 0); + if (server_fd_ < 0) { + return absl::InternalError( + absl::StrCat("socket failed: ", std::strerror(errno))); + } + + struct sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, uds_path_.c_str(), sizeof(addr.sun_path) - 1); + unlink(uds_path_.c_str()); + + if (bind(server_fd_, (struct sockaddr*)&addr, sizeof(addr)) < 0) { + close(server_fd_); + server_fd_ = -1; + return absl::InternalError( + absl::StrCat("bind failed on UDS: ", std::strerror(errno))); + } + + if (listen(server_fd_, 5) < 0) { + close(server_fd_); + server_fd_ = -1; + return absl::InternalError( + absl::StrCat("listen failed on UDS: ", std::strerror(errno))); + } + + running_ = true; + thread_ = std::thread([this]() { Run(); }); + return absl::OkStatus(); + } + + void Stop() { + running_ = false; + if (server_fd_ >= 0) { + close(server_fd_); + server_fd_ = -1; + } + if (thread_.joinable()) { + thread_.join(); + } + unlink(uds_path_.c_str()); + } + + private: + void Run() { + while (running_) { + struct timeval tv; + tv.tv_sec = 0; + tv.tv_usec = 50000; + fd_set readfds; + FD_ZERO(&readfds); + FD_SET(server_fd_, &readfds); + + int activity = select(server_fd_ + 1, &readfds, nullptr, nullptr, &tv); + if (activity <= 0) continue; + + int client_sock = accept(server_fd_, nullptr, nullptr); + if (client_sock < 0) continue; + + char control[CMSG_SPACE(sizeof(int))]; + char metadata[1024] = {0}; + + struct iovec iov[1]; + iov[0].iov_base = metadata; + iov[0].iov_len = sizeof(metadata) - 1; + + struct msghdr msg; + std::memset(&msg, 0, sizeof(msg)); + msg.msg_iov = iov; + msg.msg_iovlen = 1; + msg.msg_control = control; + msg.msg_controllen = sizeof(control); + + ssize_t recved = recvmsg(client_sock, &msg, 0); + if (recved <= 0) { + close(client_sock); + continue; + } + + struct cmsghdr* cmptr = CMSG_FIRSTHDR(&msg); + if (cmptr == nullptr || cmptr->cmsg_type != SCM_RIGHTS) { + (void)send(client_sock, "ERROR", 5, 0); + close(client_sock); + continue; + } + + int fd = *(reinterpret_cast(CMSG_DATA(cmptr))); + std::string op, key; + size_t offset = 0, size = 0; + + for (const auto& token : absl::StrSplit(metadata, ' ')) { + std::vector kv = absl::StrSplit(token, '='); + if (kv.size() == 2) { + if (kv[0] == "op") op = kv[1]; + if (kv[0] == "offset") (void)absl::SimpleAtoi(kv[1], &offset); + if (kv[0] == "size") (void)absl::SimpleAtoi(kv[1], &size); + if (kv[0] == "key") key = kv[1]; + } + } + + if (op == "write") { + void* mapped = + mmap(nullptr, offset + size, PROT_READ, MAP_SHARED, fd, 0); + if (mapped != MAP_FAILED) { + fs::create_directories(fs::path(key).parent_path()); + std::ofstream out(key, std::ios::binary); + out.write(static_cast(mapped) + offset, size); + out.close(); + munmap(mapped, offset + size); + (void)send(client_sock, "OK", 2, 0); + } else { + (void)send(client_sock, "ERROR", 5, 0); + } + } else if (op == "read") { + void* mapped = mmap(nullptr, offset + size, PROT_READ | PROT_WRITE, + MAP_SHARED, fd, 0); + if (mapped != MAP_FAILED) { + std::ifstream in(key, std::ios::binary); + if (in) { + in.read(static_cast(mapped) + offset, size); + in.close(); + (void)send(client_sock, "OK", 2, 0); + } else { + (void)send(client_sock, "ERROR", 5, 0); + } + munmap(mapped, offset + size); + } else { + (void)send(client_sock, "ERROR", 5, 0); + } + } + close(fd); + close(client_sock); + } + } + + std::string uds_path_; + std::string storage_dir_; + int server_fd_ = -1; + std::atomic running_{false}; + std::thread thread_; +}; + +TEST_F(StorageDriverTest, K5BackendUdsTransfer) { + std::string uds_path = absl::StrCat(scratch_dir_, "/k5_test.sock"); + std::string storage_dir = absl::StrCat(scratch_dir_, "/k5_storage"); + TestUdsDaemon daemon(uds_path, storage_dir); + ASSERT_TRUE(daemon.Start().ok()); + + K5Backend backend(uds_path); + EXPECT_EQ(backend.scheme(), "k5"); + + K5BlockNameMapper mapper(storage_dir, "model_test", /*tp_size=*/1, + /*rank=*/0); + BlockKey key = mapper.MapKey("hash_xyz", 0); + + const size_t kSize = 4096; + int shm_fd = memfd_create("test_k5_shm", 0); + ASSERT_GE(shm_fd, 0); + ASSERT_EQ(ftruncate(shm_fd, kSize), 0); + + void* ptr = + mmap(nullptr, kSize, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); + ASSERT_NE(ptr, MAP_FAILED); + std::memset(ptr, 0x7E, kSize); + + StorageBufferDescriptor desc; + desc.ptr = static_cast(ptr); + desc.fd = shm_fd; + desc.offset = 0; + + // Test WriteAsync + absl::Notification write_done; + absl::Status write_status; + backend.WriteAsync(key, desc, kSize, [&](const absl::Status& status) { + write_status = status; + write_done.Notify(); + }); + write_done.WaitForNotification(); + EXPECT_TRUE(write_status.ok()); + + // Test Exists + auto exists_or = backend.Exists(key); + ASSERT_TRUE(exists_or.ok()); + EXPECT_TRUE(exists_or.value()); + + // Zero buffer and test ReadAsync + std::memset(ptr, 0, kSize); + absl::Notification read_done; + absl::Status read_status; + backend.ReadAsync(key, desc, kSize, [&](const absl::Status& status) { + read_status = status; + read_done.Notify(); + }); + read_done.WaitForNotification(); + EXPECT_TRUE(read_status.ok()); + + // Verify data + uint8_t* p = static_cast(ptr); + for (size_t i = 0; i < kSize; ++i) { + EXPECT_EQ(p[i], 0x7E); + } + + munmap(ptr, kSize); + close(shm_fd); +} + +} // namespace +} // namespace storage +} // namespace kv_cache +} // namespace tpu_raiden diff --git a/tpu_sync/proto/worker_service.proto b/tpu_sync/proto/worker_service.proto index 570550f1..d9551ec6 100644 --- a/tpu_sync/proto/worker_service.proto +++ b/tpu_sync/proto/worker_service.proto @@ -51,6 +51,10 @@ service WorkerService { // Aborts an in-flight TransferProgram on this worker. rpc AbortTransfer(AbortTransferRequest) returns (AbortTransferResponse); + + // Registers storage backends dynamically on worker startup. + rpc RegisterBackends(RegisterBackendsRequest) + returns (RegisterBackendsResponse); } // Transfer endpoint descriptor for worker sub-managers. @@ -135,6 +139,59 @@ message DeleteBuffersResponse { string message = 2; } +// Simple transfer direction enum for storage operations. +enum TransferDirection { + TRANSFER_DIR_UNSPECIFIED = 0; + TRANSFER_DIR_OFFLOAD = 1; // Host RAM -> Storage (Write) + TRANSFER_DIR_RECALL = 2; // Storage -> Host RAM (Read) +} + +// Per-block storage locator descriptor. +message StorageKeyDescriptor { + // WHAT: Fully resolved physical storage locator or object identifier + // (e.g. "/tmp/raiden_proto_storage/rank_0/block_hash_a.bin" or + // "k5://block_hash_a"). WHY: Computed by BlockKeyMapper in KVCacheStore + // before issuing TransferBuffers. Low-level storage backend drivers + // (PosixKVCacheStoreBackend, K5BackendMock) need the exact file path or URI + // to execute system I/O calls (open, read, write). + string storage_key = 1; + + // WHAT: Linear byte offset from the start of the storage file, object, or + // shared memory region where block payload starts (default 0). WHY: Enables + // packed archive files or shared memory segments (e.g. K5 shared memory) + // containing multiple KV blocks to seek to payload boundaries without + // separate files. + int64 offset_bytes = 2; + + // WHAT: Optional POSIX shared memory file descriptor passed across process + // boundaries via Unix Domain Sockets (IPC) (default -1). WHY: Enables + // zero-copy inter-process communication for shared-memory daemons (e.g. K5 + // Thick Client daemon) via sendmsg(SCM_RIGHTS). + int32 file_descriptor = 3; +} + +// Batch-level storage transfer specification. +message StorageTransferSpec { + // WHAT: Simplified direction of copy (TRANSFER_DIR_OFFLOAD vs + // TRANSFER_DIR_RECALL). WHY: Storage backends only care whether they are + // writing from Host RAM to storage or reading into Host RAM. + TransferDirection direction = 1; + + // WHAT: Abstract string identifying the storage backend driver family + // (e.g. "local_disk", "lustre", "k5"). + // WHY: Used by WorkerServiceImplMock for dynamic factory driver lookup, and + // deduplicates the medium string ONCE per batch instead of repeating it per + // block. + string scheme = 2; + + // WHAT: Map keyed by host buffer index (0..N-1) containing the + // StorageKeyDescriptor for each block in the transfer batch. WHY: Explicitly + // correlates host buffer index i to its storage locator, preventing + // positional array index misalignment bugs while keeping BufferProto 100% + // pure. + map keys_by_buffer_index = 3; +} + // Specification for copying disjoint memory regions in one buffer. // The transfer applies uniformly to all buffers, i.e., across all shards and // all major dimensions (such as layers or blocks). @@ -162,6 +219,10 @@ message TransferBufferSpec { // peer's block id. Required (same length as src/dst offsets) for those // transfer types; unused otherwise. repeated BufferProto staging_host_buffers = 9; + // Option B storage transfer specification containing direction, scheme, + // and keys_by_buffer_index. Carries all storage metadata for transfers + // involving MEMORY_TYPE_STORAGE without altering core production messages. + StorageTransferSpec storage_spec = 10; } // Request to transfer data across memory spaces on a transfer worker. @@ -177,3 +238,19 @@ message TransferBuffersResponse { bool success = 1; string message = 2; } + +message BackendConfig { + string name = 1; + string scheme = 2; + uint64 capacity = 3; + map properties = 4; +} + +message RegisterBackendsRequest { + repeated BackendConfig configs = 1; +} + +message RegisterBackendsResponse { + bool success = 1; + string error_message = 2; +}