diff --git a/tpu_raiden/store_node/BUILD b/tpu_raiden/store_node/BUILD new file mode 100644 index 00000000..464f4ec7 --- /dev/null +++ b/tpu_raiden/store_node/BUILD @@ -0,0 +1,100 @@ +# 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_binary", "cc_library", "cc_test") + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "kv_transfer_spec_source", + hdrs = ["kv_transfer_spec_source.h"], + deps = [ + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + ], +) + +cc_library( + name = "kv_cache_host_store_node", + srcs = ["kv_cache_host_store_node.cc"], + hdrs = ["kv_cache_host_store_node.h"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + deps = [ + ":kv_transfer_spec_source", + "//tpu_raiden/core:kv_cache_manager_with_transfer", + "//tpu_raiden/core:kv_manager_holder", + "//tpu_raiden/core:status_macros", + "//tpu_raiden/core/controller:controller_client", + "//tpu_raiden/core/controller:worker_service_server", + "//tpu_raiden/kv_cache:kv_cache_store", + "//tpu_raiden/kv_cache:raiden_id", + "@com_google_absl//absl/log", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/random", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + ], +) + +cc_binary( + name = "kv_cache_host_store_node_main", + srcs = ["kv_cache_host_store_node_main.cc"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + deps = [ + ":kv_cache_host_store_node", + ":kv_transfer_spec_source", + "//tpu_raiden/kv_cache:raiden_id", + "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:initialize", + "@com_google_absl//absl/status:statusor", + ], +) + +cc_test( + name = "kv_cache_host_store_node_test", + srcs = ["kv_cache_host_store_node_test.cc"], + copts = [ + "-fno-strict-aliasing", + "-fexceptions", + ], + features = ["-use_header_modules"], + deps = [ + ":kv_cache_host_store_node", + ":kv_transfer_spec_source", + "//tpu_raiden/core/controller:raiden_controller", + "//tpu_raiden/kv_cache:kv_cache_store", + "//tpu_raiden/kv_cache:raiden_id", + "//tpu_raiden/kv_cache/global_registry:global_registry_client_cc", + "//tpu_raiden/kv_cache/global_registry:global_registry_server_lib", + "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_raiden/store_node/kv_cache_host_store_node.cc b/tpu_raiden/store_node/kv_cache_host_store_node.cc new file mode 100644 index 00000000..6e8c617b --- /dev/null +++ b/tpu_raiden/store_node/kv_cache_host_store_node.cc @@ -0,0 +1,210 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/store_node/kv_cache_host_store_node.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/log/log.h" +#include "absl/memory/memory.h" +#include "absl/random/random.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "tpu_raiden/core/controller/controller_client.h" +#include "tpu_raiden/core/controller/worker_service_server.h" +#include "tpu_raiden/core/kv_cache_manager_with_transfer.h" +#include "tpu_raiden/core/kv_manager_holder.h" +#include "tpu_raiden/core/status_macros.h" +#include "tpu_raiden/kv_cache/kv_cache_store.h" +#include "tpu_raiden/store_node/kv_transfer_spec_source.h" + +namespace tpu_raiden { +namespace store_node { +namespace { + +absl::Status ValidateOptions(const KVCacheHostStoreNode::Options& options) { + if (options.raiden_id.empty()) { + return absl::InvalidArgumentError("options.raiden_id must be set"); + } + if (options.store_server_ip.empty()) { + return absl::InvalidArgumentError( + "options.store_server_ip must be set: a host store node exists to be " + "reached by peers, so it always binds and advertises a concrete IP"); + } + if (options.dram_budget_bytes == 0) { + return absl::InvalidArgumentError("options.dram_budget_bytes must be set"); + } + return absl::OkStatus(); +} + +std::string ComposeEndpoint(absl::string_view ip, int port) { + return absl::StrContains(ip, ":") ? absl::StrCat("[", ip, "]:", port) + : absl::StrCat(ip, ":", port); +} + +// Transfer worker pairing id of this node's single worker. A peer pairs +// each of its workers with the worker here that has the same node_id so this +// value must mirror the peer's and is deliberately not configurable. The +// full worker topology to mirror (one worker per serving-host transfer +// rank, same node_ids) will come with the published KVTransferSpec; the +// node will then run one manager/worker per entry and this constant goes +// away. Until then a single worker under node_id 0 matches a +// single-transfer-rank serving host. +constexpr int64_t kNodeId = 0; + +} // namespace + +absl::StatusOr KVCacheHostStoreNode::WaitForSpec( + KVTransferSpecSource& source, const Options& options) { + absl::BitGen gen; + absl::Duration interval = options.spec_poll_initial; + const absl::Time deadline = absl::Now() + options.spec_wait_timeout; + while (true) { + absl::StatusOr spec = source.Get(); + if (spec.ok()) { + RETURN_IF_ERROR(ValidateSpec(*spec)); + return spec; + } + if (!absl::IsNotFound(spec.status()) && + !absl::IsUnavailable(spec.status())) { + return spec.status(); + } + if (absl::Now() >= deadline) { + return absl::DeadlineExceededError( + absl::StrCat("KV transfer spec not published within ", + absl::FormatDuration(options.spec_wait_timeout), + "; last source status: ", spec.status().ToString())); + } + absl::SleepFor(interval * absl::Uniform(gen, 0.5, 1.5)); + interval = std::min(interval * 2, options.spec_poll_max); + } +} + +absl::StatusOr KVCacheHostStoreNode::NumBlocksForBudget( + size_t dram_budget_bytes, const KVTransferSpec& spec) { + RETURN_IF_ERROR(ValidateSpec(spec)); + const size_t block_bytes = + spec.num_layers * spec.num_shards * spec.slice_byte_size; + if (dram_budget_bytes < block_bytes) { + return absl::InvalidArgumentError(absl::StrCat( + "dram_budget_bytes=", dram_budget_bytes, + " does not cover a single block of ", block_bytes, + " bytes (num_layers=", spec.num_layers, " x num_shards=", + spec.num_shards, " x slice_byte_size=", spec.slice_byte_size, ")")); + } + return dram_budget_bytes / block_bytes; +} + +absl::StatusOr> +KVCacheHostStoreNode::Create(const Options& options, + KVTransferSpecSource* kv_transfer_spec_source) { + RETURN_IF_ERROR(ValidateOptions(options)); + if (kv_transfer_spec_source == nullptr) { + return absl::InvalidArgumentError( + "kv_transfer_spec_source must not be null"); + } + + // Phase A: the only input the node cannot know on its own. + ASSIGN_OR_RETURN(const KVTransferSpec spec, + WaitForSpec(*kv_transfer_spec_source, options)); + ASSIGN_OR_RETURN(const size_t num_host_blocks, + NumBlocksForBudget(options.dram_budget_bytes, spec)); + + // Phase B: the same assembly a serving host runs, CPU-backed. + std::unique_ptr manager; + std::unique_ptr store; + try { + // 1. The pool. The CPU-only constructor allocates and zeroes the full + // num_host_blocks x num_layers x num_shards x slice_byte_size pool here, + // so the memory cost is paid at boot, never on a request path. A + // follow-up change will build one KVCacheManagerWithTransfer per entry + // of the worker topology the serving hosts publish in the global + // registry, replacing this single instance (see kNodeId above). + // local_control_port=-1: the slot protocol is engine-to-engine + // coordination and has no role on a host store node. + manager = std::make_unique( + spec.num_layers, spec.num_shards, spec.slice_byte_size, + /*local_port=*/0, + /*host_blocks_to_allocate=*/num_host_blocks, options.parallelism, + kNodeId, /*local_control_port=*/-1); + + // 2. Everything above the pool: the RaidenController coordination plane + // (num_blocks sized to the capacity, so incoming writes can allocate + // block ids through AllocateBlockIds), the LRU cache, and the + // peer-facing store server. The constructor publishes the node to the + // global registry (RegisterStore) once the store server is bound, so + // registration happens exactly when the advertised address is live. + // With no global registry configured KVCacheStore stands up no store + // server at all (the registry decides whether the peer plane exists). + store = std::make_unique( + /*capacity=*/num_host_blocks, options.global_registry_address, + options.raiden_id, static_cast(spec.num_shards), + static_cast(spec.slice_byte_size), + options.raiden_orchestrator_address, options.store_server_ip, + options.raiden_controller_port); + } catch (const std::exception& e) { + return absl::UnavailableError( + absl::StrCat("host store node assembly failed: ", e.what())); + } + + // 3. Make the manager reachable from the controller. StartServer must + // precede RegisterWorker: registration needs the worker's live control + // endpoint. Registration goes through the controller's own RPC, the same + // path a serving host's worker takes. + auto worker_server = controller::WorkerServiceServer::Create(); + RETURN_IF_ERROR(worker_server->StartServer( + /*host_allocator=*/nullptr, KVManagerHolder(manager.get()), + /*port=*/0)); + const std::string worker_endpoint = ComposeEndpoint( + options.store_server_ip, worker_server->GetRaidenWorkerPort()); + core::controller::RaidenControllerClient controller_client( + store->raiden_controller_address()); + absl::Status registered = controller_client.RegisterWorker( + "worker_0", worker_endpoint, manager->get_local_data_endpoints(), + kNodeId); + if (!registered.ok()) { + // Without a worker the node can neither receive nor serve bytes, so fail + // the whole boot rather than come up half-alive. + return absl::Status( + registered.code(), + absl::StrCat("host store node worker registration failed: ", + registered.message())); + } + + auto node = absl::WrapUnique( + new KVCacheHostStoreNode(spec, num_host_blocks, std::move(manager), + std::move(worker_server), std::move(store))); + LOG(INFO) << "KVCacheHostStoreNode up: " << node->num_host_blocks() + << " host blocks (num_layers=" << spec.num_layers + << " num_shards=" << spec.num_shards + << " slice_byte_size=" << spec.slice_byte_size << "), store server " + << node->store_server_address() << ", controller " + << node->raiden_controller_address() << ", worker " + << worker_endpoint; + return node; +} + +} // namespace store_node +} // namespace tpu_raiden diff --git a/tpu_raiden/store_node/kv_cache_host_store_node.h b/tpu_raiden/store_node/kv_cache_host_store_node.h new file mode 100644 index 00000000..bf8d877a --- /dev/null +++ b/tpu_raiden/store_node/kv_cache_host_store_node.h @@ -0,0 +1,175 @@ +// 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_STORE_NODE_KV_CACHE_HOST_STORE_NODE_H_ +#define THIRD_PARTY_TPU_RAIDEN_STORE_NODE_KV_CACHE_HOST_STORE_NODE_H_ + +#include +#include +#include + +#include "absl/status/statusor.h" +#include "absl/time/time.h" +#include "tpu_raiden/core/controller/worker_service_server.h" +#include "tpu_raiden/core/kv_cache_manager_with_transfer.h" +#include "tpu_raiden/kv_cache/kv_cache_store.h" +#include "tpu_raiden/kv_cache/raiden_id.h" +#include "tpu_raiden/store_node/kv_transfer_spec_source.h" + +namespace tpu_raiden { +namespace store_node { + +// A standalone raiden node that lends its host DRAM to the deployment as a +// KV cache pool: serving hosts evict prefix KV blocks into it and read them +// back later. It runs no engine; it is the store role of a serving host, +// deployed on its own (which is also why it can run on a machine with no +// accelerator at all). +// +// Not to be confused with kv_cache::KVCacheStoreServer, the gRPC wrapper +// hosting the peer-facing KVCacheStoreService: that is one component INSIDE +// this node, in the same way ControllerServer is one component inside +// RaidenController. +// +// Boot is two phases, split by the one input the node cannot know on its own: +// +// Phase A (no spec needed): wait on the KVTransferSpecSource until the +// deployment's KVTransferSpec is published, retrying with jittered +// backoff. Everything the node is configured with (identity, addresses, +// DRAM budget) is deployment config; the spec is the serving hosts' +// runtime truth and must come from them. +// +// Phase B (spec in hand): run the same assembly a serving host runs, +// with a CPU-backed manager: +// 1. KVCacheManagerWithTransfer (CPU-only): allocates the actual DRAM +// pool up front and opens the raw-transfer data endpoint. +// 2. KVCacheStore: builds the RaidenController, the LRU cache and the +// peer-facing store server, and publishes this node to the global +// registry (RegisterStore), so registration happens exactly when +// the node can actually serve. +// 3. WorkerService + worker registration: makes the manager reachable +// from the controller, completing the transfer path. +// +// After Create() returns, the node is discoverable and serves peers; there is +// no separate Start(). Destruction tears down in reverse: the store first +// (unregisters from the global registry, stops the store server and +// controller), then the worker server, then the manager (frees the pool). +class KVCacheHostStoreNode { + public: + struct Options { + // Identity this node registers under, in the global registry and the + // orchestrator alike. The four fields form one opaque composite key -- + // raiden never matches a field individually -- so all that matters is + // that the tuple is unique within the deployment and stable across + // restarts (re-registering the same id is how a restarted node replaces + // its stale registration). + kv_cache::RaidenId raiden_id; + + // The IP peers use to reach this node. Bind-and-advertise, same semantics + // as KVCacheStore's store_server_ip: the store server, the controller and + // the worker all bind it, and it is the host published to the global + // registry. Required: an unreachable host store node is useless. + std::string store_server_ip; + + // Controller port; 0 lets gRPC choose. + int raiden_controller_port = 0; + + // Global registry to publish this node to. Per KVCacheStore's + // construction rules the registry decides whether the peer-facing plane + // exists at all: empty means no store server is stood up -- the manager, + // controller and worker still boot (useful in tests), but peers can + // neither discover nor dial this node. + std::string global_registry_address; + + // Orchestrator to register the controller with. Empty skips registration + // (peer controller resolution is then unavailable; useful in tests). + std::string raiden_orchestrator_address; + + // Host DRAM lent to the pool. Converted to whole blocks of the received + // spec's block geometry; the remainder below one block is not allocated. + size_t dram_budget_bytes = 0; + + // Transfer parallelism of the manager. + int parallelism = 1; + + // KVTransferSpec wait: retry cadence while the source reports NotFound or + // Unavailable. The interval starts at spec_poll_initial, doubles up + // to spec_poll_max, and every sleep is jittered so a fleet of nodes + // booting together spreads its polls. spec_wait_timeout bounds the wait + // so a node that can never learn the spec (wrong registry address, dead + // deployment) fails boot visibly instead of hanging forever. + absl::Duration spec_poll_initial = absl::Seconds(1); + absl::Duration spec_poll_max = absl::Minutes(1); + absl::Duration spec_wait_timeout = absl::Minutes(120); + }; + + // Boots a host store node: blocks in WaitForSpec, then assembles the node. + // On return the node is serving and (when a global registry is configured) + // discoverable. `kv_transfer_spec_source` must outlive the call. + static absl::StatusOr> Create( + const Options& options, KVTransferSpecSource* kv_transfer_spec_source); + + KVCacheHostStoreNode(const KVCacheHostStoreNode&) = delete; + KVCacheHostStoreNode& operator=(const KVCacheHostStoreNode&) = delete; + + // Phase A alone: polls `source` until it yields a valid spec, with + // jittered exponential backoff per `options`. Returns the spec, or + // DeadlineExceeded after options.spec_wait_timeout, or the source's + // error when it is neither NotFound nor Unavailable. + static absl::StatusOr WaitForSpec( + KVTransferSpecSource& source, const Options& options); + + // Whole blocks of `spec`'s block geometry that fit in `dram_budget_bytes`. + // InvalidArgument if the budget does not cover even one block. + static absl::StatusOr NumBlocksForBudget(size_t dram_budget_bytes, + const KVTransferSpec& spec); + + const KVTransferSpec& spec() const { return spec_; } + size_t num_host_blocks() const { return num_host_blocks_; } + + kv_cache::KVCacheStore* store() const { return store_.get(); } + KVCacheManagerWithTransfer* manager() const { return manager_.get(); } + + // Advertised "host:port" of the peer-facing store server / the controller. + std::string store_server_address() const { + return store_->store_server_address(); + } + std::string raiden_controller_address() const { + return store_->raiden_controller_address(); + } + + private: + KVCacheHostStoreNode(const KVTransferSpec& spec, size_t num_host_blocks, + std::unique_ptr manager, + std::unique_ptr + worker_server, + std::unique_ptr store) + : spec_(spec), + num_host_blocks_(num_host_blocks), + manager_(std::move(manager)), + worker_server_(std::move(worker_server)), + store_(std::move(store)) {} + + KVTransferSpec spec_; + size_t num_host_blocks_ = 0; + + // Declaration order is teardown order reversed: store_ goes down first. + std::unique_ptr manager_; + std::unique_ptr worker_server_; + std::unique_ptr store_; +}; + +} // namespace store_node +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_STORE_NODE_KV_CACHE_HOST_STORE_NODE_H_ diff --git a/tpu_raiden/store_node/kv_cache_host_store_node_main.cc b/tpu_raiden/store_node/kv_cache_host_store_node_main.cc new file mode 100644 index 00000000..b8bd9e06 --- /dev/null +++ b/tpu_raiden/store_node/kv_cache_host_store_node_main.cc @@ -0,0 +1,148 @@ +// 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. + +// Standalone KV cache host store node. +// +// Lends this machine's DRAM to a raiden deployment as a KV cache pool. +// Start it alongside the serving jobs; it waits for the deployment's +// KVTransferSpec, assembles itself, publishes itself to the global registry, +// and serves until SIGTERM/SIGINT, on which it unregisters and tears down. + +#include +#include +#include +#include + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" +#include "absl/status/statusor.h" +#include "tpu_raiden/kv_cache/raiden_id.h" +#include "tpu_raiden/store_node/kv_cache_host_store_node.h" +#include "tpu_raiden/store_node/kv_transfer_spec_source.h" + +// Identity: the RaidenId this node registers under. Raiden treats the four +// fields as one opaque composite key -- no field is matched individually +// anywhere -- so the only hard requirements are that the tuple is unique +// within the deployment and stable across restarts (a restarted node replaces +// its own registration by re-registering the same id). The structure mirrors +// how serving jobs are addressed: job + replica locate the process, data +// name + replica index name the cache entity it hosts. +ABSL_FLAG(std::string, job_name, "", + "Job this host store node is deployed as, e.g. " + "'kv-host-store-node'."); +ABSL_FLAG(std::string, job_replica_id, "0", + "Replica/task index of this node within --job_name."); +ABSL_FLAG(std::string, data_name, "kv_cache", + "Name of the cache pool this node hosts. Any stable string; keep " + "the default unless one deployment hosts several distinct pools."); +ABSL_FLAG(int, data_replica_idx, 0, + "Replica index of the pool; 0 unless --data_name is replicated."); + +// Addressing. The node listens on four ports: controller (gRPC), peer-facing +// store server (gRPC), worker control plane (WorkerService gRPC), and +// transfer data plane (BlockTransport raw TCP). Only the controller port is +// configurable, because KVCacheStore's API takes it alongside +// store_server_ip; the other three always bind OS-chosen ports. That is +// safe because every address is registered or published at boot -- worker +// to controller, store and controller to the registries -- so no peer ever +// discovers this node by computing a port. +ABSL_FLAG(std::string, store_server_ip, "", + "IP peers reach this node on; bound and advertised. Required."); +ABSL_FLAG(int, raiden_controller_port, 0, + "Controller port; 0 lets gRPC choose."); +ABSL_FLAG(std::string, global_registry_address, "", + "GlobalRegistry host:port to publish this node to."); +ABSL_FLAG(std::string, raiden_orchestrator_address, "", + "Orchestrator host:port to register the controller with."); + +// Capacity. +ABSL_FLAG(size_t, dram_budget_bytes, 0, + "Host DRAM lent to the pool, in bytes. Required."); + +// Transfer plane. +ABSL_FLAG(int, parallelism, 1, + "Parallel transport streams this node uses when it initiates " + "transfers. Purely local performance tuning for this machine's " + "NIC/CPU; peers need not match it."); + +// KVTransferSpec. Stopgap flags feeding a StaticKVTransferSpecSource; +// replaced by the global-registry-backed source once spec publication lands +// there, at which point the node needs no spec flags at all. +ABSL_FLAG(size_t, num_layers, 0, "KV block geometry: layer count."); +ABSL_FLAG(size_t, num_shards, 0, "KV block geometry: shard count."); +ABSL_FLAG(size_t, slice_byte_size, 0, + "KV block geometry: bytes of one (layer, shard) slice."); + +namespace { + +using ::tpu_raiden::store_node::KVCacheHostStoreNode; +using ::tpu_raiden::store_node::KVTransferSpec; +using ::tpu_raiden::store_node::StaticKVTransferSpecSource; + +int Run() { + // Block SIGTERM/SIGINT before the node spawns any thread: threads inherit + // the mask, so delivery is confined to the sigwait below instead of + // terminating the process with no teardown. + sigset_t shutdown_signals; + sigemptyset(&shutdown_signals); + sigaddset(&shutdown_signals, SIGTERM); + sigaddset(&shutdown_signals, SIGINT); + pthread_sigmask(SIG_BLOCK, &shutdown_signals, nullptr); + + KVCacheHostStoreNode::Options options; + options.raiden_id = tpu_raiden::kv_cache::RaidenId{ + absl::GetFlag(FLAGS_job_name), absl::GetFlag(FLAGS_job_replica_id), + absl::GetFlag(FLAGS_data_name), absl::GetFlag(FLAGS_data_replica_idx)}; + options.store_server_ip = absl::GetFlag(FLAGS_store_server_ip); + options.raiden_controller_port = absl::GetFlag(FLAGS_raiden_controller_port); + options.global_registry_address = + absl::GetFlag(FLAGS_global_registry_address); + options.raiden_orchestrator_address = + absl::GetFlag(FLAGS_raiden_orchestrator_address); + options.dram_budget_bytes = absl::GetFlag(FLAGS_dram_budget_bytes); + options.parallelism = absl::GetFlag(FLAGS_parallelism); + + StaticKVTransferSpecSource kv_transfer_spec_source( + KVTransferSpec{absl::GetFlag(FLAGS_num_layers), + absl::GetFlag(FLAGS_num_shards), + absl::GetFlag(FLAGS_slice_byte_size)}); + + absl::StatusOr> node = + KVCacheHostStoreNode::Create(options, &kv_transfer_spec_source); + if (!node.ok()) { + LOG(ERROR) << "host store node failed to boot: " << node.status(); + return 1; + } + + // Park until asked to stop, then destroy the node so teardown runs in + // order: unregister from the global registry, drain and stop the servers, + // free the pool. A killed process skips all of that and leaves a stale + // registry entry behind until the next boot re-registers the same id. + int received_signal = 0; + sigwait(&shutdown_signals, &received_signal); + LOG(INFO) << "host store node shutting down on signal " << received_signal; + node->reset(); + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + absl::ParseCommandLine(argc, argv); + absl::InitializeLog(); + std::signal(SIGPIPE, SIG_IGN); + return Run(); +} diff --git a/tpu_raiden/store_node/kv_cache_host_store_node_test.cc b/tpu_raiden/store_node/kv_cache_host_store_node_test.cc new file mode 100644 index 00000000..4e632ab7 --- /dev/null +++ b/tpu_raiden/store_node/kv_cache_host_store_node_test.cc @@ -0,0 +1,262 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/store_node/kv_cache_host_store_node.h" + +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/time/time.h" +#include "grpcpp/create_channel.h" +#include "grpcpp/security/credentials.h" +#include "grpcpp/security/server_credentials.h" +#include "grpcpp/server.h" +#include "grpcpp/server_builder.h" +#include "tpu_raiden/kv_cache/raiden_id.h" +#include "tpu_raiden/store_node/kv_transfer_spec_source.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_client.h" +#include "tpu_raiden/kv_cache/global_registry/global_registry_server.h" + +namespace tpu_raiden { +namespace store_node { +namespace { + +using ::testing::HasSubstr; + +// Yields `status` until `failures` calls have been made, then the spec. +class FlakyKVTransferSpecSource : public KVTransferSpecSource { + public: + FlakyKVTransferSpecSource(KVTransferSpec spec, int failures, + absl::Status status) + : spec_(spec), failures_(failures), status_(status) {} + + absl::StatusOr Get() override { + ++calls_; + if (calls_ <= failures_) { + return status_; + } + return spec_; + } + + int calls() const { return calls_; } + + private: + KVTransferSpec spec_; + int failures_; + absl::Status status_; + int calls_ = 0; +}; + +KVTransferSpec TestSpec() { return KVTransferSpec{2, 1, 256}; } + +KVCacheHostStoreNode::Options FastPollOptions() { + KVCacheHostStoreNode::Options options; + options.spec_poll_initial = absl::Milliseconds(1); + options.spec_poll_max = absl::Milliseconds(4); + return options; +} + +TEST(ValidateSpecTest, RejectsAnyZeroField) { + EXPECT_TRUE(ValidateSpec(TestSpec()).ok()); + EXPECT_TRUE(absl::IsInvalidArgument(ValidateSpec(KVTransferSpec{0, 1, 256}))); + EXPECT_TRUE(absl::IsInvalidArgument(ValidateSpec(KVTransferSpec{2, 0, 256}))); + EXPECT_TRUE(absl::IsInvalidArgument(ValidateSpec(KVTransferSpec{2, 1, 0}))); +} + +TEST(WaitForSpecTest, RetriesNotFoundUntilPublished) { + FlakyKVTransferSpecSource source(TestSpec(), /*failures=*/3, + absl::NotFoundError("not published")); + absl::StatusOr spec = + KVCacheHostStoreNode::WaitForSpec(source, FastPollOptions()); + ASSERT_TRUE(spec.ok()) << spec.status(); + EXPECT_EQ(spec->slice_byte_size, TestSpec().slice_byte_size); + EXPECT_EQ(source.calls(), 4); +} + +TEST(WaitForSpecTest, RetriesUnavailable) { + FlakyKVTransferSpecSource source(TestSpec(), /*failures=*/2, + absl::UnavailableError("registry not up")); + absl::StatusOr spec = + KVCacheHostStoreNode::WaitForSpec(source, FastPollOptions()); + ASSERT_TRUE(spec.ok()) << spec.status(); + EXPECT_EQ(source.calls(), 3); +} + +TEST(WaitForSpecTest, PropagatesFatalErrorWithoutRetry) { + FlakyKVTransferSpecSource source(TestSpec(), /*failures=*/100, + absl::InternalError("corrupt registry")); + absl::StatusOr spec = + KVCacheHostStoreNode::WaitForSpec(source, FastPollOptions()); + EXPECT_TRUE(absl::IsInternal(spec.status())); + EXPECT_EQ(source.calls(), 1); +} + +TEST(WaitForSpecTest, RejectsInvalidPublishedSpec) { + FlakyKVTransferSpecSource source(KVTransferSpec{0, 0, 0}, /*failures=*/0, + absl::OkStatus()); + absl::StatusOr spec = + KVCacheHostStoreNode::WaitForSpec(source, FastPollOptions()); + EXPECT_TRUE(absl::IsInvalidArgument(spec.status())); +} + +TEST(WaitForSpecTest, TimesOut) { + FlakyKVTransferSpecSource source(TestSpec(), /*failures=*/1000000, + absl::NotFoundError("not published")); + KVCacheHostStoreNode::Options options = FastPollOptions(); + options.spec_wait_timeout = absl::Milliseconds(30); + absl::StatusOr spec = + KVCacheHostStoreNode::WaitForSpec(source, options); + EXPECT_TRUE(absl::IsDeadlineExceeded(spec.status())); + EXPECT_GT(source.calls(), 1); +} + +TEST(NumBlocksForBudgetTest, FloorsToWholeBlocks) { + // One block of TestSpec() is 2 * 1 * 256 = 512 bytes. + absl::StatusOr blocks = + KVCacheHostStoreNode::NumBlocksForBudget(4096, TestSpec()); + ASSERT_TRUE(blocks.ok()) << blocks.status(); + EXPECT_EQ(*blocks, 8u); + blocks = KVCacheHostStoreNode::NumBlocksForBudget(4095, TestSpec()); + ASSERT_TRUE(blocks.ok()) << blocks.status(); + EXPECT_EQ(*blocks, 7u); +} + +TEST(NumBlocksForBudgetTest, RejectsBudgetBelowOneBlock) { + absl::StatusOr blocks = + KVCacheHostStoreNode::NumBlocksForBudget(511, TestSpec()); + EXPECT_TRUE(absl::IsInvalidArgument(blocks.status())); + EXPECT_THAT(blocks.status().message(), HasSubstr("single block")); +} + +// Boot tests stand up real (loopback) gRPC servers; the env var isolates +// each node's controller/worker servers from the process-wide singletons so +// tests do not leak servers into each other. +class KVCacheHostStoreNodeBootTest : public ::testing::Test { + protected: + void SetUp() override { + setenv("RAIDEN_DISABLE_SINGLETON_WORKER", "1", /*overwrite=*/1); + } + void TearDown() override { unsetenv("RAIDEN_DISABLE_SINGLETON_WORKER"); } + + KVCacheHostStoreNode::Options BootOptions() { + KVCacheHostStoreNode::Options options = FastPollOptions(); + options.raiden_id = + kv_cache::RaidenId{"store_node_test", "0", "kv_pool", 0}; + options.store_server_ip = "localhost"; + options.dram_budget_bytes = 8 * 512; // 8 blocks of TestSpec(). + return options; + } +}; + +TEST_F(KVCacheHostStoreNodeBootTest, RequiresStoreServerIp) { + KVCacheHostStoreNode::Options options = BootOptions(); + options.store_server_ip.clear(); + StaticKVTransferSpecSource source(TestSpec()); + absl::StatusOr> node = + KVCacheHostStoreNode::Create(options, &source); + EXPECT_TRUE(absl::IsInvalidArgument(node.status())); + EXPECT_THAT(node.status().message(), HasSubstr("store_server_ip")); +} + +TEST_F(KVCacheHostStoreNodeBootTest, RequiresKVTransferSpecSource) { + absl::StatusOr> node = + KVCacheHostStoreNode::Create(BootOptions(), nullptr); + EXPECT_TRUE(absl::IsInvalidArgument(node.status())); +} + +TEST_F(KVCacheHostStoreNodeBootTest, BootsWithoutRegistryButServesNoPeers) { + StaticKVTransferSpecSource source(TestSpec()); + absl::StatusOr> node = + KVCacheHostStoreNode::Create(BootOptions(), &source); + ASSERT_TRUE(node.ok()) << node.status(); + + // The budget became whole blocks of the received spec. + EXPECT_EQ((*node)->num_host_blocks(), 8u); + EXPECT_EQ((*node)->spec().num_layers, TestSpec().num_layers); + + // KVCacheStore's construction rules make the global registry decide + // whether the peer-facing plane exists: no registry, no store server. + EXPECT_EQ((*node)->store()->store_server(), nullptr); + EXPECT_EQ((*node)->store_server_address(), ""); + + // The controller is live and knows exactly our one worker, registered + // under node_id 0 -- the pairing id the serving hosts' single worker + // uses, mirrored rather than configured. + auto workers = (*node) + ->store() + ->raiden_controller() + ->worker_registry() + ->GetRegisteredWorkers(); + ASSERT_EQ(workers.size(), 1u); + EXPECT_EQ(workers[0].node_id, 0); + + // The manager opened its raw-transfer data endpoint. + EXPECT_FALSE((*node)->manager()->get_local_endpoints().empty()); +} + +TEST_F(KVCacheHostStoreNodeBootTest, BootsServesAndPublishesWithRegistry) { + // A live in-process global registry, the same pattern + // kv_cache_store_test.cc uses. + auto service = + std::make_unique(); + grpc::ServerBuilder builder; + int port = 0; + builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), + &port); + builder.RegisterService(service.get()); + std::unique_ptr registry_server = builder.BuildAndStart(); + ASSERT_NE(registry_server, nullptr); + ASSERT_NE(port, 0); + + KVCacheHostStoreNode::Options options = BootOptions(); + options.global_registry_address = absl::StrCat("localhost:", port); + + StaticKVTransferSpecSource source(TestSpec()); + absl::StatusOr> node = + KVCacheHostStoreNode::Create(options, &source); + ASSERT_TRUE(node.ok()) << node.status(); + + // The peer-facing store server is bound and advertised. + EXPECT_THAT((*node)->store_server_address(), HasSubstr("localhost:")); + EXPECT_NE((*node)->store()->store_server(), nullptr); + + // And the node published itself: the registry resolves our RaidenId to + // the advertised store address. + auto channel = grpc::CreateChannel(options.global_registry_address, + grpc::InsecureChannelCredentials()); + kv_cache::global_registry::GlobalRegistryClient registry_client(channel); + auto store_info = registry_client.ResolveStore(options.raiden_id); + ASSERT_TRUE(store_info.ok()) << store_info.status(); + EXPECT_EQ(store_info->store_server_address(), + (*node)->store_server_address()); +} + +TEST_F(KVCacheHostStoreNodeBootTest, WaitsOutLateSpecThenBoots) { + FlakyKVTransferSpecSource source(TestSpec(), /*failures=*/2, + absl::NotFoundError("not published")); + absl::StatusOr> node = + KVCacheHostStoreNode::Create(BootOptions(), &source); + ASSERT_TRUE(node.ok()) << node.status(); + EXPECT_EQ(source.calls(), 3); + EXPECT_EQ((*node)->num_host_blocks(), 8u); +} + +} // namespace +} // namespace store_node +} // namespace tpu_raiden diff --git a/tpu_raiden/store_node/kv_transfer_spec_source.h b/tpu_raiden/store_node/kv_transfer_spec_source.h new file mode 100644 index 00000000..46676e86 --- /dev/null +++ b/tpu_raiden/store_node/kv_transfer_spec_source.h @@ -0,0 +1,90 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_STORE_NODE_KV_TRANSFER_SPEC_SOURCE_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_STORE_NODE_KV_TRANSFER_SPEC_SOURCE_H_ + +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" + +namespace tpu_raiden { +namespace store_node { + +// Everything a host store node must match about the deployment's serving hosts +// to interoperate with their KV transfers. The authority for these values is +// the serving host's manager runtime (which has the real device buffers); +// a host store node only ever receives them, it never derives them. +// +// Today this carries the block geometry: the uniform-slice block shape, +// where every (layer, shard) slice has the same byte size. Planned +// extensions -- the serving hosts' transfer worker topology, and per-pool +// block shapes for hybrid models -- grow this struct without changing the +// source interface below. +struct KVTransferSpec { + size_t num_layers = 0; + size_t num_shards = 0; + size_t slice_byte_size = 0; +}; + +// Returns InvalidArgument unless every field is positive. +inline absl::Status ValidateSpec(const KVTransferSpec& spec) { + if (spec.num_layers == 0 || spec.num_shards == 0 || + spec.slice_byte_size == 0) { + return absl::InvalidArgumentError(absl::StrCat( + "KVTransferSpec fields must all be positive, got num_layers=", + spec.num_layers, " num_shards=", spec.num_shards, + " slice_byte_size=", spec.slice_byte_size)); + } + return absl::OkStatus(); +} + +// Where a booting host store node obtains the deployment's KVTransferSpec. +// +// Get() contract: +// - OK: the spec is known. The value is fixed for the lifetime of the +// deployment; callers read it once at boot. +// - NotFound: nothing published yet. Expected during turnup, when the store +// node can come up before any serving host has published its spec; the +// caller retries. +// - Unavailable: the source itself is not reachable yet (also expected +// during turnup); the caller retries. +// - anything else: fatal. +class KVTransferSpecSource { + public: + virtual ~KVTransferSpecSource() = default; + + virtual absl::StatusOr Get() = 0; +}; + +// KVTransferSpec fixed at construction, e.g. from flags. Stopgap until the +// global-registry-backed source lands: the registry will hold the spec +// published by the serving hosts at their own registration, and a source +// implementation will poll it here through the same interface. +class StaticKVTransferSpecSource : public KVTransferSpecSource { + public: + explicit StaticKVTransferSpecSource(KVTransferSpec spec) : spec_(spec) {} + + absl::StatusOr Get() override { return spec_; } + + private: + KVTransferSpec spec_; +}; + +} // namespace store_node +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_STORE_NODE_KV_TRANSFER_SPEC_SOURCE_H_