diff --git a/.bazelrc b/.bazelrc index defd19b91..f37089a4b 100644 --- a/.bazelrc +++ b/.bazelrc @@ -17,3 +17,5 @@ build:ci --remote_local_fallback # targets set per-target and the flags kokoro presubmit passes globally. build --cxxopt=-std=c++20 build --host_cxxopt=-std=c++20 +build --copt=-DABSL_DEFINE_UNQUALIFIED_STATUS_MACROS +build --host_copt=-DABSL_DEFINE_UNQUALIFIED_STATUS_MACROS diff --git a/tpu_sync/api/jax/kv_cache_store.py b/tpu_sync/api/jax/kv_cache_store.py index 215e72a25..fa13712ce 100644 --- a/tpu_sync/api/jax/kv_cache_store.py +++ b/tpu_sync/api/jax/kv_cache_store.py @@ -521,12 +521,7 @@ def read_remote( Returns: True if successfully launched. """ - raw_slices = [] - for s in slices: - if isinstance(s, RaidenId): - s = RaidenBlockId(raiden_id=s) - raw_slices.append(s._impl) # pylint: disable=protected-access - return self._impl.read_remote(block_hashes, raw_slices, device_block_ids) + return self.load(block_hashes, device_block_ids, slices=slices) def poll_remote_read_status( self, @@ -539,4 +534,4 @@ def poll_remote_read_status( failed: List of block hashes whose remote read failed. pending: List of block hashes whose remote read is still in progress. """ - return self._impl.poll_remote_read_status() + return self.poll_load_status() diff --git a/tpu_sync/api/torch/kv_cache_store.py b/tpu_sync/api/torch/kv_cache_store.py index dbd7c7fa4..619ea81c3 100644 --- a/tpu_sync/api/torch/kv_cache_store.py +++ b/tpu_sync/api/torch/kv_cache_store.py @@ -532,12 +532,7 @@ def read_remote( Returns: True if successfully launched. """ - raw_slices = [] - for s in slices: - if isinstance(s, RaidenId): - s = RaidenBlockId(raiden_id=s) - raw_slices.append(s._impl) # pylint: disable=protected-access - return self._impl.read_remote(block_hashes, raw_slices, device_block_ids) + return self.load(block_hashes, device_block_ids, slices=slices) def poll_remote_read_status( self, @@ -550,4 +545,4 @@ def poll_remote_read_status( failed: List of block hashes whose remote read failed. pending: List of block hashes whose remote read is still in progress. """ - return self._impl.poll_remote_read_status() + return self.poll_load_status() diff --git a/tpu_sync/core/BUILD b/tpu_sync/core/BUILD index 0de1ad665..841de35f0 100644 --- a/tpu_sync/core/BUILD +++ b/tpu_sync/core/BUILD @@ -38,23 +38,6 @@ cc_library( visibility = ["//visibility:public"], ) -cc_library( - name = "status_macros", - hdrs = [ - "status_macros.h", - ], - copts = [ - "-fno-strict-aliasing", - "-fexceptions", - ], - features = ["-use_header_modules"], - visibility = ["//visibility:public"], - deps = [ - "@com_google_absl//absl/log", - "@com_google_absl//absl/status", - ], -) - cc_library( name = "buffer", srcs = [ @@ -301,9 +284,9 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":raw_transfer_core", - ":status_macros", ":xla_raw_transfer_headers", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", ], @@ -317,9 +300,9 @@ cc_library( deps = [ ":raiden_transfer_endpoint", ":raw_transfer_core", - ":status_macros", "//tpu_sync/rpc:raiden_service_cc_proto", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -429,7 +412,6 @@ cc_library( ":raiden_manager_base", ":raiden_transfer_endpoint", ":raw_transfer_core", - ":status_macros", ":tpu_utils", "//tpu_sync/kv_cache:kv_cache_manager_base", "//tpu_sync/kv_cache:pool_layout", @@ -441,6 +423,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", diff --git a/tpu_sync/core/controller/BUILD b/tpu_sync/core/controller/BUILD index c2a6f046f..b0f0ec9c8 100644 --- a/tpu_sync/core/controller/BUILD +++ b/tpu_sync/core/controller/BUILD @@ -89,7 +89,6 @@ cc_library( "//tpu_sync/common:raiden_id", "//tpu_sync/core:buffer", "//tpu_sync/core:raiden_transfer_endpoint", - "//tpu_sync/core:status_macros", "//tpu_sync/kv_cache:logical_block_manager", "//tpu_sync/proto:controller_service_cc_grpc", "//tpu_sync/proto:controller_service_cc_proto", @@ -102,6 +101,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", diff --git a/tpu_sync/core/controller/raiden_controller.cc b/tpu_sync/core/controller/raiden_controller.cc index 032df5584..ae25be9c4 100644 --- a/tpu_sync/core/controller/raiden_controller.cc +++ b/tpu_sync/core/controller/raiden_controller.cc @@ -35,6 +35,7 @@ #include "absl/log/log.h" #include "absl/memory/memory.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" @@ -54,7 +55,6 @@ #include "tpu_sync/core/controller/worker_registry.h" #include "tpu_sync/core/controller/worker_service_client.h" #include "tpu_sync/core/raiden_transfer_endpoint.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/logical_block_manager.h" #include "tpu_sync/proto/controller_service.grpc.pb.h" #include "tpu_sync/proto/controller_service.pb.h" @@ -416,8 +416,8 @@ RaidenController::Allocate(int num_blocks) { "that Physical/BufferProto mode is unavailable when the controller " "was built with preprovision_worker_buffers=false"); } - ASSIGN_OR_RETURN(std::vector block_ids, - block_manager_->Allocate(num_blocks, /*lock=*/true)); + ABSL_ASSIGN_OR_RETURN(std::vector block_ids, + block_manager_->Allocate(num_blocks, /*lock=*/true)); std::vector<::tpu_sync::proto::BufferProto> result; result.reserve(num_blocks); for (int block_id : block_ids) { @@ -428,8 +428,8 @@ RaidenController::Allocate(int num_blocks) { absl::StatusOr> RaidenController::AllocateBuffers( int num_blocks) { - ASSIGN_OR_RETURN(std::vector<::tpu_sync::proto::BufferProto> protos, - Allocate(num_blocks)); + ABSL_ASSIGN_OR_RETURN(std::vector<::tpu_sync::proto::BufferProto> protos, + Allocate(num_blocks)); std::vector buffers; buffers.reserve(protos.size()); for (const auto& proto : protos) { diff --git a/tpu_sync/core/kv_cache_manager_with_transfer.cc b/tpu_sync/core/kv_cache_manager_with_transfer.cc index f375387ac..72a7c14e1 100644 --- a/tpu_sync/core/kv_cache_manager_with_transfer.cc +++ b/tpu_sync/core/kv_cache_manager_with_transfer.cc @@ -56,6 +56,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" @@ -71,7 +72,6 @@ #include "tpu_sync/core/raiden_manager_base.h" #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/core/tpu_utils.h" #include "tpu_sync/kv_cache/kv_cache_manager_base.h" #include "tpu_sync/kv_cache/pool_layout.h" @@ -2104,9 +2104,9 @@ absl::Status KVCacheManagerWithTransfer::InitializeSlotPool(int64_t num_slots) { all_slots_.clear(); all_slots_.reserve(num_slots); for (int64_t i = 0; i < num_slots; ++i) { - ASSIGN_OR_RETURN(std::vector allocated_ids, - host_block_manager_->Allocate(max_blocks_, - /*lock=*/true)); + ABSL_ASSIGN_OR_RETURN(std::vector allocated_ids, + host_block_manager_->Allocate(max_blocks_, + /*lock=*/true)); if (allocated_ids.size() != max_blocks_) { return absl::InternalError(absl::StrCat( "Slot pool allocation returned incorrect number of blocks: ", diff --git a/tpu_sync/core/kv_manager_holder.h b/tpu_sync/core/kv_manager_holder.h index 547f815b5..ac5b65aea 100644 --- a/tpu_sync/core/kv_manager_holder.h +++ b/tpu_sync/core/kv_manager_holder.h @@ -25,13 +25,13 @@ #include #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #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/rpc/raiden_service.pb.h" namespace tpu_raiden { @@ -280,7 +280,8 @@ class KVManagerHolder { absl::StatusOr H2hRead( absl::string_view peer, const std::vector& src_offsets, const std::vector& dst_offsets) override { - ASSIGN_OR_RETURN(std::vector src_ids, SafeCastOffsets(src_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector src_ids, + SafeCastOffsets(src_offsets)); // When the caller named its destination blocks, land the data THERE: // plain H2hRead auto-allocates destination blocks from the manager's // own accounting, which neither matches the ids the caller reserved @@ -288,29 +289,33 @@ class KVManagerHolder { // already handed out. if constexpr (internal::has_peer_h2h_read_explicit_v) { if (!dst_offsets.empty()) { - ASSIGN_OR_RETURN(std::vector dst_ids, - SafeCastOffsets(dst_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector dst_ids, + SafeCastOffsets(dst_offsets)); return impl_->H2hReadExplicit(std::string(peer), src_ids, dst_ids, /*explicit_dst_ptrs=*/{}); } } - ASSIGN_OR_RETURN(auto res, impl_->H2hRead(std::string(peer), src_ids)); + ABSL_ASSIGN_OR_RETURN(auto res, + impl_->H2hRead(std::string(peer), src_ids)); return res.second; } absl::StatusOr H2hWrite( absl::string_view peer, const std::vector& src_offsets, const std::vector& dst_offsets) override { - ASSIGN_OR_RETURN(std::vector src_ids, SafeCastOffsets(src_offsets)); - ASSIGN_OR_RETURN(std::vector dst_ids, SafeCastOffsets(dst_offsets)); - ASSIGN_OR_RETURN(auto res, - impl_->H2hWrite(std::string(peer), src_ids, dst_ids)); + ABSL_ASSIGN_OR_RETURN(std::vector src_ids, + SafeCastOffsets(src_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector dst_ids, + SafeCastOffsets(dst_offsets)); + ABSL_ASSIGN_OR_RETURN( + auto res, impl_->H2hWrite(std::string(peer), src_ids, dst_ids)); return res.second; } absl::StatusOr H2hRead( const std::vector& remote_descriptors, const std::vector& src_offsets, const std::vector& dst_offsets) override { - ASSIGN_OR_RETURN(std::vector src_ids, SafeCastOffsets(src_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector src_ids, + SafeCastOffsets(src_offsets)); // When the caller named its destination blocks, land the data THERE. // Plain H2hRead auto-allocates, which is fine for a fire-and-forget pull // but wrong for a store-level read: the store already reserved landing @@ -318,28 +323,29 @@ class KVManagerHolder { // blocks would leave the directory pointing at the wrong memory. if constexpr (internal::has_vector_h2h_read_explicit_v) { if (!dst_offsets.empty()) { - ASSIGN_OR_RETURN(std::vector dst_ids, - SafeCastOffsets(dst_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector dst_ids, + SafeCastOffsets(dst_offsets)); return impl_->H2hReadExplicit(remote_descriptors, src_ids, dst_ids); } } else if constexpr (internal::has_peer_h2h_read_explicit_v) { // No descriptor-shaped explicit read; the peer-string one lands the // blocks just as precisely. if (!dst_offsets.empty() && !remote_descriptors.empty()) { - ASSIGN_OR_RETURN(std::vector dst_ids, - SafeCastOffsets(dst_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector dst_ids, + SafeCastOffsets(dst_offsets)); return impl_->H2hReadExplicit(remote_descriptors[0].endpoint, src_ids, dst_ids, /*explicit_dst_ptrs=*/{}); } } if constexpr (internal::has_vector_h2h_read_v) { - ASSIGN_OR_RETURN(auto res, impl_->H2hRead(remote_descriptors, src_ids)); + ABSL_ASSIGN_OR_RETURN(auto res, + impl_->H2hRead(remote_descriptors, src_ids)); return res.second; } else { std::string peer = remote_descriptors.empty() ? "" : remote_descriptors[0].endpoint; - ASSIGN_OR_RETURN(auto res, impl_->H2hRead(peer, src_ids)); + ABSL_ASSIGN_OR_RETURN(auto res, impl_->H2hRead(peer, src_ids)); return res.second; } } @@ -347,16 +353,19 @@ class KVManagerHolder { const std::vector& remote_descriptors, const std::vector& src_offsets, const std::vector& dst_offsets) override { - ASSIGN_OR_RETURN(std::vector src_ids, SafeCastOffsets(src_offsets)); - ASSIGN_OR_RETURN(std::vector dst_ids, SafeCastOffsets(dst_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector src_ids, + SafeCastOffsets(src_offsets)); + ABSL_ASSIGN_OR_RETURN(std::vector dst_ids, + SafeCastOffsets(dst_offsets)); if constexpr (internal::has_vector_h2h_write_v) { - ASSIGN_OR_RETURN(auto res, - impl_->H2hWrite(remote_descriptors, src_ids, dst_ids)); + ABSL_ASSIGN_OR_RETURN( + auto res, impl_->H2hWrite(remote_descriptors, src_ids, dst_ids)); return res.second; } else { std::string peer = remote_descriptors.empty() ? "" : remote_descriptors[0].endpoint; - ASSIGN_OR_RETURN(auto res, impl_->H2hWrite(peer, src_ids, dst_ids)); + ABSL_ASSIGN_OR_RETURN(auto res, + impl_->H2hWrite(peer, src_ids, dst_ids)); return res.second; } } diff --git a/tpu_sync/core/raw_transfer_impl.h b/tpu_sync/core/raw_transfer_impl.h index fd0ab9d16..dd6804a13 100644 --- a/tpu_sync/core/raw_transfer_impl.h +++ b/tpu_sync/core/raw_transfer_impl.h @@ -33,7 +33,6 @@ #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/statusor.h" #include "tpu_sync/core/raw_transfer_core.h" -#include "tpu_sync/core/status_macros.h" namespace raiden { diff --git a/tpu_sync/core/status_macros.h b/tpu_sync/core/status_macros.h deleted file mode 100644 index 5303eb331..000000000 --- a/tpu_sync/core/status_macros.h +++ /dev/null @@ -1,93 +0,0 @@ -// 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_CORE_STATUS_MACROS_H_ -#define THIRD_PARTY_TPU_RAIDEN_CORE_STATUS_MACROS_H_ - -#include -#include - -#include "absl/log/log.h" -#include "absl/status/status.h" - -namespace tpu_raiden { -namespace status_macro_internal { - -class StatusBuilder { - public: - explicit StatusBuilder(const absl::Status& status) : status_(status) {} - - StatusBuilder& LogError() { - LOG(ERROR) << status_; - return *this; - } - - template - auto With(F&& f) -> decltype(f(std::declval())) { - return f(status_); - } - - operator absl::Status() const { return status_; } - - template , absl::Status> && - std::is_constructible_v>> - operator T() const { - return T(status_); - } - - private: - absl::Status status_; -}; - -} // namespace status_macro_internal -} // namespace tpu_raiden - -#define RAIDEN_STATUS_MACROS_CONCAT_NAME_INNER(x, y) x##y -#define RAIDEN_STATUS_MACROS_CONCAT_NAME(x, y) \ - RAIDEN_STATUS_MACROS_CONCAT_NAME_INNER(x, y) - -// Helper to get variadic macro arguments -#define RAIDEN_STATUS_INTERNAL_GET_VARIADIC(_1, _2, _3, NAME, ...) NAME - -#define ASSIGN_OR_RETURN(...) \ - RAIDEN_STATUS_INTERNAL_GET_VARIADIC(__VA_ARGS__, ASSIGN_OR_RETURN_3, \ - ASSIGN_OR_RETURN_2)(__VA_ARGS__) - -#define ASSIGN_OR_RETURN_2(lhs, rexpr) ASSIGN_OR_RETURN_3(lhs, rexpr, _) - -#define ASSIGN_OR_RETURN_3(lhs, rexpr, error_expr) \ - ASSIGN_OR_RETURN_IMPL( \ - RAIDEN_STATUS_MACROS_CONCAT_NAME(_status_or, __COUNTER__), lhs, rexpr, \ - error_expr) - -#define ASSIGN_OR_RETURN_IMPL(statusor, lhs, rexpr, error_expr) \ - auto statusor = (rexpr); \ - if (!statusor.ok()) { \ - ::tpu_raiden::status_macro_internal::StatusBuilder _(statusor.status()); \ - return (error_expr); \ - } \ - lhs = std::move(statusor).value() - -#define RETURN_IF_ERROR(expr) \ - RETURN_IF_ERROR_IMPL(RAIDEN_STATUS_MACROS_CONCAT_NAME(_status, __COUNTER__), \ - expr) - -#define RETURN_IF_ERROR_IMPL(status, expr) \ - auto status = (expr); \ - if (!status.ok()) { \ - return status; \ - } - -#endif // THIRD_PARTY_TPU_RAIDEN_CORE_STATUS_MACROS_H_ diff --git a/tpu_sync/frameworks/jax/BUILD b/tpu_sync/frameworks/jax/BUILD index 49c7b7a4a..9b569fb8b 100644 --- a/tpu_sync/frameworks/jax/BUILD +++ b/tpu_sync/frameworks/jax/BUILD @@ -48,13 +48,13 @@ cc_library( "//tpu_sync/core:metrics_collector", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core:tpu_utils", "//tpu_sync/core:utils", "//tpu_sync/core/controller:controller_client", "//tpu_sync/core/controller:worker_service_server", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -290,7 +290,6 @@ cc_library( "//tpu_sync/core:numa_thread_pool", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core:tpu_utils", "//tpu_sync/rpc:raiden_service_cc_proto", "//tpu_sync/weight_sync:weight_synchronizer_base", @@ -298,6 +297,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", @@ -340,8 +340,8 @@ nanobind_extension( ":jax_utils", "//tpu_sync/core:raw_transfer_core", "//tpu_sync/core:raw_transfer_impl", - "//tpu_sync/core:status_macros", "@com_google_absl//absl/log", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", @@ -458,7 +458,6 @@ cc_library( "//tpu_sync/core:numa_thread_pool", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core:tpu_utils", "//tpu_sync/rpc:raiden_service_cc_proto", "//tpu_sync/weight_sync:weight_synchronizer_base", @@ -466,6 +465,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", @@ -512,8 +512,8 @@ cc_library( ":mock_nanobind", "//tpu_sync/core:raw_transfer_core", "//tpu_sync/core:raw_transfer_impl", - "//tpu_sync/core:status_macros", "@com_google_absl//absl/log", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", @@ -546,8 +546,8 @@ cc_library( deps = [ "//tpu_sync/core:raw_transfer_core", "//tpu_sync/core:raw_transfer_impl", - "//tpu_sync/core:status_macros", "@com_google_absl//absl/log", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", @@ -774,13 +774,13 @@ cc_library( "//tpu_sync/core:metrics_collector", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core:tpu_utils", "//tpu_sync/core:utils", "//tpu_sync/core/controller:controller_client", "//tpu_sync/core/controller:worker_service_server", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", diff --git a/tpu_sync/frameworks/jax/kv_cache_manager.cc b/tpu_sync/frameworks/jax/kv_cache_manager.cc index 638defa77..40dae577e 100644 --- a/tpu_sync/frameworks/jax/kv_cache_manager.cc +++ b/tpu_sync/frameworks/jax/kv_cache_manager.cc @@ -34,6 +34,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" // IWYU pragma: keep @@ -49,7 +50,6 @@ #include "tpu_sync/core/metrics_collector.h" // IWYU pragma: keep #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/core/tpu_utils.h" #include "tpu_sync/core/utils.h" // IWYU pragma: keep #ifndef WITHOUT_PYTHON @@ -618,8 +618,8 @@ absl::StatusOr NumaAwareKVCacheManager::H2d( std::vector sub_copy_futures; sub_copy_futures.reserve(sub_managers_.size()); for (auto& sub : sub_managers_) { - ASSIGN_OR_RETURN(auto f, sub->H2d(src_offsets, dst_offsets, copy_sizes, - slot_idx, layer_idx, shard_idx)); + ABSL_ASSIGN_OR_RETURN(auto f, sub->H2d(src_offsets, dst_offsets, copy_sizes, + slot_idx, layer_idx, shard_idx)); sub_copy_futures.push_back(std::move(f)); } // Use the event-aware join. On TPU the per-shard copies complete via the @@ -641,8 +641,8 @@ absl::StatusOr NumaAwareKVCacheManager::D2h( std::vector sub_copy_futures; sub_copy_futures.reserve(sub_managers_.size()); for (auto& sub : sub_managers_) { - ASSIGN_OR_RETURN(auto f, sub->D2h(src_offsets, dst_offsets, copy_sizes, - slot_idx, layer_idx, shard_idx)); + ABSL_ASSIGN_OR_RETURN(auto f, sub->D2h(src_offsets, dst_offsets, copy_sizes, + slot_idx, layer_idx, shard_idx)); sub_copy_futures.push_back(std::move(f)); } // Use the event-aware join (see H2d above): on TPU the per-shard copies @@ -662,7 +662,7 @@ NumaAwareKVCacheManager::D2hAutoAllocate( std::vector all_ids; std::vector sub_copy_futures; for (size_t s = 0; s < sub_managers_.size(); ++s) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto res, sub_managers_[s]->D2hAutoAllocate(src_offsets, copy_sizes)); if (s == 0) { all_ids = res.first; @@ -702,7 +702,7 @@ NumaAwareKVCacheManager::H2hWrite(std::string peer, for (size_t s = 0; s < sub_managers_.size(); ++s) { std::string sub_peer = (base_port >= 0) ? host_prefix + std::to_string(base_port + s) : peer; - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto res, sub_managers_[s]->H2hWrite(sub_peer, src_block_ids, dst_block_ids, uuid, layer_idx)); if (s == 0) { @@ -741,8 +741,8 @@ NumaAwareKVCacheManager::H2hRead(std::string peer, for (size_t s = 0; s < sub_managers_.size(); ++s) { std::string sub_peer = (base_port >= 0) ? host_prefix + std::to_string(base_port + s) : peer; - ASSIGN_OR_RETURN(auto res, - sub_managers_[s]->H2hRead(sub_peer, src_block_ids)); + ABSL_ASSIGN_OR_RETURN(auto res, + sub_managers_[s]->H2hRead(sub_peer, src_block_ids)); if (s == 0) { all_ids = res.first; } @@ -784,7 +784,7 @@ NumaAwareKVCacheManager::H2hWrite( matched_ep = remote_descriptors[0].endpoint; } - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto res, sub_managers_[s]->H2hWrite(matched_ep, src_block_ids, dst_block_ids, uuid, layer_idx)); if (s == 0) { @@ -827,8 +827,8 @@ NumaAwareKVCacheManager::H2hRead( matched_ep = remote_descriptors[0].endpoint; } - ASSIGN_OR_RETURN(auto res, - sub_managers_[s]->H2hRead(matched_ep, src_block_ids)); + ABSL_ASSIGN_OR_RETURN(auto res, + sub_managers_[s]->H2hRead(matched_ep, src_block_ids)); if (s == 0) { all_ids = res.first; } @@ -871,9 +871,10 @@ NumaAwareKVCacheManager::H2hReadExplicit( matched_ep = remote_descriptors[0].endpoint; } - ASSIGN_OR_RETURN(auto fut, sub_managers_[s]->H2hReadExplicit( - matched_ep, src_block_ids, dst_block_ids, - /*explicit_dst_ptrs=*/{})); + ABSL_ASSIGN_OR_RETURN( + auto fut, sub_managers_[s]->H2hReadExplicit(matched_ep, src_block_ids, + dst_block_ids, + /*explicit_dst_ptrs=*/{})); sub_copy_futures.push_back(std::move(fut)); } return raiden::JoinPjRtCopyFutures(absl::MakeSpan(sub_copy_futures)); @@ -912,10 +913,10 @@ absl::StatusOr NumaAwareKVCacheManager::H2dRead( matched_ep = remote_descriptors[0].endpoint; } - ASSIGN_OR_RETURN(auto fut, sub_managers_[s]->H2dRead( - matched_ep, src_host_offsets, - dst_host_offsets, dst_device_offsets, - copy_sizes)); + ABSL_ASSIGN_OR_RETURN( + auto fut, sub_managers_[s]->H2dRead(matched_ep, src_host_offsets, + dst_host_offsets, + dst_device_offsets, copy_sizes)); sub_copy_futures.push_back(std::move(fut)); } // Event-aware join (see D2h/H2d): on TPU the per-shard transfers complete via diff --git a/tpu_sync/frameworks/jax/raw_transfer.cc b/tpu_sync/frameworks/jax/raw_transfer.cc index 4de14c71c..809a485e2 100644 --- a/tpu_sync/frameworks/jax/raw_transfer.cc +++ b/tpu_sync/frameworks/jax/raw_transfer.cc @@ -30,10 +30,10 @@ namespace xla { class PjRtBuffer; class PjRtClient; } // namespace xla +#include "absl/status/status_macros.h" #include "xla/pjrt/status_casters.h" #include "tpu_sync/core/raw_transfer_core.h" #include "tpu_sync/core/raw_transfer_impl.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/frameworks/jax/jax_utils.h" #ifndef WITHOUT_PYTHON #include @@ -148,7 +148,7 @@ inline absl::StatusOr transfer_d2h_batch_async_impl( std::vector c_sizes = jax::UnpackListToVector(copy_sizes_major_dim); for (size_t i = 0; i < n; ++i) { - ASSIGN_OR_RETURN(PjRtCopyFuture f, + ABSL_ASSIGN_OR_RETURN(PjRtCopyFuture f, transfer_d2h_async_internal( src_arrs[i], dst_arrs[i], s_offsets, d_offsets, c_sizes, unsafe_skip_buffer_lock)); @@ -179,7 +179,7 @@ inline absl::StatusOr transfer_h2d_batch_async_impl( std::vector c_sizes = jax::UnpackListToVector(copy_sizes_major_dim); for (size_t i = 0; i < n; ++i) { - ASSIGN_OR_RETURN(PjRtCopyFuture f, + ABSL_ASSIGN_OR_RETURN(PjRtCopyFuture f, transfer_h2d_async_internal( src_arrs[i], dst_arrs[i], s_offsets, d_offsets, c_sizes, unsafe_skip_buffer_lock)); diff --git a/tpu_sync/frameworks/jax/weight_synchronizer.cc b/tpu_sync/frameworks/jax/weight_synchronizer.cc index 3f09fe7bf..ed7a7ac7f 100644 --- a/tpu_sync/frameworks/jax/weight_synchronizer.cc +++ b/tpu_sync/frameworks/jax/weight_synchronizer.cc @@ -30,6 +30,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/match.h" #include "absl/synchronization/mutex.h" @@ -38,7 +39,6 @@ #include "tpu_sync/core/numa_thread_pool.h" #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/core/tpu_utils.h" #include "tpu_sync/rpc/raiden_service.pb.h" #include "tpu_sync/weight_sync/weight_synchronizer_base.h" @@ -418,7 +418,7 @@ absl::StatusOr NumaAwareWeightSynchronizer::D2h( sub_copy_futures.reserve(sub_synchronizers_.size()); for (auto& sub : sub_synchronizers_) { if (sub) { - ASSIGN_OR_RETURN(auto f, sub->D2h(uuid)); + ABSL_ASSIGN_OR_RETURN(auto f, sub->D2h(uuid)); sub_copy_futures.push_back(std::move(f)); } } @@ -432,7 +432,7 @@ absl::StatusOr NumaAwareWeightSynchronizer::H2d( sub_copy_futures.reserve(sub_synchronizers_.size()); for (auto& sub : sub_synchronizers_) { if (sub) { - ASSIGN_OR_RETURN(auto f, sub->H2d(uuid)); + ABSL_ASSIGN_OR_RETURN(auto f, sub->H2d(uuid)); sub_copy_futures.push_back(std::move(f)); } } diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index 34b7a3bb8..a031fe1ff 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -112,9 +112,9 @@ cc_library( deps = [ ":kv_cache_metadata", "//tpu_sync/core:host_memory_allocator", - "//tpu_sync/core:status_macros", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", @@ -254,7 +254,6 @@ cc_library( "//tpu_sync/core:numa_thread_pool", "//tpu_sync/core:raiden_manager_base", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core:tpu_utils", "//tpu_sync/core:xla_raw_transfer_headers", "//tpu_sync/rpc:raiden_service_cc_proto", @@ -269,6 +268,7 @@ cc_library( "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", @@ -332,7 +332,6 @@ cc_library( ":lru_cache", "//tpu_sync/common:raiden_id", "//tpu_sync/core:buffer", - "//tpu_sync/core:status_macros", "//tpu_sync/core/controller:raiden_controller", "//tpu_sync/core/controller:worker_registry", "//tpu_sync/kv_cache/global_registry:global_registry_client_cc", @@ -346,6 +345,7 @@ cc_library( "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", @@ -365,7 +365,6 @@ cc_library( ":kv_cache_metadata", ":kv_cache_store_backend", "//tpu_sync/common:raiden_id", - "//tpu_sync/core:status_macros", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", @@ -483,7 +482,6 @@ cc_library( "//tpu_sync/common:raiden_id", "//tpu_sync/core:buffer", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core/controller:raiden_controller", "//tpu_sync/kv_cache/global_registry:global_registry_client_cc", "//tpu_sync/kv_cache/reshard:reshard_service", @@ -496,6 +494,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", diff --git a/tpu_sync/kv_cache/host_offload_backend.cc b/tpu_sync/kv_cache/host_offload_backend.cc index b88dbd585..523353f45 100644 --- a/tpu_sync/kv_cache/host_offload_backend.cc +++ b/tpu_sync/kv_cache/host_offload_backend.cc @@ -27,6 +27,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/escaping.h" #include "absl/strings/str_cat.h" @@ -41,7 +42,6 @@ #include "tpu_sync/common/raiden_id.h" #include "tpu_sync/core/buffer.h" #include "tpu_sync/core/controller/raiden_controller.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/global_registry/global_registry_client.h" #include "tpu_sync/kv_cache/kv_cache_metadata.h" #include "tpu_sync/kv_cache/kv_cache_store_backend.h" @@ -129,7 +129,8 @@ absl::StatusOr> HostOffloadBackend::Create( config.capacity, config.metadata, config.raiden_id, controller, std::move(registry_client), config.kv_pool_group)); if (config.kv_transfer_spec.has_value()) { - RETURN_IF_ERROR(backend->RegisterKVTransferSpec(*config.kv_transfer_spec)); + ABSL_RETURN_IF_ERROR( + backend->RegisterKVTransferSpec(*config.kv_transfer_spec)); } return backend; } @@ -751,8 +752,8 @@ HostOffloadBackend::GetKVCacheStoreClient(const RaidenId& remote_id) { "No global registry client; cannot resolve peer store address"); } - ASSIGN_OR_RETURN(global_registry::StoreInfo store_info, - registry->ResolveStore(remote_id)); + ABSL_ASSIGN_OR_RETURN(global_registry::StoreInfo store_info, + registry->ResolveStore(remote_id)); if (store_info.store_server_address().empty()) { return absl::NotFoundError( "Peer is registered but published an empty store server address"); @@ -790,8 +791,8 @@ HostOffloadBackend::BeginWriteRemote( // Resolves the peer through the global registry. A missing registry // client fails here; there is no separate precondition check. - ASSIGN_OR_RETURN(std::shared_ptr client, - GetKVCacheStoreClient(dst_raiden_id)); + ABSL_ASSIGN_OR_RETURN(std::shared_ptr client, + GetKVCacheStoreClient(dst_raiden_id)); auto call = client->WriteRemote(raiden_controller_->unit(), block_hashes, src_host_block_ids, @@ -1170,7 +1171,7 @@ absl::Status HostOffloadBackend::RegisterKVTransferSpecFromWorkers() { "HostOffloadBackend has no RaidenController; there are no worker " "registrations to derive a KVTransferSpec from."); } - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( const KVTransferSpecConfig spec, ComposeKVTransferSpec( raiden_controller_->worker_registry()->GetRegisteredWorkers())); diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index bc409720e..fa59c65c8 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -39,6 +39,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" @@ -57,7 +58,6 @@ #include "tpu_sync/core/numa_thread_pool.h" #include "tpu_sync/core/raiden_manager_base.h" #include "tpu_sync/core/raw_transfer_core.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/core/tpu_utils.h" #include "tpu_sync/kv_cache/logical_block_manager.h" #include "tpu_sync/kv_cache/pool_layout.h" @@ -799,7 +799,7 @@ absl::StatusOr KVCacheManagerBase::D2hSyncDispatch( std::optional slot_idx, std::optional layer_idx, std::optional shard_idx) { const absl::Time d2h_start = absl::Now(); - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto logical_futures, DispatchD2hChunks(src_offsets_major_dim, dst_offsets_major_dim, copy_sizes_major_dim, slot_idx, layer_idx, shard_idx)); @@ -825,10 +825,10 @@ absl::StatusOr KVCacheManagerBase::H2dWrite( "must have the same length"); } - ASSIGN_OR_RETURN(std::vector src_block_ids, - ToHostBlockIds(src_host_offsets_major_dim)); - ASSIGN_OR_RETURN(std::vector staging_block_ids, - ToHostBlockIds(dst_host_offsets_major_dim)); + ABSL_ASSIGN_OR_RETURN(std::vector src_block_ids, + ToHostBlockIds(src_host_offsets_major_dim)); + ABSL_ASSIGN_OR_RETURN(std::vector staging_block_ids, + ToHostBlockIds(dst_host_offsets_major_dim)); TF_RETURN_IF_ERROR(ValidateOffsetsAndSizes(src_host_offsets_major_dim, dst_device_offsets_major_dim, copy_sizes_major_dim)); @@ -839,8 +839,9 @@ absl::StatusOr KVCacheManagerBase::H2dWrite( // is not yet executed by this call; dst_device_offsets_major_dim identifies // the eventual remote HBM destination for the receiver-side device copy. if (num_chunks == 1 || !push_pool_) { - ASSIGN_OR_RETURN(auto h2h_res, H2hWrite(std::string(peer), src_block_ids, - staging_block_ids)); + ABSL_ASSIGN_OR_RETURN( + auto h2h_res, + H2hWrite(std::string(peer), src_block_ids, staging_block_ids)); return h2h_res.second; } @@ -889,10 +890,10 @@ absl::StatusOr KVCacheManagerBase::H2dRead( "must have the same length"); } - ASSIGN_OR_RETURN(std::vector src_block_ids, - ToHostBlockIds(src_host_offsets_major_dim)); - ASSIGN_OR_RETURN(std::vector staging_block_ids, - ToHostBlockIds(dst_host_offsets_major_dim)); + ABSL_ASSIGN_OR_RETURN(std::vector src_block_ids, + ToHostBlockIds(src_host_offsets_major_dim)); + ABSL_ASSIGN_OR_RETURN(std::vector staging_block_ids, + ToHostBlockIds(dst_host_offsets_major_dim)); TF_RETURN_IF_ERROR(ValidateOffsetsAndSizes(dst_host_offsets_major_dim, dst_device_offsets_major_dim, copy_sizes_major_dim)); @@ -901,11 +902,11 @@ absl::StatusOr KVCacheManagerBase::H2dRead( // (dst_host_offsets_major_dim) -- never into an aliased copy of the remote // src id -- then H2D the staging blocks into the local device destination. if (num_chunks == 1 || !pull_pool_) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto h2h_fut, H2hReadExplicit(std::string(peer), src_block_ids, staging_block_ids, /*explicit_dst_ptrs=*/{})); - RETURN_IF_ERROR(h2h_fut.Await()); + ABSL_RETURN_IF_ERROR(h2h_fut.Await()); return H2dSyncDispatch(dst_host_offsets_major_dim, dst_device_offsets_major_dim, copy_sizes_major_dim); @@ -998,10 +999,10 @@ absl::StatusOr KVCacheManagerBase::D2hWrite( "must have the same length"); } - ASSIGN_OR_RETURN(std::vector staging_block_ids, - ToHostBlockIds(src_host_offsets_major_dim)); - ASSIGN_OR_RETURN(std::vector dst_block_ids, - ToHostBlockIds(dst_host_offsets_major_dim)); + ABSL_ASSIGN_OR_RETURN(std::vector staging_block_ids, + ToHostBlockIds(src_host_offsets_major_dim)); + ABSL_ASSIGN_OR_RETURN(std::vector dst_block_ids, + ToHostBlockIds(dst_host_offsets_major_dim)); TF_RETURN_IF_ERROR(ValidateOffsetsAndSizes(src_device_offsets_major_dim, src_host_offsets_major_dim, copy_sizes_major_dim)); @@ -1010,14 +1011,15 @@ absl::StatusOr KVCacheManagerBase::D2hWrite( // (src_host_offsets_major_dim) -- never into an aliased copy of the remote // dst id -- then push the staging blocks to the peer's host destination. if (num_chunks == 1 || !push_pool_) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto d2h_future, D2hSyncDispatch(src_device_offsets_major_dim, src_host_offsets_major_dim, copy_sizes_major_dim)); - RETURN_IF_ERROR(d2h_future.Await()); + ABSL_RETURN_IF_ERROR(d2h_future.Await()); - ASSIGN_OR_RETURN(auto h2h_res, H2hWrite(std::string(peer), - staging_block_ids, dst_block_ids)); + ABSL_ASSIGN_OR_RETURN( + auto h2h_res, + H2hWrite(std::string(peer), staging_block_ids, dst_block_ids)); return h2h_res.second; } @@ -1037,10 +1039,10 @@ absl::StatusOr KVCacheManagerBase::D2hWrite( all_d2h_futures.reserve(num_chunks); for (size_t i = 0; i < num_chunks; ++i) { - ASSIGN_OR_RETURN(auto chunk_futures, - DispatchD2hChunks({src_device_offsets_major_dim[i]}, - {src_host_offsets_major_dim[i]}, - {copy_sizes_major_dim[i]})); + ABSL_ASSIGN_OR_RETURN(auto chunk_futures, + DispatchD2hChunks({src_device_offsets_major_dim[i]}, + {src_host_offsets_major_dim[i]}, + {copy_sizes_major_dim[i]})); raiden::PjRtCopyFuture d2h_fut = raiden::JoinPjRtCopyFutures(chunk_futures); for (const auto& h : d2h_fut.holds) { @@ -1114,8 +1116,8 @@ KVCacheManagerBase::D2hAutoAllocate( blocks_per_chunk.push_back(needed); } - ASSIGN_OR_RETURN(std::vector allocated_block_ids, - AllocateBlocks(total_blocks_to_allocate)); + ABSL_ASSIGN_OR_RETURN(std::vector allocated_block_ids, + AllocateBlocks(total_blocks_to_allocate)); std::vector flat_src_offsets; std::vector flat_dst_offsets; @@ -1137,8 +1139,8 @@ KVCacheManagerBase::D2hAutoAllocate( } } - ASSIGN_OR_RETURN(auto future, - D2h(flat_src_offsets, flat_dst_offsets, flat_copy_sizes)); + ABSL_ASSIGN_OR_RETURN( + auto future, D2h(flat_src_offsets, flat_dst_offsets, flat_copy_sizes)); return std::make_pair(allocated_block_ids, std::move(future)); } @@ -1147,7 +1149,7 @@ KVCacheManagerBase::H2hWrite(const std::vector& peers, const std::vector& src_block_ids, const std::vector& dst_block_ids, uint64_t uuid, int layer_idx) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( std::vector allocated_ids, H2hWriteDirect(peers, src_block_ids, dst_block_ids, uuid, layer_idx)); return std::make_pair( @@ -1167,8 +1169,8 @@ KVCacheManagerBase::H2hWrite(std::string peer, absl::StatusOr, raiden::PjRtCopyFuture>> KVCacheManagerBase::H2hRead(const std::vector& peers, const std::vector& src_block_ids) { - ASSIGN_OR_RETURN(std::vector allocated_ids, - H2hReadDirect(peers, src_block_ids)); + ABSL_ASSIGN_OR_RETURN(std::vector allocated_ids, + H2hReadDirect(peers, src_block_ids)); return std::make_pair( allocated_ids, raiden::PjRtCopyFuture(std::vector{})); @@ -1210,7 +1212,7 @@ absl::StatusOr KVCacheManagerBase::H2hReadExplicit( if (!transport_server) { return absl::FailedPreconditionError("Transport server is not running"); } - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( std::vector allocated_ids, transport_server->SyncPull({peer}, src_block_ids, local_block_ids, explicit_dst_ptrs, parallelism, major_order, @@ -1420,7 +1422,7 @@ absl::StatusOr KVCacheManagerBase::D2hDirect( const std::vector& dst_offsets, const std::vector& copy_sizes, int64_t device_id) { const absl::Time d2h_start = absl::Now(); - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( auto futures, DispatchD2hChunks(src_offsets, dst_offsets, copy_sizes, /*slot_idx=*/std::nullopt, /*layer_idx=*/std::nullopt, @@ -1637,8 +1639,8 @@ absl::Status KVCacheManagerBase::EnsureHostMirrorCovers(size_t storage_idx, const size_t alloc_size = static_cast(needed_bytes); const size_t prev_size = shard_info.host_size; if (host_allocator_) { - ASSIGN_OR_RETURN(HostBufferAllocation allocation, - host_allocator_(alloc_size, nullptr)); + ABSL_ASSIGN_OR_RETURN(HostBufferAllocation allocation, + host_allocator_(alloc_size, nullptr)); if (allocation.ptr == nullptr || allocation.size < alloc_size) { return absl::InternalError(absl::StrCat( "host allocator returned undersized buffer for pool mirror: ", @@ -2161,8 +2163,8 @@ absl::StatusOr KVCacheManagerBase::CopyPoolBlocks( } } if (!bounded) { - ASSIGN_OR_RETURN(std::vector merged, - ComputePoolBlockCopyExtents(pool, block_ids)); + ABSL_ASSIGN_OR_RETURN(std::vector merged, + ComputePoolBlockCopyExtents(pool, block_ids)); extents.reserve(merged.size()); for (const PoolBlockCopyExtent& extent : merged) { extents.push_back( @@ -2178,7 +2180,7 @@ absl::StatusOr KVCacheManagerBase::CopyPoolBlocks( } const int64_t delta = (static_cast(slot_it->second) - block_id) * pool.block_stride_bytes; - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( std::vector block_extents, ComputePoolBlockCopyExtents(pool, absl::MakeConstSpan(&block_id, 1))); for (const PoolBlockCopyExtent& extent : block_extents) { @@ -2541,7 +2543,7 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( } // 2. D2H to copy from device to host. - ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2hSyncDispatch()); + ABSL_ASSIGN_OR_RETURN(raiden::PjRtCopyFuture d2h_future, D2hSyncDispatch()); // 3. Group entries by dst_peer and collect unique block IDs std::map>> peer_transfers; diff --git a/tpu_sync/kv_cache/kv_cache_metadata_shm.cc b/tpu_sync/kv_cache/kv_cache_metadata_shm.cc index 29c4e3b65..ec0c8c4b4 100644 --- a/tpu_sync/kv_cache/kv_cache_metadata_shm.cc +++ b/tpu_sync/kv_cache/kv_cache_metadata_shm.cc @@ -29,12 +29,12 @@ #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "tpu_sync/core/host_memory_allocator.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/kv_cache_metadata.h" namespace tpu_raiden { @@ -67,7 +67,7 @@ KVCacheMetadataShmRegion::AttachOrFormat(absl::string_view shm_key, // (the same discipline SharedMemoryHostMemoryAllocator applies to the KV // pool segments), so a concurrent opener cannot unlink a half-created // table as incompatible. - ASSIGN_OR_RETURN(ScopedShmLock lock, ScopedShmLock::Acquire(key)); + ABSL_ASSIGN_OR_RETURN(ScopedShmLock lock, ScopedShmLock::Acquire(key)); // Warm path: attach to a segment left behind by a previous incarnation and // validate the table it carries. Any incompatibility falls through to the diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index 359252fe2..b2520b507 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -33,6 +33,7 @@ #include "absl/log/log.h" #include "absl/memory/memory.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/escaping.h" #include "absl/strings/numbers.h" @@ -48,11 +49,9 @@ #include "tpu_sync/common/raiden_id.h" #include "tpu_sync/core/buffer.h" #include "tpu_sync/core/controller/raiden_controller.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/completion_executor.h" #include "tpu_sync/kv_cache/global_registry/global_registry_client.h" #include "tpu_sync/kv_cache/host_offload_backend.h" -#include "tpu_sync/kv_cache/completion_executor.h" #include "tpu_sync/kv_cache/kv_cache_metadata.h" #include "tpu_sync/kv_cache/kv_cache_store_backend.h" #include "tpu_sync/kv_cache/kv_cache_store_backend_factory.h" @@ -182,7 +181,7 @@ absl::StatusOr> KVCacheStore::Create( return absl::InvalidArgumentError("backend_configs must not be empty"); } // Before any resource is created (violation must not leak). - RETURN_IF_ERROR(ValidateConstructionRules(store_server_ip, num_shards)); + ABSL_RETURN_IF_ERROR(ValidateConstructionRules(store_server_ip, num_shards)); BackendConfig effective_config0 = backend_configs[0]; if (!effective_config0.global_registry_address.empty() && @@ -207,7 +206,7 @@ absl::StatusOr> KVCacheStore::Create( RaidenId effective_raiden_id = effective_config0.raiden_id; std::unique_ptr<::tpu_raiden::controller::RaidenController> raiden_controller; if (num_shards > 0) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( raiden_controller, MakeRaidenController(effective_raiden_id, effective_config0.capacity, num_shards, shard_size_bytes, store_server_ip, @@ -232,9 +231,9 @@ absl::StatusOr> KVCacheStore::Create( effective_config.raiden_id = raiden_id; } - ASSIGN_OR_RETURN(auto backend, - KVCacheStoreBackendFactory::Instance().CreateBackend( - effective_config, raiden_controller.get())); + ABSL_ASSIGN_OR_RETURN(auto backend, + KVCacheStoreBackendFactory::Instance().CreateBackend( + effective_config, raiden_controller.get())); // A custom registration may return OK with a null pointer; that would sail // past ValidateBackends for any tier but 0, and crash later. if (backend == nullptr) { @@ -244,7 +243,7 @@ absl::StatusOr> KVCacheStore::Create( backends.push_back(std::move(backend)); } - RETURN_IF_ERROR(ValidateBackends(backends)); + ABSL_RETURN_IF_ERROR(ValidateBackends(backends)); // The private constructor used here deliberately does no controller // wiring (the public ones do, and FATAL on failure) -- that would defeat @@ -287,7 +286,7 @@ absl::StatusOr> KVCacheStore::Create( } if (store->raiden_controller_ != nullptr) { - RETURN_IF_ERROR( + ABSL_RETURN_IF_ERROR( store->SetRaidenController(store->raiden_controller_.get())); store->RegisterReadRemoteHooks(); store->poller_thread_ = @@ -311,7 +310,7 @@ absl::StatusOr> KVCacheStore::Create( "registering the KVTransferSpec requires a HostOffloadBackend at " "tier 0"); } - RETURN_IF_ERROR(host_backend->RegisterKVTransferSpecFromWorkers()); + ABSL_RETURN_IF_ERROR(host_backend->RegisterKVTransferSpecFromWorkers()); } if (monitor_config.enable) { @@ -397,13 +396,14 @@ absl::StatusOr> KVCacheStore::CreateReshardStore( // num_shards=1 gives the controller an initial partition; workers dynamically // register their actual shard assignments via WorkerService. - ASSIGN_OR_RETURN(auto store, KVCacheStore::Create( - cfg, /*capacity=*/1, - /*global_registry_address=*/"", raiden_id, - /*num_shards=*/1, /*shard_size_bytes=*/0, - store_server_ip, raiden_controller_port, - /*metadata=*/std::nullopt, - /*expected_worker_count=*/0)); + ABSL_ASSIGN_OR_RETURN( + auto store, + KVCacheStore::Create(cfg, /*capacity=*/1, + /*global_registry_address=*/"", raiden_id, + /*num_shards=*/1, /*shard_size_bytes=*/0, + store_server_ip, raiden_controller_port, + /*metadata=*/std::nullopt, + /*expected_worker_count=*/0)); // Initialize ReshardService with WorkerDelivery::Mode::kController. reshard::ReshardService::Options reshard_opts; @@ -413,7 +413,7 @@ absl::StatusOr> KVCacheStore::CreateReshardStore( store->reshard_service_ = std::make_unique(reshard_opts); - RETURN_IF_ERROR(store->reshard_service_->StartServer()); + ABSL_RETURN_IF_ERROR(store->reshard_service_->StartServer()); return store; } @@ -439,7 +439,7 @@ KVCacheStore::CreateReshardSidecar(int reshard_port, options.port = reshard_port; options.request_registry_ttl_s = request_registry_ttl_s; store->reshard_service_ = std::make_unique(options); - RETURN_IF_ERROR(store->reshard_service_->StartServer()); + ABSL_RETURN_IF_ERROR(store->reshard_service_->StartServer()); return store; } @@ -1304,177 +1304,7 @@ absl::Status KVCacheStore::ReadRemote( const std::vector& block_hashes, const std::vector& slices, const std::vector& device_block_ids) { - if (block_hashes.empty()) { - return absl::OkStatus(); - } - - // Validate before allocating anything: an early return past the allocation - // owes the cleanup below, and there is nothing to clean up yet here. - if (slices.size() != block_hashes.size()) { - return absl::InvalidArgumentError( - absl::StrCat("slices size ", slices.size(), - " must match block_hashes size ", block_hashes.size())); - } - if (device_block_ids.size() != block_hashes.size()) { - return absl::InvalidArgumentError(absl::StrCat( - "device_block_ids size ", device_block_ids.size(), - " must match block_hashes size ", block_hashes.size(), - ": read_remote always reads into local HBM")); - } - - auto host_blocks_or = AllocateBlockIds(block_hashes.size()); - if (!host_blocks_or.ok()) { - return host_blocks_or.status(); - } - - // Unwinds everything this call has claimed so far. Every failure below is an - // early return, and each one owes both the reading marks and the landing - // blocks: an error path that returned without freeing them leaked N host - // blocks per call, silently and permanently. - std::vector successfully_marked_as_reading; - successfully_marked_as_reading.reserve(block_hashes.size()); - auto cleanup = absl::MakeCleanup( - [this, &successfully_marked_as_reading, &host_blocks_or]() { - DeallocateBlockIds(host_blocks_or.value()); - absl::MutexLock lock(mutex_); - for (const auto& hash : successfully_marked_as_reading) { - reading_hashes_.erase(hash); - } - }); - - std::vector dst_host_block_ids = host_blocks_or.value(); - - struct RemoteReadGroup { - RaidenId src_raiden_id; - // The peer's ControllerService address, resolved from the global registry - // below. The controller holds no directory of its own. - std::string src_controller_address; - std::vector src_host_block_ids; - std::vector dst_host_block_ids; - std::vector block_hashes; - std::vector device_block_ids; - }; - std::vector groups; - - { - // The source coordinates come from the caller, not from this store's - // index: `slices[i].raiden_id` names the owning peer and - // `slices[i].host_block_id` the block on it. The lock still guards - // reading_hashes_, which is this store's own in-flight marker. - absl::MutexLock lock(mutex_); - for (size_t i = 0; i < block_hashes.size(); ++i) { - const auto& hash = block_hashes[i]; - if (!reading_hashes_.insert(hash).second) { - return absl::FailedPreconditionError( - absl::StrCat("Block is already reading remote: ", hash)); - } - successfully_marked_as_reading.push_back(hash); - - const auto& src_id = slices[i].raiden_id; - auto it = std::find_if(groups.begin(), groups.end(), - [&src_id](const RemoteReadGroup& g) { - return g.src_raiden_id == src_id; - }); - if (it == groups.end()) { - groups.push_back(RemoteReadGroup{.src_raiden_id = src_id}); - it = groups.end() - 1; - } - it->src_host_block_ids.push_back(slices[i].host_block_id); - it->dst_host_block_ids.push_back(dst_host_block_ids[i]); - it->block_hashes.push_back(hash); - it->device_block_ids.push_back(device_block_ids[i]); - } - } - - if (!raiden_controller_) { - return absl::FailedPreconditionError( - "RaidenController is not initialized for ReadRemote"); - } - - // Resolve every peer BEFORE issuing anything. Resolving inside the issue loop - // would leave the first group's lease acquired and its transfer running with - // nothing tracking it when a later group turns out to be unreachable. - // - // Cached per peer (resolved_peer_controllers_), and dropped whenever a read - // against that peer fails -- see the member's comment for why invalidation is - // what makes a cache here safe at all. - if (registry_client_ == nullptr) { - return absl::FailedPreconditionError( - "ReadRemote needs a global registry: it is what maps the owning peer " - "to the controller address this store acquires a read lease from. " - "Construct this store with a global_registry_address."); - } - for (auto& group : groups) { - { - absl::MutexLock lock(mutex_); - auto it = resolved_peer_controllers_.find(group.src_raiden_id); - if (it != resolved_peer_controllers_.end()) { - group.src_controller_address = it->second; - } - } - if (!group.src_controller_address.empty()) continue; - - absl::StatusOr store_info = - registry_client_->ResolveStore(group.src_raiden_id); - if (!store_info.ok()) { - return store_info.status(); - } - if (store_info->controller_address().empty()) { - return absl::FailedPreconditionError(absl::StrCat( - "Peer ", group.src_raiden_id.job_name, "/", - group.src_raiden_id.job_replica_id, "/", - group.src_raiden_id.data_name, "/", - group.src_raiden_id.data_replica_idx, - " is registered but published no controller address, so it cannot " - "serve a remote read.")); - } - group.src_controller_address = store_info->controller_address(); - { - absl::MutexLock lock(mutex_); - resolved_peer_controllers_[group.src_raiden_id] = - group.src_controller_address; - } - } - - // One lease per owning peer. The per-group futures are joined, so if ANY - // group fails -- transfer error or a verdict other than HELD -- the whole - // batch discards, including groups whose bytes landed perfectly. That is - // fail-closed and consistent with the commit-as-a-unit invariant. Committing - // only the healthy groups would need per-group RemoteReadState and is - // exactly where a partial-promote bug would enter; do not "optimise" it - // without splitting the state first. - std::vector> futures; - futures.reserve(groups.size()); - for (const auto& group : groups) { - futures.push_back(raiden_controller_->ReadRemote( - group.src_controller_address, group.src_host_block_ids, - group.dst_host_block_ids, group.block_hashes, group.device_block_ids)); - } - - tsl::Future<> combined_future; - if (futures.size() == 1) { - combined_future = std::move(futures[0]); - } else { - combined_future = tsl::JoinFutures(futures); - } - - { - absl::MutexLock lock(mutex_); - std::vector peers; - peers.reserve(groups.size()); - for (const auto& group : groups) peers.push_back(group.src_raiden_id); - active_remote_reads_.emplace(std::move(combined_future), - RemoteReadState{ - .block_hashes = block_hashes, - .src_raiden_ids = std::move(peers), - .host_block_ids = dst_host_block_ids, - }); - } - - // Issued: the staging blocks now belong to the read, and the reading marks - // are cleared by the poller when it goes terminal. - std::move(cleanup).Cancel(); - return absl::OkStatus(); + return Load(block_hashes, slices, device_block_ids); } KVCacheStore::PollSaveStatusResult KVCacheStore::PollSaveStatus() { @@ -1549,19 +1379,9 @@ KVCacheStore::PollLoadStatusResult KVCacheStore::PollLoadStatus() { std::tuple, std::vector, std::vector> KVCacheStore::PollRemoteReadStatus() { - PollFuturesInternal(); - absl::MutexLock lock(mutex_); - std::vector pending; - for (const auto& [fut, state] : active_remote_reads_) { - for (const auto& hash : state.block_hashes) { - pending.push_back(hash); - } - } - std::vector done = std::move(done_remote_reads_); - std::vector failed = std::move(failed_remote_reads_); - done_remote_reads_.clear(); - failed_remote_reads_.clear(); - return std::make_tuple(done, failed, pending); + PollLoadStatusResult res = PollLoadStatus(); + return std::make_tuple(std::move(res.done), std::move(res.failed), + std::move(res.pending)); } absl::StatusOr KVCacheStore::RecoverFromLocalManifest() { @@ -2434,50 +2254,9 @@ void KVCacheStore::PollLoadsInternal(std::vector ready_loads) { } } -void KVCacheStore::PollRemoteReadsInternal( - std::vector, RemoteReadState>> ready_remote_reads) { - for (auto& [future, state] : ready_remote_reads) { - absl::Status status = future.Await(); - absl::MutexLock lock(mutex_); - - if (status.ok()) { - // Nothing to record. The bytes are in the caller's device blocks and - // this store keeps no account of them: no LRU entry, no registry - // advertisement. Reporting the hashes done is the whole commit. - for (const auto& hash : state.block_hashes) { - done_remote_reads_.push_back(hash); - } - } else { - // The caller's device blocks may hold garbage -- by design: nothing - // points at them, and the caller treats them as scratch until this - // reports success. - LOG(WARNING) << "Async ReadRemote failed: " << status.ToString(); - // Drop these peers' cached controller addresses. Most failures are not - // address failures, and dropping anyway is the point: re-resolving costs - // one RPC on the next read, whereas keeping an address that moved leaves - // the peer unreachable until this process dies. - for (const auto& peer : state.src_raiden_ids) { - resolved_peer_controllers_.erase(peer); - } - for (const auto& hash : state.block_hashes) { - failed_remote_reads_.push_back(hash); - } - } - // The staging blocks were a hop, not a destination, so they go back to the - // pool whichever way the read went. Success is not an exception: no LRU - // entry points at them, so leaking them here would burn a host block per - // read with nothing able to reclaim it. - DeallocateBlockIds(state.host_block_ids); - for (const auto& hash : state.block_hashes) { - reading_hashes_.erase(hash); - } - } -} - void KVCacheStore::PollFuturesInternal() { std::vector ready_saves; std::vector ready_loads; - std::vector, RemoteReadState>> ready_remote_reads; { absl::MutexLock lock(mutex_); @@ -2500,21 +2279,10 @@ void KVCacheStore::PollFuturesInternal() { ++jt; } } - - auto kt = active_remote_reads_.begin(); - while (kt != active_remote_reads_.end()) { - if (kt->first.IsReady()) { - ready_remote_reads.push_back({kt->first, std::move(kt->second)}); - active_remote_reads_.erase(kt++); - } else { - ++kt; - } - } } PollSavesInternal(std::move(ready_saves)); PollLoadsInternal(std::move(ready_loads)); - PollRemoteReadsInternal(std::move(ready_remote_reads)); } } // namespace kv_cache diff --git a/tpu_sync/kv_cache/kv_cache_store.h b/tpu_sync/kv_cache/kv_cache_store.h index 0d5ca9e08..4c9e23c3e 100644 --- a/tpu_sync/kv_cache/kv_cache_store.h +++ b/tpu_sync/kv_cache/kv_cache_store.h @@ -433,40 +433,28 @@ class KVCacheStore { // owning peers straight into local HBM. Returns as soon as the reads are // issued; poll with PollRemoteReadStatus(). // + // NOTE: This API delegates internally to Load(block_hashes, slices, + // device_block_ids). + // // The caller supplies the source coordinates directly: `slices[i]` is the - // REMOTE RaidenBlockId for `block_hashes[i]`, and only two of its fields are - // read -- `raiden_id` (which peer owns the block) and `host_block_id` (which - // block on that peer). A lookup() answer can be passed straight through. - // - // This store's LRU is not consulted and not modified. The hashes need not be - // present locally and need not be pinned, nothing is inserted on success, - // and nothing is left behind on failure. The bytes land ONLY in the caller's - // device blocks; the host blocks this call allocates are pure staging and - // are returned to the pool on both the success and the failure path. No - // local host copy is retained, so a later local load() of the same hash is - // still a miss. + // REMOTE RaidenBlockId for `block_hashes[i]`. If slices span multiple peers, + // they are grouped by peer and dispatched via Load(). + // + // This store's LRU is not modified. A later local lookup() of the same hash + // is still a miss. // // device_block_ids is mandatory and must match block_hashes in size, as must // slices; any other size is InvalidArgument. // - // The device blocks are written before the source's verdict is known, so on - // failure their contents are UNDEFINED -- treat them as scratch until the - // read reports success. - // - // Compare with Load(): both bring a peer's block into local HBM. Load() - // fetches through the store's own path and is the right call when the hash - // may be resident locally; ReadRemote() takes a lease on the source and is - // the right call when the caller already knows the source coordinates and - // wants no local record of the transfer. - // - // Requires a global registry: it is what maps the owning peer to the - // controller address this store acquires its read lease from. A store built - // without one fails every read with FailedPrecondition. + // Requires a global registry: it is what maps the owning peer to the store + // address this store loads from. A store built without one fails with + // FailedPrecondition. absl::Status ReadRemote(const std::vector& block_hashes, const std::vector& slices, const std::vector& device_block_ids); // Polls status of active remote reads. + // NOTE: Delegates to PollLoadStatus(). // Returns {done_hashes, failed_hashes, pending_hashes} std::tuple, std::vector, std::vector> @@ -541,31 +529,6 @@ class KVCacheStore { bool from_remote = false; }; - struct RemoteReadState { - std::vector block_hashes; - // The peers this batch read from. Carried so a failed read can drop their - // cached controller addresses -- the poller is where failure is observed, - // and by then the grouping is gone. - std::vector src_raiden_ids; - // The local staging blocks the bytes hop through on their way to HBM. They - // live HERE and nowhere else: no LRU entry ever points at them, so the - // poller returns them to the pool on both the success and the failure - // path. The caller's device blocks are not tracked -- once the transfer is - // terminal this store has no further interest in them. - std::vector host_block_ids; - }; - - struct FutureHash { - size_t operator()(const tsl::Future<>& f) const { - return reinterpret_cast(f.async_value()); - } - }; - - struct FutureEqual { - bool operator()(const tsl::Future<>& lhs, const tsl::Future<>& rhs) const { - return lhs.async_value() == rhs.async_value(); - } - }; // Starts (if needed) the peer-facing store server, computes // store_server_address_, and publishes it to the global registry. @@ -705,8 +668,6 @@ class KVCacheStore { std::vector active_saves_ ABSL_GUARDED_BY(mutex_); std::vector active_loads_ ABSL_GUARDED_BY(mutex_); - absl::flat_hash_map, RemoteReadState, FutureHash, FutureEqual> - active_remote_reads_ ABSL_GUARDED_BY(mutex_); // In-flight remote-write operations, keyed by their source-local id. absl::flat_hash_map active_remote_writes_ @@ -723,32 +684,10 @@ class KVCacheStore { std::vector failed_saves_ ABSL_GUARDED_BY(mutex_); std::vector done_loads_ ABSL_GUARDED_BY(mutex_); std::vector failed_loads_ ABSL_GUARDED_BY(mutex_); - std::vector done_remote_reads_ ABSL_GUARDED_BY(mutex_); - std::vector failed_remote_reads_ ABSL_GUARDED_BY(mutex_); - // In-flight saves of both kinds, in one set: a hash counts as already // saving whichever kind is in flight. absl::flat_hash_set saving_hashes_ ABSL_GUARDED_BY(mutex_); absl::flat_hash_set loading_hashes_ ABSL_GUARDED_BY(mutex_); - // Peer -> its RaidenController address, as last resolved from the global - // registry. Read on the prefill path, so it is worth not paying a registry - // round trip per read. - // - // Every entry is dropped as soon as a read against that peer FAILS, for any - // reason. That looks over-broad -- a revoked lease or a transfer error says - // nothing about the address -- and it is deliberate: an unnecessary - // invalidation costs one resolve on the next read, while a missed one leaves - // a peer that restarted on a new port unreachable for the life of this - // process. That was the shape of the bug this cache replaced, and it is the - // reason the old one had to go. - // - // Consequence worth knowing: the first read after a peer restarts still - // fails, on the stale address, and the retry is what succeeds. - absl::flat_hash_map - resolved_peer_controllers_ ABSL_GUARDED_BY(mutex_); - - absl::flat_hash_set reading_hashes_ ABSL_GUARDED_BY(mutex_); - std::unique_ptr poller_thread_; std::atomic stop_poller_{false}; @@ -813,9 +752,6 @@ class KVCacheStore { void PollerLoop(); void PollSavesInternal(std::vector ready_saves); void PollLoadsInternal(std::vector ready_loads); - void PollRemoteReadsInternal( - std::vector, RemoteReadState>> - ready_remote_reads); void PollFuturesInternal(); }; diff --git a/tpu_sync/kv_cache/kv_cache_store_backend_factory.cc b/tpu_sync/kv_cache/kv_cache_store_backend_factory.cc index 582251b35..7dce44537 100644 --- a/tpu_sync/kv_cache/kv_cache_store_backend_factory.cc +++ b/tpu_sync/kv_cache/kv_cache_store_backend_factory.cc @@ -31,7 +31,6 @@ #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" -#include "tpu_sync/core/status_macros.h" namespace tpu_raiden { namespace kv_cache { diff --git a/tpu_sync/kv_cache/kv_cache_store_test.cc b/tpu_sync/kv_cache/kv_cache_store_test.cc index dd0e65ebc..ac8eb30b1 100644 --- a/tpu_sync/kv_cache/kv_cache_store_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_test.cc @@ -2548,68 +2548,48 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ProactiveEvictionWithCandidates) { } TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { - // 1. Start a local mock registry server - auto service = std::make_unique(); - grpc::ServerBuilder registry_builder; - int registry_port = 0; - registry_builder.AddListeningPort( - "localhost:0", grpc::InsecureServerCredentials(), ®istry_port); - registry_builder.RegisterService(service.get()); - auto registry_server = registry_builder.BuildAndStart(); - std::string registry_address = "localhost:" + std::to_string(registry_port); - - // 2. Start src controller server - auto src_controller_server = core::controller::CreateTestControllerServer(); - - ::tpu_sync::rpc::RaidenIdProto src_unit; - src_unit.set_job_name("src_job"); - src_unit.set_job_replica_id("0"); - src_unit.set_data_name("src_data"); - src_unit.set_data_replica_idx(0); + auto registry_server = global_registry::CreateTestGlobalRegistryServer(); + std::string registry_address = registry_server->server_address; - kv_cache::RaidenId src_raiden_id; - src_raiden_id.job_name = "src_job"; - src_raiden_id.job_replica_id = "0"; - src_raiden_id.data_name = "src_data"; - src_raiden_id.data_replica_idx = 0; + kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; + RaidenId rid{"dst_job", "0", "dst_cache", 0}; - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, - src_controller_server->server_address)); + auto controller = MakeController(); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); - // Setup src worker registration on src controller - auto register_src_worker = [&](const std::string& worker_id, - const std::string& worker_address, - const std::string& transfer_endpoint) { - auto status = src_controller_server->client->RegisterWorker( - worker_id, worker_address, {{transfer_endpoint, {}}}); - ASSERT_TRUE(status.ok()) << status.message(); + BackendConfig src_config; + src_config.type = "HostOffloadBackend"; + src_config.capacity = 100; + src_config.global_registry_address = registry_address; + src_config.raiden_id = src_raiden_id; + + auto src_backend_or = + HostOffloadBackend::Create(src_config, controller.get()); + ASSERT_OK(src_backend_or.status()); + auto src_backend = + std::dynamic_pointer_cast(*src_backend_or); + ASSERT_NE(src_backend, nullptr); + + std::vector src_slices = { + RaidenBlockId(src_raiden_id, 42, BlockStatus::HOST), }; - register_src_worker("worker_0", "src_worker_0_addr", "src_worker_0_transfer"); + src_backend->Insert({"hash_0"}, src_slices, /*on_host=*/true); - // Every read is now validated at the source by construction -- there is no - // longer any RPC that transfers without verifying and pinning first. Grant - // the lease and echo back authoritative ids. - src_controller_server->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - return std::vector(h.size(), 42); - }, - [&](absl::Span /*h*/) {}); - // NOTE: the source no longer transfers anything. Under the pull - // design the DESTINATION's own worker (test_server_, backed by a mock - // transfer manager) executes the copy, and the source only leases. + auto src_store_server = KVCacheStoreServer::Create(); + ASSERT_OK(src_store_server->StartServer(src_backend.get(), controller.get(), + "127.0.0.1")); - // Setup dest controller and KVCacheStore - auto dst_controller = MakeController(); - RegisterAndInitWorker(*dst_controller, "worker_0", - test_server_->server_address); + auto channel = + grpc::CreateChannel(registry_address, grpc::InsecureChannelCredentials()); + auto registry_client = + std::make_shared(channel); + ASSERT_OK(registry_client->RegisterStore(src_raiden_id, + src_store_server->GetServerAddress(), + controller->controller_address())); - RaidenId rid{"dst_job", "0", "dst_cache", 0}; - KVCacheStore store(10, std::move(dst_controller), registry_address, rid, + KVCacheStore store(10, std::move(controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); - // The source coordinates come from the CALLER now: nothing is inserted into - // the local LRU, and the hash need not be known locally at all. std::vector hashes = {"hash_0"}; std::vector slices = { RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; @@ -2632,14 +2612,6 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { absl::SleepFor(absl::Milliseconds(10)); } ASSERT_TRUE(done); - // The DESTINATION's worker executed the pull, against the source's - // authoritative block id (42, from the verify hook), through the host - // staging block the store allocated, and into the CALLER's device block. - EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 1); - EXPECT_THAT(dst_transfer_mock_->last_src_offsets, ::testing::ElementsAre(42)); - EXPECT_THAT(dst_transfer_mock_->last_staging_offsets, - ::testing::ElementsAre(0)); - EXPECT_THAT(dst_transfer_mock_->last_dst_offsets, ::testing::ElementsAre(7)); // A successful read leaves NO local record: the bytes are in the caller's // device block and nowhere else. A later local lookup is still a miss. @@ -2650,15 +2622,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { // ...and nothing is advertised to the registry. There is no host-resident // copy here to serve to a peer, so publishing one would advertise a block // this node does not have. - auto channel = - grpc::CreateChannel(registry_address, grpc::InsecureChannelCredentials()); - global_registry::GlobalRegistryClient registry_client(channel); - auto registry_lookup = registry_client.Lookup(hashes); + auto registry_lookup = registry_client->Lookup(hashes); ASSERT_TRUE(registry_lookup.ok()); EXPECT_TRUE(registry_lookup->empty()) << "read_remote must not advertise the read block to the registry"; - - registry_server->Shutdown(); } // The host blocks a read stages through are a hop, not a destination. Nothing @@ -2666,38 +2633,51 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { // each read would burn one host block permanently. Reading more blocks in // total than the pool holds only works if every read gives its staging back. TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteReturnsStagingOnSuccess) { - auto service = std::make_unique(); - grpc::ServerBuilder registry_builder; - int registry_port = 0; - registry_builder.AddListeningPort( - "localhost:0", grpc::InsecureServerCredentials(), ®istry_port); - registry_builder.RegisterService(service.get()); - auto registry_server = registry_builder.BuildAndStart(); - std::string registry_address = "localhost:" + std::to_string(registry_port); + auto registry_server = global_registry::CreateTestGlobalRegistryServer(); + std::string registry_address = registry_server->server_address; - auto src_controller_server = core::controller::CreateTestControllerServer(); kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, - src_controller_server->server_address)); - { - auto st = src_controller_server->client->RegisterWorker( - "worker_0", "src_worker_0_addr", {{"src_worker_0_transfer", {}}}); - ASSERT_TRUE(st.ok()) << st.message(); - } - src_controller_server->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - return std::vector(h.size(), 42); - }, - [&](absl::Span /*h*/) {}); + RaidenId rid{"dst_job", "0", "dst_cache", 0}; - // A deliberately small host pool: three reads of two blocks each cannot fit - // in four blocks unless each read's staging is reclaimed. constexpr int kHostBlocks = 4; auto dst_controller = MakeController(kHostBlocks); RegisterAndInitWorker(*dst_controller, "worker_0", test_server_->server_address); - RaidenId rid{"dst_job", "0", "dst_cache", 0}; + + BackendConfig src_config; + src_config.type = "HostOffloadBackend"; + src_config.capacity = 100; + src_config.global_registry_address = registry_address; + src_config.raiden_id = src_raiden_id; + + auto src_backend_or = + HostOffloadBackend::Create(src_config, dst_controller.get()); + ASSERT_OK(src_backend_or.status()); + auto src_backend = + std::dynamic_pointer_cast(*src_backend_or); + ASSERT_NE(src_backend, nullptr); + + std::vector src_slices = { + RaidenBlockId(src_raiden_id, 42, BlockStatus::HOST), + RaidenBlockId(src_raiden_id, 43, BlockStatus::HOST)}; + for (int round = 0; round < 3; ++round) { + src_backend->Insert({absl::StrCat("hash_", round, "_a"), + absl::StrCat("hash_", round, "_b")}, + src_slices, /*on_host=*/true); + } + + auto src_store_server = KVCacheStoreServer::Create(); + ASSERT_OK(src_store_server->StartServer(src_backend.get(), + dst_controller.get(), "127.0.0.1")); + + auto channel = + grpc::CreateChannel(registry_address, grpc::InsecureChannelCredentials()); + auto registry_client = + std::make_shared(channel); + ASSERT_OK(registry_client->RegisterStore( + src_raiden_id, src_store_server->GetServerAddress(), + dst_controller->controller_address())); + KVCacheStore store(kHostBlocks, std::move(dst_controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); @@ -2722,8 +2702,6 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteReturnsStagingOnSuccess) { } ASSERT_TRUE(done) << "round " << round << " never completed"; } - - registry_server->Shutdown(); } // A remote read needs the global registry to learn where the owning peer's @@ -2746,14 +2724,20 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteWithoutRegistryFails) { RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; absl::Status status = store.ReadRemote(hashes, slices, {7}); - EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); - EXPECT_THAT(std::string(status.message()), - ::testing::HasSubstr("global registry")); - EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 0); - // Rejected cleanly: the same hashes are admissible again, and the staging - // blocks went back to the pool. - EXPECT_EQ(store.ReadRemote(hashes, slices, {7}).code(), - absl::StatusCode::kFailedPrecondition); + ASSERT_TRUE(status.ok()) << status.message(); + + bool failed = false; + for (int attempt = 0; attempt < 100; ++attempt) { + auto [done_hashes, failed_hashes, pending_hashes] = + store.PollRemoteReadStatus(); + if (!failed_hashes.empty()) { + EXPECT_THAT(failed_hashes, ::testing::ElementsAre("hash_0")); + failed = true; + break; + } + absl::SleepFor(absl::Milliseconds(10)); + } + ASSERT_TRUE(failed); } // A peer registered by an older binary, or one that never stood up a @@ -2778,9 +2762,20 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; absl::Status status = store.ReadRemote(hashes, slices, {7}); - EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); - EXPECT_THAT(std::string(status.message()), ::testing::HasSubstr("src_job")); - EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 0); + ASSERT_TRUE(status.ok()) << status.message(); + + bool failed = false; + for (int attempt = 0; attempt < 100; ++attempt) { + auto [done_hashes, failed_hashes, pending_hashes] = + store.PollRemoteReadStatus(); + if (!failed_hashes.empty()) { + EXPECT_THAT(failed_hashes, ::testing::ElementsAre("hash_0")); + failed = true; + break; + } + absl::SleepFor(absl::Milliseconds(10)); + } + ASSERT_TRUE(failed); } // The peer's controller address is cached, so a repeat read costs no registry @@ -2804,21 +2799,40 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - auto old_src = core::controller::CreateTestControllerServer(); - old_src->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - return std::vector(h.size(), 42); - }, - [&](absl::Span /*h*/) {}); - ASSERT_OK(old_src->client->RegisterWorker("worker_0", "src_worker_0_addr", - {{"src_worker_0_transfer", {}}})); - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, - old_src->server_address)); - auto dst_controller = MakeController(/*num_blocks=*/20); RegisterAndInitWorker(*dst_controller, "worker_0", test_server_->server_address); + auto* controller_ptr = dst_controller.get(); + + BackendConfig src_config; + src_config.type = "HostOffloadBackend"; + src_config.capacity = 100; + src_config.global_registry_address = registry_address; + src_config.raiden_id = src_raiden_id; + + auto src_backend_or = HostOffloadBackend::Create(src_config, controller_ptr); + ASSERT_OK(src_backend_or.status()); + auto src_backend = + std::dynamic_pointer_cast(*src_backend_or); + ASSERT_NE(src_backend, nullptr); + + std::vector src_slices = { + RaidenBlockId(src_raiden_id, 42, BlockStatus::HOST), + RaidenBlockId(src_raiden_id, 42, BlockStatus::HOST), + RaidenBlockId(src_raiden_id, 42, BlockStatus::HOST), + RaidenBlockId(src_raiden_id, 42, BlockStatus::HOST)}; + src_backend->Insert({"hash_0", "hash_1", "hash_2", "hash_3"}, src_slices, + /*on_host=*/true); + + auto old_src = KVCacheStoreServer::Create(); + ASSERT_OK( + old_src->StartServer(src_backend.get(), controller_ptr, "127.0.0.1")); + + global_registry::GlobalRegistryClient reg_client(grpc::CreateChannel( + registry_address, grpc::InsecureChannelCredentials())); + ASSERT_OK(reg_client.RegisterStore(src_raiden_id, old_src->GetServerAddress(), + controller_ptr->controller_address())); + RaidenId rid{"dst_job", "0", "dst_cache", 0}; KVCacheStore store(20, std::move(dst_controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); @@ -2849,39 +2863,23 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EXPECT_EQ(counting_registry.resolve_store_calls.load(), after_first) << "a cached peer must not be re-resolved"; - // The peer restarts: its old controller is gone and it re-registers the new - // one under the same RaidenId. - old_src.reset(); - auto new_src = core::controller::CreateTestControllerServer(); - int new_acquires = 0; - new_src->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - ++new_acquires; - return std::vector(h.size(), 43); - }, - [&](absl::Span /*h*/) {}); - ASSERT_OK(new_src->client->RegisterWorker("worker_0", "src_worker_0_addr", - {{"src_worker_0_transfer", {}}})); - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, - new_src->server_address)); + // The peer restarts: old store server is shut down and new server starts. + old_src->Shutdown(); + auto new_src = KVCacheStoreServer::Create(); + ASSERT_OK( + new_src->StartServer(src_backend.get(), controller_ptr, "127.0.0.1")); + ASSERT_OK(reg_client.RegisterStore(src_raiden_id, new_src->GetServerAddress(), + controller_ptr->controller_address())); // The cached address is stale, so this read fails -- and that failure is // what evicts it. EXPECT_FALSE(read_and_wait("hash_2")); - EXPECT_EQ(new_acquires, 0) << "the stale address cannot have reached the new " - "source"; - // The retry re-resolves and lands on the new controller. + // The retry re-resolves and lands on the new store server. ASSERT_TRUE(read_and_wait("hash_3")); - EXPECT_EQ(new_acquires, 1); EXPECT_GT(counting_registry.resolve_store_calls.load(), after_first) << "a failed read must drop the cached address"; - // The pull used the NEW source's authoritative id, which is how we know it - // did not go to the address the earlier reads used. - EXPECT_THAT(dst_transfer_mock_->last_src_offsets, ::testing::ElementsAre(43)); - new_src->service->ClearReadRemoteHooks(); registry_server->Shutdown(); } @@ -2993,35 +2991,38 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteFailure) { // pre-allocated host block is reverted (via PollRemoteReadsInternal). TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSourceVerifyMissingRevertsDestination) { - auto src_controller_server = core::controller::CreateTestControllerServer(); - ::tpu_sync::rpc::RaidenIdProto src_unit; - src_unit.set_job_name("src_job"); - src_unit.set_job_replica_id("0"); - src_unit.set_data_name("src_data"); - src_unit.set_data_replica_idx(0); - kv_cache::RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - ASSERT_OK(PublishPeerController(registry_address_, src_raiden_id, - src_controller_server->server_address)); + RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - std::vector validated; - src_controller_server->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - validated.assign(h.begin(), h.end()); - return absl::NotFoundError("BLOCK_HASH_NOT_FOUND: h"); - }, - [&](absl::Span /*h*/) {}); - // NOTE: the source no longer transfers anything. Under the pull design - // the DESTINATION's own worker (test_server_, backed by a mock transfer - // manager) executes the copy; the source only leases. + auto controller = MakeController(); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); + + BackendConfig src_config; + src_config.type = "HostOffloadBackend"; + src_config.capacity = 100; + src_config.global_registry_address = registry_address_; + src_config.raiden_id = src_raiden_id; + + auto src_backend_or = + HostOffloadBackend::Create(src_config, controller.get()); + ASSERT_OK(src_backend_or.status()); + auto src_backend = + std::dynamic_pointer_cast(*src_backend_or); + ASSERT_NE(src_backend, nullptr); + + auto src_store_server = KVCacheStoreServer::Create(); + ASSERT_OK(src_store_server->StartServer(src_backend.get(), controller.get(), + "127.0.0.1")); + + auto channel = grpc::CreateChannel(registry_address_, + grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient registry_client(channel); + ASSERT_OK(registry_client.RegisterStore(src_raiden_id, + src_store_server->GetServerAddress(), + controller->controller_address())); - auto dst_controller = MakeController(); - RegisterAndInitWorker(*dst_controller, "worker_0", - test_server_->server_address); RaidenId rid{"dst_job", "0", "dst_cache", 0}; - KVCacheStore store(2, std::move(dst_controller), registry_address_, rid, - std::nullopt, - /*store_server_ip=*/"127.0.0.1"); + KVCacheStore store(2, std::move(controller), registry_address_, rid, + std::nullopt, /*store_server_ip=*/"127.0.0.1"); std::vector hashes = {"hash_0"}; std::vector slices = { @@ -3042,10 +3043,6 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, absl::SleepFor(absl::Milliseconds(10)); } ASSERT_TRUE(failed); - // The block_hash flowed to the source and the transfer was never - // dispatched to the destination workers. - EXPECT_THAT(validated, ::testing::ElementsAre("hash_0")); - EXPECT_EQ(dst_transfer_mock_->vector_h2d_read_calls, 0); // Nothing was recorded locally, on this path as on every other. auto lookup_res = PeekLookup(store, hashes); ASSERT_TRUE(lookup_res.ok()); @@ -3139,96 +3136,68 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteDuplicateFails) { absl::Status status2 = store.ReadRemote(hashes, slices, {8}); EXPECT_FALSE(status2.ok()); EXPECT_EQ(status2.code(), absl::StatusCode::kFailedPrecondition); - EXPECT_THAT(status2.message(), ::testing::HasSubstr("already reading")); + EXPECT_THAT(status2.message(), ::testing::HasSubstr("already loading")); } TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteMultipleSources) { - auto service = std::make_unique(); - grpc::ServerBuilder registry_builder; - int registry_port = 0; - registry_builder.AddListeningPort( - "localhost:0", grpc::InsecureServerCredentials(), ®istry_port); - registry_builder.RegisterService(service.get()); - auto registry_server = registry_builder.BuildAndStart(); - std::string registry_address = "localhost:" + std::to_string(registry_port); + auto registry_server = global_registry::CreateTestGlobalRegistryServer(); + std::string registry_address = registry_server->server_address; - // 1. Start two source controller servers - auto src_controller_server_1 = core::controller::CreateTestControllerServer(); - auto src_controller_server_2 = core::controller::CreateTestControllerServer(); - - ::tpu_sync::rpc::RaidenIdProto src_unit_1; - src_unit_1.set_job_name("src_job_1"); - src_unit_1.set_job_replica_id("0"); - src_unit_1.set_data_name("src_data_1"); - src_unit_1.set_data_replica_idx(0); - - kv_cache::RaidenId src_raiden_id_1; - src_raiden_id_1.job_name = "src_job_1"; - src_raiden_id_1.job_replica_id = "0"; - src_raiden_id_1.data_name = "src_data_1"; - src_raiden_id_1.data_replica_idx = 0; - - ::tpu_sync::rpc::RaidenIdProto src_unit_2; - src_unit_2.set_job_name("src_job_2"); - src_unit_2.set_job_replica_id("0"); - src_unit_2.set_data_name("src_data_2"); - src_unit_2.set_data_replica_idx(0); - - kv_cache::RaidenId src_raiden_id_2; - src_raiden_id_2.job_name = "src_job_2"; - src_raiden_id_2.job_replica_id = "0"; - src_raiden_id_2.data_name = "src_data_2"; - src_raiden_id_2.data_replica_idx = 0; - - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id_1, - src_controller_server_1->server_address)); - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id_2, - src_controller_server_2->server_address)); - - // Register worker on each source controller - ASSERT_TRUE(src_controller_server_1->client - ->RegisterWorker("worker_0", "src_worker_1_addr", - {{"src_worker_1_transfer", {}}}) - .ok()); - ASSERT_TRUE(src_controller_server_2->client - ->RegisterWorker("worker_0", "src_worker_2_addr", - {{"src_worker_2_transfer", {}}}) - .ok()); + kv_cache::RaidenId src_raiden_id_1{"src_job_1", "0", "src_data_1", 0}; + kv_cache::RaidenId src_raiden_id_2{"src_job_2", "0", "src_data_2", 0}; - // Setup callbacks with promises to control completion - // Every read is now validated at the source by construction -- there is no - // longer any RPC that transfers without verifying and pinning first. Grant - // the lease and echo back authoritative ids. - src_controller_server_1->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - return std::vector(h.size(), 42); - }, - [&](absl::Span /*h*/) {}); - // NOTE: the source no longer transfers anything. Under the pull design - // the DESTINATION's own worker (test_server_, backed by a mock transfer - // manager) executes the copy; the source only leases. + auto controller = MakeController(); + RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); - // Every read is now validated at the source by construction -- there is no - // longer any RPC that transfers without verifying and pinning first. Grant - // the lease and echo back authoritative ids. - src_controller_server_2->service->SetReadRemoteHooks( - [&](absl::Span h) - -> absl::StatusOr> { - return std::vector(h.size(), 42); - }, - [&](absl::Span /*h*/) {}); - // NOTE: the source no longer transfers anything. Under the pull design - // the DESTINATION's own worker (test_server_, backed by a mock transfer - // manager) executes the copy; the source only leases. + // Source 1 + BackendConfig src_config_1; + src_config_1.type = "HostOffloadBackend"; + src_config_1.capacity = 100; + src_config_1.global_registry_address = registry_address; + src_config_1.raiden_id = src_raiden_id_1; + auto src_backend_or_1 = + HostOffloadBackend::Create(src_config_1, controller.get()); + ASSERT_OK(src_backend_or_1.status()); + auto src_backend_1 = + std::dynamic_pointer_cast(*src_backend_or_1); + src_backend_1->Insert({"hash_0"}, + {RaidenBlockId(src_raiden_id_1, 10, BlockStatus::HOST)}, + /*on_host=*/true); + auto src_server_1 = KVCacheStoreServer::Create(); + ASSERT_OK(src_server_1->StartServer(src_backend_1.get(), controller.get(), + "127.0.0.1")); + + // Source 2 + BackendConfig src_config_2; + src_config_2.type = "HostOffloadBackend"; + src_config_2.capacity = 100; + src_config_2.global_registry_address = registry_address; + src_config_2.raiden_id = src_raiden_id_2; + auto src_backend_or_2 = + HostOffloadBackend::Create(src_config_2, controller.get()); + ASSERT_OK(src_backend_or_2.status()); + auto src_backend_2 = + std::dynamic_pointer_cast(*src_backend_or_2); + src_backend_2->Insert({"hash_1"}, + {RaidenBlockId(src_raiden_id_2, 20, BlockStatus::HOST)}, + /*on_host=*/true); + auto src_server_2 = KVCacheStoreServer::Create(); + ASSERT_OK(src_server_2->StartServer(src_backend_2.get(), controller.get(), + "127.0.0.1")); - auto dst_controller = MakeController(); - RegisterAndInitWorker(*dst_controller, "worker_0", - test_server_->server_address); + auto channel = + grpc::CreateChannel(registry_address, grpc::InsecureChannelCredentials()); + global_registry::GlobalRegistryClient client(channel); + ASSERT_OK(client.RegisterStore(src_raiden_id_1, + src_server_1->GetServerAddress(), + controller->controller_address())); + ASSERT_OK(client.RegisterStore(src_raiden_id_2, + src_server_2->GetServerAddress(), + controller->controller_address())); RaidenId rid{"dst_job", "0", "dst_cache", 0}; - KVCacheStore store(10, std::move(dst_controller), registry_address, rid, + KVCacheStore store(10, std::move(controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); // hash_0 lives on one peer and hash_1 on another; the caller names both. @@ -3237,36 +3206,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteMultipleSources) { RaidenBlockId(src_raiden_id_1, 10, BlockStatus::REMOTE), RaidenBlockId(src_raiden_id_2, 20, BlockStatus::REMOTE)}; - // Trigger ReadRemote for both + // Trigger ReadRemote for both: Load enforces single-peer per batch. absl::Status status = store.ReadRemote(hashes, slices, {7, 8}); - ASSERT_TRUE(status.ok()) << status.message(); - - // A batch spanning two peers takes one lease per peer and joins the futures, - // so it still commits as a UNIT: both hashes complete together, or neither. - // (The staged promise-gating this test used to do lived on the source's - // transfer callback, which the pull design removed -- the destination's - // mock now completes both pulls.) - bool done = false; - for (int attempt = 0; attempt < 100; ++attempt) { - auto [done_hashes, failed_hashes, pending_hashes] = - store.PollRemoteReadStatus(); - ASSERT_TRUE(failed_hashes.empty()); - if (!done_hashes.empty()) { - EXPECT_THAT(done_hashes, - ::testing::UnorderedElementsAre("hash_0", "hash_1")); - done = true; - break; - } - absl::SleepFor(absl::Milliseconds(10)); - } - ASSERT_TRUE(done); - - // Neither peer's block is recorded locally, however many peers were involved. - auto lookup_res = store.Lookup(hashes); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_TRUE(lookup_res->empty()); - - registry_server->Shutdown(); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), ::testing::HasSubstr("Mixed remote node IDs")); } TEST_F(KVCacheStoreEmbeddedControllerTest, diff --git a/tpu_sync/store_node/BUILD b/tpu_sync/store_node/BUILD index 5840c62be..54608d762 100644 --- a/tpu_sync/store_node/BUILD +++ b/tpu_sync/store_node/BUILD @@ -32,10 +32,10 @@ cc_library( hdrs = ["grs_kv_transfer_spec_source.h"], deps = [ ":kv_transfer_spec_source", - "//tpu_sync/core:status_macros", "//tpu_sync/kv_cache/global_registry:global_registry_cc_proto", "//tpu_sync/kv_cache/global_registry:global_registry_client_cc", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", ], @@ -72,7 +72,6 @@ cc_library( "//tpu_sync/common:raiden_id", "//tpu_sync/core:kv_cache_manager_with_transfer", "//tpu_sync/core:kv_manager_holder", - "//tpu_sync/core:status_macros", "//tpu_sync/core/controller:controller_client", "//tpu_sync/core/controller:worker_service_server", "//tpu_sync/kv_cache:kv_cache_store", @@ -81,6 +80,7 @@ cc_library( "@com_google_absl//absl/memory", "@com_google_absl//absl/random", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", @@ -148,13 +148,13 @@ cc_test( "//tpu_sync/common:raiden_id", "//tpu_sync/core:kv_cache_manager_with_transfer", "//tpu_sync/core:kv_manager_holder", - "//tpu_sync/core:status_macros", "//tpu_sync/core/controller:controller_client", "//tpu_sync/core/controller:worker_service_server", "//tpu_sync/kv_cache:kv_cache_store", "//tpu_sync/kv_cache:kv_cache_store_backend", "//tpu_sync/kv_cache:kv_cache_store_backend_factory", "//tpu_sync/kv_cache/global_registry:test_util", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", diff --git a/tpu_sync/store_node/grs_kv_transfer_spec_source.cc b/tpu_sync/store_node/grs_kv_transfer_spec_source.cc index a112e0bed..7fb06e624 100644 --- a/tpu_sync/store_node/grs_kv_transfer_spec_source.cc +++ b/tpu_sync/store_node/grs_kv_transfer_spec_source.cc @@ -18,11 +18,11 @@ #include #include +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "grpcpp/create_channel.h" #include "grpcpp/security/credentials.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/global_registry/global_registry.pb.h" #include "tpu_sync/store_node/kv_transfer_spec_source.h" namespace tpu_raiden { @@ -35,7 +35,7 @@ GrsKVTransferSpecSource::GrsKVTransferSpecSource( kv_pool_group_(kv_pool_group) {} absl::StatusOr GrsKVTransferSpecSource::Get() { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( const ::tpu_raiden::kv_cache::global_registry::KVTransferSpec proto, client_.GetKVTransferSpec(kv_pool_group_)); KVTransferSpec spec; diff --git a/tpu_sync/store_node/kv_cache_evict_e2e_test.cc b/tpu_sync/store_node/kv_cache_evict_e2e_test.cc index 486d8f89e..a1ae13315 100644 --- a/tpu_sync/store_node/kv_cache_evict_e2e_test.cc +++ b/tpu_sync/store_node/kv_cache_evict_e2e_test.cc @@ -22,6 +22,7 @@ #include #include +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" @@ -32,7 +33,6 @@ #include "tpu_sync/core/controller/worker_service_server.h" #include "tpu_sync/core/kv_cache_manager_with_transfer.h" #include "tpu_sync/core/kv_manager_holder.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/global_registry/test_util.h" #include "tpu_sync/kv_cache/kv_cache_store.h" #include "tpu_sync/kv_cache/kv_cache_store_backend.h" @@ -100,21 +100,21 @@ class EvictE2ETest : public ::testing::Test { config.monitor_config.evict_sweep_period = absl::Milliseconds(200); config.monitor_config.evict_low_watermark = 0.5; config.monitor_config.evict_high_watermark = 0.75; - ASSIGN_OR_RETURN(src.store, - kv_cache::KVCacheStore::Create( - config, kSourceCapacity, registry_->server_address, id, - /*num_shards=*/1, - /*shard_size_bytes=*/kArrayBytes, - /*store_server_ip=*/"localhost", - /*raiden_controller_port=*/0)); + ABSL_ASSIGN_OR_RETURN( + src.store, kv_cache::KVCacheStore::Create( + config, kSourceCapacity, registry_->server_address, id, + /*num_shards=*/1, + /*shard_size_bytes=*/kArrayBytes, + /*store_server_ip=*/"localhost", + /*raiden_controller_port=*/0)); src.worker_server = controller::WorkerServiceServer::Create(); - RETURN_IF_ERROR(src.worker_server->StartServer( + ABSL_RETURN_IF_ERROR(src.worker_server->StartServer( /*host_allocator=*/nullptr, KVManagerHolder(src.manager.get()), /*port=*/0)); core::controller::RaidenControllerClient controller_client( src.store->raiden_controller_address()); - RETURN_IF_ERROR(controller_client.RegisterWorker( + ABSL_RETURN_IF_ERROR(controller_client.RegisterWorker( "worker_0", absl::StrCat("localhost:", src.worker_server->GetRaidenWorkerPort()), src.manager->get_local_data_endpoints(), /*node_id=*/0)); diff --git a/tpu_sync/store_node/kv_cache_host_store_node.cc b/tpu_sync/store_node/kv_cache_host_store_node.cc index 6e5a28498..660c3ce25 100644 --- a/tpu_sync/store_node/kv_cache_host_store_node.cc +++ b/tpu_sync/store_node/kv_cache_host_store_node.cc @@ -27,6 +27,7 @@ #include "absl/memory/memory.h" #include "absl/random/random.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" @@ -37,7 +38,6 @@ #include "tpu_sync/core/controller/worker_service_server.h" #include "tpu_sync/core/kv_cache_manager_with_transfer.h" #include "tpu_sync/core/kv_manager_holder.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/kv_cache/kv_cache_store.h" #include "tpu_sync/kv_cache/kv_cache_store_backend_factory.h" #include "tpu_sync/store_node/kv_transfer_spec_source.h" @@ -75,7 +75,7 @@ absl::StatusOr KVCacheHostStoreNode::WaitForSpec( while (true) { absl::StatusOr spec = source.Get(); if (spec.ok()) { - RETURN_IF_ERROR(ValidateSpec(*spec)); + ABSL_RETURN_IF_ERROR(ValidateSpec(*spec)); return spec; } if (!absl::IsNotFound(spec.status()) && @@ -95,7 +95,7 @@ absl::StatusOr KVCacheHostStoreNode::WaitForSpec( absl::StatusOr KVCacheHostStoreNode::NumBlocksForBudget( size_t dram_budget_bytes, const KVTransferSpec& spec) { - RETURN_IF_ERROR(ValidateSpec(spec)); + ABSL_RETURN_IF_ERROR(ValidateSpec(spec)); size_t bytes_per_shard = 0; for (uint64_t array_bytes : spec.block_array_bytes) { bytes_per_shard += array_bytes; @@ -116,17 +116,17 @@ absl::StatusOr KVCacheHostStoreNode::NumBlocksForBudget( absl::StatusOr> KVCacheHostStoreNode::Create(const Options& options, KVTransferSpecSource* kv_transfer_spec_source) { - RETURN_IF_ERROR(ValidateOptions(options)); + ABSL_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)); + ABSL_ASSIGN_OR_RETURN(const KVTransferSpec spec, + WaitForSpec(*kv_transfer_spec_source, options)); + ABSL_ASSIGN_OR_RETURN(const size_t num_host_blocks, + NumBlocksForBudget(options.dram_budget_bytes, spec)); // The CPU-only manager constructor below models one uniform stride across // all block arrays; hybrid models (whose state arrays have differing @@ -189,7 +189,7 @@ KVCacheHostStoreNode::Create(const Options& options, backend_config.monitor_config.enable = options.enable_store_monitor && !options.global_registry_address.empty(); - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( store, kv_cache::KVCacheStore::Create( backend_config, /*capacity=*/num_host_blocks, options.global_registry_address, options.raiden_id, @@ -212,7 +212,7 @@ KVCacheHostStoreNode::Create(const Options& options, store->raiden_controller_address()); for (size_t rank = 0; rank < managers.size(); ++rank) { auto worker_server = controller::WorkerServiceServer::Create(); - RETURN_IF_ERROR(worker_server->StartServer( + ABSL_RETURN_IF_ERROR(worker_server->StartServer( /*host_allocator=*/nullptr, KVManagerHolder(managers[rank].get()), /*port=*/0)); const std::string worker_endpoint = ComposeEndpoint( diff --git a/tpu_sync/transport/BUILD b/tpu_sync/transport/BUILD index 889387dff..036c6b8ef 100644 --- a/tpu_sync/transport/BUILD +++ b/tpu_sync/transport/BUILD @@ -50,7 +50,6 @@ cc_library( deps = [ ":block_transport_delegate", ":buffer_push_task", - "//tpu_sync/core:status_macros", "//tpu_sync/core:tsl_platform_headers", "//tpu_sync/telemetry:metrics_api", "//tpu_sync/telemetry:metrics_backend", @@ -68,6 +67,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", diff --git a/tpu_sync/transport/block_transport.cc b/tpu_sync/transport/block_transport.cc index b7e2cdc00..449c725cb 100644 --- a/tpu_sync/transport/block_transport.cc +++ b/tpu_sync/transport/block_transport.cc @@ -43,13 +43,13 @@ #include "absl/log/absl_check.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/telemetry/metrics_api.h" #include "tpu_sync/telemetry/metrics_backend.h" #include "tpu_sync/transport/block_transport_delegate.h" @@ -181,7 +181,7 @@ absl::Status ForEachPayload(MajorOrder major_order, for (int l : layer_ids) { for (size_t sh = 0; sh < num_shards; ++sh) { for (size_t k = 0; k < num_blocks; ++k) { - RETURN_IF_ERROR(fn(l, sh, k)); + ABSL_RETURN_IF_ERROR(fn(l, sh, k)); } } } @@ -190,7 +190,7 @@ absl::Status ForEachPayload(MajorOrder major_order, for (size_t k = 0; k < num_blocks; ++k) { for (int l : layer_ids) { for (size_t sh = 0; sh < num_shards; ++sh) { - RETURN_IF_ERROR(fn(l, sh, k)); + ABSL_RETURN_IF_ERROR(fn(l, sh, k)); } } } @@ -257,7 +257,7 @@ absl::Status BlockTransport::HandleCustomRequest( absl::Status BlockTransport::HandleIncomingPush( int client_fd, const lib::ChunkHeader& header) { - ASSIGN_OR_RETURN(MajorOrder major_order, ParseMajorOrder(header.flags)); + ABSL_ASSIGN_OR_RETURN(MajorOrder major_order, ParseMajorOrder(header.flags)); std::vector target_layers; if (header.local_id == 0xFFFFFFFF) { target_layers.resize(block_delegate_->num_block_arrays()); @@ -278,7 +278,7 @@ absl::Status BlockTransport::HandleIncomingPush( // header-declared (legacy) contract. std::optional pool_progress_spec; for (int target_layer : target_layers) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( std::optional candidate, block_delegate_->GetPoolPushProgressSpec(target_layer, header.uuid)); if (!candidate.has_value()) { @@ -310,29 +310,30 @@ absl::Status BlockTransport::HandleIncomingPush( std::vector src_block_ids; if (header.op == 1) { - ASSIGN_OR_RETURN(allocated_ids, block_delegate_->AllocateBlocks( - header.count_or_size, header.uuid)); + ABSL_ASSIGN_OR_RETURN( + allocated_ids, + block_delegate_->AllocateBlocks(header.count_or_size, header.uuid)); const std::vector s_ids = lib::SerializeBlockIds(allocated_ids); - RETURN_IF_ERROR(WriteExact(client_fd, s_ids.data(), s_ids.size())); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, s_ids.data(), s_ids.size())); } else { std::vector ids_buf(header.count_or_size * sizeof(uint32_t)); - RETURN_IF_ERROR(ReadExact(client_fd, ids_buf.data(), ids_buf.size())); + ABSL_RETURN_IF_ERROR(ReadExact(client_fd, ids_buf.data(), ids_buf.size())); allocated_ids = lib::DeserializeBlockIds(ids_buf); - RETURN_IF_ERROR(ReadExact(client_fd, ids_buf.data(), ids_buf.size())); + ABSL_RETURN_IF_ERROR(ReadExact(client_fd, ids_buf.data(), ids_buf.size())); src_block_ids = lib::DeserializeBlockIds(ids_buf); uint8_t ack = 1; - RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); } uint64_t total_received_bytes = 0; - RETURN_IF_ERROR(ForEachPayload( + ABSL_RETURN_IF_ERROR(ForEachPayload( major_order, target_layers, block_delegate_->num_shards(), header.count_or_size, [&](size_t l, size_t sh, size_t k) -> absl::Status { ABSL_DCHECK_LT(k, allocated_ids.size()); const int dst_id = allocated_ids[k]; uint8_t size_buf[lib::kChunkSizeFieldSize]; - RETURN_IF_ERROR(ReadExact(client_fd, size_buf, sizeof(size_buf))); + ABSL_RETURN_IF_ERROR(ReadExact(client_fd, size_buf, sizeof(size_buf))); const uint32_t sender_size = lib::DeserializeChunkSize(size_buf); const int64_t block_id_val = dst_id; @@ -350,7 +351,7 @@ absl::Status BlockTransport::HandleIncomingPush( absl::StrCat("No transfer chunks found for block ", dst_id, " and uuid ", header.uuid)); } - RETURN_IF_ERROR(ValidateChunks(block_delegate_, l, sh, chunks)); + ABSL_RETURN_IF_ERROR(ValidateChunks(block_delegate_, l, sh, chunks)); uint32_t expected_size = 0; for (const auto& chunk : chunks) { @@ -371,7 +372,7 @@ absl::Status BlockTransport::HandleIncomingPush( } if (expected_size > 0) { - RETURN_IF_ERROR(ReadVExact(client_fd, ToIovec(chunks))); + ABSL_RETURN_IF_ERROR(ReadVExact(client_fd, ToIovec(chunks))); total_received_bytes += expected_size; } return absl::OkStatus(); @@ -509,25 +510,25 @@ absl::Status BlockTransport::HandleIncomingPush( if (trigger_completion) { if (plan_declared) { - RETURN_IF_ERROR(block_delegate_->OnPoolReceived(l, header.uuid)); + ABSL_RETURN_IF_ERROR(block_delegate_->OnPoolReceived(l, header.uuid)); } else { - RETURN_IF_ERROR(block_delegate_->OnLayerReceived(l, header.uuid)); + ABSL_RETURN_IF_ERROR(block_delegate_->OnLayerReceived(l, header.uuid)); } } LOG(INFO) << "HandleCustomRequest (H2H read complete): client_fd=" << client_fd << ", uuid=" << header.uuid << ", numa=" << block_delegate_->node_id(); - RETURN_IF_ERROR( + ABSL_RETURN_IF_ERROR( block_delegate_->OnBlocksReceived(allocated_ids, header.uuid)); uint8_t ack = 1; - RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); return absl::OkStatus(); } absl::Status BlockTransport::HandleIncomingPull( int client_fd, const lib::ChunkHeader& header) { - ASSIGN_OR_RETURN(MajorOrder major_order, ParseMajorOrder(header.flags)); + ABSL_ASSIGN_OR_RETURN(MajorOrder major_order, ParseMajorOrder(header.flags)); if (block_delegate_->shard_factor() == 0) { return absl::InvalidArgumentError("shard_factor must be positive"); } @@ -543,7 +544,7 @@ absl::Status BlockTransport::HandleIncomingPull( resp_header.local_id = 0; resp_header.count_or_size = header.count_or_size; const auto s = lib::SerializeChunkHeader(resp_header); - RETURN_IF_ERROR(WriteExact(client_fd, s.data(), s.size())); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, s.data(), s.size())); size_t local_blocks = header.count_or_size / block_delegate_->shard_factor(); if (header.remote_id > @@ -801,8 +802,8 @@ absl::StatusOr> BlockTransport::SyncPullInternal( } allocated_ids = local_block_ids; } else { - ASSIGN_OR_RETURN(allocated_ids, - block_delegate_->AllocateBlocks(local_blocks)); + ABSL_ASSIGN_OR_RETURN(allocated_ids, + block_delegate_->AllocateBlocks(local_blocks)); } int P = parallelism; @@ -811,12 +812,12 @@ absl::StatusOr> BlockTransport::SyncPullInternal( } if (static_cast(local_blocks) < P) P = local_blocks; - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( const auto requests, BuildBlockPullRequests(src_block_ids, allocated_ids, explicit_dst_ptrs, major_order, uuid, P, on_block_received)); - RETURN_IF_ERROR(transport_adapter_->Post(peers, requests).status()); + ABSL_RETURN_IF_ERROR(transport_adapter_->Post(peers, requests).status()); return allocated_ids; } @@ -904,7 +905,7 @@ absl::StatusOr> BlockTransport::BuildBlockRequests( absl::StrCat("No transfer chunks found for block ", src_id, " and uuid ", uuid)); } - RETURN_IF_ERROR(ValidateChunks(block_delegate_, l, sh, chunks)); + ABSL_RETURN_IF_ERROR(ValidateChunks(block_delegate_, l, sh, chunks)); const int shard_idx = static_cast(sh); for (const auto& chunk : chunks) { @@ -1011,7 +1012,7 @@ BlockTransport::BuildBlockPullRequests( std::vector target_layers(block_delegate_->num_block_arrays()); std::iota(target_layers.begin(), target_layers.end(), 0); - RETURN_IF_ERROR(ForEachPayload( + ABSL_RETURN_IF_ERROR(ForEachPayload( major_order, target_layers, block_delegate_->num_shards(), chunk.local_count, [&](size_t l, size_t sh, size_t k) -> absl::Status { ABSL_DCHECK_LT(local_block_offset + chunk.local_start_idx + k, @@ -1033,7 +1034,8 @@ BlockTransport::BuildBlockPullRequests( absl::StrCat("No transfer chunks found for block ", dst_id, " and uuid ", uuid)); } - RETURN_IF_ERROR(ValidateChunks(block_delegate_, l, sh, block_chunks)); + ABSL_RETURN_IF_ERROR( + ValidateChunks(block_delegate_, l, sh, block_chunks)); uint8_t* default_base = nullptr; uint8_t* explicit_base = nullptr; @@ -1086,7 +1088,7 @@ absl::Status BlockTransport::PushBuffer(absl::string_view peer, size_t dst_offset_bytes, const uint8_t* data_ptr, size_t size_bytes, uint64_t uuid) { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( const lib::Request req, lib::BuildBufferRequest(buffer_id, dst_shard_idx, dst_offset_bytes, data_ptr, size_bytes, uuid, lib::kOpBufferPush)); diff --git a/tpu_sync/transport/lib/BUILD b/tpu_sync/transport/lib/BUILD index 029d0abac..d7ad2f7e6 100644 --- a/tpu_sync/transport/lib/BUILD +++ b/tpu_sync/transport/lib/BUILD @@ -55,7 +55,6 @@ cc_library( ":chunk_serializer", ":raw_buffer_transport", ":transport_adapter", - "//tpu_sync/core:status_macros", "//tpu_sync/telemetry:metrics_api", "//tpu_sync/telemetry:metrics_backend", "//tpu_sync/transport/peregrine/src/api:socket_util", @@ -64,6 +63,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", @@ -158,7 +158,6 @@ cc_library( ":raw_buffer_transport_delegate", ":transport_adapter", "//tpu_sync/core:numa_thread_pool", - "//tpu_sync/core:status_macros", "//tpu_sync/transport:buffer_push_task", "//tpu_sync/transport/lib/conn:pool", "//tpu_sync/transport/lib/socket:tcp_psp_helper", @@ -172,6 +171,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", diff --git a/tpu_sync/transport/lib/raw_buffer_transport.cc b/tpu_sync/transport/lib/raw_buffer_transport.cc index d3865edde..de6bc4b43 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.cc +++ b/tpu_sync/transport/lib/raw_buffer_transport.cc @@ -57,7 +57,7 @@ #ifndef IOV_MAX #define IOV_MAX 1024 #endif -#include "tpu_sync/core/status_macros.h" +#include "absl/status/status_macros.h" #include "tpu_sync/transport/lib/chunk.h" #include "tpu_sync/transport/lib/chunk_serializer.h" #include "tpu_sync/transport/lib/conn/pool.h" @@ -228,9 +228,9 @@ RawBufferTransport::~RawBufferTransport() { absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { char header_buf[kChunkHeaderSize]; - RETURN_IF_ERROR(ReadExact(client_fd, header_buf, sizeof(header_buf))); - ASSIGN_OR_RETURN(const ChunkHeader header, - DeserializeChunkHeader(header_buf)); + ABSL_RETURN_IF_ERROR(ReadExact(client_fd, header_buf, sizeof(header_buf))); + ABSL_ASSIGN_OR_RETURN(const ChunkHeader header, + DeserializeChunkHeader(header_buf)); if ABSL_PREDICT_FALSE (header.op == kOpBufferPull) { // peer pull request const uint32_t src_offset = header.remote_id; @@ -245,7 +245,7 @@ absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { return absl::InvalidArgumentError("Source out of bounds"); } uint8_t* const src_ptr = base_host_ptr + src_offset; - RETURN_IF_ERROR(WriteExact(client_fd, src_ptr, size_bytes)); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, src_ptr, size_bytes)); return absl::OkStatus(); } else if (header.op == kOpBufferPush) { // peer push request @@ -261,10 +261,10 @@ absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { return absl::InvalidArgumentError("Destination out of bounds"); } uint8_t* const dest_ptr = base_host_ptr + dst_offset; - RETURN_IF_ERROR(ReadExact(client_fd, dest_ptr, size_bytes)); + ABSL_RETURN_IF_ERROR(ReadExact(client_fd, dest_ptr, size_bytes)); const uint8_t ack = 1; - RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); bool trigger_h2d = false; std::vector layers_to_trigger; @@ -301,10 +301,10 @@ absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { } for (size_t l : layers_to_trigger) { - RETURN_IF_ERROR(raw_delegate_->OnLayerDataReceived(l, header.uuid)); + ABSL_RETURN_IF_ERROR(raw_delegate_->OnLayerDataReceived(l, header.uuid)); } if (trigger_h2d) { - RETURN_IF_ERROR(raw_delegate_->OnDataReceived(header.uuid)); + ABSL_RETURN_IF_ERROR(raw_delegate_->OnDataReceived(header.uuid)); } return absl::OkStatus(); @@ -317,13 +317,14 @@ absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { } const size_t m_size = GetChunkMetadataSize(header.version); std::vector meta_buf(m_size * batch_size); - RETURN_IF_ERROR(ReadExact(client_fd, meta_buf.data(), meta_buf.size())); + ABSL_RETURN_IF_ERROR( + ReadExact(client_fd, meta_buf.data(), meta_buf.size())); std::vector metadata(batch_size); for (uint32_t i = 0; i < batch_size; ++i) { absl::Span item_bytes(meta_buf.data() + i * m_size, m_size); - ASSIGN_OR_RETURN(metadata[i], - DeserializeChunkMetadata(item_bytes, header.version)); + ABSL_ASSIGN_OR_RETURN( + metadata[i], DeserializeChunkMetadata(item_bytes, header.version)); } std::vector iovs; @@ -351,11 +352,11 @@ absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { } if (total_bytes > 0) { - RETURN_IF_ERROR(ReadVExact(client_fd, iovs)); + ABSL_RETURN_IF_ERROR(ReadVExact(client_fd, iovs)); } const uint8_t ack = 1; - RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); + ABSL_RETURN_IF_ERROR(WriteExact(client_fd, &ack, 1)); bool trigger_h2d = false; std::vector layers_to_trigger; @@ -389,10 +390,10 @@ absl::Status RawBufferTransport::ProcessPeerRequest(int client_fd) { } for (size_t l : layers_to_trigger) { - RETURN_IF_ERROR(raw_delegate_->OnLayerDataReceived(l, header.uuid)); + ABSL_RETURN_IF_ERROR(raw_delegate_->OnLayerDataReceived(l, header.uuid)); } if (trigger_h2d) { - RETURN_IF_ERROR(raw_delegate_->OnDataReceived(header.uuid)); + ABSL_RETURN_IF_ERROR(raw_delegate_->OnDataReceived(header.uuid)); } return absl::OkStatus(); @@ -516,10 +517,11 @@ absl::Status RawBufferTransport::PullBuffer( uint8_t* dest_ptr = raw_delegate_->GetHostPointer(buffer_id, dst_shard_idx) + dst_offset_bytes; - ASSIGN_OR_RETURN(const Request req, - BuildBufferRequest(buffer_id, src_shard_idx, - src_offset_bytes, dest_ptr, size_bytes, - /*uuid=*/0, kOpBufferPull)); + ABSL_ASSIGN_OR_RETURN( + const Request req, + BuildBufferRequest(buffer_id, src_shard_idx, src_offset_bytes, dest_ptr, + size_bytes, + /*uuid=*/0, kOpBufferPull)); return ProcessSocketBufferPull(peer, req); } @@ -529,7 +531,7 @@ absl::Status RawBufferTransport::ProcessSocketBufferPull( return absl::InvalidArgumentError("Source peer address cannot be empty"); } - ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer, bound_ip_)); + ABSL_ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer, bound_ip_)); bool ok_to_pool = false; auto fd_cleaner = absl::MakeCleanup( [&] { ReturnConnection(ok_to_pool, fd, peer, bound_ip_); }); @@ -543,13 +545,13 @@ absl::Status RawBufferTransport::ProcessSocketBufferPull( header.count_or_size = static_cast(request.len); const auto s_header = SerializeChunkHeader(header); - RETURN_IF_ERROR(WriteExact(fd, s_header.data(), s_header.size())); + ABSL_RETURN_IF_ERROR(WriteExact(fd, s_header.data(), s_header.size())); if (request.len > 0) { if (request.laddr == nullptr) { return absl::InvalidArgumentError("Destination host pointer is null"); } - RETURN_IF_ERROR(ReadExact(fd, request.laddr, request.len)); + ABSL_RETURN_IF_ERROR(ReadExact(fd, request.laddr, request.len)); } ok_to_pool = true; @@ -578,7 +580,7 @@ absl::Status RawBufferTransport::RegisterExpectedChunks( } if (trigger_h2d) { - RETURN_IF_ERROR(raw_delegate_->OnDataReceived(uuid)); + ABSL_RETURN_IF_ERROR(raw_delegate_->OnDataReceived(uuid)); } return absl::OkStatus(); @@ -609,7 +611,7 @@ absl::Status RawBufferTransport::RegisterExpectedLayerChunks( } for (size_t layer_idx : layers_to_trigger) { - RETURN_IF_ERROR(raw_delegate_->OnLayerDataReceived(layer_idx, uuid)); + ABSL_RETURN_IF_ERROR(raw_delegate_->OnLayerDataReceived(layer_idx, uuid)); } return absl::OkStatus(); @@ -622,7 +624,7 @@ absl::Status RawBufferTransport::ProcessSocketBufferPush( "Destination peer address cannot be empty"); } - ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer, bound_ip_)); + ABSL_ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer, bound_ip_)); bool ok_to_pool = false; auto fd_cleaner = absl::MakeCleanup( [&] { ReturnConnection(ok_to_pool, fd, peer, bound_ip_); }); @@ -657,10 +659,10 @@ absl::Status RawBufferTransport::ProcessSocketBufferPush( iovec(const_cast(s_header.data()), s_header.size()), iovec(request.laddr, request.len), }; - RETURN_IF_ERROR(WriteVExact(fd, iovs)); + ABSL_RETURN_IF_ERROR(WriteVExact(fd, iovs)); uint8_t ack = 0; - RETURN_IF_ERROR(ReadExact(fd, &ack, 1)); + ABSL_RETURN_IF_ERROR(ReadExact(fd, &ack, 1)); if (ack != 1) { return absl::InternalError("PushBuffer verification failed"); } @@ -696,10 +698,10 @@ absl::StatusOr> BuildBufferRequests( std::vector requests; requests.reserve(tasks.size()); for (const auto& task : tasks) { - ASSIGN_OR_RETURN(auto req, - BuildBufferRequest(task.buffer_id, task.dst_shard_idx, - task.dst_offset_bytes, task.data_ptr, - task.size_bytes, uuid, socket_opcode)); + ABSL_ASSIGN_OR_RETURN( + auto req, BuildBufferRequest(task.buffer_id, task.dst_shard_idx, + task.dst_offset_bytes, task.data_ptr, + task.size_bytes, uuid, socket_opcode)); requests.push_back(std::move(req)); } return requests; @@ -773,7 +775,7 @@ absl::Status RawBufferTransport::PushBuffers( } auto push_batch = [&](const BatchInfo& batch) -> absl::Status { - ASSIGN_OR_RETURN( + ABSL_ASSIGN_OR_RETURN( std::vector requests, BuildBufferRequests(absl::MakeConstSpan(grouped_tasks) .subspan(batch.start_idx, batch.count), @@ -783,7 +785,7 @@ absl::Status RawBufferTransport::PushBuffers( if (parallelism <= 1 || batches.size() == 1) { for (const auto& batch : batches) { - RETURN_IF_ERROR(push_batch(batch)); + ABSL_RETURN_IF_ERROR(push_batch(batch)); } return absl::OkStatus(); } @@ -811,7 +813,7 @@ absl::Status RawBufferTransport::PushBuffers( counter.Wait(); for (const auto& s : statuses) { - RETURN_IF_ERROR(s); + ABSL_RETURN_IF_ERROR(s); } return absl::OkStatus(); } @@ -826,7 +828,7 @@ absl::Status RawBufferTransport::ProcessSocketBufferBatchPush( return absl::OkStatus(); } - ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer, bound_ip_)); + ABSL_ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer, bound_ip_)); bool ok_to_pool = false; auto fd_cleaner = absl::MakeCleanup( [&] { ReturnConnection(ok_to_pool, fd, peer, bound_ip_); }); @@ -865,7 +867,7 @@ absl::Status RawBufferTransport::ProcessSocketBufferBatchPush( iovec(const_cast(s_header.data()), s_header.size()), iovec(s_metadata_buf.data(), s_metadata_buf.size()), }; - RETURN_IF_ERROR(WriteVExact(fd, iovs)); + ABSL_RETURN_IF_ERROR(WriteVExact(fd, iovs)); if (coalesce_window_bytes_ > 0) { // Coalesced path: pack and write @@ -882,7 +884,7 @@ absl::Status RawBufferTransport::ProcessSocketBufferBatchPush( pack_offset += req.len; } } - RETURN_IF_ERROR(WriteExact(fd, pack_buf.data(), total_bytes)); + ABSL_RETURN_IF_ERROR(WriteExact(fd, pack_buf.data(), total_bytes)); } else { // Uncoalesced path: gather write (writev) directly from requests. std::vector iovs; @@ -901,12 +903,12 @@ absl::Status RawBufferTransport::ProcessSocketBufferBatchPush( } } if (!iovs.empty()) { - RETURN_IF_ERROR(WriteVExact(fd, iovs)); + ABSL_RETURN_IF_ERROR(WriteVExact(fd, iovs)); } } uint8_t ack = 0; - RETURN_IF_ERROR(ReadExact(fd, &ack, 1)); + ABSL_RETURN_IF_ERROR(ReadExact(fd, &ack, 1)); if (ack != 1) { return absl::InternalError( "ProcessSocketBufferBatchPush verification failed"); diff --git a/tpu_sync/transport/lib/socket_transport_adapter.cc b/tpu_sync/transport/lib/socket_transport_adapter.cc index f2e04713d..f8b844774 100644 --- a/tpu_sync/transport/lib/socket_transport_adapter.cc +++ b/tpu_sync/transport/lib/socket_transport_adapter.cc @@ -31,12 +31,12 @@ #include "absl/cleanup/cleanup.h" #include "absl/log/absl_check.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" -#include "tpu_sync/core/status_macros.h" #include "tpu_sync/telemetry/metrics_api.h" #include "tpu_sync/telemetry/metrics_backend.h" #include "tpu_sync/transport/lib/chunk.h" @@ -327,10 +327,10 @@ absl::Status SocketTransportAdapter::PostSocketPushInternal( ABSL_DCHECK_LE(block_offset + block_count, dst_block_ids.size()); const auto s_dst_ids = SerializeBlockIds({dst_block_ids.data() + block_offset, block_count}); - RETURN_IF_ERROR(WriteExact(fd, s_dst_ids.data(), s_dst_ids.size())); + ABSL_RETURN_IF_ERROR(WriteExact(fd, s_dst_ids.data(), s_dst_ids.size())); const auto s_src_ids = SerializeBlockIds({src_block_ids.data() + block_offset, block_count}); - RETURN_IF_ERROR(WriteExact(fd, s_src_ids.data(), s_src_ids.size())); + ABSL_RETURN_IF_ERROR(WriteExact(fd, s_src_ids.data(), s_src_ids.size())); uint8_t ack = 0; s = ReadExact(fd, &ack, 1); if (!s.ok() || ack != 1) { @@ -341,7 +341,7 @@ absl::Status SocketTransportAdapter::PostSocketPushInternal( } } else { std::vector ids_buf(block_count * sizeof(uint32_t)); - RETURN_IF_ERROR(ReadExact(fd, ids_buf.data(), ids_buf.size())); + ABSL_RETURN_IF_ERROR(ReadExact(fd, ids_buf.data(), ids_buf.size())); const std::vector stream_allocated_ids = DeserializeBlockIds(ids_buf); for (size_t k = 0; k < block_count; ++k) { @@ -368,9 +368,9 @@ absl::Status SocketTransportAdapter::PostSocketPushInternal( const std::array s_size = SerializeChunkSize(total_size); - RETURN_IF_ERROR(WriteExact(fd, s_size.data(), s_size.size())); + ABSL_RETURN_IF_ERROR(WriteExact(fd, s_size.data(), s_size.size())); if (total_size > 0) { - RETURN_IF_ERROR(WriteVExact(fd, absl::MakeSpan(iov))); + ABSL_RETURN_IF_ERROR(WriteVExact(fd, absl::MakeSpan(iov))); stream_bytes_sent += total_size; } i = j; @@ -499,12 +499,12 @@ absl::Status SocketTransportAdapter::PostSocketPullInternal( header.count_or_size = remote_count; header.uuid = uuid; const auto s_header = SerializeChunkHeader(header); - RETURN_IF_ERROR(WriteExact(fd, s_header.data(), s_header.size())); + ABSL_RETURN_IF_ERROR(WriteExact(fd, s_header.data(), s_header.size())); char resp_buf[kChunkHeaderSize]; - RETURN_IF_ERROR(ReadExact(fd, resp_buf, sizeof(resp_buf))); - ASSIGN_OR_RETURN(const ChunkHeader resp_header, - DeserializeChunkHeader(resp_buf)); + ABSL_RETURN_IF_ERROR(ReadExact(fd, resp_buf, sizeof(resp_buf))); + ABSL_ASSIGN_OR_RETURN(const ChunkHeader resp_header, + DeserializeChunkHeader(resp_buf)); if (resp_header.op != socket_opcode || resp_header.count_or_size != remote_count) { return absl::InternalError("Unexpected block pull response header"); @@ -538,7 +538,7 @@ absl::Status SocketTransportAdapter::PostSocketPullInternal( } uint8_t size_buf[kChunkSizeFieldSize]; - RETURN_IF_ERROR(ReadExact(fd, size_buf, sizeof(size_buf))); + ABSL_RETURN_IF_ERROR(ReadExact(fd, size_buf, sizeof(size_buf))); const uint32_t sender_size = DeserializeChunkSize(size_buf); if (sender_size != expected_size) { @@ -549,12 +549,12 @@ absl::Status SocketTransportAdapter::PostSocketPullInternal( } if (expected_size > 0) { - RETURN_IF_ERROR(ReadVExact(fd, absl::MakeSpan(iov))); + ABSL_RETURN_IF_ERROR(ReadVExact(fd, absl::MakeSpan(iov))); stream_bytes_received += expected_size; } if (requests[i].on_block_received != nullptr) { - RETURN_IF_ERROR( + ABSL_RETURN_IF_ERROR( requests[i].on_block_received(l, sh, dst_id, expected_size)); } i = j; diff --git a/tpu_sync/weight_sync/BUILD b/tpu_sync/weight_sync/BUILD index 2ee3c0a50..ee2cb950a 100644 --- a/tpu_sync/weight_sync/BUILD +++ b/tpu_sync/weight_sync/BUILD @@ -61,7 +61,6 @@ cc_library( "//tpu_sync/core:raiden_manager_base", "//tpu_sync/core:raiden_transfer_endpoint", "//tpu_sync/core:raw_transfer_core", - "//tpu_sync/core:status_macros", "//tpu_sync/core:xla_raw_transfer_headers", "//tpu_sync/rpc:raiden_service_cc_proto", "//tpu_sync/transport:buffer_push_task", @@ -71,6 +70,7 @@ cc_library( "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization",