From b8bdb91291db319f2af1a4ecd05f92d4b2c70027 Mon Sep 17 00:00:00 2001 From: justinlu Date: Sun, 6 Sep 2026 20:11:14 -0700 Subject: [PATCH] Refactor kv_cache to conform to TPU Raiden status handling conventions. PiperOrigin-RevId: 977239376 --- .bazelrc | 4 +- tpu_sync/kv_cache/BUILD | 19 + tpu_sync/kv_cache/host_offload_backend.cc | 77 +- .../kv_cache/host_offload_backend_test.cc | 206 ++-- tpu_sync/kv_cache/kv_cache_manager_base.cc | 132 +-- tpu_sync/kv_cache/kv_cache_manager_test.cc | 415 ++++---- tpu_sync/kv_cache/kv_cache_metadata.cc | 11 +- tpu_sync/kv_cache/kv_cache_metadata_shm.cc | 27 +- .../kv_cache/kv_cache_metadata_shm_test.cc | 111 +- tpu_sync/kv_cache/kv_cache_metadata_test.cc | 166 ++- tpu_sync/kv_cache/kv_cache_store.cc | 83 +- .../kv_cache/kv_cache_store_client_test.cc | 28 +- .../kv_cache/kv_cache_store_server_test.cc | 34 +- tpu_sync/kv_cache/kv_cache_store_service.cc | 18 +- .../kv_cache/kv_cache_store_service_test.cc | 211 ++-- tpu_sync/kv_cache/kv_cache_store_test.cc | 990 +++++++++--------- .../kv_cache/kv_cache_store_wrapper_test.cc | 34 +- .../kv_cache/logical_block_manager_test.cc | 122 +-- tpu_sync/kv_cache/pool_layout.cc | 133 ++- tpu_sync/kv_cache/pool_layout_test.cc | 95 +- tpu_sync/kv_cache/pool_staging_test.cc | 54 +- tpu_sync/kv_cache/store_monitor_test.cc | 46 +- 22 files changed, 1449 insertions(+), 1567 deletions(-) diff --git a/.bazelrc b/.bazelrc index defd19b9..ca7be8ce 100644 --- a/.bazelrc +++ b/.bazelrc @@ -1,3 +1,6 @@ +# Allow empty globs across all Bazel commands (required for external deps e.g. Eigen / XLA under Bazel 8). +common --incompatible_disallow_empty_glob=false + build:oss --experimental_repo_remote_exec build:oss --nocheck_visibility build:oss --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir. @@ -5,7 +8,6 @@ build:oss --override_module=torch_tpu=%workspace%/third_party/torch_tpu_stub build:oss --action_env=CC=clang-18 build:oss --action_env=CXX=clang++-18 build:oss --repo_env=CC=clang-18 -build:oss --incompatible_disallow_empty_glob=false build:ci --google_default_credentials=true build:ci --remote_cache=grpcs://remotebuildexecution.googleapis.com diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index 7e45bf83..da928365 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -70,9 +70,11 @@ cc_test( deps = [ ":logical_block_manager", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -83,6 +85,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ "@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", @@ -96,10 +99,12 @@ cc_test( deps = [ ":kv_cache_metadata", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -155,10 +160,12 @@ cc_test( ":kv_cache_store_wrapper", "//tpu_sync/common:raiden_id", "//tpu_sync/kv_cache/global_registry:test_util", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -168,10 +175,12 @@ cc_test( deps = [ ":kv_cache_metadata", ":kv_cache_metadata_shm", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -187,6 +196,7 @@ cc_library( deps = [ "//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", @@ -207,9 +217,11 @@ cc_test( "//tpu_sync/rpc:raiden_service_cc_proto", "//tpu_sync/transport:block_transport_delegate", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/time", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -229,6 +241,7 @@ cc_test( "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -427,10 +440,12 @@ cc_test( "//tpu_sync/kv_cache/global_registry:global_registry_server_lib", "//tpu_sync/kv_cache/global_registry:test_util", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@xla//xla/tsl/platform:statusor", ], ) @@ -575,6 +590,7 @@ cc_test( "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -743,6 +759,7 @@ cc_test( "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", "@xla//xla/tsl/concurrency:future", + "@xla//xla/tsl/platform:statusor", ], ) @@ -771,6 +788,7 @@ cc_test( "//tpu_sync/kv_cache/global_registry:test_util", "//tpu_sync/proto:kv_cache_store_service_cc_grpc", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", @@ -778,6 +796,7 @@ cc_test( "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", "@xla//xla/tsl/concurrency:future", + "@xla//xla/tsl/platform:statusor", ], ) diff --git a/tpu_sync/kv_cache/host_offload_backend.cc b/tpu_sync/kv_cache/host_offload_backend.cc index 523353f4..879935cc 100644 --- a/tpu_sync/kv_cache/host_offload_backend.cc +++ b/tpu_sync/kv_cache/host_offload_backend.cc @@ -217,13 +217,13 @@ absl::StatusOr HostOffloadBackend::Lookup( // at the first one. std::vector remote_hits; if (!missing_hashes.empty() && options.enable_global && client != nullptr) { - auto global_res_or = client->Lookup(missing_hashes, local_id); - if (global_res_or.ok()) { - remote_hits = std::move(global_res_or).value(); - } else { - LOG(WARNING) << "Global registry lookup failed: " - << global_res_or.status().message(); - } + auto lookup_remote = [&]() -> absl::Status { + ABSL_ASSIGN_OR_RETURN( + remote_hits, client->Lookup(missing_hashes, local_id), + _.LogWarning() << "Global registry lookup failed: "); + return absl::OkStatus(); + }; + (void)lookup_remote(); } // Phases 3 and 4, under one lock: walk the request in order taking each hash @@ -580,11 +580,7 @@ absl::StatusOr HostOffloadBackend::RecoverFromLocalManifest() { for (const KVCacheMetadata::Entry* entry : recoverable) { block_ids.push_back(entry->block_id); } - absl::Status allocate_status = - raiden_controller_->AllocateTargetBlockIds(block_ids); - if (!allocate_status.ok()) { - return allocate_status; - } + ABSL_RETURN_IF_ERROR(raiden_controller_->AllocateTargetBlockIds(block_ids)); uint64_t max_seq = 0; for (const KVCacheMetadata::Entry* entry : recoverable) { @@ -799,29 +795,29 @@ HostOffloadBackend::BeginWriteRemote( BuildLocalWorkerEndpoints(raiden_controller_), absl::ToInt64Milliseconds(requested_deadline), hold_window, std::move(on_verdict)); - auto response_or = call.ack.Await(); - if (!response_or.ok()) { + auto response = call.ack.Await(); + if (!response.ok()) { // On a transport error the peer may have restarted on a new port; drop // the cached client so the next attempt re-resolves it. An application // error (e.g. RESOURCE_EXHAUSTED) means the peer is alive, so the // channel is kept. - if (IsTransportError(response_or.status())) { + if (IsTransportError(response.status())) { InvalidateStoreClient(dst_raiden_id); } - return response_or.status(); + return response.status(); } RemoteWriteAck ack; ack.cancel = std::move(call.cancel); - ack.operation_id = response_or->operation_id(); - ack.granted_deadline = absl::Milliseconds(response_or->granted_deadline_ms()); - switch (response_or->exist_state()) { + ack.operation_id = response->operation_id(); + ack.granted_deadline = absl::Milliseconds(response->granted_deadline_ms()); + switch (response->exist_state()) { case ::tpu_raiden::kv_cache::proto::WRITE_ALL_EXIST: ack.all_exist = true; return ack; case ::tpu_raiden::kv_cache::proto::WRITE_PARTIAL_EXIST: - ack.existing_hashes.assign(response_or->existing_hashes().begin(), - response_or->existing_hashes().end()); + ack.existing_hashes.assign(response->existing_hashes().begin(), + response->existing_hashes().end()); return ack; default: break; @@ -839,11 +835,11 @@ tsl::Future HostOffloadBackend::PollWriteRemoteAsync(const RaidenId& dst_raiden_id, uint64_t operation_id, int64_t wait_ms) { - auto client_or = GetKVCacheStoreClient(dst_raiden_id); - if (!client_or.ok()) { - return tsl::Future(client_or.status()); + auto client = GetKVCacheStoreClient(dst_raiden_id); + if (!client.ok()) { + return tsl::Future(client.status()); } - return (*client_or)->PollWriteRemote(operation_id, wait_ms); + return (*client)->PollWriteRemote(operation_id, wait_ms); } std::vector HostOffloadBackend::AlreadyPresentHostResident( @@ -991,19 +987,18 @@ tsl::Future<> HostOffloadBackend::Load( tsl::Future<> HostOffloadBackend::LoadRemoteBlocks( const RaidenId& remote_id, absl::Span block_hashes, absl::Span device_block_ids) { - auto client_or = GetKVCacheStoreClient(remote_id); - if (!client_or.ok()) { - return tsl::Future<>(client_or.status()); + auto client = GetKVCacheStoreClient(remote_id); + if (!client.ok()) { + return tsl::Future<>(client.status()); } - std::shared_ptr client = std::move(client_or.value()); + std::shared_ptr client_ptr = *std::move(client); - auto host_blocks_or = - raiden_controller_->AllocateBlockIds(block_hashes.size()); - if (!host_blocks_or.ok()) { - return tsl::Future<>(host_blocks_or.status()); + auto host_blocks = raiden_controller_->AllocateBlockIds(block_hashes.size()); + if (!host_blocks.ok()) { + return tsl::Future<>(host_blocks.status()); } - std::vector dst_host_block_ids(host_blocks_or.value().begin(), - host_blocks_or.value().end()); + std::vector dst_host_block_ids(host_blocks->begin(), + host_blocks->end()); auto [load_promise, load_future] = tsl::MakePromise<>(); @@ -1011,8 +1006,8 @@ tsl::Future<> HostOffloadBackend::LoadRemoteBlocks( std::vector<::tpu_sync::proto::RaidenWorkerEndpointsProto> client_worker_endpoints = BuildLocalWorkerEndpoints(raiden_controller_); tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> fetch_future = - client->Fetch(block_hashes, device_block_ids, dst_host_block_ids, - client_raiden_id, client_worker_endpoints); + client_ptr->Fetch(block_hashes, device_block_ids, dst_host_block_ids, + client_raiden_id, client_worker_endpoints); fetch_future.OnReady( [this, remote_id, dst_host_block_ids, @@ -1020,17 +1015,17 @@ tsl::Future<> HostOffloadBackend::LoadRemoteBlocks( device_block_ids.end()), load_promise = std::move(load_promise)]( const absl::StatusOr<::tpu_raiden::kv_cache::proto::FetchResponse>& - response_or) mutable { - if (!response_or.ok()) { + fetch_response) mutable { + if (!fetch_response.ok()) { (void)raiden_controller_->DeallocateBlockIds(dst_host_block_ids); // The peer may have restarted on a new port; drop the cached client // so the next attempt re-resolves instead of redialling a dead one. InvalidateStoreClient(remote_id); - load_promise.Set(response_or.status()); + load_promise.Set(fetch_response.status()); return; } - const auto& response = response_or.value(); + const auto& response = *fetch_response; if (!response.failed_block_hashes().empty()) { (void)raiden_controller_->DeallocateBlockIds(dst_host_block_ids); std::string err_msg = response.error_message().empty() diff --git a/tpu_sync/kv_cache/host_offload_backend_test.cc b/tpu_sync/kv_cache/host_offload_backend_test.cc index 93dc11f5..2ed9563a 100644 --- a/tpu_sync/kv_cache/host_offload_backend_test.cc +++ b/tpu_sync/kv_cache/host_offload_backend_test.cc @@ -88,23 +88,20 @@ TEST(HostOffloadBackendTest, BasicInsertAndLookup) { EXPECT_EQ(backend.GetSize(), 2); // Lookup both - auto lookup_res = backend.Lookup({"h1", "h2"}); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_EQ(lookup_res->size(), 2); - EXPECT_EQ((*lookup_res)[0].first, "h1"); - EXPECT_EQ((*lookup_res)[0].second.host_block_id, 10); - EXPECT_EQ((*lookup_res)[1].first, "h2"); - EXPECT_EQ((*lookup_res)[1].second.host_block_id, 11); + TF_ASSERT_OK_AND_ASSIGN(auto lookup_res, backend.Lookup({"h1", "h2"})); + EXPECT_EQ(lookup_res.size(), 2); + EXPECT_EQ(lookup_res[0].first, "h1"); + EXPECT_EQ(lookup_res[0].second.host_block_id, 10); + EXPECT_EQ(lookup_res[1].first, "h2"); + EXPECT_EQ(lookup_res[1].second.host_block_id, 11); // Partial miss at end - auto partial_res = backend.Lookup({"h1", "h2", "h3"}); - ASSERT_TRUE(partial_res.ok()); - EXPECT_EQ(partial_res->size(), 2); + TF_ASSERT_OK_AND_ASSIGN(auto partial_res, backend.Lookup({"h1", "h2", "h3"})); + EXPECT_EQ(partial_res.size(), 2); // Miss at start - auto miss_res = backend.Lookup({"h3", "h1"}); - ASSERT_TRUE(miss_res.ok()); - EXPECT_TRUE(miss_res->empty()); + TF_ASSERT_OK_AND_ASSIGN(auto miss_res, backend.Lookup({"h3", "h1"})); + EXPECT_TRUE(miss_res.empty()); } TEST(HostOffloadBackendTest, SnapshotAndPinHostResidentForRepublish) { @@ -166,9 +163,8 @@ TEST(HostOffloadBackendTest, LookupUnboundedByAvailableSpace) { EXPECT_EQ(backend.GetAvailableSpace(), 0); // Lookup still succeeds completely despite available_space() == 0 - auto lookup_res = backend.Lookup({"h1", "h2"}); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_EQ(lookup_res->size(), 2); + TF_ASSERT_OK_AND_ASSIGN(auto lookup_res, backend.Lookup({"h1", "h2"})); + EXPECT_EQ(lookup_res.size(), 2); } TEST(HostOffloadBackendTest, InsertAndLockRollbackOnCapacityExceeded) { @@ -214,11 +210,10 @@ TEST(HostOffloadBackendTest, InsertAndLockRebindsStaleHbmEntry) { BlockStatus::HBM); ASSERT_TRUE(backend.InsertAndLock({"h1"}, {second}, /*on_host=*/false)); - auto lookup_res = backend.Lookup({"h1"}); - ASSERT_TRUE(lookup_res.ok()); - ASSERT_EQ(lookup_res->size(), 1); - EXPECT_EQ((*lookup_res)[0].second.device_block_id, 9); - EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HBM); + TF_ASSERT_OK_AND_ASSIGN(auto lookup_res, backend.Lookup({"h1"})); + ASSERT_EQ(lookup_res.size(), 1); + EXPECT_EQ(lookup_res[0].second.device_block_id, 9); + EXPECT_EQ(lookup_res[0].second.status, BlockStatus::HBM); } TEST(HostOffloadBackendTest, InsertAndLockKeepsHostResidentEntry) { @@ -235,11 +230,10 @@ TEST(HostOffloadBackendTest, InsertAndLockKeepsHostResidentEntry) { BlockStatus::HBM); ASSERT_TRUE(backend.InsertAndLock({"h1"}, {reoffer}, /*on_host=*/false)); - auto lookup_res = backend.Lookup({"h1"}); - ASSERT_TRUE(lookup_res.ok()); - ASSERT_EQ(lookup_res->size(), 1); - EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST); - EXPECT_EQ((*lookup_res)[0].second.host_block_id, 3); + TF_ASSERT_OK_AND_ASSIGN(auto lookup_res, backend.Lookup({"h1"})); + ASSERT_EQ(lookup_res.size(), 1); + EXPECT_EQ(lookup_res[0].second.status, BlockStatus::HOST); + EXPECT_EQ(lookup_res[0].second.host_block_id, 3); } TEST(HostOffloadBackendTest, LookupReturnsRemoteDescriptors) { @@ -253,7 +247,7 @@ TEST(HostOffloadBackendTest, LookupReturnsRemoteDescriptors) { {.prefix_hash = "r_hash1", .raiden_id = remote_node_id, .block_id = 42}, {.prefix_hash = "r_hash2", .raiden_id = remote_node_id, .block_id = 43}, }; - ASSERT_TRUE(registry_client->Register(regs).ok()); + ABSL_ASSERT_OK(registry_client->Register(regs)); RaidenId local_node_id{"local_job", "0", "data", 0}; ::tpu_sync::rpc::RaidenIdProto unit_proto; @@ -277,22 +271,22 @@ TEST(HostOffloadBackendTest, LookupReturnsRemoteDescriptors) { HostOffloadBackend::Create(config, controller.get())); EXPECT_EQ(backend->name(), "HostOffloadBackend"); - auto lookup_res = backend->Lookup({"r_hash1", "r_hash2"}); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_EQ(lookup_res->size(), 2); - EXPECT_EQ((*lookup_res)[0].first, "r_hash1"); - EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::REMOTE); - EXPECT_EQ((*lookup_res)[0].second.host_block_id, 42); - EXPECT_EQ((*lookup_res)[0].second.raiden_id, remote_node_id); + TF_ASSERT_OK_AND_ASSIGN(auto lookup_res, + backend->Lookup({"r_hash1", "r_hash2"})); + EXPECT_EQ(lookup_res.size(), 2); + EXPECT_EQ(lookup_res[0].first, "r_hash1"); + EXPECT_EQ(lookup_res[0].second.status, BlockStatus::REMOTE); + EXPECT_EQ(lookup_res[0].second.host_block_id, 42); + EXPECT_EQ(lookup_res[0].second.raiden_id, remote_node_id); - EXPECT_EQ((*lookup_res)[1].first, "r_hash2"); - EXPECT_EQ((*lookup_res)[1].second.status, BlockStatus::REMOTE); - EXPECT_EQ((*lookup_res)[1].second.host_block_id, 43); + EXPECT_EQ(lookup_res[1].first, "r_hash2"); + EXPECT_EQ(lookup_res[1].second.status, BlockStatus::REMOTE); + EXPECT_EQ(lookup_res[1].second.host_block_id, 43); // Lookup with miss stops at miss - auto partial_res = backend->Lookup({"r_hash1", "missing_hash"}); - ASSERT_TRUE(partial_res.ok()); - EXPECT_EQ(partial_res->size(), 1); + TF_ASSERT_OK_AND_ASSIGN(auto partial_res, + backend->Lookup({"r_hash1", "missing_hash"})); + EXPECT_EQ(partial_res.size(), 1); } TEST(HostOffloadBackendTest, @@ -308,7 +302,7 @@ TEST(HostOffloadBackendTest, .raiden_id = local_node_id, .block_id = 99}, }; - ASSERT_TRUE(registry_client->Register(regs).ok()); + ABSL_ASSERT_OK(registry_client->Register(regs)); ::tpu_sync::rpc::RaidenIdProto unit_proto; unit_proto.set_job_name(local_node_id.job_name); @@ -330,9 +324,8 @@ TEST(HostOffloadBackendTest, TF_ASSERT_OK_AND_ASSIGN(auto backend, HostOffloadBackend::Create(config, controller.get())); - auto lookup_res = backend->Lookup({"local_g_hash"}); - ASSERT_TRUE(lookup_res.ok()); - EXPECT_EQ(lookup_res->size(), 0); + TF_ASSERT_OK_AND_ASSIGN(auto lookup_res, backend->Lookup({"local_g_hash"})); + EXPECT_EQ(lookup_res.size(), 0); } TEST(HostOffloadBackendTest, CreateRegistersKVTransferSpecFromConfig) { @@ -367,17 +360,17 @@ TEST(HostOffloadBackendTest, CreateRegistersKVTransferSpecFromConfig) { auto channel = grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); global_registry::GlobalRegistryClient registry_client(channel); - auto spec = registry_client.GetKVTransferSpec("node_group"); - ASSERT_TRUE(spec.ok()) << spec.status(); - ASSERT_EQ(spec->block_arrays_size(), 2); - EXPECT_EQ(spec->block_arrays(0).block_bytes(), 4096); - EXPECT_EQ(spec->block_arrays(1).block_bytes(), 512); - EXPECT_EQ(spec->num_kv_shards(), 2); - EXPECT_EQ(spec->num_workers(), 2); + TF_ASSERT_OK_AND_ASSIGN(auto spec, + registry_client.GetKVTransferSpec("node_group")); + ASSERT_EQ(spec.block_arrays_size(), 2); + EXPECT_EQ(spec.block_arrays(0).block_bytes(), 4096); + EXPECT_EQ(spec.block_arrays(1).block_bytes(), 512); + EXPECT_EQ(spec.num_kv_shards(), 2); + EXPECT_EQ(spec.num_workers(), 2); // Creating another backend with the identical spec is a no-op validation; // a differing spec fails creation. - EXPECT_TRUE(HostOffloadBackend::Create(config, controller.get()).ok()); + ABSL_EXPECT_OK(HostOffloadBackend::Create(config, controller.get())); config.kv_transfer_spec->num_kv_shards = 4; EXPECT_TRUE(absl::IsInvalidArgument( HostOffloadBackend::Create(config, controller.get()).status())); @@ -422,7 +415,7 @@ TEST(HostOffloadBackendTest, ComposeKVTransferSpecFromUniformWorkers) { workers.push_back(GeometryWorker("w0", /*node_id=*/0, {4096, 512}, 2)); auto spec = HostOffloadBackendTest::ComposeKVTransferSpec(workers); - ASSERT_OK(spec.status()); + ABSL_ASSERT_OK(spec.status()); EXPECT_EQ(spec->block_array_bytes, (std::vector{4096, 512})); EXPECT_EQ(spec->num_kv_shards, 2); EXPECT_EQ(spec->num_workers, 2); @@ -502,7 +495,7 @@ TEST(HostOffloadBackendTest, ServerLifecycleAndControllerInitialization) { auto store_server = KVCacheStoreServer::Create(); // A wildcard bind reports no publishable address, // so bind a real, dialable host. - ASSERT_OK( + ABSL_ASSERT_OK( store_server->StartServer(backend.get(), controller.get(), "127.0.0.1")); EXPECT_GT(store_server->GetGrpcPort(), 0); EXPECT_FALSE(store_server->GetServerAddress().empty()); @@ -537,7 +530,7 @@ TEST(HostOffloadBackendTest, StartServerStripsControllerPort) { auto store_server = KVCacheStoreServer::Create(); std::string ctrl_addr = controller->controller_address(); std::string target_host = ctrl_addr.substr(0, ctrl_addr.rfind(':')); - ASSERT_OK( + ABSL_ASSERT_OK( store_server->StartServer(backend.get(), controller.get(), target_host)); EXPECT_GT(store_server->GetGrpcPort(), 0); EXPECT_NE(store_server->GetGrpcPort(), 12345); @@ -563,7 +556,7 @@ TEST(HostOffloadBackendTest, EndToEndFetchRPC) { RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; RaidenId dst_raiden_id{"dst_job", "0", "dst_data", 0}; - ASSERT_OK(src_controller_server->client->RegisterWorker( + ABSL_ASSERT_OK(src_controller_server->client->RegisterWorker( "worker_0", test_worker_server->server_address, {{test_worker_server->server_address, {}}})); @@ -583,7 +576,7 @@ TEST(HostOffloadBackendTest, EndToEndFetchRPC) { .raiden_id = dst_raiden_id, .block_id = 102}, }; - ASSERT_OK(registry_client->Register(registrations)); + ABSL_ASSERT_OK(registry_client->Register(registrations)); // 6. Create destination HostOffloadBackend & RaidenController ::tpu_sync::rpc::RaidenIdProto dst_unit_proto; @@ -619,15 +612,15 @@ TEST(HostOffloadBackendTest, EndToEndFetchRPC) { core::controller::RaidenControllerClient dst_controller_client( dst_controller->controller_address()); - ASSERT_OK(dst_controller_client.RegisterWorker( + ABSL_ASSERT_OK(dst_controller_client.RegisterWorker( "dst_worker_0", test_worker_server->server_address, {{test_worker_server->server_address, {}}})); auto store_server = KVCacheStoreServer::Create(); // A wildcard bind reports no publishable address, // so bind a real, dialable host. - ASSERT_OK(store_server->StartServer(backend.get(), dst_controller.get(), - "127.0.0.1")); + ABSL_ASSERT_OK(store_server->StartServer(backend.get(), dst_controller.get(), + "127.0.0.1")); EXPECT_GT(store_server->GetGrpcPort(), 0); // 7. Issue Fetch RPC using KVCacheStoreClient @@ -641,7 +634,7 @@ TEST(HostOffloadBackendTest, EndToEndFetchRPC) { .Fetch(hashes, /*device_block_ids=*/{}, host_ids, dst_controller->unit()) .Await(); - ASSERT_OK(fetch_res.status()); + ABSL_ASSERT_OK(fetch_res.status()); EXPECT_THAT(fetch_res->done_block_hashes(), UnorderedElementsAre("fetch_hash_1", "fetch_hash_2")); @@ -699,7 +692,7 @@ TEST(HostOffloadBackendTest, LoadSuccess) { .raiden_id = remote_node_id, .block_id = 42}, }; - ASSERT_OK(registry_client->Register(regs)); + ABSL_ASSERT_OK(registry_client->Register(regs)); // Setup local RaidenController ::tpu_sync::rpc::RaidenIdProto local_unit; @@ -737,12 +730,12 @@ TEST(HostOffloadBackendTest, LoadSuccess) { auto remote_server = KVCacheStoreServer::Create(); // A wildcard bind reports no publishable address, // so bind a real, dialable host -- this test publishes it below. - ASSERT_OK(remote_server->StartServer(remote_backend.get(), controller.get(), - "127.0.0.1")); + ABSL_ASSERT_OK(remote_server->StartServer(remote_backend.get(), + controller.get(), "127.0.0.1")); - ASSERT_OK(registry_client->RegisterStore(remote_node_id, - remote_server->GetServerAddress(), - controller->controller_address())); + ABSL_ASSERT_OK(registry_client->RegisterStore( + remote_node_id, remote_server->GetServerAddress(), + controller->controller_address())); BackendConfig local_config; local_config.type = "HostOffloadBackend"; @@ -766,7 +759,7 @@ TEST(HostOffloadBackendTest, LoadSuccess) { core::controller::RaidenControllerClient controller_client( controller->controller_address()); - ASSERT_OK(controller_client.RegisterWorker( + ABSL_ASSERT_OK(controller_client.RegisterWorker( "worker_0", test_worker_server->server_address, {{test_worker_server->server_address, {}}})); @@ -774,7 +767,7 @@ TEST(HostOffloadBackendTest, LoadSuccess) { std::vector hashes = {"load_hash_1"}; std::vector dev_ids = {5}; auto load_future = backend->Load(remote_node_id, hashes, dev_ids); - EXPECT_OK(load_future.Await()); + ABSL_EXPECT_OK(load_future.Await()); remote_server->Shutdown(); } @@ -799,7 +792,7 @@ TEST(HostOffloadBackendTest, LoadLocalSuccess) { core::controller::RaidenControllerClient controller_client( controller->controller_address()); - ASSERT_OK(controller_client.RegisterWorker( + ABSL_ASSERT_OK(controller_client.RegisterWorker( "worker_0", test_worker_server->server_address, {{test_worker_server->server_address, {}}})); @@ -818,7 +811,7 @@ TEST(HostOffloadBackendTest, LoadLocalSuccess) { /*on_host=*/true); auto load_future = backend->Load(RaidenId{}, {"local_hash_1"}, {5}); - EXPECT_OK(load_future.Await()); + ABSL_EXPECT_OK(load_future.Await()); } TEST(HostOffloadBackendTest, LoadLocalMissingBlockError) { @@ -903,7 +896,7 @@ TEST(HostOffloadBackendTest, StoreServerOverride) { auto backend = std::dynamic_pointer_cast(backend_base); ASSERT_NE(backend, nullptr); EXPECT_EQ(backend->store_server(), nullptr); - ASSERT_TRUE(backend->StartServer("127.0.0.1").ok()); + ABSL_ASSERT_OK(backend->StartServer("127.0.0.1")); EXPECT_NE(backend->store_server(), nullptr); backend->store_server()->Shutdown(); } @@ -1000,7 +993,7 @@ TEST(HostOffloadBackendWriteRemoteTest, InsertAllOrNothingRespectsPinnedSpace) { TEST(HostOffloadBackendWriteRemoteTest, RegisterBlocksAsyncIsOkWithoutARegistry) { HostOffloadBackendTest::Backend backend(/*capacity=*/8); - EXPECT_TRUE(backend.RegisterBlocksAsync({"a"}, {1}).Await().ok()); + ABSL_EXPECT_OK(backend.RegisterBlocksAsync({"a"}, {1}).Await()); } TEST(HostOffloadBackendWriteRemoteTest, @@ -1013,13 +1006,13 @@ TEST(HostOffloadBackendWriteRemoteTest, std::make_shared( reg_server->channel)); - ASSERT_TRUE(backend.RegisterBlocksAsync({"a", "b"}, {7, 8}).Await().ok()); + ABSL_ASSERT_OK(backend.RegisterBlocksAsync({"a", "b"}, {7, 8}).Await()); - auto looked_up = reg_server->client->Lookup({"a", "b"}); - ASSERT_TRUE(looked_up.ok()) << looked_up.status().ToString(); - ASSERT_EQ(looked_up->size(), 2); - EXPECT_EQ((*looked_up)[0].block_id(), 7); - EXPECT_EQ((*looked_up)[1].block_id(), 8); + TF_ASSERT_OK_AND_ASSIGN(auto looked_up, + reg_server->client->Lookup({"a", "b"})); + ASSERT_EQ(looked_up.size(), 2); + EXPECT_EQ(looked_up[0].block_id(), 7); + EXPECT_EQ(looked_up[1].block_id(), 8); } // COMMITTED is only allowed to mean "globally reachable", so a publish that @@ -1054,9 +1047,9 @@ TEST(HostOffloadBackendTest, LookupAndPinBasicHits) { backend.Insert(hashes, slices, /*on_host=*/true); - auto res = backend.Lookup(hashes, LookupOptions{.pin_found = true}); - ASSERT_TRUE(res.ok()); - EXPECT_EQ(res->size(), 3); + TF_ASSERT_OK_AND_ASSIGN( + auto res, backend.Lookup(hashes, LookupOptions{.pin_found = true})); + EXPECT_EQ(res.size(), 3); EXPECT_EQ(backend.GetPinCount("h1"), 1); EXPECT_EQ(backend.GetPinCount("h2"), 1); EXPECT_EQ(backend.GetPinCount("h3"), 1); @@ -1072,10 +1065,10 @@ TEST(HostOffloadBackendTest, LookupAndPinPartialMiss) { backend.Insert(hashes, slices, /*on_host=*/true); - auto res = - backend.Lookup({"h1", "h2", "h3"}, LookupOptions{.pin_found = true}); - ASSERT_TRUE(res.ok()); - EXPECT_EQ(res->size(), 2); + TF_ASSERT_OK_AND_ASSIGN( + auto res, + backend.Lookup({"h1", "h2", "h3"}, LookupOptions{.pin_found = true})); + EXPECT_EQ(res.size(), 2); EXPECT_EQ(backend.GetPinCount("h1"), 1); EXPECT_EQ(backend.GetPinCount("h2"), 1); EXPECT_EQ(backend.GetPinCount("h3"), 0); @@ -1090,7 +1083,7 @@ TEST(HostOffloadBackendTest, LookupAndPinRemoteDescriptorsUnpinnedLocally) { std::vector regs = { {.prefix_hash = "r1", .raiden_id = remote_node_id, .block_id = 42}, }; - ASSERT_TRUE(registry_client->Register(regs).ok()); + ABSL_ASSERT_OK(registry_client->Register(regs)); RaidenId local_node_id{"local_job", "0", "data", 0}; ::tpu_sync::rpc::RaidenIdProto unit_proto; @@ -1113,12 +1106,13 @@ TEST(HostOffloadBackendTest, LookupAndPinRemoteDescriptorsUnpinnedLocally) { TF_ASSERT_OK_AND_ASSIGN(auto backend, HostOffloadBackend::Create(config, controller.get())); - auto lookup_res = backend->Lookup( - {"r1"}, LookupOptions{.enable_global = true, .pin_found = true}); - ASSERT_TRUE(lookup_res.ok()); - ASSERT_EQ(lookup_res->size(), 1); - EXPECT_EQ((*lookup_res)[0].first, "r1"); - EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::REMOTE); + TF_ASSERT_OK_AND_ASSIGN( + auto lookup_res, + backend->Lookup({"r1"}, + LookupOptions{.enable_global = true, .pin_found = true})); + ASSERT_EQ(lookup_res.size(), 1); + EXPECT_EQ(lookup_res[0].first, "r1"); + EXPECT_EQ(lookup_res[0].second.status, BlockStatus::REMOTE); EXPECT_EQ(backend->GetPinCount("r1"), 0); } @@ -1177,7 +1171,7 @@ void InsertLocal(KVCacheStoreBackend* backend, const RaidenId& local_id, void RegisterGlobal(global_registry::GlobalRegistryClient* client, const RaidenId& owner, const std::string& hash, int block_id) { - ASSERT_OK(client->Register( + ABSL_ASSERT_OK(client->Register( {{.prefix_hash = hash, .raiden_id = owner, .block_id = block_id}})); } @@ -1190,7 +1184,7 @@ TEST(HostOffloadBackendTest, LookupInterleavesLocalAndRemoteHits) { InsertLocal(f->backend.get(), f->local_id, "l2", 12); auto res = f->backend->Lookup({"r1", "l1", "r2", "l2", "nowhere"}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); ASSERT_EQ(res->size(), 4) << "the answer must run to the first hash that " "neither source can resolve"; @@ -1224,7 +1218,7 @@ TEST(HostOffloadBackendTest, InsertLocal(f->backend.get(), f->local_id, "l2", 12); auto res = f->backend->Lookup({"r1", "r2", "l1", "l2"}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); ASSERT_EQ(res->size(), 4); EXPECT_EQ((*res)[2].second.status, BlockStatus::HOST); EXPECT_EQ((*res)[2].second.host_block_id, 11); @@ -1243,7 +1237,7 @@ TEST(HostOffloadBackendTest, RegisterGlobal(f->registry->client.get(), f->peer_id, "l1", 91); auto res = f->backend->Lookup({"r1", "l1"}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); ASSERT_EQ(res->size(), 2); EXPECT_EQ((*res)[1].first, "l1"); EXPECT_EQ((*res)[1].second.status, BlockStatus::HOST); @@ -1259,7 +1253,7 @@ TEST(HostOffloadBackendTest, LookupInterleavedStopsAtTheFirstAbsoluteMiss) { // may be reported however reachable it is. auto res = f->backend->Lookup({"l1", "gap", "l2"}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); ASSERT_EQ(res->size(), 1); EXPECT_EQ((*res)[0].first, "l1"); } @@ -1275,7 +1269,7 @@ TEST(HostOffloadBackendTest, LookupInterleavedPinsOnlyTheReturnedPrefix) { // never release it. auto res = f->backend->Lookup({"r1", "l1", "gap", "l2"}, LookupOptions{.pin_found = true}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); ASSERT_EQ(res->size(), 2); EXPECT_EQ(f->backend->GetPinCount("l1"), 1); EXPECT_EQ(f->backend->GetPinCount("l2"), 0); @@ -1290,12 +1284,12 @@ TEST(HostOffloadBackendTest, LookupInterleavedDisabledStopsAtFirstLocalMiss) { auto legacy = f->backend->Lookup( {"r1", "l1"}, LookupOptions{.enable_interleaved_lookup = false}); - ASSERT_OK(legacy.status()); + ABSL_ASSERT_OK(legacy.status()); ASSERT_EQ(legacy->size(), 1); EXPECT_EQ((*legacy)[0].first, "r1"); auto interleaved = f->backend->Lookup({"r1", "l1"}); - ASSERT_OK(interleaved.status()); + ABSL_ASSERT_OK(interleaved.status()); ASSERT_EQ(interleaved->size(), 2); EXPECT_EQ((*interleaved)[1].first, "l1"); EXPECT_EQ((*interleaved)[1].second.status, BlockStatus::HOST); @@ -1309,13 +1303,13 @@ TEST(HostOffloadBackendTest, LookupInterleavedWithoutGlobalStopsAtLocalMiss) { auto res = f->backend->Lookup({"r1", "l1"}, LookupOptions{.enable_global = false}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); EXPECT_TRUE(res->empty()) << "with no registry to consult, a local miss " "still ends the answer"; auto res2 = f->backend->Lookup({"l1", "r1"}, LookupOptions{.enable_global = false}); - ASSERT_OK(res2.status()); + ABSL_ASSERT_OK(res2.status()); ASSERT_EQ(res2->size(), 1); EXPECT_EQ((*res2)[0].first, "l1"); } @@ -1332,7 +1326,7 @@ TEST(HostOffloadBackendTest, LookupInterleavedWithoutRegistryClient) { auto res = backend.Lookup({"l1", "gap", "l2"}, LookupOptions{.pin_found = true}); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); ASSERT_EQ(res->size(), 1); EXPECT_EQ((*res)[0].first, "l1"); EXPECT_EQ(backend.GetPinCount("l1"), 1); @@ -1350,7 +1344,7 @@ TEST(HostOffloadBackendTest, DeleteSkipsPinnedBlocks) { // Deleting a pinned block must skip it without erasing the block. backend.Delete({"h1"}, {}); auto lookup_res = backend.Lookup({"h1"}); - ASSERT_OK(lookup_res.status()); + ABSL_ASSERT_OK(lookup_res.status()); EXPECT_EQ(lookup_res->size(), 1); backend.Release({"h1"}); @@ -1359,7 +1353,7 @@ TEST(HostOffloadBackendTest, DeleteSkipsPinnedBlocks) { // Now that it's unpinned, Delete removes it. backend.Delete({"h1"}, {}); auto lookup_after = backend.Lookup({"h1"}); - ASSERT_OK(lookup_after.status()); + ABSL_ASSERT_OK(lookup_after.status()); EXPECT_TRUE(lookup_after->empty()); } diff --git a/tpu_sync/kv_cache/kv_cache_manager_base.cc b/tpu_sync/kv_cache/kv_cache_manager_base.cc index 90b7527c..d8dfae8c 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_base.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_base.cc @@ -210,8 +210,8 @@ raiden::PjRtCopyFuture JoinAndRecordTelemetry( auto joined_future = raiden::JoinPjRtCopyFutures(futures); if (telemetry::RaidenMetricStore::GetGlobalMetricStore().HasBackends()) { joined_future.OnReady([start_time, metric = std::string(metric_name)]( - auto status_or) { - if (status_or.ok()) { + const auto& result) { + if (result.ok()) { telemetry::RaidenMetricStore::GetGlobalMetricStore().ObserveHistogram( metric, {}, absl::ToDoubleMilliseconds(absl::Now() - start_time)); } @@ -337,14 +337,13 @@ KVCacheManagerBase::KVCacheManagerBase( size_t alloc_size = num_host_blocks * layer_block_byte_size(layer_idx); if (host_allocator) { const xla::PjRtDevice* target_dev = dst_buffer.device; - auto status_or_allocation = host_allocator(alloc_size, target_dev); - if (!status_or_allocation.ok()) { - throw std::runtime_error(absl::StrCat( - "Host allocator failed for size: ", alloc_size, - ", error: ", status_or_allocation.status().ToString())); + auto host_alloc = host_allocator(alloc_size, target_dev); + if (!host_alloc.ok()) { + throw std::runtime_error( + absl::StrCat("Host allocator failed for size: ", alloc_size, + ", error: ", host_alloc.status().ToString())); } - HostBufferAllocation allocation = - std::move(status_or_allocation).value(); + HostBufferAllocation allocation = *std::move(host_alloc); if (alloc_size > 0 && allocation.ptr == nullptr) { throw std::runtime_error(absl::StrCat( "Host allocator returned null buffer for size: ", alloc_size)); @@ -445,14 +444,13 @@ KVCacheManagerBase::KVCacheManagerBase( int num_host_blocks = host_blocks_to_allocate.value_or(0); size_t alloc_size = num_host_blocks * layer_block_byte_size(layer_idx); if (host_allocator) { - auto status_or_allocation = host_allocator(alloc_size, nullptr); - if (!status_or_allocation.ok()) { - throw std::runtime_error(absl::StrCat( - "Host allocator failed for size: ", alloc_size, - ", error: ", status_or_allocation.status().ToString())); + auto host_alloc = host_allocator(alloc_size, nullptr); + if (!host_alloc.ok()) { + throw std::runtime_error( + absl::StrCat("Host allocator failed for size: ", alloc_size, + ", error: ", host_alloc.status().ToString())); } - HostBufferAllocation allocation = - std::move(status_or_allocation).value(); + HostBufferAllocation allocation = *std::move(host_alloc); if (alloc_size > 0 && allocation.ptr == nullptr) { throw std::runtime_error(absl::StrCat( "Host allocator returned null buffer for size: ", alloc_size)); @@ -543,21 +541,20 @@ void KVCacheManagerBase::WorkerLoop() { task = std::move(task_queue_.front()); task_queue_.pop(); } - auto status_or_future = std::move(task.work)(); - if (!status_or_future.ok()) { - task.promise.Set(status_or_future.status()); + auto future = std::move(task.work)(); + if (!future.ok()) { + task.promise.Set(future.status()); } else { // TODO(b/539581381): Research whether using PJRT_Event_OnReady in a // PJRT-owned thread or a dedicated completion polling thread is // preferable to per-call callback threads in OnReady. - status_or_future->OnReady( - [promise = std::move(task.promise)](auto status_or_holds) mutable { - if (status_or_holds.ok()) { - promise.Set(); - } else { - promise.Set(status_or_holds.status()); - } - }); + future->OnReady([promise = std::move(task.promise)](auto holds) mutable { + if (holds.ok()) { + promise.Set(); + } else { + promise.Set(holds.status()); + } + }); } } } @@ -646,14 +643,10 @@ absl::StatusOr KVCacheManagerBase::H2dSyncDispatch( for (const auto& [node, works] : grouped_work) { VLOG(1) << "H2d: Executing inline dispatch for NUMA node " << node << ", works count: " << works.size(); - auto status_or_local_futures = - DispatchH2dWork(works, slot_idx, is_partial, src_c, dst_c, sizes_c); - if (!status_or_local_futures.ok()) { - VLOG(1) << "H2d: Inline dispatch failed: " - << status_or_local_futures.status().ToString(); - return status_or_local_futures.status(); - } - auto local_futures = std::move(status_or_local_futures).value(); + ABSL_ASSIGN_OR_RETURN( + auto local_futures, + DispatchH2dWork(works, slot_idx, is_partial, src_c, dst_c, sizes_c), + _.VLog(1) << "H2d: Inline dispatch failed: "); for (size_t i = 0; i < works.size(); ++i) { const auto& work = works[i]; logical_futures[work.layer_idx * num_shards_ + work.shard_idx] = @@ -679,14 +672,9 @@ absl::StatusOr KVCacheManagerBase::H2dSyncDispatch( VLOG(1) << "H2d: Awaiting scheduled dispatches..."; for (auto& pf : pending_futures) { - auto status_or_local_futures = pf.future.get(); - if (!status_or_local_futures.ok()) { - VLOG(1) << "H2d: Scheduled dispatch failed: " - << status_or_local_futures.status().ToString(); - return status_or_local_futures.status(); - } + ABSL_ASSIGN_OR_RETURN(auto local_futures, pf.future.get(), + _.VLog(1) << "H2d: Scheduled dispatch failed: "); VLOG(1) << "H2d: Scheduled dispatch completed successfully."; - auto local_futures = std::move(status_or_local_futures).value(); for (size_t i = 0; i < pf.works.size(); ++i) { const auto& work = pf.works[i]; logical_futures[work.layer_idx * num_shards_ + work.shard_idx] = @@ -784,14 +772,10 @@ KVCacheManagerBase::DispatchD2hChunks(const std::vector& src_offsets, for (const auto& [node, works] : grouped_work) { VLOG(1) << "DispatchD2hChunks: Executing inline dispatch for NUMA node " << node << ", works count: " << works.size(); - auto status_or_local_futures = - DispatchD2hWork(works, slot_idx, is_partial, src_c, dst_c, sizes_c); - if (!status_or_local_futures.ok()) { - VLOG(1) << "DispatchD2hChunks: Inline dispatch failed: " - << status_or_local_futures.status().ToString(); - return status_or_local_futures.status(); - } - auto local_futures = std::move(status_or_local_futures).value(); + ABSL_ASSIGN_OR_RETURN( + auto local_futures, + DispatchD2hWork(works, slot_idx, is_partial, src_c, dst_c, sizes_c), + _.VLog(1) << "DispatchD2hChunks: Inline dispatch failed: "); for (size_t i = 0; i < works.size(); ++i) { const auto& work = works[i]; logical_futures[work.layer_idx * num_shards_ + work.shard_idx] = @@ -818,15 +802,11 @@ KVCacheManagerBase::DispatchD2hChunks(const std::vector& src_offsets, VLOG(1) << "DispatchD2hChunks: Awaiting scheduled dispatches..."; for (auto& pf : pending_futures) { - auto status_or_local_futures = pf.future.get(); - if (!status_or_local_futures.ok()) { - VLOG(1) << "DispatchD2hChunks: Scheduled dispatch failed: " - << status_or_local_futures.status().ToString(); - return status_or_local_futures.status(); - } + ABSL_ASSIGN_OR_RETURN( + auto local_futures, pf.future.get(), + _.VLog(1) << "DispatchD2hChunks: Scheduled dispatch failed: "); VLOG(1) << "DispatchD2hChunks: Scheduled dispatch completed successfully."; - auto local_futures = std::move(status_or_local_futures).value(); for (size_t i = 0; i < pf.works.size(); ++i) { const auto& work = pf.works[i]; logical_futures[work.layer_idx * num_shards_ + work.shard_idx] = @@ -987,16 +967,16 @@ absl::StatusOr KVCacheManagerBase::H2dRead( return; } - auto h2h_fut_or = + auto h2h_fut = H2hReadExplicit(peer_str, {src_block_id}, {staging_block_id}, /*explicit_dst_ptrs=*/{}); - if (!h2h_fut_or.ok()) { - state->SetError(h2h_fut_or.status()); + if (!h2h_fut.ok()) { + state->SetError(h2h_fut.status()); state->MarkChunkComplete(); return; } - absl::Status h2h_status = h2h_fut_or->Await(); + absl::Status h2h_status = h2h_fut->Await(); if (!h2h_status.ok()) { state->SetError(h2h_status); state->MarkChunkComplete(); @@ -1008,17 +988,17 @@ absl::StatusOr KVCacheManagerBase::H2dRead( return; } - auto h2d_fut_or = + auto h2d_fut = H2dSyncDispatch({staging_offset}, {dst_device_offset}, {size}); - if (!h2d_fut_or.ok()) { - state->SetError(h2d_fut_or.status()); + if (!h2d_fut.ok()) { + state->SetError(h2d_fut.status()); state->MarkChunkComplete(); return; } - h2d_fut_or->OnReady([state](auto status_or) { - if (!status_or.ok()) { - state->SetError(status_or.status()); + h2d_fut->OnReady([state](const auto& result) { + if (!result.ok()) { + state->SetError(result.status()); } state->MarkChunkComplete(); }); @@ -1117,9 +1097,9 @@ absl::StatusOr KVCacheManagerBase::D2hWrite( int staging_block_id = chunks[i].staging_block_id; int dst_block_id = chunks[i].dst_block_id; chunks[i].d2h_fut.OnReady([this, pool, state, peer_str, staging_block_id, - dst_block_id](auto status_or) { - if (!status_or.ok()) { - state->SetError(status_or.status()); + dst_block_id](const auto& result) { + if (!result.ok()) { + state->SetError(result.status()); state->MarkChunkComplete(); return; } @@ -2597,10 +2577,10 @@ absl::Status KVCacheManagerBase::PushKVCacheResharded( } } - d2h_future.OnReady([this, request, peer_transfers, numa](auto status_or) { - if (!status_or.ok()) { + d2h_future.OnReady([this, request, peer_transfers, numa](const auto& result) { + if (!result.ok()) { LOG(ERROR) << "D2H copy failed for resharded push uuid " << request.uuid() - << ": " << status_or.status().ToString(); + << ": " << result.status().ToString(); return; } @@ -2883,12 +2863,12 @@ KVCacheManagerBase::GetBlockChunks(size_t layer_idx, size_t shard_idx, const PoolSpec& pool = pools_[layer_idx]; for (int64_t block_id : block_ids) { if (accumulated_bytes >= total_bytes) break; - auto extents_or = ComputePoolBlockCopyExtents( + auto extents_status = ComputePoolBlockCopyExtents( pool, absl::MakeConstSpan(&block_id, 1)); - if (!extents_or.ok()) return {}; + if (!extents_status.ok()) return {}; uint8_t* storage_base = GetHostPointer(pool.storage_index, shard_idx); if (storage_base == nullptr) return {}; - for (const PoolBlockCopyExtent& extent : *extents_or) { + for (const PoolBlockCopyExtent& extent : *extents_status) { if (accumulated_bytes >= total_bytes) break; const size_t size = std::min(static_cast(extent.size_bytes), total_bytes - accumulated_bytes); diff --git a/tpu_sync/kv_cache/kv_cache_manager_test.cc b/tpu_sync/kv_cache/kv_cache_manager_test.cc index 7ef1f3f4..90b8e39a 100644 --- a/tpu_sync/kv_cache/kv_cache_manager_test.cc +++ b/tpu_sync/kv_cache/kv_cache_manager_test.cc @@ -27,11 +27,13 @@ #include #include "absl/base/thread_annotations.h" #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/strings/str_cat.h" -#include "absl/synchronization/mutex.h" #include "absl/strings/string_view.h" -#include "tpu_sync/core/raw_transfer_core.h" +#include "absl/synchronization/mutex.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/core/raiden_manager_base.h" +#include "tpu_sync/core/raw_transfer_core.h" #include "tpu_sync/kv_cache/kv_cache_manager_base.h" #include "tpu_sync/rpc/raiden_service.pb.h" #include "tpu_sync/telemetry/metrics_api.h" @@ -122,17 +124,17 @@ TEST(KVCacheManagerTest, RegisterPoolsValidatesAgainstStorage) { /*slice_byte_size=*/128, /*host_blocks=*/2); // Overlapping pools on one storage are allowed (aliased-raw pattern). - absl::Status status = manager.RegisterPools({ + ABSL_ASSERT_OK(manager.RegisterPools({ DensePool("kind_a", 0, 0, 128, 2), StridedPool("kind_b", 0, 0, 128, 2), DensePool("kind_a", 1, 64, 64, 3), - }); - ASSERT_TRUE(status.ok()) << status.ToString(); + })); EXPECT_TRUE(manager.has_explicit_pools()); EXPECT_EQ(manager.num_pools(), 3); // storage_index out of range. - status = manager.RegisterPools({DensePool("kind_a", 5, 0, 128, 2)}); + absl::Status status = + manager.RegisterPools({DensePool("kind_a", 5, 0, 128, 2)}); EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); EXPECT_THAT(status.message(), testing::HasSubstr("storage_index")); @@ -147,8 +149,8 @@ TEST(KVCacheManagerTest, RegisterPoolsValidatesAgainstStorage) { device_backed.SetLayerPhysicalSizeForTest(/*layer_idx=*/0, /*physical_size=*/128, /*major_dim_size=*/1); - EXPECT_TRUE( - device_backed.RegisterPools({DensePool("kind_a", 0, 0, 128, 1)}).ok()); + ABSL_EXPECT_OK( + device_backed.RegisterPools({DensePool("kind_a", 0, 0, 128, 1)})); status = device_backed.RegisterPools({DensePool("kind_a", 0, 0, 128, 2)}); EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); EXPECT_THAT(status.message(), testing::HasSubstr("exceeds storage bytes")); @@ -158,32 +160,30 @@ TEST(KVCacheManagerTest, RegisterPoolsValidatesAgainstStorage) { TEST(KVCacheManagerTest, GetPoolBlockRefPointerMath) { TestKVCacheManager manager(/*num_layers=*/1, /*num_shards=*/1, /*slice_byte_size=*/128, /*host_blocks=*/4); - ASSERT_TRUE(manager - .RegisterPools({ - DensePool("kind_a", 0, 0, 128, 2), - DensePool("kind_b", 0, 256, 64, 4), - }) - .ok()); + ABSL_ASSERT_OK(manager.RegisterPools({ + DensePool("kind_a", 0, 0, 128, 2), + DensePool("kind_b", 0, 256, 64, 4), + })); uint8_t* base = manager.GetHostPointer(/*layer_idx=*/0, /*shard_idx=*/0); - auto ref = manager.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, - /*block_id=*/1); - ASSERT_TRUE(ref.ok()) << ref.status().ToString(); - EXPECT_EQ(ref->ptr, base + 128); - EXPECT_EQ(ref->block_stride_bytes, 128); - EXPECT_EQ(ref->pool->tag, "kind_a"); - EXPECT_EQ(ref->pool_idx, 0); - EXPECT_EQ(ref->block_id, 1); - - ref = manager.GetPoolBlockRef(/*pool_idx=*/1, /*shard_idx=*/0, - /*block_id=*/0); - ASSERT_TRUE(ref.ok()) << ref.status().ToString(); - EXPECT_EQ(ref->ptr, base + 256); - ref = manager.GetPoolBlockRef(/*pool_idx=*/1, /*shard_idx=*/0, - /*block_id=*/3); - ASSERT_TRUE(ref.ok()) << ref.status().ToString(); - EXPECT_EQ(ref->ptr, base + 256 + 3 * 64); + TF_ASSERT_OK_AND_ASSIGN( + auto ref, + manager.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, /*block_id=*/1)); + EXPECT_EQ(ref.ptr, base + 128); + EXPECT_EQ(ref.block_stride_bytes, 128); + EXPECT_EQ(ref.pool->tag, "kind_a"); + EXPECT_EQ(ref.pool_idx, 0); + EXPECT_EQ(ref.block_id, 1); + + TF_ASSERT_OK_AND_ASSIGN( + ref, + manager.GetPoolBlockRef(/*pool_idx=*/1, /*shard_idx=*/0, /*block_id=*/0)); + EXPECT_EQ(ref.ptr, base + 256); + TF_ASSERT_OK_AND_ASSIGN( + ref, + manager.GetPoolBlockRef(/*pool_idx=*/1, /*shard_idx=*/0, /*block_id=*/3)); + EXPECT_EQ(ref.ptr, base + 256 + 3 * 64); EXPECT_EQ(manager.GetPoolBlockRef(0, 0, 2).status().code(), absl::StatusCode::kOutOfRange); @@ -207,9 +207,7 @@ TEST(KVCacheManagerTest, RegisterPoolsGrowsDeviceBackedHostMirror) { manager.SetLayerPhysicalSizeForTest(/*layer_idx=*/0, /*physical_size=*/256, /*major_dim_size=*/1); - absl::Status status = - manager.RegisterPools({DensePool("kind_a", 0, 0, 64, 4)}); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(manager.RegisterPools({DensePool("kind_a", 0, 0, 64, 4)})); EXPECT_GE(manager.GetHostSize(/*layer_idx=*/0, /*shard_idx=*/0), 256); uint8_t* grown = manager.GetHostPointer(/*layer_idx=*/0, /*shard_idx=*/0); @@ -217,10 +215,10 @@ TEST(KVCacheManagerTest, RegisterPoolsGrowsDeviceBackedHostMirror) { for (size_t i = 0; i < 64; ++i) { EXPECT_EQ(grown[i], static_cast(i)); } - auto last_ref = manager.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, - /*block_id=*/3); - ASSERT_TRUE(last_ref.ok()) << last_ref.status().ToString(); - EXPECT_EQ(last_ref->ptr, grown + 3 * 64); + TF_ASSERT_OK_AND_ASSIGN(auto last_ref, manager.GetPoolBlockRef( + /*pool_idx=*/0, /*shard_idx=*/0, + /*block_id=*/3)); + EXPECT_EQ(last_ref.ptr, grown + 3 * 64); } // The pool table is frozen while plans are active. @@ -231,11 +229,11 @@ TEST(KVCacheManagerTest, RegisterPoolsFailsAfterActivePlanRegistered) { request.set_uuid(445566); request.set_is_sender(true); - absl::Status status = - manager.RegisterActivePlan(445566, request, /*is_sender=*/true); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK( + manager.RegisterActivePlan(445566, request, /*is_sender=*/true)); - status = manager.RegisterPools({DensePool("kind_a", 0, 0, 128, 1)}); + absl::Status status = + manager.RegisterPools({DensePool("kind_a", 0, 0, 128, 1)}); EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); EXPECT_THAT(status.message(), testing::HasSubstr("active plans")); } @@ -245,13 +243,11 @@ TEST(KVCacheManagerTest, PoolIndicesWithTag) { /*slice_byte_size=*/128, /*host_blocks=*/2); PoolSpec pool_b = DensePool("kind_b", 1, 0, 128, 2); pool_b.dtype_tag = "dtype_b"; - ASSERT_TRUE(manager - .RegisterPools({ - DensePool("kind_a", 0, 0, 128, 2), - pool_b, - StridedPool("kind_a", 1, 0, 128, 2), - }) - .ok()); + ABSL_ASSERT_OK(manager.RegisterPools({ + DensePool("kind_a", 0, 0, 128, 2), + pool_b, + StridedPool("kind_a", 1, 0, 128, 2), + })); EXPECT_THAT(manager.PoolIndicesWithTag("kind_a"), testing::ElementsAre(0, 2)); EXPECT_THAT(manager.PoolIndicesWithTag("kind_b"), testing::ElementsAre(1)); @@ -263,13 +259,11 @@ TEST(KVCacheManagerTest, RegisterActivePlanChecksPoolDtypeTags) { /*slice_byte_size=*/128, /*host_blocks=*/2); PoolSpec pool_b = DensePool("kind_b", 1, 0, 128, 2); pool_b.dtype_tag = "dtype_b"; - ASSERT_TRUE(manager - .RegisterPools({ - DensePool("kind_a", 0, 0, 128, 2), - pool_b, - StridedPool("kind_a", 1, 0, 128, 2), - }) - .ok()); + ABSL_ASSERT_OK(manager.RegisterPools({ + DensePool("kind_a", 0, 0, 128, 2), + pool_b, + StridedPool("kind_a", 1, 0, 128, 2), + })); tpu_sync::rpc::StartTransferRequest request; request.set_uuid(777); @@ -287,21 +281,18 @@ TEST(KVCacheManagerTest, RegisterActivePlanChecksPoolDtypeTags) { EXPECT_THAT(status.message(), testing::HasSubstr("dtype tag mismatch")); request.set_pool_dtype_tags(1, "dtype_b"); - status = manager.RegisterActivePlan(777, request, /*is_sender=*/true); - ASSERT_TRUE(status.ok()) << status.ToString(); - EXPECT_TRUE(manager.UnregisterActivePlan(777).ok()); + ABSL_ASSERT_OK(manager.RegisterActivePlan(777, request, /*is_sender=*/true)); + ABSL_EXPECT_OK(manager.UnregisterActivePlan(777)); } TEST(KVCacheManagerTest, ExplicitPoolAddressingUsesPoolBaseAndStride) { TestKVCacheManager manager(/*num_layers=*/2, /*num_shards=*/1, /*slice_byte_size=*/128, /*host_blocks=*/4); - ASSERT_TRUE(manager - .RegisterPools({ - DensePool("kind_a", /*storage_index=*/1, - /*base_offset=*/128, /*stride=*/64, - /*num_blocks=*/4), - }) - .ok()); + ABSL_ASSERT_OK(manager.RegisterPools({ + DensePool("kind_a", /*storage_index=*/1, + /*base_offset=*/128, /*stride=*/64, + /*num_blocks=*/4), + })); uint8_t* storage_base = manager.GetHostPointer(/*layer_idx=*/1, /*shard_idx=*/0); @@ -333,32 +324,31 @@ TEST(KVCacheManagerTest, DensePool("kind_a", /*storage_index=*/1, /*base_offset=*/64, /*stride=*/128, /*num_blocks=*/2), }; - ASSERT_TRUE(sender.RegisterPools(pools).ok()); - ASSERT_TRUE(receiver.RegisterPools(pools).ok()); - - auto src_ref = sender.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, - /*block_id=*/0); - auto dst_ref = receiver.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, - /*block_id=*/1); - ASSERT_TRUE(src_ref.ok()) << src_ref.status().ToString(); - ASSERT_TRUE(dst_ref.ok()) << dst_ref.status().ToString(); + ABSL_ASSERT_OK(sender.RegisterPools(pools)); + ABSL_ASSERT_OK(receiver.RegisterPools(pools)); + + TF_ASSERT_OK_AND_ASSIGN( + auto src_ref, + sender.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, /*block_id=*/0)); + TF_ASSERT_OK_AND_ASSIGN(auto dst_ref, receiver.GetPoolBlockRef( + /*pool_idx=*/0, /*shard_idx=*/0, + /*block_id=*/1)); std::vector pattern(128); for (size_t i = 0; i < pattern.size(); ++i) { pattern[i] = static_cast((i * 17 + 3) % 251); } - std::memcpy(src_ref->ptr, pattern.data(), pattern.size()); + std::memcpy(src_ref.ptr, pattern.data(), pattern.size()); const std::optional receiver_port = receiver.local_port(); ASSERT_TRUE(receiver_port.has_value()); tpu_sync::rpc::StartTransferRequest dummy_plan; - ASSERT_TRUE( - receiver.RegisterActivePlan(0, dummy_plan, /*is_sender=*/false).ok()); - auto pushed = sender.H2hWriteDirect( + ABSL_ASSERT_OK( + receiver.RegisterActivePlan(0, dummy_plan, /*is_sender=*/false)); + ABSL_ASSERT_OK(sender.H2hWriteDirect( absl::StrCat(receiver.local_ip(), ":", *receiver_port), /*src_block_ids=*/{0}, /*dst_block_ids=*/{1}, /*uuid=*/0, - /*layer_idx=*/0); - ASSERT_TRUE(pushed.ok()) << pushed.status().ToString(); - EXPECT_EQ(std::memcmp(dst_ref->ptr, pattern.data(), pattern.size()), 0); + /*layer_idx=*/0)); + EXPECT_EQ(std::memcmp(dst_ref.ptr, pattern.data(), pattern.size()), 0); const uint8_t* receiver_storage0 = receiver.GetHostPointer(/*layer_idx=*/0, /*shard_idx=*/0); @@ -395,8 +385,8 @@ TEST(KVCacheManagerTest, StridedPool("aliased", /*storage_index=*/0, /*base_offset=*/32, /*stride=*/128, /*num_blocks=*/2), }; - ASSERT_TRUE(sender.RegisterPools(pools).ok()); - ASSERT_TRUE(receiver.RegisterPools(pools).ok()); + ABSL_ASSERT_OK(sender.RegisterPools(pools)); + ABSL_ASSERT_OK(receiver.RegisterPools(pools)); uint8_t* src = sender.GetHostPointer(/*layer_idx=*/0, /*shard_idx=*/0); uint8_t* dst = receiver.GetHostPointer(/*layer_idx=*/0, /*shard_idx=*/0); @@ -409,22 +399,21 @@ TEST(KVCacheManagerTest, // The logical two-stride array would end at byte 288. Its last live byte is // exactly byte 256, so admission and the last block reference are valid. - auto last_ref = sender.GetPoolBlockRef(/*pool_idx=*/0, /*shard_idx=*/0, - /*block_id=*/1); - ASSERT_TRUE(last_ref.ok()) << last_ref.status().ToString(); - EXPECT_EQ(last_ref->ptr, src + 160); + TF_ASSERT_OK_AND_ASSIGN(auto last_ref, sender.GetPoolBlockRef( + /*pool_idx=*/0, /*shard_idx=*/0, + /*block_id=*/1)); + EXPECT_EQ(last_ref.ptr, src + 160); EXPECT_EQ(sender.GetBlockArrayHostSize(/*pool_idx=*/0, /*shard_idx=*/0), 224); const std::optional receiver_port = receiver.local_port(); ASSERT_TRUE(receiver_port.has_value()); tpu_sync::rpc::StartTransferRequest dummy_plan; - ASSERT_TRUE( - receiver.RegisterActivePlan(0, dummy_plan, /*is_sender=*/false).ok()); - auto pushed = sender.H2hWriteDirect( + ABSL_ASSERT_OK( + receiver.RegisterActivePlan(0, dummy_plan, /*is_sender=*/false)); + ABSL_ASSERT_OK(sender.H2hWriteDirect( absl::StrCat(receiver.local_ip(), ":", *receiver_port), /*src_block_ids=*/{0}, /*dst_block_ids=*/{1}, /*uuid=*/0, - /*layer_idx=*/0); - ASSERT_TRUE(pushed.ok()) << pushed.status().ToString(); + /*layer_idx=*/0)); EXPECT_TRUE(std::all_of(dst + 160, dst + 192, [](uint8_t value) { return value == 0x11; })); EXPECT_TRUE(std::all_of(dst + 192, dst + 224, @@ -444,63 +433,60 @@ TEST(KVCacheManagerTest, ExplicitPoolTransportEnumeratesAllPools) { DensePool("kind_b", /*storage_index=*/0, /*base_offset=*/256, /*stride=*/64, /*num_blocks=*/2), }; - ASSERT_TRUE(sender.RegisterPools(pools).ok()); - ASSERT_TRUE(receiver.RegisterPools(pools).ok()); + ABSL_ASSERT_OK(sender.RegisterPools(pools)); + ABSL_ASSERT_OK(receiver.RegisterPools(pools)); ASSERT_EQ(sender.num_block_arrays(), 2); ASSERT_EQ(receiver.num_block_arrays(), 2); - auto sender_a = sender.GetPoolBlockRef(0, 0, 0); - auto sender_b = sender.GetPoolBlockRef(1, 0, 0); - auto receiver_a = receiver.GetPoolBlockRef(0, 0, 1); - auto receiver_b = receiver.GetPoolBlockRef(1, 0, 1); - ASSERT_TRUE(sender_a.ok() && sender_b.ok() && receiver_a.ok() && - receiver_b.ok()); - std::memset(sender_a->ptr, 0xA1, 64); - std::memset(sender_b->ptr, 0xB2, 64); + TF_ASSERT_OK_AND_ASSIGN(auto sender_a, sender.GetPoolBlockRef(0, 0, 0)); + TF_ASSERT_OK_AND_ASSIGN(auto sender_b, sender.GetPoolBlockRef(1, 0, 0)); + TF_ASSERT_OK_AND_ASSIGN(auto receiver_a, receiver.GetPoolBlockRef(0, 0, 1)); + TF_ASSERT_OK_AND_ASSIGN(auto receiver_b, receiver.GetPoolBlockRef(1, 0, 1)); + std::memset(sender_a.ptr, 0xA1, 64); + std::memset(sender_b.ptr, 0xB2, 64); const std::optional receiver_port = receiver.local_port(); ASSERT_TRUE(receiver_port.has_value()); tpu_sync::rpc::StartTransferRequest dummy_plan; - ASSERT_TRUE( - receiver.RegisterActivePlan(0, dummy_plan, /*is_sender=*/false).ok()); - auto pushed = sender.H2hWriteDirect( + ABSL_ASSERT_OK( + receiver.RegisterActivePlan(0, dummy_plan, /*is_sender=*/false)); + ABSL_ASSERT_OK(sender.H2hWriteDirect( absl::StrCat(receiver.local_ip(), ":", *receiver_port), /*src_block_ids=*/{0}, /*dst_block_ids=*/{1}, /*uuid=*/0, - /*layer_idx=*/-1); - ASSERT_TRUE(pushed.ok()) << pushed.status().ToString(); - EXPECT_TRUE(std::all_of(receiver_a->ptr, receiver_a->ptr + 64, + /*layer_idx=*/-1)); + EXPECT_TRUE(std::all_of(receiver_a.ptr, receiver_a.ptr + 64, [](uint8_t value) { return value == 0xA1; })); - EXPECT_TRUE(std::all_of(receiver_b->ptr, receiver_b->ptr + 64, + EXPECT_TRUE(std::all_of(receiver_b.ptr, receiver_b.ptr + 64, [](uint8_t value) { return value == 0xB2; })); - auto receiver_a0 = receiver.GetPoolBlockRef(0, 0, 0); - auto receiver_b0 = receiver.GetPoolBlockRef(1, 0, 0); - ASSERT_TRUE(receiver_a0.ok() && receiver_b0.ok()); - std::memset(receiver_a0->ptr, 0, 64); - std::memset(receiver_b0->ptr, 0, 64); + TF_ASSERT_OK_AND_ASSIGN(auto receiver_a0, receiver.GetPoolBlockRef(0, 0, 0)); + TF_ASSERT_OK_AND_ASSIGN(auto receiver_b0, receiver.GetPoolBlockRef(1, 0, 0)); + std::memset(receiver_a0.ptr, 0, 64); + std::memset(receiver_b0.ptr, 0, 64); const std::optional sender_port = sender.local_port(); ASSERT_TRUE(sender_port.has_value()); - auto pulled = + TF_ASSERT_OK_AND_ASSIGN( + auto pulled, receiver.H2hReadDirect(absl::StrCat(sender.local_ip(), ":", *sender_port), - /*src_block_ids=*/{0}); - ASSERT_TRUE(pulled.ok()) << pulled.status().ToString(); - ASSERT_EQ(*pulled, std::vector({0})); - EXPECT_TRUE(std::all_of(receiver_a0->ptr, receiver_a0->ptr + 64, + /*src_block_ids=*/{0})); + ASSERT_EQ(pulled, std::vector({0})); + EXPECT_TRUE(std::all_of(receiver_a0.ptr, receiver_a0.ptr + 64, [](uint8_t value) { return value == 0xA1; })); - EXPECT_TRUE(std::all_of(receiver_b0->ptr, receiver_b0->ptr + 64, + EXPECT_TRUE(std::all_of(receiver_b0.ptr, receiver_b0.ptr + 64, [](uint8_t value) { return value == 0xB2; })); std::vector external_a(2 * 64, 0); std::vector external_b(2 * 64, 0); std::vector explicit_pool_bases = {external_a.data(), external_b.data()}; - auto explicit_pull = receiver.H2hReadExplicit( - absl::StrCat(sender.local_ip(), ":", *sender_port), - /*src_block_ids=*/{0}, /*local_block_ids=*/{1}, explicit_pool_bases, - /*parallelism=*/1, transport::MajorOrder::kLayerMajor, - /*on_block_received=*/nullptr); - ASSERT_TRUE(explicit_pull.ok()) << explicit_pull.status().ToString(); - ASSERT_TRUE(explicit_pull->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto explicit_pull, + receiver.H2hReadExplicit( + absl::StrCat(sender.local_ip(), ":", *sender_port), + /*src_block_ids=*/{0}, /*local_block_ids=*/{1}, explicit_pool_bases, + /*parallelism=*/1, transport::MajorOrder::kLayerMajor, + /*on_block_received=*/nullptr)); + ABSL_ASSERT_OK(explicit_pull.Await()); EXPECT_TRUE(std::all_of(external_a.begin(), external_a.begin() + 64, [](uint8_t value) { return value == 0; })); EXPECT_TRUE(std::all_of(external_a.begin() + 64, external_a.end(), @@ -516,7 +502,7 @@ TEST(KVCacheManagerTest, PoolBlockCopiesRejectHostOnlyManager) { /*slice_byte_size=*/64, /*local_port=*/std::nullopt, /*host_blocks_to_allocate=*/2); - ASSERT_TRUE(manager.RegisterPools({DensePool("kind_a", 0, 0, 64, 2)}).ok()); + ABSL_ASSERT_OK(manager.RegisterPools({DensePool("kind_a", 0, 0, 64, 2)})); auto d2h = manager.D2hPoolBlocks(/*pool_idx=*/0, /*block_ids=*/{0}); ASSERT_FALSE(d2h.ok()); @@ -545,12 +531,12 @@ TEST(KVCacheManagerTest, ImplicitPoolsMirrorStorages) { EXPECT_EQ(pool->block_stride_bytes, 128); EXPECT_EQ(pool->num_blocks, 2); - auto ref = manager.GetPoolBlockRef(/*pool_idx=*/1, /*shard_idx=*/0, - /*block_id=*/1); - ASSERT_TRUE(ref.ok()) << ref.status().ToString(); - EXPECT_EQ(ref->ptr, manager.GetHostPointer(/*layer_idx=*/1, - /*shard_idx=*/0) + - 128); + TF_ASSERT_OK_AND_ASSIGN( + auto ref, manager.GetPoolBlockRef(/*pool_idx=*/1, /*shard_idx=*/0, + /*block_id=*/1)); + EXPECT_EQ(ref.ptr, manager.GetHostPointer(/*layer_idx=*/1, + /*shard_idx=*/0) + + 128); } TEST(KVCacheManagerTest, UnregisterActivePlanAllowsUuidReuse) { @@ -563,17 +549,16 @@ TEST(KVCacheManagerTest, UnregisterActivePlanAllowsUuidReuse) { absl::Status status = manager.UnregisterActivePlan(112233); EXPECT_EQ(status.code(), absl::StatusCode::kNotFound); - status = manager.RegisterActivePlan(112233, request, /*is_sender=*/true); - EXPECT_TRUE(status.ok()) << status.ToString(); + ABSL_EXPECT_OK( + manager.RegisterActivePlan(112233, request, /*is_sender=*/true)); status = manager.RegisterActivePlan(112233, request, /*is_sender=*/true); EXPECT_EQ(status.code(), absl::StatusCode::kAlreadyExists); - status = manager.UnregisterActivePlan(112233); - EXPECT_TRUE(status.ok()) << status.ToString(); + ABSL_EXPECT_OK(manager.UnregisterActivePlan(112233)); - status = manager.RegisterActivePlan(112233, request, /*is_sender=*/true); - EXPECT_TRUE(status.ok()) << status.ToString(); + ABSL_EXPECT_OK( + manager.RegisterActivePlan(112233, request, /*is_sender=*/true)); } TEST(KVCacheManagerTest, D2hFailsWithMismatchedCopySpecLengths) { @@ -685,14 +670,13 @@ TEST(KVCacheManagerTest, AsymmetricBlockSizesGetBlockChunks) { entry->set_dst_block_id(0); // Register active plan on both sides - absl::Status status = - sender.RegisterActivePlan(112233, request, /*is_sender=*/true); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK( + sender.RegisterActivePlan(112233, request, /*is_sender=*/true)); // Receiver schedule should be the same request but marked as is_sender=false request.set_is_sender(false); - status = receiver.RegisterActivePlan(112233, request, /*is_sender=*/false); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK( + receiver.RegisterActivePlan(112233, request, /*is_sender=*/false)); // 3. Resolve chunks on Sender (should return offset 64 from block 0 base) std::vector src_block_ids = {0}; @@ -836,9 +820,9 @@ TEST(KVCacheManagerTest, D2hWriteSuccessWithMockD2h) { std::vector dst_host_offsets = {1}; // remote destination std::vector copy_sizes = {1}; - auto res = sender.D2hWrite(receiver_peer, src_device_offsets, - src_host_offsets, dst_host_offsets, copy_sizes); - ASSERT_TRUE(res.ok()) << res.status().ToString(); + ABSL_ASSERT_OK(sender.D2hWrite(receiver_peer, src_device_offsets, + src_host_offsets, dst_host_offsets, + copy_sizes)); EXPECT_TRUE(sender.d2h_called_); EXPECT_EQ(sender.last_src_offsets_, src_device_offsets); // The D2H stage lands in the EXPLICIT local staging blocks, not in a local @@ -873,10 +857,11 @@ TEST(KVCacheManagerTest, D2hWritePipelinedSuccess) { std::vector dst_host_offsets = {0, 1}; // remote destination std::vector copy_sizes = {1, 1}; - auto res = sender.D2hWrite(receiver_peer, src_device_offsets, - src_host_offsets, dst_host_offsets, copy_sizes); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, + sender.D2hWrite(receiver_peer, src_device_offsets, src_host_offsets, + dst_host_offsets, copy_sizes)); + ABSL_EXPECT_OK(res.Await()); EXPECT_TRUE(sender.dispatch_d2h_chunks_called_); EXPECT_EQ(sender.dispatched_src_offsets_.size(), 2); @@ -909,17 +894,17 @@ TEST(KVCacheManagerTest, H2dReadSuccess) { std::memset(receiver_buf, 0, 256); // Test empty src_offsets returns OK empty future - auto empty_res = receiver.H2dRead(sender_peer, {}, {}, {}, {}); - ASSERT_TRUE(empty_res.ok()) << empty_res.status().ToString(); - EXPECT_TRUE(empty_res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN(auto empty_res, + receiver.H2dRead(sender_peer, {}, {}, {}, {})); + ABSL_EXPECT_OK(empty_res.Await()); // Test H2dRead reading sender block 0 via local staging block 0 into // receiver device block 0. - auto res = receiver.H2dRead(sender_peer, /*src_host=*/{0}, - /*dst_host(staging)=*/{0}, /*dst_device=*/{0}, - /*copy_sizes=*/{1}); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, receiver.H2dRead(sender_peer, /*src_host=*/{0}, + /*dst_host(staging)=*/{0}, /*dst_device=*/{0}, + /*copy_sizes=*/{1})); + ABSL_EXPECT_OK(res.Await()); EXPECT_TRUE(std::all_of(receiver_buf, receiver_buf + 128, [](uint8_t v) { return v == 0xEF; })); @@ -950,10 +935,11 @@ TEST(KVCacheManagerTest, H2dReadPipelinedSuccess) { std::vector dst_device_offsets = {0, 1}; std::vector copy_sizes = {1, 1}; - auto res = receiver.H2dRead(sender_peer, src_host_offsets, dst_host_offsets, - dst_device_offsets, copy_sizes); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, + receiver.H2dRead(sender_peer, src_host_offsets, dst_host_offsets, + dst_device_offsets, copy_sizes)); + ABSL_EXPECT_OK(res.Await()); EXPECT_TRUE(std::all_of(receiver_buf, receiver_buf + 128, [](uint8_t v) { return v == 0x55; })); @@ -975,11 +961,11 @@ TEST(KVCacheManagerTest, H2dReadCallsH2dForTpuHbmDestination) { ASSERT_NE(sender_buf, nullptr); std::memset(sender_buf, 0x77, 128); - auto res = receiver.H2dRead(sender_peer, /*src_host=*/{0}, - /*dst_host(staging)=*/{0}, /*dst_device=*/{1}, - /*copy_sizes=*/{1}); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, receiver.H2dRead(sender_peer, /*src_host=*/{0}, + /*dst_host(staging)=*/{0}, /*dst_device=*/{1}, + /*copy_sizes=*/{1})); + ABSL_EXPECT_OK(res.Await()); // H2dRead MUST trigger Stage 2 H2d DMA from the explicit staging block {0} // into TPU HBM destination offset {1}. @@ -1014,10 +1000,11 @@ TEST(KVCacheManagerTest, H2dWriteSuccess) { std::vector dst_device_offsets = {0}; std::vector copy_sizes = {1}; - auto res = sender.H2dWrite(receiver_peer, src_host_offsets, dst_host_offsets, - dst_device_offsets, copy_sizes); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, + sender.H2dWrite(receiver_peer, src_host_offsets, dst_host_offsets, + dst_device_offsets, copy_sizes)); + ABSL_EXPECT_OK(res.Await()); EXPECT_TRUE(std::all_of(receiver_buf, receiver_buf + 128, [](uint8_t v) { return v == 0; })); @@ -1051,10 +1038,11 @@ TEST(KVCacheManagerTest, H2dWritePipelinedSuccess) { std::vector dst_device_offsets = {0, 1}; std::vector copy_sizes = {1, 1}; - auto res = sender.H2dWrite(receiver_peer, src_host_offsets, dst_host_offsets, - dst_device_offsets, copy_sizes); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, + sender.H2dWrite(receiver_peer, src_host_offsets, dst_host_offsets, + dst_device_offsets, copy_sizes)); + ABSL_EXPECT_OK(res.Await()); EXPECT_TRUE(std::all_of(receiver_buf, receiver_buf + 128, [](uint8_t v) { return v == 0x33; })); @@ -1088,11 +1076,11 @@ TEST(KVCacheManagerTest, H2dReadExplicitStagingDoesNotClobberAliasedBlock) { std::memset(receiver_buf, 0x99, 128); std::memset(receiver_buf + 128, 0, 128); - auto res = receiver.H2dRead(sender_peer, /*src_host=*/{0}, - /*dst_host(staging)=*/{1}, /*dst_device=*/{0}, - /*copy_sizes=*/{1}); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, receiver.H2dRead(sender_peer, /*src_host=*/{0}, + /*dst_host(staging)=*/{1}, /*dst_device=*/{0}, + /*copy_sizes=*/{1})); + ABSL_EXPECT_OK(res.Await()); // Data staged into the explicit staging block 1. EXPECT_TRUE(std::all_of(receiver_buf + 128, receiver_buf + 256, @@ -1131,11 +1119,11 @@ TEST(KVCacheManagerTest, D2hWriteExplicitStagingIsPushedNotAliasedBlock) { std::memset(sender_buf + 128, 0x99, 128); std::memset(receiver_buf, 0, 256); - auto res = sender.D2hWrite(receiver_peer, /*src_device=*/{0}, - /*src_host(staging)=*/{0}, /*dst_host=*/{1}, - /*copy_sizes=*/{1}); - ASSERT_TRUE(res.ok()) << res.status().ToString(); - EXPECT_TRUE(res->Await().ok()); + TF_ASSERT_OK_AND_ASSIGN( + auto res, sender.D2hWrite(receiver_peer, /*src_device=*/{0}, + /*src_host(staging)=*/{0}, /*dst_host=*/{1}, + /*copy_sizes=*/{1})); + ABSL_EXPECT_OK(res.Await()); // The peer received the STAGING block's payload, not the sentinel from the // sender's local block 1 (the would-be alias of the remote dst id). @@ -1215,17 +1203,14 @@ TEST(KVCacheManagerTest, BackgroundWorkerThreadExecutesInFifoOrder) { /*slice_byte_size=*/128, /*host_blocks=*/2); // Queue H2D, D2H, H2D sequentially - auto f1 = manager.H2d({0}, {0}, {1}); - auto f2 = manager.D2h({0}, {0}, {1}); - auto f3 = manager.H2d({0}, {0}, {1}); - ASSERT_TRUE(f1.ok()); - ASSERT_TRUE(f2.ok()); - ASSERT_TRUE(f3.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto f1, manager.H2d({0}, {0}, {1})); + TF_ASSERT_OK_AND_ASSIGN(auto f2, manager.D2h({0}, {0}, {1})); + TF_ASSERT_OK_AND_ASSIGN(auto f3, manager.H2d({0}, {0}, {1})); // Await all futures - EXPECT_TRUE(f1->Await().ok()); - EXPECT_TRUE(f2->Await().ok()); - EXPECT_TRUE(f3->Await().ok()); + ABSL_EXPECT_OK(f1.Await()); + ABSL_EXPECT_OK(f2.Await()); + ABSL_EXPECT_OK(f3.Await()); absl::MutexLock lock(manager.mu_); EXPECT_EQ(manager.h2d_count_, 2); @@ -1242,8 +1227,7 @@ TEST(KVCacheManagerTest, BackgroundWorkerThreadDisabledByDefault) { TestBackgroundKVCacheManager manager(/*num_layers=*/1, /*num_shards=*/1, /*slice_byte_size=*/128, /*host_blocks=*/2); - auto f1 = manager.H2d({0}, {0}, {1}); - ASSERT_TRUE(f1.ok()); + ABSL_ASSERT_OK(manager.H2d({0}, {0}, {1})); absl::MutexLock lock(manager.mu_); EXPECT_EQ(manager.h2d_count_, 1); } @@ -1270,13 +1254,13 @@ TEST(KVCacheManagerTest, TelemetryMetricsObservedWhenEnabled) { std::vector offsets = {0}; std::vector sizes = {1}; - ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture h2d_res, - manager.H2d(offsets, offsets, sizes)); - EXPECT_OK(h2d_res.Await()); + TF_ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture h2d_res, + manager.H2d(offsets, offsets, sizes)); + ABSL_EXPECT_OK(h2d_res.Await()); - ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture d2h_res, - manager.D2h(offsets, offsets, sizes)); - EXPECT_OK(d2h_res.Await()); + TF_ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture d2h_res, + manager.D2h(offsets, offsets, sizes)); + ABSL_EXPECT_OK(d2h_res.Await()); } TEST(KVCacheManagerTest, TelemetryMetricsSkippedWhenDisabled) { @@ -1290,13 +1274,13 @@ TEST(KVCacheManagerTest, TelemetryMetricsSkippedWhenDisabled) { std::vector offsets = {0}; std::vector sizes = {1}; - ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture h2d_res, - manager.H2d(offsets, offsets, sizes)); - EXPECT_OK(h2d_res.Await()); + TF_ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture h2d_res, + manager.H2d(offsets, offsets, sizes)); + ABSL_EXPECT_OK(h2d_res.Await()); - ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture d2h_res, - manager.D2h(offsets, offsets, sizes)); - EXPECT_OK(d2h_res.Await()); + TF_ASSERT_OK_AND_ASSIGN(raiden::PjRtCopyFuture d2h_res, + manager.D2h(offsets, offsets, sizes)); + ABSL_EXPECT_OK(d2h_res.Await()); } TEST(KVCacheManagerTest, D2hWritePipelinedTelemetryBatchObservation) { @@ -1327,11 +1311,11 @@ TEST(KVCacheManagerTest, D2hWritePipelinedTelemetryBatchObservation) { std::vector dst_host_offsets = {0, 1}; std::vector copy_sizes = {1, 1}; - ASSERT_OK_AND_ASSIGN( + TF_ASSERT_OK_AND_ASSIGN( raiden::PjRtCopyFuture res, sender.D2hWrite(receiver_peer, src_device_offsets, src_host_offsets, dst_host_offsets, copy_sizes)); - EXPECT_OK(res.Await()); + ABSL_EXPECT_OK(res.Await()); } TEST(KVCacheManagerTest, BufferAllocatedHostDramTelemetry) { @@ -1383,10 +1367,9 @@ TEST(KVCacheManagerTest, BufferAllocatedHostDramMaintenance) { manager.SetLayerPhysicalSizeForTest(/*layer_idx=*/0, /*physical_size=*/1024, /*major_dim_size=*/1); - absl::Status status = manager.RegisterPools( + ABSL_ASSERT_OK(manager.RegisterPools( {DensePool("kind_a", /*storage_index=*/0, /*base_offset=*/0, - /*stride=*/128, /*num_blocks=*/8)}); - ASSERT_TRUE(status.ok()) << status.ToString(); + /*stride=*/128, /*num_blocks=*/8)})); EXPECT_EQ(manager.GetAllocatedHostDramBytes(), 3072); // Subsequent call to UpdateAllocatedOccupancyMetric maintains and reports diff --git a/tpu_sync/kv_cache/kv_cache_metadata.cc b/tpu_sync/kv_cache/kv_cache_metadata.cc index a85754fd..5ee5d391 100644 --- a/tpu_sync/kv_cache/kv_cache_metadata.cc +++ b/tpu_sync/kv_cache/kv_cache_metadata.cc @@ -22,6 +22,7 @@ #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" @@ -70,10 +71,7 @@ size_t KVCacheMetadata::RequiredSizeBytes(int num_blocks) { absl::StatusOr KVCacheMetadata::Format( absl::Span region, int num_blocks, absl::string_view model_uid) { - absl::Status status = ValidateRegion(region, num_blocks, model_uid); - if (!status.ok()) { - return status; - } + ABSL_RETURN_IF_ERROR(ValidateRegion(region, num_blocks, model_uid)); std::memset(region.data(), 0, RequiredSizeBytes(num_blocks)); auto* header = reinterpret_cast(region.data()); @@ -90,10 +88,7 @@ absl::StatusOr KVCacheMetadata::Format( absl::StatusOr KVCacheMetadata::Attach( absl::Span region, int num_blocks, absl::string_view model_uid) { - absl::Status status = ValidateRegion(region, num_blocks, model_uid); - if (!status.ok()) { - return status; - } + ABSL_RETURN_IF_ERROR(ValidateRegion(region, num_blocks, model_uid)); auto* header = reinterpret_cast(region.data()); const KVCacheMetadataHeader expected; diff --git a/tpu_sync/kv_cache/kv_cache_metadata_shm.cc b/tpu_sync/kv_cache/kv_cache_metadata_shm.cc index 79e60deb..9ae5cd00 100644 --- a/tpu_sync/kv_cache/kv_cache_metadata_shm.cc +++ b/tpu_sync/kv_cache/kv_cache_metadata_shm.cc @@ -118,17 +118,17 @@ KVCacheMetadataShmRegion::AttachOrFormat(absl::string_view shm_key, } if (mapped != MAP_FAILED) { auto span = absl::MakeSpan(static_cast(mapped), size); - auto metadata_or = KVCacheMetadata::Attach(span, num_blocks, model_uid); - if (metadata_or.ok()) { + auto metadata = KVCacheMetadata::Attach(span, num_blocks, model_uid); + if (metadata.ok()) { VLOG(1) << "Attached to the surviving KV metadata table " << key; return std::unique_ptr( - new KVCacheMetadataShmRegion(fd, mapped, size, *metadata_or, + new KVCacheMetadataShmRegion(fd, mapped, size, *metadata, /*warm=*/true, num_blocks, std::string(model_uid))); } LOG(WARNING) << "Surviving KV metadata table " << key << " failed validation, re-creating: " - << metadata_or.status().message(); + << metadata.status().message(); munmap(mapped, size); } else { LOG(WARNING) << "Surviving KV metadata segment " << key @@ -165,17 +165,17 @@ KVCacheMetadataShmRegion::AttachOrFormat(absl::string_view shm_key, return status; } auto span = absl::MakeSpan(static_cast(mapped), size); - auto metadata_or = KVCacheMetadata::Format(span, num_blocks, model_uid); - if (!metadata_or.ok()) { + auto metadata = KVCacheMetadata::Format(span, num_blocks, model_uid); + if (!metadata.ok()) { munmap(mapped, size); close(fd); shm_unlink(key.c_str()); - return metadata_or.status(); + return metadata.status(); } VLOG(1) << "Formatted a fresh KV metadata table " << key; - return std::unique_ptr(new KVCacheMetadataShmRegion( - fd, mapped, size, *metadata_or, /*warm=*/false, num_blocks, - std::string(model_uid))); + return std::unique_ptr( + new KVCacheMetadataShmRegion(fd, mapped, size, *metadata, /*warm=*/false, + num_blocks, std::string(model_uid))); } KVCacheMetadataShmRegion::KVCacheMetadataShmRegion( @@ -196,11 +196,8 @@ KVCacheMetadataShmRegion::~KVCacheMetadataShmRegion() { absl::Status KVCacheMetadataShmRegion::Reformat() { auto span = absl::MakeSpan(static_cast(mapped_), mapped_size_); - auto metadata_or = KVCacheMetadata::Format(span, num_blocks_, model_uid_); - if (!metadata_or.ok()) { - return metadata_or.status(); - } - metadata_ = *metadata_or; + ABSL_ASSIGN_OR_RETURN(metadata_, + KVCacheMetadata::Format(span, num_blocks_, model_uid_)); warm_ = false; return absl::OkStatus(); } diff --git a/tpu_sync/kv_cache/kv_cache_metadata_shm_test.cc b/tpu_sync/kv_cache/kv_cache_metadata_shm_test.cc index b8e1b5e6..2d7a28ef 100644 --- a/tpu_sync/kv_cache/kv_cache_metadata_shm_test.cc +++ b/tpu_sync/kv_cache/kv_cache_metadata_shm_test.cc @@ -23,8 +23,10 @@ #include #include +#include "absl/status/status_matchers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/kv_cache/kv_cache_metadata.h" namespace tpu_raiden { @@ -53,92 +55,85 @@ class ScopedShmKey { TEST(KVCacheMetadataShmTest, ColdStartFormatsEmptyTable) { ScopedShmKey key("cold"); - auto region_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(region_or.ok()) << region_or.status().ToString(); - EXPECT_FALSE((*region_or)->warm()); - EXPECT_THAT((*region_or)->metadata().ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto region, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 8, "model_a")); + EXPECT_FALSE(region->warm()); + EXPECT_THAT(region->metadata().ValidEntries(), IsEmpty()); } TEST(KVCacheMetadataShmTest, TableSurvivesProcessRestart) { ScopedShmKey key("restart"); - auto region_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(region_or.ok()) << region_or.status().ToString(); - KVCacheMetadata metadata = (*region_or)->metadata(); - ASSERT_TRUE(metadata.Set(1, "hash_b", /*seq=*/2).ok()); - ASSERT_TRUE(metadata.Set(3, "hash_d", /*seq=*/1).ok()); + TF_ASSERT_OK_AND_ASSIGN(auto region, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 8, "model_a")); + KVCacheMetadata metadata = region->metadata(); + ABSL_ASSERT_OK(metadata.Set(1, "hash_b", /*seq=*/2)); + ABSL_ASSERT_OK(metadata.Set(3, "hash_d", /*seq=*/1)); // Dropping the region simulates the engine dying: the mapping goes away, // the segment survives. - region_or->reset(); + region.reset(); - auto revived_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(revived_or.ok()) << revived_or.status().ToString(); - EXPECT_TRUE((*revived_or)->warm()); + TF_ASSERT_OK_AND_ASSIGN( + auto revived, + KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a")); + EXPECT_TRUE(revived->warm()); EXPECT_THAT( - (*revived_or)->metadata().ValidEntries(), + revived->metadata().ValidEntries(), ElementsAre(FieldsAre(1, "hash_b", 2), FieldsAre(3, "hash_d", 1))); } TEST(KVCacheMetadataShmTest, NumBlocksMismatchRecreates) { ScopedShmKey key("num_blocks"); - auto region_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(region_or.ok()) << region_or.status().ToString(); - KVCacheMetadata metadata = (*region_or)->metadata(); - ASSERT_TRUE(metadata.Set(1, "hash_b", /*seq=*/2).ok()); - region_or->reset(); + TF_ASSERT_OK_AND_ASSIGN(auto region, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 8, "model_a")); + KVCacheMetadata metadata = region->metadata(); + ABSL_ASSERT_OK(metadata.Set(1, "hash_b", /*seq=*/2)); + region.reset(); // A different table geometry must not resurrect the old entries, whether // the surviving segment is too small (grown table) or its header disagrees // (shrunk table). - auto grown_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 64, "model_a"); - ASSERT_TRUE(grown_or.ok()) << grown_or.status().ToString(); - EXPECT_FALSE((*grown_or)->warm()); - EXPECT_THAT((*grown_or)->metadata().ValidEntries(), IsEmpty()); - region_or = std::move(grown_or); - KVCacheMetadata grown_metadata = (*region_or)->metadata(); - ASSERT_TRUE(grown_metadata.Set(1, "hash_b", /*seq=*/2).ok()); - region_or->reset(); - - auto shrunk_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(shrunk_or.ok()) << shrunk_or.status().ToString(); - EXPECT_FALSE((*shrunk_or)->warm()); - EXPECT_THAT((*shrunk_or)->metadata().ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto grown, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 64, "model_a")); + EXPECT_FALSE(grown->warm()); + EXPECT_THAT(grown->metadata().ValidEntries(), IsEmpty()); + region = std::move(grown); + KVCacheMetadata grown_metadata = region->metadata(); + ABSL_ASSERT_OK(grown_metadata.Set(1, "hash_b", /*seq=*/2)); + region.reset(); + + TF_ASSERT_OK_AND_ASSIGN(auto shrunk, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 8, "model_a")); + EXPECT_FALSE(shrunk->warm()); + EXPECT_THAT(shrunk->metadata().ValidEntries(), IsEmpty()); } TEST(KVCacheMetadataShmTest, ModelUidMismatchRecreates) { ScopedShmKey key("model_uid"); - auto region_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(region_or.ok()) << region_or.status().ToString(); - KVCacheMetadata metadata = (*region_or)->metadata(); - ASSERT_TRUE(metadata.Set(1, "hash_b", /*seq=*/2).ok()); - region_or->reset(); + TF_ASSERT_OK_AND_ASSIGN(auto region, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 8, "model_a")); + KVCacheMetadata metadata = region->metadata(); + ABSL_ASSERT_OK(metadata.Set(1, "hash_b", /*seq=*/2)); + region.reset(); // A table recorded under another model must not resurrect its bindings. - auto other_model_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_b"); - ASSERT_TRUE(other_model_or.ok()) << other_model_or.status().ToString(); - EXPECT_FALSE((*other_model_or)->warm()); - EXPECT_THAT((*other_model_or)->metadata().ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN( + auto other_model, + KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_b")); + EXPECT_FALSE(other_model->warm()); + EXPECT_THAT(other_model->metadata().ValidEntries(), IsEmpty()); } TEST(KVCacheMetadataShmTest, ReformatErasesEntries) { ScopedShmKey key("reformat"); - auto region_or = - KVCacheMetadataShmRegion::AttachOrFormat(key.key(), 8, "model_a"); - ASSERT_TRUE(region_or.ok()) << region_or.status().ToString(); - KVCacheMetadata metadata = (*region_or)->metadata(); - ASSERT_TRUE(metadata.Set(1, "hash_b", /*seq=*/2).ok()); - - ASSERT_TRUE((*region_or)->Reformat().ok()); - EXPECT_FALSE((*region_or)->warm()); - EXPECT_THAT((*region_or)->metadata().ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto region, KVCacheMetadataShmRegion::AttachOrFormat( + key.key(), 8, "model_a")); + KVCacheMetadata metadata = region->metadata(); + ABSL_ASSERT_OK(metadata.Set(1, "hash_b", /*seq=*/2)); + + ABSL_ASSERT_OK(region->Reformat()); + EXPECT_FALSE(region->warm()); + EXPECT_THAT(region->metadata().ValidEntries(), IsEmpty()); // Copies attached before the reformat observe the emptied table. EXPECT_THAT(metadata.ValidEntries(), IsEmpty()); } diff --git a/tpu_sync/kv_cache/kv_cache_metadata_test.cc b/tpu_sync/kv_cache/kv_cache_metadata_test.cc index 06e41b12..989beaf3 100644 --- a/tpu_sync/kv_cache/kv_cache_metadata_test.cc +++ b/tpu_sync/kv_cache/kv_cache_metadata_test.cc @@ -22,13 +22,16 @@ #include #include #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/tsl/platform/statusor.h" namespace tpu_raiden { namespace kv_cache { namespace { +using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::FieldsAre; using ::testing::IsEmpty; @@ -51,103 +54,98 @@ class Region { TEST(KVCacheMetadataTest, FormatCreatesEmptyTable) { Region region(4); - auto metadata_or = KVCacheMetadata::Format(region.span(), 4); - ASSERT_TRUE(metadata_or.ok()); - EXPECT_EQ(metadata_or->num_blocks(), 4); - EXPECT_THAT(metadata_or->ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 4)); + EXPECT_EQ(metadata.num_blocks(), 4); + EXPECT_THAT(metadata.ValidEntries(), IsEmpty()); } TEST(KVCacheMetadataTest, FormatRejectsInvalidRegions) { Region region(4); - EXPECT_EQ(KVCacheMetadata::Format(region.span(), 0).status().code(), - absl::StatusCode::kInvalidArgument); - EXPECT_EQ(KVCacheMetadata::Format(absl::Span(), 4).status().code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(KVCacheMetadata::Format(region.span(), 0), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(KVCacheMetadata::Format(absl::Span(), 4), + StatusIs(absl::StatusCode::kInvalidArgument)); // Too small for 8 blocks. - EXPECT_EQ( - KVCacheMetadata::Format(region.span().subspan(0, 128), 8).status().code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(KVCacheMetadata::Format(region.span().subspan(0, 128), 8), + StatusIs(absl::StatusCode::kInvalidArgument)); // Misaligned. - EXPECT_EQ( - KVCacheMetadata::Format(region.span().subspan(1), 2).status().code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(KVCacheMetadata::Format(region.span().subspan(1), 2), + StatusIs(absl::StatusCode::kInvalidArgument)); } TEST(KVCacheMetadataTest, SetAndClearRoundTrip) { Region region(4); - auto metadata_or = KVCacheMetadata::Format(region.span(), 4); - ASSERT_TRUE(metadata_or.ok()); - - ASSERT_TRUE(metadata_or->Set(1, "hash_b", /*seq=*/2).ok()); - ASSERT_TRUE(metadata_or->Set(3, "hash_d", /*seq=*/1).ok()); - EXPECT_THAT( - metadata_or->ValidEntries(), - ElementsAre(FieldsAre(1, "hash_b", 2), FieldsAre(3, "hash_d", 1))); - - ASSERT_TRUE(metadata_or->Clear(1).ok()); - EXPECT_THAT(metadata_or->ValidEntries(), - ElementsAre(FieldsAre(3, "hash_d", 1))); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 4)); + + ABSL_ASSERT_OK(metadata.Set(1, "hash_b", /*seq=*/2)); + ABSL_ASSERT_OK(metadata.Set(3, "hash_d", /*seq=*/1)); + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(FieldsAre(1, "hash_b", 2), + FieldsAre(3, "hash_d", 1))); + + ABSL_ASSERT_OK(metadata.Clear(1)); + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(FieldsAre(3, "hash_d", 1))); } TEST(KVCacheMetadataTest, SetOverwritesPreviousBinding) { Region region(2); - auto metadata_or = KVCacheMetadata::Format(region.span(), 2); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 2)); - ASSERT_TRUE(metadata_or->Set(0, "old_hash", /*seq=*/1).ok()); - ASSERT_TRUE(metadata_or->Set(0, "new", /*seq=*/2).ok()); - EXPECT_THAT(metadata_or->ValidEntries(), ElementsAre(FieldsAre(0, "new", 2))); + ABSL_ASSERT_OK(metadata.Set(0, "old_hash", /*seq=*/1)); + ABSL_ASSERT_OK(metadata.Set(0, "new", /*seq=*/2)); + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(FieldsAre(0, "new", 2))); } TEST(KVCacheMetadataTest, SetValidatesArguments) { Region region(2); - auto metadata_or = KVCacheMetadata::Format(region.span(), 2); - ASSERT_TRUE(metadata_or.ok()); - - EXPECT_EQ(metadata_or->Set(-1, "h", 0).code(), - absl::StatusCode::kInvalidArgument); - EXPECT_EQ(metadata_or->Set(2, "h", 0).code(), - absl::StatusCode::kInvalidArgument); - EXPECT_EQ(metadata_or->Set(0, "", 0).code(), - absl::StatusCode::kInvalidArgument); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 2)); + + EXPECT_THAT(metadata.Set(-1, "h", 0), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(metadata.Set(2, "h", 0), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(metadata.Set(0, "", 0), + StatusIs(absl::StatusCode::kInvalidArgument)); std::string too_long(KVCacheMetadata::kMaxHashLength + 1, 'x'); - EXPECT_EQ(metadata_or->Set(0, too_long, 0).code(), - absl::StatusCode::kInvalidArgument); - EXPECT_EQ(metadata_or->Clear(2).code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(metadata.Set(0, too_long, 0), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(metadata.Clear(2), StatusIs(absl::StatusCode::kInvalidArgument)); std::string max_length(KVCacheMetadata::kMaxHashLength, 'y'); - EXPECT_TRUE(metadata_or->Set(0, max_length, 0).ok()); - EXPECT_THAT(metadata_or->ValidEntries(), + ABSL_EXPECT_OK(metadata.Set(0, max_length, 0)); + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(FieldsAre(0, max_length, 0))); } TEST(KVCacheMetadataTest, HashesAreOpaqueBytes) { Region region(1); - auto metadata_or = KVCacheMetadata::Format(region.span(), 1); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 1)); std::string binary_hash("\x00\xff\x00raiden\x01", 10); - ASSERT_TRUE(metadata_or->Set(0, binary_hash, /*seq=*/7).ok()); - EXPECT_THAT(metadata_or->ValidEntries(), + ABSL_ASSERT_OK(metadata.Set(0, binary_hash, /*seq=*/7)); + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(FieldsAre(0, binary_hash, 7))); } TEST(KVCacheMetadataTest, AttachRecoversEntriesFromSurvivingRegion) { Region region(4); { - auto metadata_or = KVCacheMetadata::Format(region.span(), 4); - ASSERT_TRUE(metadata_or.ok()); - ASSERT_TRUE(metadata_or->Set(0, "hash_a", /*seq=*/3).ok()); - ASSERT_TRUE(metadata_or->Set(2, "hash_c", /*seq=*/4).ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 4)); + ABSL_ASSERT_OK(metadata.Set(0, "hash_a", /*seq=*/3)); + ABSL_ASSERT_OK(metadata.Set(2, "hash_c", /*seq=*/4)); // The view is dropped here; the region survives, as shared memory would // across an engine crash. } - auto recovered_or = KVCacheMetadata::Attach(region.span(), 4); - ASSERT_TRUE(recovered_or.ok()); - EXPECT_THAT( - recovered_or->ValidEntries(), - ElementsAre(FieldsAre(0, "hash_a", 3), FieldsAre(2, "hash_c", 4))); + TF_ASSERT_OK_AND_ASSIGN(auto recovered, + KVCacheMetadata::Attach(region.span(), 4)); + EXPECT_THAT(recovered.ValidEntries(), ElementsAre(FieldsAre(0, "hash_a", 3), + FieldsAre(2, "hash_c", 4))); } TEST(KVCacheMetadataTest, AttachRejectsUnformattedOrMismatchedRegions) { @@ -157,52 +155,48 @@ TEST(KVCacheMetadataTest, AttachRejectsUnformattedOrMismatchedRegions) { // Never formatted. std::memset(region.span().data(), 0, KVCacheMetadata::RequiredSizeBytes(4)); - EXPECT_EQ(KVCacheMetadata::Attach(region.span(), 4).status().code(), - absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(KVCacheMetadata::Attach(region.span(), 4), + StatusIs(absl::StatusCode::kFailedPrecondition)); - ASSERT_TRUE(KVCacheMetadata::Format(region.span(), 4).ok()); + ABSL_ASSERT_OK(KVCacheMetadata::Format(region.span(), 4)); // Formatted for 4 blocks, attached expecting 8 (block pool resized across // the restart): the table no longer matches the pool, treat as cold start. - EXPECT_EQ(KVCacheMetadata::Attach(region.span(), 8).status().code(), - absl::StatusCode::kFailedPrecondition); - EXPECT_TRUE(KVCacheMetadata::Attach(region.span(), 4).ok()); + EXPECT_THAT(KVCacheMetadata::Attach(region.span(), 8), + StatusIs(absl::StatusCode::kFailedPrecondition)); + ABSL_EXPECT_OK(KVCacheMetadata::Attach(region.span(), 4)); } TEST(KVCacheMetadataTest, AttachValidatesModelUid) { Region region(4); - ASSERT_TRUE(KVCacheMetadata::Format(region.span(), 4, "model_a").ok()); + ABSL_ASSERT_OK(KVCacheMetadata::Format(region.span(), 4, "model_a")); // A table recorded under another model (or none) must not attach: its // bindings describe another model's blocks. - EXPECT_TRUE(KVCacheMetadata::Attach(region.span(), 4, "model_a").ok()); - EXPECT_EQ( - KVCacheMetadata::Attach(region.span(), 4, "model_b").status().code(), - absl::StatusCode::kFailedPrecondition); - EXPECT_EQ(KVCacheMetadata::Attach(region.span(), 4).status().code(), - absl::StatusCode::kFailedPrecondition); + ABSL_EXPECT_OK(KVCacheMetadata::Attach(region.span(), 4, "model_a")); + EXPECT_THAT(KVCacheMetadata::Attach(region.span(), 4, "model_b"), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_THAT(KVCacheMetadata::Attach(region.span(), 4), + StatusIs(absl::StatusCode::kFailedPrecondition)); // Longer than the header field can record. - EXPECT_EQ(KVCacheMetadata::Format(region.span(), 4, std::string(64, 'x')) - .status() - .code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(KVCacheMetadata::Format(region.span(), 4, std::string(64, 'x')), + StatusIs(absl::StatusCode::kInvalidArgument)); } TEST(KVCacheMetadataTest, FormatWipesSurvivingEntries) { Region region(2); - auto metadata_or = KVCacheMetadata::Format(region.span(), 2); - ASSERT_TRUE(metadata_or.ok()); - ASSERT_TRUE(metadata_or->Set(0, "stale", /*seq=*/1).ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 2)); + ABSL_ASSERT_OK(metadata.Set(0, "stale", /*seq=*/1)); - auto reformatted_or = KVCacheMetadata::Format(region.span(), 2); - ASSERT_TRUE(reformatted_or.ok()); - EXPECT_THAT(reformatted_or->ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto reformatted, + KVCacheMetadata::Format(region.span(), 2)); + EXPECT_THAT(reformatted.ValidEntries(), IsEmpty()); } TEST(KVCacheMetadataTest, UncommittedEntryIsInvisible) { Region region(2); - auto metadata_or = KVCacheMetadata::Format(region.span(), 2); - ASSERT_TRUE(metadata_or.ok()); + ABSL_ASSERT_OK(KVCacheMetadata::Format(region.span(), 2)); // Simulate a crash after the hash bytes landed but before the entry was // committed: write the fields directly and leave `valid` unset. @@ -212,9 +206,9 @@ TEST(KVCacheMetadataTest, UncommittedEntryIsInvisible) { entry->hash_len = 4; std::memcpy(entry->hash, "torn", 4); - auto recovered_or = KVCacheMetadata::Attach(region.span(), 2); - ASSERT_TRUE(recovered_or.ok()); - EXPECT_THAT(recovered_or->ValidEntries(), IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto recovered, + KVCacheMetadata::Attach(region.span(), 2)); + EXPECT_THAT(recovered.ValidEntries(), IsEmpty()); } } // namespace diff --git a/tpu_sync/kv_cache/kv_cache_store.cc b/tpu_sync/kv_cache/kv_cache_store.cc index edd7ba7a..cf3af1f2 100644 --- a/tpu_sync/kv_cache/kv_cache_store.cc +++ b/tpu_sync/kv_cache/kv_cache_store.cc @@ -214,20 +214,20 @@ absl::StatusOr> KVCacheStore::Create( if (num_shards > 0 && !effective_config0.metadata.has_value()) { const char* shm_key_env = std::getenv("RAIDEN_SHM_KEY"); if (shm_key_env != nullptr && std::strlen(shm_key_env) > 0) { - RETURN_IF_ERROR( + ABSL_RETURN_IF_ERROR( SharedMemoryHostMemoryAllocator::ValidateShmNameParts(shm_key_env)); const char* model_uid = std::getenv("RAIDEN_SHM_MODEL_UID"); - auto region_or = KVCacheMetadataShmRegion::AttachOrFormat( + auto region = KVCacheMetadataShmRegion::AttachOrFormat( MetadataShmKey(effective_config0.raiden_id), static_cast(effective_config0.capacity), model_uid != nullptr ? model_uid : "default_model"); - if (region_or.ok()) { - metadata_region = *std::move(region_or); + if (region.ok()) { + metadata_region = *std::move(region); effective_config0.metadata = metadata_region->metadata(); } else { LOG(WARNING) << "KV metadata table unavailable, serving without " "crash recovery: " - << region_or.status().message(); + << region.status().message(); } } } @@ -941,10 +941,11 @@ absl::StatusOr KVCacheStore::Lookup( if (start_idx >= block_hashes.size()) break; if (!backend) continue; - auto res_or = backend->Lookup( + auto lookup_result = backend->Lookup( absl::MakeSpan(block_hashes).subspan(start_idx), options); - if (!res_or.ok()) { - if (!accumulated_results.empty() && absl::IsNotFound(res_or.status())) { + if (!lookup_result.ok()) { + if (!accumulated_results.empty() && + absl::IsNotFound(lookup_result.status())) { break; } if (options.pin_found && !accumulated_results.empty()) { @@ -955,10 +956,10 @@ absl::StatusOr KVCacheStore::Lookup( } Release(matched_hashes); } - return res_or.status(); + return lookup_result.status(); } - const auto& res = res_or.value(); + const auto& res = *lookup_result; for (const auto& pair : res) { accumulated_results.push_back(pair); ++start_idx; @@ -1063,9 +1064,7 @@ absl::Status KVCacheStore::SaveLocal( { absl::MutexLock lock(mutex_); - auto lookup_or = backend()->Lookup(block_hashes); - if (!lookup_or.ok()) return lookup_or.status(); - const auto& slices = lookup_or.value(); + ABSL_ASSIGN_OR_RETURN(const auto& slices, backend()->Lookup(block_hashes)); if (slices.size() < block_hashes.size()) { return absl::NotFoundError( absl::StrCat("Block hash not found: ", block_hashes[slices.size()])); @@ -1094,12 +1093,12 @@ absl::Status KVCacheStore::SaveLocal( save_tracker_.AddPending(block_hashes); } - auto host_blocks_or = AllocateBlockIds(block_hashes.size()); - if (!host_blocks_or.ok()) { + auto host_blocks = AllocateBlockIds(block_hashes.size()); + if (!host_blocks.ok()) { save_tracker_.RemovePending(block_hashes); - return host_blocks_or.status(); + return host_blocks.status(); } - const auto& host_block_ids = host_blocks_or.value(); + const auto& host_block_ids = *host_blocks; std::vector src_buffers; src_buffers.reserve(src_device_block_ids.size()); @@ -1145,9 +1144,7 @@ absl::Status KVCacheStore::Load(absl::Span block_hashes, { absl::MutexLock lock(mutex_); - auto lookup_or = backend()->Lookup(block_hashes); - if (!lookup_or.ok()) return lookup_or.status(); - const auto& slices = lookup_or.value(); + ABSL_ASSIGN_OR_RETURN(const auto& slices, backend()->Lookup(block_hashes)); if (slices.size() < block_hashes.size()) { return absl::NotFoundError( absl::StrCat("Block hash not found: ", block_hashes[slices.size()])); @@ -1310,9 +1307,7 @@ void KVCacheStore::RegisterReadRemoteHooks() { absl::StatusOr> KVCacheStore::ValidateAndPinHostBlocks( absl::Span block_hashes) { absl::MutexLock lock(mutex_); - auto lookup_or = backend()->Lookup(block_hashes); - if (!lookup_or.ok()) return lookup_or.status(); - const auto& slices = lookup_or.value(); + ABSL_ASSIGN_OR_RETURN(const auto& slices, backend()->Lookup(block_hashes)); if (slices.size() < block_hashes.size()) { return absl::NotFoundError( absl::StrCat("BLOCK_HASH_NOT_FOUND: ", block_hashes[slices.size()])); @@ -1446,10 +1441,10 @@ bool KVCacheStore::SweepOnce() { // staler placement data. sweep_active_ = true; placement_targets_.clear(); - auto targets_or = + auto targets = registry_client_->GetPlacementTargets(raiden_id_, kMaxPlacementTargets); - if (targets_or.ok()) { - for (const auto& info : *targets_or) { + if (targets.ok()) { + for (const auto& info : *targets) { placement_targets_.push_back(RaidenId{ info.raiden_id().job_name(), info.raiden_id().job_replica_id(), info.raiden_id().data_name(), @@ -1457,7 +1452,7 @@ bool KVCacheStore::SweepOnce() { } } else { LOG(WARNING) << "Evict sweep could not fetch placement targets: " - << targets_or.status() + << targets.status() << ". Dropping cold blocks locally instead."; } } else if (free_ratio >= monitor_config_.evict_high_watermark) { @@ -1763,9 +1758,7 @@ absl::Status KVCacheStore::SaveRemote( std::vector src_host_block_ids; { absl::MutexLock lock(mutex_); - auto lookup_or = backend->Lookup(block_hashes); - if (!lookup_or.ok()) return lookup_or.status(); - const auto& slices = lookup_or.value(); + ABSL_ASSIGN_OR_RETURN(const auto& slices, backend->Lookup(block_hashes)); if (slices.size() < block_hashes.size()) { return absl::NotFoundError( absl::StrCat("Block hash not found: ", block_hashes[slices.size()])); @@ -1947,11 +1940,10 @@ absl::Status KVCacheStore::SaveRemote( } }; - auto ack_or = - backend->BeginWriteRemote(dst_raiden_id, block_hashes, src_host_block_ids, - hold - kRemoteWriteMargin, hold, - std::move(on_verdict)); - if (!ack_or.ok()) { + auto ack_res = backend->BeginWriteRemote( + dst_raiden_id, block_hashes, src_host_block_ids, + hold - kRemoteWriteMargin, hold, std::move(on_verdict)); + if (!ack_res.ok()) { // The offer failed before any ack. Undo it completely: release the pin, // clear the marks, and report only through the return status. // @@ -1968,11 +1960,10 @@ absl::Status KVCacheStore::SaveRemote( sweep_tracker_.RemovePending(taken->block_hashes); } } - return ack_or.status(); + return ack_res.status(); } - - const auto& ack = *ack_or; + const auto& ack = *ack_res; if (ack.all_exist) { // SUCCESS with nothing to wait for. auto taken = TakeRemoteWrite(op_key); @@ -2077,10 +2068,10 @@ void KVCacheStore::PollSavesInternal(std::vector ready_saves) { // Hoisted out of the lookup scope: the caller's pins on these hashes // are consumed below, on a path that may run after this scope ends. std::vector update_hashes; - auto lookup_or = backend()->Lookup(state.block_hashes, - LookupOptions{.enable_global = false}); - if (lookup_or.ok()) { - const auto& slices = lookup_or.value(); + auto lookup = backend()->Lookup(state.block_hashes, + LookupOptions{.enable_global = false}); + if (lookup.ok()) { + const auto& slices = *lookup; std::vector update_slices; for (size_t i = 0; i < state.block_hashes.size(); ++i) { const auto& hash = state.block_hashes[i]; @@ -2193,10 +2184,10 @@ void KVCacheStore::PollLoadsInternal(std::vector ready_loads) { // Local source: the entry exists here by construction, so this lookup is // purely local -- no registry fallback, which would otherwise put a // blocking RPC inside the poller while it holds mutex_. - auto lookup_or = backend()->Lookup(state.block_hashes, - LookupOptions{.enable_global = false}); - if (lookup_or.ok()) { - const auto& slices = lookup_or.value(); + auto lookup = backend()->Lookup(state.block_hashes, + LookupOptions{.enable_global = false}); + if (lookup.ok()) { + const auto& slices = *lookup; std::vector update_hashes; std::vector update_slices; for (size_t i = 0; i < state.block_hashes.size(); ++i) { diff --git a/tpu_sync/kv_cache/kv_cache_store_client_test.cc b/tpu_sync/kv_cache/kv_cache_store_client_test.cc index a43725b9..d4c53d8b 100644 --- a/tpu_sync/kv_cache/kv_cache_store_client_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_client_test.cc @@ -28,6 +28,7 @@ #include "grpcpp/security/server_credentials.h" #include "grpcpp/support/status.h" #include "xla/tsl/concurrency/future.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/core/raiden_future.h" #include "tpu_sync/proto/kv_cache_store_service.grpc.pb.h" #include "tpu_sync/proto/kv_cache_store_service.pb.h" @@ -106,9 +107,8 @@ TEST_F(KVCacheStoreClientTest, FetchReturnsFutureWithFetchResponseSuccess) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_ids, client_id); - auto response_or = future.Await(); - ASSERT_OK(response_or.status()); - EXPECT_THAT(response_or->done_block_hashes(), + TF_ASSERT_OK_AND_ASSIGN(auto response, future.Await()); + EXPECT_THAT(response.done_block_hashes(), UnorderedElementsAre("hash_1", "hash_2")); } @@ -119,8 +119,8 @@ TEST_F(KVCacheStoreClientTest, FetchReturnsFutureWithErrorStatusOnRPCFailure) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), StatusIs(absl::StatusCode::kInternal)); + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kInternal)); } TEST_F(KVCacheStoreClientTest, FetchPopulatesRequestFieldsCorrectly) { @@ -135,8 +135,7 @@ TEST_F(KVCacheStoreClientTest, FetchPopulatesRequestFieldsCorrectly) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, dev_ids, host_ids, client_id); - auto response_or = future.Await(); - ASSERT_OK(response_or.status()); + TF_ASSERT_OK_AND_ASSIGN(auto response, future.Await()); const auto& req = service_->last_request(); EXPECT_THAT(req.block_hashes(), UnorderedElementsAre("hash_a", "hash_b")); @@ -156,9 +155,8 @@ TEST_F(KVCacheStoreClientTest, FetchValidatesMismatchedDeviceBlockIds) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, dev_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), - StatusIs(absl::StatusCode::kInvalidArgument)); + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kInvalidArgument)); } TEST_F(KVCacheStoreClientTest, FetchValidatesMismatchedHostBlockIds) { @@ -167,18 +165,16 @@ TEST_F(KVCacheStoreClientTest, FetchValidatesMismatchedHostBlockIds) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), - StatusIs(absl::StatusCode::kInvalidArgument)); + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kInvalidArgument)); } TEST_F(KVCacheStoreClientTest, FetchEmptyHashesReturnsEmptyResponse) { std::vector hashes; tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes); - auto response_or = future.Await(); - ASSERT_OK(response_or.status()); - EXPECT_EQ(response_or->done_block_hashes_size(), 0); + TF_ASSERT_OK_AND_ASSIGN(auto response, future.Await()); + EXPECT_EQ(response.done_block_hashes_size(), 0); } } // namespace diff --git a/tpu_sync/kv_cache/kv_cache_store_server_test.cc b/tpu_sync/kv_cache/kv_cache_store_server_test.cc index e98b108f..a291d41f 100644 --- a/tpu_sync/kv_cache/kv_cache_store_server_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_server_test.cc @@ -57,7 +57,7 @@ class KVCacheStoreServerTest : public ::testing::Test { RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; RaidenId dst_raiden_id{"dst_job", "0", "dst_data", 0}; - ASSERT_OK(src_controller_server_->client->RegisterWorker( + ABSL_ASSERT_OK(src_controller_server_->client->RegisterWorker( "worker_0", test_worker_server_->server_address, {{test_worker_server_->server_address, {}}})); @@ -78,7 +78,7 @@ class KVCacheStoreServerTest : public ::testing::Test { ::tpu_raiden::core::controller::RaidenControllerClient dst_controller_client(store_->raiden_controller_address()); - ASSERT_OK(dst_controller_client.RegisterWorker( + ABSL_ASSERT_OK(dst_controller_client.RegisterWorker( "dst_worker_0", test_worker_server_->server_address, {{test_worker_server_->server_address, {}}})); @@ -88,7 +88,7 @@ class KVCacheStoreServerTest : public ::testing::Test { RaidenBlockId(src_raiden_id, 10, BlockStatus::HOST), RaidenBlockId(src_raiden_id, 11, BlockStatus::HOST), }; - ASSERT_TRUE(store_->Insert(test_hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store_->Insert(test_hashes, slices, /*on_host=*/true)); } void TearDown() override { @@ -111,8 +111,8 @@ TEST_F(KVCacheStoreServerTest, StartServerWithRawPointerStore) { server_ = KVCacheStoreServer::Create(); // A wildcard bind reports no publishable address (no in-tree caller ever // wildcard-binds), so use a real, dialable host. - ASSERT_OK(server_->StartServer(store_->backend().get(), - store_->raiden_controller(), "127.0.0.1:0")); + ABSL_ASSERT_OK(server_->StartServer( + store_->backend().get(), store_->raiden_controller(), "127.0.0.1:0")); int port = server_->GetGrpcPort(); EXPECT_GT(port, 0); @@ -127,7 +127,7 @@ TEST_F(KVCacheStoreServerTest, StartServerWithRawPointerStore) { std::vector host_ids = {100, 101}; auto fetch_res = client.Fetch(hashes, /*device_block_ids=*/{}, host_ids).Await(); - ASSERT_OK(fetch_res.status()); + ABSL_ASSERT_OK(fetch_res.status()); EXPECT_THAT(fetch_res->done_block_hashes(), UnorderedElementsAre("block_hash_1", "block_hash_2")); @@ -138,8 +138,8 @@ TEST_F(KVCacheStoreServerTest, StartServerWithRawPointerStore) { TEST_F(KVCacheStoreServerTest, RestartServerAfterShutdown) { server_ = KVCacheStoreServer::Create(); - ASSERT_OK(server_->StartServer(store_->backend().get(), - store_->raiden_controller(), "127.0.0.1:0")); + ABSL_ASSERT_OK(server_->StartServer( + store_->backend().get(), store_->raiden_controller(), "127.0.0.1:0")); int first_port = server_->GetGrpcPort(); EXPECT_GT(first_port, 0); @@ -148,8 +148,8 @@ TEST_F(KVCacheStoreServerTest, RestartServerAfterShutdown) { EXPECT_TRUE(server_->GetServerAddress().empty()); // Restart server on a new ephemeral port - ASSERT_OK(server_->StartServer(store_->backend().get(), - store_->raiden_controller(), "127.0.0.1:0")); + ABSL_ASSERT_OK(server_->StartServer( + store_->backend().get(), store_->raiden_controller(), "127.0.0.1:0")); int second_port = server_->GetGrpcPort(); EXPECT_GT(second_port, 0); @@ -162,7 +162,7 @@ TEST_F(KVCacheStoreServerTest, RestartServerAfterShutdown) { std::vector host_ids = {100, 101}; auto fetch_res = client.Fetch(hashes, /*device_block_ids=*/{}, host_ids).Await(); - ASSERT_OK(fetch_res.status()); + ABSL_ASSERT_OK(fetch_res.status()); EXPECT_THAT(fetch_res->done_block_hashes(), UnorderedElementsAre("block_hash_1", "block_hash_2")); @@ -171,13 +171,13 @@ TEST_F(KVCacheStoreServerTest, RestartServerAfterShutdown) { TEST_F(KVCacheStoreServerTest, MultipleServersCanRunConcurrently) { auto server1 = KVCacheStoreServer::Create(); - ASSERT_OK(server1->StartServer(store_->backend().get(), - store_->raiden_controller(), "[::]:0")); + ABSL_ASSERT_OK(server1->StartServer(store_->backend().get(), + store_->raiden_controller(), "[::]:0")); int port1 = server1->GetGrpcPort(); auto server2 = KVCacheStoreServer::Create(); - ASSERT_OK(server2->StartServer(store_->backend().get(), - store_->raiden_controller(), "[::]:0")); + ABSL_ASSERT_OK(server2->StartServer(store_->backend().get(), + store_->raiden_controller(), "[::]:0")); int port2 = server2->GetGrpcPort(); EXPECT_GT(port1, 0); @@ -199,8 +199,8 @@ TEST_F(KVCacheStoreServerTest, StartServerWithInvalidAddressFails) { TEST_F(KVCacheStoreServerTest, StartServerWithHostOnlySucceeds) { auto server = KVCacheStoreServer::Create(); - ASSERT_OK(server->StartServer(store_->backend().get(), - store_->raiden_controller(), "127.0.0.1")); + ABSL_ASSERT_OK(server->StartServer(store_->backend().get(), + store_->raiden_controller(), "127.0.0.1")); EXPECT_GT(server->GetGrpcPort(), 0); server->Shutdown(); } diff --git a/tpu_sync/kv_cache/kv_cache_store_service.cc b/tpu_sync/kv_cache/kv_cache_store_service.cc index cab078e8..7d23d989 100644 --- a/tpu_sync/kv_cache/kv_cache_store_service.cc +++ b/tpu_sync/kv_cache/kv_cache_store_service.cc @@ -456,14 +456,14 @@ ::grpc::ServerUnaryReactor* KVCacheStoreServiceImpl::Fetch( LookupOptions options; options.enable_global = false; options.pin_found = true; - auto lookup_or = backend_->Lookup(block_hashes, options); - if (!lookup_or.ok()) { + auto lookup = backend_->Lookup(block_hashes, options); + if (!lookup.ok()) { reactor->Finish(::grpc::Status( ::grpc::StatusCode::NOT_FOUND, - absl::StrCat("Validation failed: ", lookup_or.status().message()))); + absl::StrCat("Validation failed: ", lookup.status().message()))); return reactor; } - const auto& lookup_slices = lookup_or.value(); + const auto& lookup_slices = *lookup; // The lookup pinned everything it returned. Release that on every exit, // including the refusals below. @@ -696,17 +696,17 @@ KVCacheStoreServiceImpl::WriteRemote( std::min(absl::Milliseconds(request->deadline_ms()), DeadlineCap()); // Allocate landing blocks in destination host DRAM for the transfer. - auto allocated_ids_or = controller_->AllocateBlockIds(block_hashes.size()); - if (!allocated_ids_or.ok()) { + auto allocated_ids = controller_->AllocateBlockIds(block_hashes.size()); + if (!allocated_ids.ok()) { auto* reactor = new WriteRemoteServerReactor(/*gate=*/nullptr); reactor->Finish(::grpc::Status( ::grpc::StatusCode::RESOURCE_EXHAUSTED, absl::StrCat("Failed to allocate destination landing blocks: ", - allocated_ids_or.status().message()))); + allocated_ids.status().message()))); return reactor; } - std::vector landing_block_ids(allocated_ids_or->begin(), - allocated_ids_or->end()); + std::vector landing_block_ids(allocated_ids->begin(), + allocated_ids->end()); const absl::Time now = absl::Now(); const absl::Time deadline = now + granted_deadline; diff --git a/tpu_sync/kv_cache/kv_cache_store_service_test.cc b/tpu_sync/kv_cache/kv_cache_store_service_test.cc index fbe16379..520d9250 100644 --- a/tpu_sync/kv_cache/kv_cache_store_service_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_service_test.cc @@ -28,6 +28,7 @@ #include #include +#include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/synchronization/mutex.h" #include "absl/time/clock.h" @@ -38,6 +39,7 @@ #include "grpcpp/security/server_credentials.h" #include "grpcpp/support/status.h" #include "xla/tsl/concurrency/future.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/common/raiden_id.h" #include "tpu_sync/core/buffer.h" #include "tpu_sync/core/controller/controller_client.h" @@ -76,7 +78,7 @@ class KVCacheStoreServiceTest : public ::testing::Test { RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; RaidenId dst_raiden_id{"dst_job", "0", "dst_data", 0}; - ASSERT_OK(src_controller_server_->client->RegisterWorker( + ABSL_ASSERT_OK(src_controller_server_->client->RegisterWorker( "worker_0", test_worker_server_->server_address, {{test_worker_server_->server_address, {}}})); @@ -97,7 +99,7 @@ class KVCacheStoreServiceTest : public ::testing::Test { ::tpu_raiden::core::controller::RaidenControllerClient dst_controller_client(store_->raiden_controller_address()); - ASSERT_OK(dst_controller_client.RegisterWorker( + ABSL_ASSERT_OK(dst_controller_client.RegisterWorker( "dst_worker_0", test_worker_server_->server_address, {{test_worker_server_->server_address, {}}})); @@ -110,7 +112,7 @@ class KVCacheStoreServiceTest : public ::testing::Test { RaidenBlockId(src_raiden_id, 12, BlockStatus::HOST), RaidenBlockId(src_raiden_id, 13, BlockStatus::HOST), }; - ASSERT_TRUE(store_->Insert(test_hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store_->Insert(test_hashes, slices, /*on_host=*/true)); // Setup KVCacheStoreServiceImpl & gRPC server service_ = std::make_unique( @@ -152,9 +154,8 @@ TEST_F(KVCacheStoreServiceTest, FetchEmptyRequest) { std::vector empty_hashes; tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(empty_hashes); - auto response_or = future.Await(); - ASSERT_OK(response_or.status()); - EXPECT_EQ(response_or->done_block_hashes_size(), 0); + TF_ASSERT_OK_AND_ASSIGN(auto response, future.Await()); + EXPECT_EQ(response.done_block_hashes_size(), 0); } TEST_F(KVCacheStoreServiceTest, Fetch5StepWorkflowSuccess) { @@ -173,11 +174,10 @@ TEST_F(KVCacheStoreServiceTest, Fetch5StepWorkflowSuccess) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids, client_id, {client_ep}); - auto response_or = future.Await(); - ASSERT_OK(response_or.status()); - EXPECT_THAT(response_or->done_block_hashes(), + TF_ASSERT_OK_AND_ASSIGN(auto response, future.Await()); + EXPECT_THAT(response.done_block_hashes(), UnorderedElementsAre("block_hash_1", "block_hash_2")); - EXPECT_EQ(response_or->failed_block_hashes_size(), 0); + EXPECT_EQ(response.failed_block_hashes_size(), 0); } TEST_F(KVCacheStoreServiceTest, FetchCrossNodeMissingEndpointsFails) { @@ -190,9 +190,8 @@ TEST_F(KVCacheStoreServiceTest, FetchCrossNodeMissingEndpointsFails) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids, client_id); - auto response_or = future.Await(); - EXPECT_FALSE(response_or.status().ok()); - EXPECT_EQ(response_or.status().code(), absl::StatusCode::kInvalidArgument); + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kInvalidArgument)); } TEST_F(KVCacheStoreServiceTest, FetchValidationFailsForMissingHash) { @@ -201,8 +200,8 @@ TEST_F(KVCacheStoreServiceTest, FetchValidationFailsForMissingHash) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), StatusIs(absl::StatusCode::kNotFound)); + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kNotFound)); } // Real block hashes are raw digests, not text. Every other case in this file @@ -219,10 +218,9 @@ TEST_F(KVCacheStoreServiceTest, FetchValidationFailsForMissingHash) { TEST_F(KVCacheStoreServiceTest, FetchRoundTripsNonUtf8Hash) { const std::string binary_hash("\xff\xfe\x80\x00\x01\xc0\xaf\xed\xa0\x80", 10); RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - ASSERT_TRUE(store_->Insert( + ABSL_ASSERT_OK(store_->Insert( {binary_hash}, {RaidenBlockId(src_raiden_id, 20, BlockStatus::HOST)}, - /*on_host=*/true) - .ok()); + /*on_host=*/true)); ::tpu_sync::rpc::RaidenIdProto client_id; client_id.set_job_name("client_job"); @@ -236,11 +234,9 @@ TEST_F(KVCacheStoreServiceTest, FetchRoundTripsNonUtf8Hash) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch({binary_hash}, /*device_block_ids=*/{}, /*host_block_ids=*/{102}, client_id, {client_ep}); - auto response_or = future.Await(); - ASSERT_OK(response_or.status()); - EXPECT_THAT(response_or->done_block_hashes(), - UnorderedElementsAre(binary_hash)); - EXPECT_EQ(response_or->failed_block_hashes_size(), 0); + TF_ASSERT_OK_AND_ASSIGN(auto response, future.Await()); + EXPECT_THAT(response.done_block_hashes(), UnorderedElementsAre(binary_hash)); + EXPECT_EQ(response.failed_block_hashes_size(), 0); } TEST_F(KVCacheStoreServiceTest, FetchValidationFailsForNonHostBlock) { @@ -252,13 +248,13 @@ TEST_F(KVCacheStoreServiceTest, FetchValidationFailsForNonHostBlock) { RaidenBlockId(src_raiden_id, /*host_block_id=*/-1, /*device_block_id=*/50, BlockStatus::HBM), }; - ASSERT_TRUE(store_->Insert(hashes, slices, /*on_host=*/false).ok()); + ABSL_ASSERT_OK(store_->Insert(hashes, slices, /*on_host=*/false)); std::vector host_block_ids = {100}; tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kFailedPrecondition)); } @@ -274,9 +270,9 @@ TEST_F(KVCacheStoreServiceTest, ARefusedFetchLeavesNoPinBehind) { std::vector hashes = {"block_hash_1", "no_such_hash"}; std::vector host_block_ids = {100, 101}; - auto response_or = + auto response = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids).Await(); - EXPECT_THAT(response_or.status(), StatusIs(absl::StatusCode::kNotFound)); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kNotFound)); EXPECT_EQ(store_->backend()->GetPinCount("block_hash_1"), before); } @@ -293,9 +289,9 @@ TEST_F(KVCacheStoreServiceTest, ARefusedFetchDoesNotReleaseAPinItDidNotTake) { "block_hash_1"}; std::vector host_block_ids = {100, 101, 102}; - auto response_or = + auto response = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids).Await(); - EXPECT_THAT(response_or.status(), StatusIs(absl::StatusCode::kNotFound)); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kNotFound)); EXPECT_EQ(store_->backend()->GetPinCount("block_hash_1"), before); } @@ -309,9 +305,9 @@ TEST_F(KVCacheStoreServiceTest, ACompletedFetchLeavesNoPinBehind) { std::vector hashes = {"block_hash_1", "block_hash_2"}; std::vector host_block_ids = {100, 101}; - auto response_or = - client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids).Await(); - ASSERT_OK(response_or.status()); + TF_ASSERT_OK_AND_ASSIGN( + auto response, + client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids).Await()); EXPECT_EQ(store_->backend()->GetPinCount("block_hash_1"), before_1); EXPECT_EQ(store_->backend()->GetPinCount("block_hash_2"), before_2); @@ -323,9 +319,8 @@ TEST_F(KVCacheStoreServiceTest, FetchMismatchedHostBlockCount) { tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = client_->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), - StatusIs(absl::StatusCode::kInvalidArgument)); + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kInvalidArgument)); } TEST_F(KVCacheStoreServiceTest, FetchNullStoreHandling) { @@ -346,8 +341,8 @@ TEST_F(KVCacheStoreServiceTest, FetchNullStoreHandling) { std::vector host_block_ids = {100}; tsl::Future<::tpu_raiden::kv_cache::proto::FetchResponse> future = null_client->Fetch(hashes, /*device_block_ids=*/{}, host_block_ids); - auto response_or = future.Await(); - EXPECT_THAT(response_or.status(), + auto response = future.Await(); + EXPECT_THAT(response.status(), StatusIs(absl::StatusCode::kFailedPrecondition)); null_server->Shutdown(); } @@ -390,7 +385,7 @@ TEST_F(KVCacheStoreServiceTest, ConcurrentFetchRPCs) { } for (int i = 0; i < kNumThreads; ++i) { - ASSERT_OK(results[i].status()); + ABSL_ASSERT_OK(results[i].status()); EXPECT_THAT( results[i]->done_block_hashes(), UnorderedElementsAre("concurrent_hash_" + std::to_string(2 * i), @@ -426,7 +421,7 @@ TEST_F(KVCacheStoreServiceTest, FetchRoutesToClientAdvertisedEndpoints) { client_ ->Fetch(hashes, /*device_block_ids=*/{}, host_ids, client_id, groups) .Await(); - ASSERT_TRUE(res.ok()) << res.status(); + ABSL_ASSERT_OK(res); ASSERT_EQ(dst_transfer_mock_->last_write_descriptors.size(), 1); EXPECT_EQ(dst_transfer_mock_->last_write_descriptors[0].endpoint, @@ -457,7 +452,7 @@ TEST_F(KVCacheStoreServiceTest, SameNodeFetchNeedsNoEndpoints) { std::vector hashes = {"block_hash_1"}; std::vector host_ids = {201}; auto res = client_->Fetch(hashes, /*device_block_ids=*/{}, host_ids).Await(); - EXPECT_TRUE(res.ok()) << res.status(); + ABSL_EXPECT_OK(res); } // Two workers on this node, each with its own transfer manager, and two client @@ -488,17 +483,17 @@ TEST_F(KVCacheStoreServiceTest, FetchWithMultiWorkerEndpointsRoutesPerWorker) { /*store_server_ip=*/"127.0.0.1"); ::tpu_raiden::core::controller::RaidenControllerClient ctrl_client( store.raiden_controller_address()); - ASSERT_OK(ctrl_client.RegisterWorker("w_a", worker_a->server_address, - {{worker_a->server_address, {}}}, - /*node_id=*/10)); - ASSERT_OK(ctrl_client.RegisterWorker("w_b", worker_b->server_address, - {{worker_b->server_address, {}}}, - /*node_id=*/20)); + ABSL_ASSERT_OK(ctrl_client.RegisterWorker("w_a", worker_a->server_address, + {{worker_a->server_address, {}}}, + /*node_id=*/10)); + ABSL_ASSERT_OK(ctrl_client.RegisterWorker("w_b", worker_b->server_address, + {{worker_b->server_address, {}}}, + /*node_id=*/20)); std::vector hashes = {"multi_hash"}; std::vector slices = { RaidenBlockId(multi_id, 7, BlockStatus::HOST)}; - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); KVCacheStoreServiceImpl service(store.backend().get(), store.raiden_controller()); @@ -524,7 +519,7 @@ TEST_F(KVCacheStoreServiceTest, FetchWithMultiWorkerEndpointsRoutesPerWorker) { .Fetch(hashes, /*device_block_ids=*/{}, /*host_block_ids=*/{301}, client_id, groups) .Await(); - ASSERT_OK(res.status()); + ABSL_ASSERT_OK(res.status()); // Each worker saw its own peer -- not the other's, and not both. ASSERT_EQ(mock_a->last_write_descriptors.size(), 1); @@ -722,7 +717,7 @@ TEST_F(WriteRemoteTest, AllExistNeedsNoSourceEndpoints) { ->WriteRemote(SrcIdProto(), {"a"}, src_ids, {}, 5000, absl::Seconds(30)) .ack.Await(); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); EXPECT_EQ(response->exist_state(), ::tpu_raiden::kv_cache::proto::WRITE_ALL_EXIST); } @@ -764,7 +759,7 @@ TEST_F(WriteRemoteTest, AllExistIsAnImmediateSuccess) { RaidenBlockId(dst_id_, 2, BlockStatus::HOST)})); auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); EXPECT_EQ(response->operation_id(), 0); EXPECT_EQ(response->exist_state(), ::tpu_raiden::kv_cache::proto::WRITE_ALL_EXIST); @@ -777,7 +772,7 @@ TEST_F(WriteRemoteTest, PartialExistIsRefusedAndNamesWhatItHas) { {"a"}, {RaidenBlockId(dst_id_, 1, BlockStatus::HOST)})); auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); EXPECT_EQ(response->operation_id(), 0); EXPECT_EQ(response->exist_state(), ::tpu_raiden::kv_cache::proto::WRITE_PARTIAL_EXIST); @@ -788,7 +783,7 @@ TEST_F(WriteRemoteTest, PartialExistIsRefusedAndNamesWhatItHas) { TEST_F(WriteRemoteTest, AcceptsAndAnswersWithoutWaitingForBytes) { auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); EXPECT_NE(response->operation_id(), 0); EXPECT_EQ(response->exist_state(), ::tpu_raiden::kv_cache::proto::WRITE_EXIST_STATE_UNSPECIFIED); @@ -808,7 +803,7 @@ TEST_F(WriteRemoteTest, AcceptsAndAnswersWithoutWaitingForBytes) { TEST_F(WriteRemoteTest, GrantedDeadlineIsClampedToTheLocalCap) { auto response = Offer({"a"}, /*deadline_ms=*/absl::ToInt64Milliseconds(absl::Hours(1))); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); EXPECT_LE(response->granted_deadline_ms(), absl::ToInt64Milliseconds(absl::Seconds(25))); latch_.Release(absl::CancelledError("done with this test")); @@ -817,7 +812,7 @@ TEST_F(WriteRemoteTest, GrantedDeadlineIsClampedToTheLocalCap) { TEST_F(WriteRemoteTest, CommitInsertsTheBlocksAndReportsThem) { auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); latch_.Release(absl::OkStatus()); ASSERT_EQ(AwaitTerminal(response->operation_id()), @@ -829,7 +824,7 @@ TEST_F(WriteRemoteTest, CommitInsertsTheBlocksAndReportsThem) { TEST_F(WriteRemoteTest, TransferFailureFreesTheLandingBlocks) { auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); latch_.Release(absl::InternalError("pull failed")); ASSERT_EQ(AwaitTerminal(response->operation_id()), @@ -838,15 +833,14 @@ TEST_F(WriteRemoteTest, TransferFailureFreesTheLandingBlocks) { EXPECT_TRUE( store_->backend()->AlreadyPresentHostResident({"a", "b"}).empty()); // The blocks came back: the whole pool is allocatable again. - auto reallocated = store_->raiden_controller()->AllocateBlockIds(kCapacity); - EXPECT_TRUE(reallocated.ok()) << reallocated.status().ToString(); + ABSL_EXPECT_OK(store_->raiden_controller()->AllocateBlockIds(kCapacity)); } // A transfer that resolves after the deadline must not insert or register // anything. The bytes are discarded and the blocks come back. TEST_F(WriteRemoteTest, ATransferThatResolvesPastTheDeadlineNeverCommits) { auto response = Offer({"a", "b"}, /*deadline_ms=*/100); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); const uint64_t op_id = response->operation_id(); absl::SleepFor(absl::Milliseconds(300)); @@ -869,14 +863,13 @@ TEST_F(WriteRemoteTest, ATransferThatResolvesPastTheDeadlineNeverCommits) { ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::FAILED); EXPECT_TRUE(store_->backend()->AlreadyPresentHostResident({"a", "b"}).empty()) << "a post-deadline transfer inserted its bytes anyway"; - auto reallocated = store_->raiden_controller()->AllocateBlockIds(kCapacity); - EXPECT_TRUE(reallocated.ok()) << "the deferred free never happened: " - << reallocated.status().ToString(); + ABSL_EXPECT_OK(store_->raiden_controller()->AllocateBlockIds(kCapacity)) + << "the deferred free never happened"; } TEST_F(WriteRemoteTest, WriteOpsShrinksWithoutPollingTraffic) { auto response = Offer({"a", "b"}, /*deadline_ms=*/100); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); const uint64_t op_id = response->operation_id(); ASSERT_EQ(service_->InFlightWriteOpsCountForTesting(), 1); @@ -903,7 +896,7 @@ TEST_F(WriteRemoteTest, TheClaimRefusesALateTransferEvenIfNoThreadFiredIt) { service_->PauseDeadlineFiringForTesting(); auto response = Offer({"a", "b"}, /*deadline_ms=*/100); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); const uint64_t op_id = response->operation_id(); absl::SleepFor(absl::Milliseconds(300)); @@ -917,15 +910,14 @@ TEST_F(WriteRemoteTest, TheClaimRefusesALateTransferEvenIfNoThreadFiredIt) { << "a transfer that resolved past its deadline was allowed to commit"; EXPECT_TRUE( store_->backend()->AlreadyPresentHostResident({"a", "b"}).empty()); - auto reallocated = store_->raiden_controller()->AllocateBlockIds(kCapacity); - EXPECT_TRUE(reallocated.ok()) << reallocated.status().ToString(); + ABSL_EXPECT_OK(store_->raiden_controller()->AllocateBlockIds(kCapacity)); } // A concurrent writer landed every hash mid-flight: the claimed path reports // ALL_EXIST and still frees its landing blocks. TEST_F(WriteRemoteTest, LosingTheRaceAtInsertTimeReportsAllExistAndFrees) { auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); ASSERT_TRUE(store_->backend()->InsertAllOrNothing( {"a", "b"}, {RaidenBlockId(dst_id_, 6, BlockStatus::HOST), @@ -934,16 +926,15 @@ TEST_F(WriteRemoteTest, LosingTheRaceAtInsertTimeReportsAllExistAndFrees) { latch_.Release(absl::OkStatus()); ASSERT_EQ(AwaitTerminal(response->operation_id()), ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::ALL_EXIST); - auto reallocated = store_->raiden_controller()->AllocateBlockIds(kCapacity); - EXPECT_TRUE(reallocated.ok()) << "the landing blocks were never returned: " - << reallocated.status().ToString(); + ABSL_EXPECT_OK(store_->raiden_controller()->AllocateBlockIds(kCapacity)) + << "the landing blocks were never returned"; } // A PARTIAL_EXIST found at insert time reaches the poll with the same list // as an ack-time answer. TEST_F(WriteRemoteTest, PartialExistDiscoveredAtInsertTimeReachesThePoll) { auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); ASSERT_TRUE(store_->backend()->InsertAllOrNothing( {"a"}, {RaidenBlockId(dst_id_, 6, BlockStatus::HOST)})); @@ -955,8 +946,7 @@ TEST_F(WriteRemoteTest, PartialExistDiscoveredAtInsertTimeReachesThePoll) { EXPECT_THAT(last_poll_.existing_hashes(), UnorderedElementsAre("a")); // "b" was never inserted: this destination does not do partial writes. EXPECT_TRUE(store_->backend()->AlreadyPresentHostResident({"b"}).empty()); - auto reallocated = store_->raiden_controller()->AllocateBlockIds(kCapacity); - EXPECT_TRUE(reallocated.ok()) << reallocated.status().ToString(); + ABSL_EXPECT_OK(store_->raiden_controller()->AllocateBlockIds(kCapacity)); } // Landing blocks come from free blocks only; a full destination refuses the @@ -964,8 +954,7 @@ TEST_F(WriteRemoteTest, PartialExistDiscoveredAtInsertTimeReachesThePoll) { TEST_F(WriteRemoteTest, RefusesWhenThereAreNoFreeBlocksAndEvictsNothing) { ASSERT_TRUE(store_->backend()->InsertAllOrNothing( {"victim"}, {RaidenBlockId(dst_id_, 0, BlockStatus::HOST)})); - auto drained = store_->raiden_controller()->AllocateBlockIds(kCapacity); - ASSERT_TRUE(drained.ok()) << drained.status().ToString(); + ABSL_ASSERT_OK(store_->raiden_controller()->AllocateBlockIds(kCapacity)); auto response = Offer({"a"}); EXPECT_THAT(response.status(), @@ -977,7 +966,7 @@ TEST_F(WriteRemoteTest, RefusesWhenThereAreNoFreeBlocksAndEvictsNothing) { TEST_F(WriteRemoteTest, PollOfAnUnknownOperationIsUnknown) { auto poll = Poll(999999); - ASSERT_OK(poll.status()); + ABSL_ASSERT_OK(poll.status()); EXPECT_EQ(poll->state(), ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::UNKNOWN); } @@ -1039,7 +1028,7 @@ TEST(WriteRemotePublishTest, PublishDoesNotBlockTheTransferCompletion) { client.WriteRemote(src_id, hashes, src_ids, {group}, 5000, absl::Seconds(30)) .ack.Await(); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); const uint64_t op_id = response->operation_id(); ASSERT_NE(op_id, 0); @@ -1065,7 +1054,7 @@ TEST(WriteRemotePublishTest, PublishDoesNotBlockTheTransferCompletion) { // ... and the operation has not settled, which is what proves the wait that // did not happen was a real one rather than a registry that answered fast. auto polled = client.PollWriteRemote(op_id).Await(); - ASSERT_OK(polled.status()); + ABSL_ASSERT_OK(polled.status()); EXPECT_EQ(polled->state(), ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::PENDING); @@ -1074,7 +1063,7 @@ TEST(WriteRemotePublishTest, PublishDoesNotBlockTheTransferCompletion) { ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse final_poll; for (int i = 0; i < 300; ++i) { auto p = client.PollWriteRemote(op_id).Await(); - ASSERT_OK(p.status()); + ABSL_ASSERT_OK(p.status()); if (p->state() != ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::PENDING) { final_poll = *p; @@ -1090,7 +1079,7 @@ TEST(WriteRemotePublishTest, PublishDoesNotBlockTheTransferCompletion) { UnorderedElementsAre("publish_a", "publish_b")); auto looked_up = registry_server->client->Lookup(hashes); - ASSERT_OK(looked_up.status()); + ABSL_ASSERT_OK(looked_up.status()); EXPECT_EQ(looked_up->size(), 2) << "the landed blocks were never advertised"; @@ -1142,7 +1131,7 @@ TEST(WriteRemoteRegistryFailureTest, StoredButUnregisteredIsReportedAsSuch) { client.WriteRemote(src_id, hashes, src_ids, {group}, 5000, absl::Seconds(30)) .ack.Await(); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); const uint64_t op_id = response->operation_id(); ASSERT_NE(op_id, 0); @@ -1153,7 +1142,7 @@ TEST(WriteRemoteRegistryFailureTest, StoredButUnregisteredIsReportedAsSuch) { ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse poll; for (int i = 0; i < 300; ++i) { auto polled = client.PollWriteRemote(op_id).Await(); - ASSERT_OK(polled.status()); + ABSL_ASSERT_OK(polled.status()); if (polled->state() != ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::PENDING) { poll = *polled; @@ -1185,7 +1174,7 @@ TEST(WriteRemoteRegistryFailureTest, StoredButUnregisteredIsReportedAsSuch) { // and must not leave a callback pointing at freed memory. TEST_F(WriteRemoteTest, TeardownWithAnOutstandingTransferIsBoundedAndSafe) { auto response = Offer({"a", "b"}); - ASSERT_OK(response.status()); + ABSL_ASSERT_OK(response.status()); ASSERT_EQ(latch_.issued(), 1); server_->Shutdown(); @@ -1244,15 +1233,14 @@ TEST(FetchWithdrawTest, ARefusedFetchWithdrawsOnlyTheEntriesItCannotBack) { /*store_server_ip=*/"127.0.0.1"); // Two blocks this store really holds... - ASSERT_OK(store.Insert({"held_a", "held_b"}, - {RaidenBlockId(src_id, 10, BlockStatus::HOST), - RaidenBlockId(src_id, 11, BlockStatus::HOST)}, - /*on_host=*/true)); + ABSL_ASSERT_OK(store.Insert({"held_a", "held_b"}, + {RaidenBlockId(src_id, 10, BlockStatus::HOST), + RaidenBlockId(src_id, 11, BlockStatus::HOST)}, + /*on_host=*/true)); // ...and three entries advertising it, one of which names a block that was // never inserted. That is what a peer follows to get here. - ASSERT_OK(registry->Register({{"gone", src_id, 12}, - {"held_a", src_id, 10}, - {"held_b", src_id, 11}})); + ABSL_ASSERT_OK(registry->Register( + {{"gone", src_id, 12}, {"held_a", src_id, 10}, {"held_b", src_id, 11}})); KVCacheStoreServiceImpl service(store.backend().get(), store.raiden_controller()); @@ -1284,18 +1272,18 @@ TEST(FetchWithdrawTest, ARefusedFetchWithdrawsOnlyTheEntriesItCannotBack) { const absl::Time give_up = absl::Now() + absl::Seconds(5); while (absl::Now() < give_up) { auto gone = registry->Lookup({"gone"}, peer); - ASSERT_OK(gone.status()); + ABSL_ASSERT_OK(gone.status()); if (gone->empty()) break; absl::SleepFor(absl::Milliseconds(20)); } // Asked as a peer, so the caller filter does not hide the answer. auto gone = registry->Lookup({"gone"}, peer); - ASSERT_OK(gone.status()); + ABSL_ASSERT_OK(gone.status()); EXPECT_TRUE(gone->empty()) << "the entry this store cannot back survived"; auto held = registry->Lookup({"held_a", "held_b"}, peer); - ASSERT_OK(held.status()); + ABSL_ASSERT_OK(held.status()); EXPECT_EQ(held->size(), 2) << "a refused fetch withdrew blocks this store is holding"; @@ -1316,7 +1304,7 @@ TEST(FetchWithdrawTest, TheWithdrawDoesNotDelayTheRefusal) { KVCacheStore store(/*capacity=*/8, counting.server->server_address, src_id, /*num_shards=*/1, /*shard_size_bytes=*/1024, /*store_server_ip=*/"127.0.0.1"); - ASSERT_OK(registry->Register({{"gone", src_id, 12}})); + ABSL_ASSERT_OK(registry->Register({{"gone", src_id, 12}})); KVCacheStoreServiceImpl service(store.backend().get(), store.raiden_controller()); @@ -1378,7 +1366,7 @@ TEST(FetchWithdrawTest, AFetchRefusedOnTheWrongTierKeepsTheEntry) { BlockStatus::HBM)}, /*on_host=*/false) .first); - ASSERT_OK(registry->Register({{"staged", src_id, 10}})); + ABSL_ASSERT_OK(registry->Register({{"staged", src_id, 10}})); KVCacheStoreServiceImpl service(backend, store.raiden_controller()); ::grpc::ServerBuilder builder; @@ -1399,7 +1387,7 @@ TEST(FetchWithdrawTest, AFetchRefusedOnTheWrongTierKeepsTheEntry) { ExpectNoWithdrawWithin(registry_service, absl::Seconds(2)); const RaidenId peer{"peer_job_tier", "0", "peer_data", 0}; auto after = registry->Lookup({"staged"}, peer); - ASSERT_OK(after.status()); + ABSL_ASSERT_OK(after.status()); EXPECT_EQ(after->size(), 1) << "a fetch refused on tier withdrew an entry for a block this store " "still holds"; @@ -1431,7 +1419,7 @@ TEST(FetchWithdrawTest, ARefusedFetchKeepsAnEvictionCandidatesEntry) { {RaidenBlockId(src_id, 10, BlockStatus::HOST)}, /*on_host=*/true) .first); - ASSERT_OK(registry->Register({{"demoted", src_id, 10}})); + ABSL_ASSERT_OK(registry->Register({{"demoted", src_id, 10}})); // InsertAndLock, not Insert: it discards what Put displaces, on purpose -- // the displaced entry becomes a candidate that still holds its host block. @@ -1463,7 +1451,7 @@ TEST(FetchWithdrawTest, ARefusedFetchKeepsAnEvictionCandidatesEntry) { ExpectNoWithdrawWithin(registry_service, absl::Seconds(2)); const RaidenId peer{"peer_job_candidate", "0", "peer_data", 0}; auto after = registry->Lookup({"demoted"}, peer); - ASSERT_OK(after.status()); + ABSL_ASSERT_OK(after.status()); EXPECT_EQ(after->size(), 1) << "a refused fetch withdrew an eviction candidate, which still holds " "its host block"; @@ -1486,7 +1474,7 @@ TEST(FetchWithdrawTest, AHashOnlyAPeerHoldsIsAMissAndThePeerKeepsItsEntry) { /*store_server_ip=*/"127.0.0.1"); // Registered, but on somebody else. This store never held it. - ASSERT_OK(registry->Register({{"peer_only", peer_id, 7}})); + ABSL_ASSERT_OK(registry->Register({{"peer_only", peer_id, 7}})); KVCacheStoreServiceImpl service(store.backend().get(), store.raiden_controller()); @@ -1520,7 +1508,7 @@ TEST(FetchWithdrawTest, AHashOnlyAPeerHoldsIsAMissAndThePeerKeepsItsEntry) { const RaidenId asker{"asker_job", "0", "asker_data", 0}; auto after = registry->Lookup({"peer_only"}, asker); - ASSERT_OK(after.status()); + ABSL_ASSERT_OK(after.status()); ASSERT_EQ(after->size(), 1); EXPECT_EQ((*after)[0].raiden_id().job_name(), peer_id.job_name); @@ -1530,17 +1518,16 @@ TEST(FetchWithdrawTest, AHashOnlyAPeerHoldsIsAMissAndThePeerKeepsItsEntry) { // A poll with wait_ms parks while the operation is pending and returns as // soon as it goes terminal. TEST_F(WriteRemoteTest, PollWriteRemoteWithWaitMsAwaitsUntilTerminal) { - auto ack_or = Offer({"wait_a", "wait_b"}); - ASSERT_OK(ack_or.status()); - const uint64_t op_id = ack_or->operation_id(); + TF_ASSERT_OK_AND_ASSIGN(auto ack, Offer({"wait_a", "wait_b"})); + const uint64_t op_id = ack.operation_id(); ASSERT_EQ(latch_.issued(), 1); std::atomic poll_done = false; proto::PollWriteRemoteResponse poll_resp; std::thread poll_thread([&]() { - auto resp_or = client_->PollWriteRemote(op_id, /*wait_ms=*/5000).Await(); - if (resp_or.ok()) { - poll_resp = *resp_or; + auto poll_res = client_->PollWriteRemote(op_id, /*wait_ms=*/5000).Await(); + if (poll_res.ok()) { + poll_resp = *poll_res; } poll_done = true; }); @@ -1621,14 +1608,12 @@ TEST_F(WriteRemoteTest, AWaitingPollIsCappedByTheServersOwnClock) { setenv("RAIDEN_REMOTE_WRITE_DEADLINE_S", "1", /*overwrite=*/1); service_->PauseDeadlineFiringForTesting(); - auto ack_or = Offer({"cap_a"}); - ASSERT_OK(ack_or.status()); + TF_ASSERT_OK_AND_ASSIGN(auto ack, Offer({"cap_a"})); ASSERT_EQ(latch_.issued(), 1); const absl::Time before = absl::Now(); - auto resp_or = - client_->PollWriteRemote(ack_or->operation_id(), /*wait_ms=*/60000) - .Await(); + auto poll_res = + client_->PollWriteRemote(ack.operation_id(), /*wait_ms=*/60000).Await(); const absl::Duration waited = absl::Now() - before; if (previous != nullptr) { @@ -1637,8 +1622,8 @@ TEST_F(WriteRemoteTest, AWaitingPollIsCappedByTheServersOwnClock) { unsetenv("RAIDEN_REMOTE_WRITE_DEADLINE_S"); } - ASSERT_OK(resp_or.status()); - EXPECT_EQ(resp_or->state(), proto::PollWriteRemoteResponse::PENDING); + TF_ASSERT_OK_AND_ASSIGN(auto resp, poll_res); + EXPECT_EQ(resp.state(), proto::PollWriteRemoteResponse::PENDING); EXPECT_GE(waited, absl::Seconds(4)) << "the poll did not wait at all"; EXPECT_LT(waited, absl::Seconds(30)) << "the requested one-minute wait was honoured instead of clamped"; diff --git a/tpu_sync/kv_cache/kv_cache_store_test.cc b/tpu_sync/kv_cache/kv_cache_store_test.cc index ac8eb30b..fbaa8de7 100644 --- a/tpu_sync/kv_cache/kv_cache_store_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_test.cc @@ -215,7 +215,7 @@ TEST(KVCacheStoreTest, EvictionTracking) { RaidenId{"inference_server", "1", "kv_cache", 1}}; // 1. Insert 101 and 102, filling the cache. Both come back PINNED. - EXPECT_TRUE(controller.Insert(hashes_1_2, slices_1_2, true).ok()); + ABSL_EXPECT_OK(controller.Insert(hashes_1_2, slices_1_2, true)); // 2. Insert 103. available_space() is capacity minus the PINNED entries, so // with both slots pinned there is nothing insert is allowed to reclaim and @@ -237,7 +237,7 @@ TEST(KVCacheStoreTest, EvictionTracking) { // prefix caches: dropping the tail of a shared prefix costs less than // dropping its head. controller.Release(hashes_1_2); - EXPECT_TRUE(controller.Insert(hash_3, slice_3, true).ok()); + ABSL_EXPECT_OK(controller.Insert(hash_3, slice_3, true)); // 4. Verify that 102 is in candidates. EXPECT_THAT(KVCacheStoreTest::GetEvictCandidateKeys(controller), @@ -245,7 +245,7 @@ TEST(KVCacheStoreTest, EvictionTracking) { // 5. Verify that lookup for 102 misses (candidate invisible with Peek). auto lookup_res = PeekLookup(controller, {"102"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_TRUE(lookup_res->empty()); // 102 should still be in candidates. EXPECT_THAT(KVCacheStoreTest::GetEvictCandidateKeys(controller), @@ -299,7 +299,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { { auto lookup_res = store.Lookup({"local_only_hash"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); EXPECT_EQ((*lookup_res)[0].second.raiden_id.job_name, "local_job"); @@ -311,7 +311,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { { auto lookup_res = store.Lookup({"shared_hash"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "shared_hash"); // Should return local info, not remote info from registry @@ -323,7 +323,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { { auto lookup_res = store.Lookup({"global_hash_1", "global_hash_2"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); @@ -343,7 +343,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { // local_only_hash and stop. { auto lookup_res = PeekLookup(store, {"local_only_hash", "global_hash_1"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); } @@ -354,7 +354,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { auto lookup_res = store.Lookup({"local_only_hash", "global_hash_1", "global_hash_2"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 3); EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); @@ -377,7 +377,7 @@ TEST(KVCacheStoreTest, GlobalLookupFallback) { auto lookup_res = store.Lookup( {"local_only_hash", "global_hash_1", "missing_hash", "global_hash_2"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); // local_only_hash, global_hash_1 EXPECT_EQ((*lookup_res)[0].first, "local_only_hash"); EXPECT_EQ((*lookup_res)[1].first, "global_hash_1"); @@ -523,7 +523,7 @@ TEST(KVCacheStoreTest, GlobalLookupRegistryDown) { // should return the local hit. Observation only, so no pin is taken. auto lookup_res = store.Lookup({"local_hash", "missing_hash"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "local_hash"); EXPECT_EQ((*lookup_res)[0].second.raiden_id.job_name, "local_job"); @@ -546,10 +546,9 @@ TEST(KVCacheStoreTest, ValidateAndPinHostBlocksSuccessReDerivesIdsAndPins) { BlockStatus::HOST_AND_HBM)}; ASSERT_TRUE(InsertResident(store, hashes, slices, /*on_host=*/true)); - auto ids_or = store.ValidateAndPinHostBlocks(hashes); - ASSERT_TRUE(ids_or.ok()) << ids_or.status().message(); + TF_ASSERT_OK_AND_ASSIGN(auto ids, store.ValidateAndPinHostBlocks(hashes)); // Source ids are re-derived from the LRU (not from the request). - EXPECT_THAT(*ids_or, ::testing::ElementsAre(5, 7)); + EXPECT_THAT(ids, ::testing::ElementsAre(5, 7)); EXPECT_EQ(store.GetPinCount("h0"), 1); EXPECT_EQ(store.GetPinCount("h1"), 1); @@ -561,9 +560,9 @@ TEST(KVCacheStoreTest, ValidateAndPinHostBlocksSuccessReDerivesIdsAndPins) { TEST(KVCacheStoreTest, ValidateAndPinHostBlocksMissingReturnsNotFound) { KVCacheStore store(4, "", {}, /*num_shards=*/1, /*shard_size_bytes=*/512, /*store_server_ip=*/"127.0.0.1"); - auto ids_or = + auto ids = store.ValidateAndPinHostBlocks(std::vector{"missing"}); - EXPECT_TRUE(absl::IsNotFound(ids_or.status())) << ids_or.status(); + EXPECT_TRUE(absl::IsNotFound(ids.status())) << ids.status(); } TEST(KVCacheStoreTest, ValidateAndPinHostBlocksWrongStatusFailedPrecondition) { @@ -575,8 +574,8 @@ TEST(KVCacheStoreTest, ValidateAndPinHostBlocksWrongStatusFailedPrecondition) { rid, /*host_block_id=*/-1, /*device_block_id=*/0, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, hashes, slices, /*on_host=*/false)); - auto ids_or = store.ValidateAndPinHostBlocks(hashes); - EXPECT_TRUE(absl::IsFailedPrecondition(ids_or.status())) << ids_or.status(); + auto ids = store.ValidateAndPinHostBlocks(hashes); + EXPECT_TRUE(absl::IsFailedPrecondition(ids.status())) << ids.status(); EXPECT_EQ(store.GetPinCount("hbm_h"), 0); } @@ -594,8 +593,8 @@ TEST(KVCacheStoreTest, ValidateAndPinHostBlocksAtomicRollbackOnPartialMiss) { BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, hashes, slices, /*on_host=*/false)); - auto ids_or = store.ValidateAndPinHostBlocks(hashes); - EXPECT_FALSE(ids_or.ok()); + auto ids = store.ValidateAndPinHostBlocks(hashes); + EXPECT_FALSE(ids.ok()); EXPECT_EQ(store.GetPinCount("ok"), 0); EXPECT_EQ(store.GetPinCount("bad"), 0); } @@ -603,9 +602,9 @@ TEST(KVCacheStoreTest, ValidateAndPinHostBlocksAtomicRollbackOnPartialMiss) { TEST(KVCacheStoreTest, ValidateAndPinHostBlocksEmptyInputIsOk) { KVCacheStore store(4, "", {}, /*num_shards=*/1, /*shard_size_bytes=*/512, /*store_server_ip=*/"127.0.0.1"); - auto ids_or = store.ValidateAndPinHostBlocks(std::vector{}); - ASSERT_TRUE(ids_or.ok()); - EXPECT_TRUE(ids_or->empty()); + TF_ASSERT_OK_AND_ASSIGN( + auto ids, store.ValidateAndPinHostBlocks(std::vector{})); + EXPECT_TRUE(ids.empty()); } TEST(KVCacheStoreTest, @@ -617,12 +616,11 @@ TEST(KVCacheStoreTest, std::vector slices = {RaidenBlockId( rid, /*host_block_id=*/9, /*device_block_id=*/-1, BlockStatus::HOST)}; // Insert pins once, and this case counts from that pin. - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); EXPECT_EQ(store.GetPinCount("h0"), 1); - auto ids_or = store.ValidateAndPinHostBlocks(hashes); - ASSERT_TRUE(ids_or.ok()); - EXPECT_THAT(*ids_or, ::testing::ElementsAre(9)); + TF_ASSERT_OK_AND_ASSIGN(auto ids, store.ValidateAndPinHostBlocks(hashes)); + EXPECT_THAT(ids, ::testing::ElementsAre(9)); EXPECT_EQ(store.GetPinCount("h0"), 2); // verify added a second pin. store.UnpinHostBlocks(hashes); @@ -643,7 +641,7 @@ TEST(KVCacheStoreTest, LookupCapLimit) { // Lookup 3 hashes, but capacity is 2. It should only return 2. std::vector lookup_hashes = {"101", "102", "103"}; auto lookup_res = PeekLookup(store, lookup_hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].first, "101"); EXPECT_EQ((*lookup_res)[1].first, "102"); @@ -669,11 +667,9 @@ TEST(KVCacheStoreTest, LookupCapLimitWithGlobal) { RaidenId host3{"job3", "0", "kv_cache", 0}; int32_t block3 = 44; - ASSERT_TRUE(registry_client - .Register({{hash1, host1, block1}, - {hash2, host2, block2}, - {hash3, host3, block3}}) - .ok()); + ABSL_ASSERT_OK(registry_client.Register({{hash1, host1, block1}, + {hash2, host2, block2}, + {hash3, host3, block3}})); // 3. Create KVCacheStore with capacity 2 RaidenId store_id{"store_job", "0", "kv_cache", 0}; @@ -685,7 +681,7 @@ TEST(KVCacheStoreTest, LookupCapLimitWithGlobal) { std::vector lookup_hashes = {"global_hash_1", "global_hash_2", "global_hash_3"}; auto lookup_res = store.Lookup(lookup_hashes, /*enable_global=*/true); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].first, "global_hash_1"); EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); @@ -707,9 +703,8 @@ TEST(KVCacheStoreTest, LookupCapLimitMixed) { RaidenId host3{"job3", "0", "kv_cache", 0}; int32_t block3 = 44; - ASSERT_TRUE( - registry_client.Register({{hash2, host2, block2}, {hash3, host3, block3}}) - .ok()); + ABSL_ASSERT_OK(registry_client.Register( + {{hash2, host2, block2}, {hash3, host3, block3}})); // 3. Create KVCacheStore with capacity 2 RaidenId store_id{"store_job", "0", "kv_cache", 0}; @@ -729,7 +724,7 @@ TEST(KVCacheStoreTest, LookupCapLimitMixed) { "global_hash_3"}; auto lookup_res = store.Lookup(lookup_hashes, LookupOptions{.enable_global = true}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].first, "local_hash_1"); EXPECT_EQ((*lookup_res)[1].first, "global_hash_2"); @@ -745,7 +740,7 @@ TEST(KVCacheStoreTest, LookupDefaultSkipsTheRegistry) { auto& registry_client = *reg_server->client; RaidenId owner{"peer_job", "0", "kv_cache", 0}; - ASSERT_TRUE(registry_client.Register({{"registry_only", owner, 42}}).ok()); + ABSL_ASSERT_OK(registry_client.Register({{"registry_only", owner, 42}})); RaidenId store_id{"store_job", "0", "kv_cache", 0}; KVCacheStore store(4, reg_server->server_address, store_id, @@ -754,13 +749,13 @@ TEST(KVCacheStoreTest, LookupDefaultSkipsTheRegistry) { // Defaulted: local miss, registry never consulted. auto defaulted = store.Lookup({"registry_only"}); - ASSERT_TRUE(defaulted.ok()); + ABSL_ASSERT_OK(defaulted); EXPECT_TRUE(defaulted->empty()); // Asked globally, the same hash resolves. auto global = store.Lookup({"registry_only"}, LookupOptions{.enable_global = true}); - ASSERT_TRUE(global.ok()); + ABSL_ASSERT_OK(global); ASSERT_EQ(global->size(), 1); EXPECT_EQ((*global)[0].second.status, BlockStatus::REMOTE); EXPECT_EQ((*global)[0].second.raiden_id, owner); @@ -779,13 +774,13 @@ TEST(KVCacheStoreTest, LookupAvailableSpaceLimit) { ASSERT_TRUE(InsertResident(store, hashes, slices, true)); // Pin 101. Pinned count = 1. Available space = 3 - 1 = 2. - EXPECT_TRUE(store.Lookup({"101"}).ok()); + ABSL_EXPECT_OK(store.Lookup({"101"})); // Lookup 4 hashes. Lookup is non-mutating and unbounded by available space, // returning all 3 cached blocks up to the first miss ("104"). std::vector lookup_hashes = {"101", "102", "103", "104"}; auto lookup_res = PeekLookup(store, lookup_hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 3); EXPECT_EQ((*lookup_res)[0].first, "101"); EXPECT_EQ((*lookup_res)[1].first, "102"); @@ -810,7 +805,7 @@ TEST(KVCacheStoreTest, InsertPinsExistingAndNewAlike) { std::vector slices = { RaidenId{"local_job", "0", "kv_cache", 0}, RaidenId{"remote_job", "0", "kv_cache", 42}}; - EXPECT_TRUE(store.Insert({"local_1", "remote_1"}, slices, true).ok()); + ABSL_EXPECT_OK(store.Insert({"local_1", "remote_1"}, slices, true)); EXPECT_EQ(store.GetPinCount("local_1"), 1); EXPECT_EQ(store.GetPinCount("remote_1"), 1); @@ -860,7 +855,7 @@ TEST(KVCacheStoreTest, EvictRaceCondition) { ASSERT_TRUE(InsertResident(store, local_hashes, local_slices, true)); // Pin local_1 - ASSERT_TRUE(store.Lookup({"local_1"}).ok()); + ABSL_ASSERT_OK(store.Lookup({"local_1"})); EXPECT_EQ(store.GetPinCount("local_1"), 1); // Attempt Evict on local_1 (which is pinned) @@ -871,7 +866,7 @@ TEST(KVCacheStoreTest, EvictRaceCondition) { // count above is the subject, so the check must not add one of its own. EXPECT_EQ(store.GetPinCount("local_1"), 1); auto lookup_res = PeekLookup(store, {"local_1"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_EQ(lookup_res->size(), 1); } @@ -911,12 +906,12 @@ MakeRecoveryController(const RaidenId& rid, int num_blocks) { TEST(KVCacheStoreTest, MetadataKeepsEvictionCandidates) { MetadataRegion region(4); - auto metadata_or = KVCacheMetadata::Format(region.span(), 4); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 4)); RaidenId rid{"local_job", "0", "kv_cache", 0}; KVCacheStore store(2, MakeRecoveryController(rid, 4), - /*global_registry_address=*/"", rid, *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); ASSERT_TRUE(InsertResident(store, {"host_1", "host_2"}, @@ -936,7 +931,7 @@ TEST(KVCacheStoreTest, MetadataKeepsEvictionCandidates) { {RaidenBlockId(rid, 2, BlockStatus::HOST)}, true)); EXPECT_THAT(KVCacheStoreTest::GetEvictCandidateKeys(store), ElementsAre("host_2")); - EXPECT_THAT(metadata_or->ValidEntries(), + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(::testing::FieldsAre(0, "host_1", 1), ::testing::FieldsAre(1, "host_2", 0), ::testing::FieldsAre(2, "host_3", 2))); @@ -948,7 +943,7 @@ TEST(KVCacheStoreTest, MetadataKeepsEvictionCandidates) { {RaidenBlockId(rid, 3, BlockStatus::HOST)}, true)); EXPECT_THAT(KVCacheStoreTest::GetEvictCandidateKeys(store), ::testing::IsEmpty()); - EXPECT_THAT(metadata_or->ValidEntries(), + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(::testing::FieldsAre(0, "host_1", 1), ::testing::FieldsAre(2, "host_3", 2), ::testing::FieldsAre(3, "host_2", 3))); @@ -956,22 +951,22 @@ TEST(KVCacheStoreTest, MetadataKeepsEvictionCandidates) { TEST(KVCacheStoreTest, EvictClearsMetadataEntries) { MetadataRegion region(2); - auto metadata_or = KVCacheMetadata::Format(region.span(), 2); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 2)); RaidenId rid{"local_job", "0", "kv_cache", 0}; KVCacheStore store(2, MakeRecoveryController(rid, 2), - /*global_registry_address=*/"", rid, *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); ASSERT_TRUE(InsertResident(store, {"host_1", "host_2"}, {RaidenBlockId(rid, 0, BlockStatus::HOST), RaidenBlockId(rid, 1, BlockStatus::HOST)}, true)); - ASSERT_EQ(metadata_or->ValidEntries().size(), 2); + ASSERT_EQ(metadata.ValidEntries().size(), 2); EXPECT_EQ(KVCacheStoreTest::Evict(store, {"host_1"}), 1); - EXPECT_THAT(metadata_or->ValidEntries(), + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(::testing::FieldsAre(1, "host_2", 0))); } @@ -1018,7 +1013,7 @@ class KVCacheStoreEmbeddedControllerTest : public ::testing::Test { controller.controller_address()); auto status = client.RegisterWorker(worker_id, worker_address, {{worker_address, {}}}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); } std::unique_ptr<::tpu_raiden::controller::RaidenController> MakeController( @@ -1072,8 +1067,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveReusesFreedBlocksAfterEvict) { RaidenBlockId(rid, -1, 0, BlockStatus::HBM), RaidenBlockId(rid, -1, 1, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, first, first_slices, false)); - ASSERT_TRUE(store.Lookup(first).ok()); - ASSERT_TRUE(save_and_wait(first).ok()); + ABSL_ASSERT_OK(store.Lookup(first)); + ABSL_ASSERT_OK(save_and_wait(first)); EXPECT_EQ(controller_ptr->block_manager()->num_locked_blocks(), 2); // Evict everything: the successful save consumed the pins, so the @@ -1088,9 +1083,9 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveReusesFreedBlocksAfterEvict) { std::vector second_slices = { RaidenBlockId(rid, -1, 0, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, second, second_slices, false)); - ASSERT_TRUE(store.Lookup(second).ok()); + ABSL_ASSERT_OK(store.Lookup(second)); absl::Status status = save_and_wait(second); - EXPECT_TRUE(status.ok()) << status.message(); + ABSL_EXPECT_OK(status); EXPECT_EQ(controller_ptr->block_manager()->num_locked_blocks(), 1); } @@ -1115,11 +1110,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveSuccess) { ASSERT_TRUE(InsertResident(store, hashes, slices, false)); // Pin them - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); // Save them absl::Status status = store.Save(hashes); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for completion bool done = false; @@ -1151,7 +1146,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveSuccess) { // Verify status in store is updated to HOST_AND_HBM auto lookup_res = PeekLookup(store, hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_EQ((*lookup_res)[0].second.host_block_id, 0); @@ -1182,11 +1177,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadSuccess) { ASSERT_TRUE(InsertResident(store, hashes, slices, true)); // Pin them - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); // Load them to device block 2 and 3 absl::Status status = store.Load(hashes, {2, 3}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for completion bool done = false; @@ -1213,7 +1208,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadSuccess) { // Verify status in store is updated to HOST_AND_HBM auto lookup_res = PeekLookup(store, hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_EQ((*lookup_res)[0].second.host_block_id, 0); @@ -1241,10 +1236,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesSuccess) { RaidenBlockId(rid, 1, -1, BlockStatus::HOST)}; ASSERT_TRUE(InsertResident(store, hashes, slices, true)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); absl::Status status = store.Load(hashes, slices, {2, 3}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); bool done = false; while (!done) { @@ -1314,8 +1309,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesUnpinnedFails) { // With the pin the caller was supposed to hold, it goes through -- and a // successful load CONSUMES that pin, same as the no-slices form. - ASSERT_TRUE(store.Lookup(hashes).ok()); - ASSERT_TRUE(store.Load(hashes, slices, {2}).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); + ABSL_ASSERT_OK(store.Load(hashes, slices, {2})); bool done = false; for (int attempt = 0; attempt < 100 && !done; ++attempt) { auto [load_done, load_failed, load_pending] = store.PollLoadStatus(); @@ -1350,10 +1345,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LocalSaveConsumesTheCallerPin) { std::vector slices = { RaidenBlockId(rid, -1, 3, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, hashes, slices, /*on_host=*/false)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); ASSERT_EQ(store.GetPinCount("hash_1"), 1); - ASSERT_OK(store.Save(hashes)); + ABSL_ASSERT_OK(store.Save(hashes)); bool done = false; for (int attempt = 0; attempt < 200 && !done; ++attempt) { @@ -1378,7 +1373,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LocalSaveConsumesTheCallerPin) { // The entry survives the unpin, carrying the save's result. auto after = PeekLookup(store, hashes); - ASSERT_TRUE(after.ok()); + ABSL_ASSERT_OK(after); ASSERT_EQ(after->size(), 1); EXPECT_EQ((*after)[0].second.status, BlockStatus::HOST_AND_HBM); } @@ -1403,10 +1398,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LocalLoadConsumesTheCallerPin) { std::vector slices = { RaidenBlockId(rid, 0, -1, BlockStatus::HOST)}; ASSERT_TRUE(InsertResident(store, hashes, slices, true)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); ASSERT_EQ(store.GetPinCount("hash_1"), 1); - ASSERT_OK(store.Load(hashes, {2})); + ABSL_ASSERT_OK(store.Load(hashes, {2})); bool done = false; for (int attempt = 0; attempt < 100 && !done; ++attempt) { @@ -1421,7 +1416,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LocalLoadConsumesTheCallerPin) { << "a successful local load must consume the caller's pin"; // The entry survives the unpin and carries the load's result. auto after = PeekLookup(store, hashes); - ASSERT_TRUE(after.ok()); + ABSL_ASSERT_OK(after); ASSERT_EQ(after->size(), 1); EXPECT_EQ((*after)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_EQ((*after)[0].second.device_block_id, 2); @@ -1448,10 +1443,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, FailedLocalLoadKeepsTheCallerPin) { std::vector slices = { RaidenBlockId(rid, 0, -1, BlockStatus::HOST)}; ASSERT_TRUE(InsertResident(store, hashes, slices, true)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); ASSERT_EQ(store.GetPinCount("hash_1"), 1); - ASSERT_OK(store.Load(hashes, {2})); + ABSL_ASSERT_OK(store.Load(hashes, {2})); bool failed = false; for (int attempt = 0; attempt < 100 && !failed; ++attempt) { @@ -1469,7 +1464,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, FailedLocalLoadKeepsTheCallerPin) { << "a failed load must leave the caller's pin so a retry can hold on"; // The entry itself is untouched: still HOST, still the same block. auto after = PeekLookup(store, hashes); - ASSERT_TRUE(after.ok()); + ABSL_ASSERT_OK(after); ASSERT_EQ(after->size(), 1); EXPECT_EQ((*after)[0].second.status, BlockStatus::HOST); // The caller decides: retry (the pin is still good) or give up. @@ -1500,11 +1495,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, std::vector slices = { RaidenBlockId(rid, -1, 3, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, hashes, slices, /*on_host=*/false)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); ASSERT_EQ(store.GetPinCount("hash_1"), 1); const size_t free_before = controller_ptr->block_manager()->num_free_blocks(); - ASSERT_OK(store.Save(hashes)); + ABSL_ASSERT_OK(store.Save(hashes)); bool failed = false; for (int attempt = 0; attempt < 200 && !failed; ++attempt) { @@ -1525,7 +1520,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, << "the host blocks a failed save allocated must return to the pool"; // Still an HBM-only entry: the failed save recorded no host residency. auto after = PeekLookup(store, hashes); - ASSERT_TRUE(after.ok()); + ABSL_ASSERT_OK(after); ASSERT_EQ(after->size(), 1); EXPECT_EQ((*after)[0].second.status, BlockStatus::HBM); store.Release(hashes); @@ -1590,7 +1585,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveRequiresHbmResidency) { std::vector hashes = {"hash_1"}; std::vector slices = { RaidenBlockId(rid, 0, -1, BlockStatus::HOST)}; - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); absl::Status status = store.Save(hashes); EXPECT_TRUE(absl::IsFailedPrecondition(status)) << status; @@ -1637,10 +1632,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesAlreadyLoadingFails) { RaidenBlockId(rid, 0, -1, BlockStatus::HOST)}; ASSERT_TRUE(InsertResident(store, hashes, slices, true)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); absl::Status status1 = store.Load(hashes, slices, {2}); - ASSERT_TRUE(status1.ok()); + ABSL_ASSERT_OK(status1); absl::Status status2 = store.Load(hashes, slices, {3}); EXPECT_TRUE(absl::IsFailedPrecondition(status2)); @@ -1665,7 +1660,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesMixedStatusesFails) { KVCacheStoreTest::PlantIndexEntry(store, {"hash_1"}, {slices[0]}, /*on_host=*/false); ASSERT_TRUE(InsertResident(store, {"hash_2"}, {slices[1]}, /*on_host=*/true)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); absl::Status status = store.Load(hashes, slices, {2, 3}); EXPECT_TRUE(absl::IsInvalidArgument(status)); @@ -1689,11 +1684,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesRemoteSuccess) { remote_config.global_registry_address = registry_address; remote_config.raiden_id = remote_rid; - auto remote_backend_or = - HostOffloadBackend::Create(remote_config, controller.get()); - ASSERT_OK(remote_backend_or.status()); + TF_ASSERT_OK_AND_ASSIGN( + auto remote_backend_raw, + HostOffloadBackend::Create(remote_config, controller.get())); auto remote_backend = - std::dynamic_pointer_cast(*remote_backend_or); + std::dynamic_pointer_cast(remote_backend_raw); ASSERT_NE(remote_backend, nullptr); std::vector remote_slices = { @@ -1703,16 +1698,16 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesRemoteSuccess) { /*on_host=*/true); auto remote_server = KVCacheStoreServer::Create(); - ASSERT_OK(remote_server->StartServer(remote_backend.get(), controller.get(), - "127.0.0.1")); + ABSL_ASSERT_OK(remote_server->StartServer(remote_backend.get(), + 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(remote_rid, - remote_server->GetServerAddress(), - controller->controller_address())); + ABSL_ASSERT_OK(registry_client->RegisterStore( + remote_rid, remote_server->GetServerAddress(), + controller->controller_address())); KVCacheStore store(10, std::move(controller), registry_address, local_rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); @@ -1726,7 +1721,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesRemoteSuccess) { // No pin: a load from a peer requires none and consumes none. absl::Status status = store.Load(hashes, slices, {5}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); bool done = false; for (int attempt = 0; attempt < 100; ++attempt) { @@ -1745,7 +1740,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadWithSlicesRemoteSuccess) { // Nothing is recorded for a peer source, so the entry this test inserted up // front is left exactly as it was: still REMOTE, still naming the peer. auto lookup_res = PeekLookup(store, hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::REMOTE); EXPECT_EQ((*lookup_res)[0].second.host_block_id, 42); @@ -1772,29 +1767,30 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadRemoteWithSlicesRecordsNothing) { remote_config.global_registry_address = registry_address; remote_config.raiden_id = remote_rid; - auto remote_backend_or = - HostOffloadBackend::Create(remote_config, controller.get()); - ASSERT_OK(remote_backend_or.status()); + TF_ASSERT_OK_AND_ASSIGN( + auto remote_backend_raw, + HostOffloadBackend::Create(remote_config, controller.get())); auto remote_backend = - std::dynamic_pointer_cast(*remote_backend_or); + std::dynamic_pointer_cast(remote_backend_raw); ASSERT_NE(remote_backend, nullptr); remote_backend->Insert({"slice_load_hash"}, {RaidenBlockId(remote_rid, 42, BlockStatus::HOST)}, /*on_host=*/true); auto remote_server = KVCacheStoreServer::Create(); - ASSERT_OK(remote_server->StartServer(remote_backend.get(), controller.get(), - "127.0.0.1")); + ABSL_ASSERT_OK(remote_server->StartServer(remote_backend.get(), + 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(remote_rid, - remote_server->GetServerAddress(), - controller->controller_address())); + ABSL_ASSERT_OK(registry_client->RegisterStore( + remote_rid, remote_server->GetServerAddress(), + controller->controller_address())); // And advertise the block: indexing it in the remote backend does not, and // this case needs the registry to answer for a hash the store lacks. - ASSERT_OK(registry_client->Register({{"slice_load_hash", remote_rid, 42}})); + ABSL_ASSERT_OK( + registry_client->Register({{"slice_load_hash", remote_rid, 42}})); KVCacheStore store(10, std::move(controller), registry_address, local_rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); @@ -1804,14 +1800,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadRemoteWithSlicesRecordsNothing) { // The registry answers, but a registry-only hit never enters the local index // and is never pinned -- so this load has no caller pin to consume either. auto resolved = store.Lookup(hashes, /*enable_global=*/true); - ASSERT_TRUE(resolved.ok()); + ABSL_ASSERT_OK(resolved); ASSERT_EQ(resolved->size(), 1); EXPECT_EQ((*resolved)[0].second.status, BlockStatus::REMOTE); EXPECT_TRUE(PeekLookup(store, hashes)->empty()) << "a registry-only hit must not have entered the local index"; std::vector slices = {(*resolved)[0].second}; - ASSERT_OK(store.Load(hashes, slices, {5})); + ABSL_ASSERT_OK(store.Load(hashes, slices, {5})); bool done = false; for (int attempt = 0; attempt < 100 && !done; ++attempt) { @@ -1887,10 +1883,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveMultiWorkerSuccess) { ASSERT_TRUE(InsertResident(store, hashes, slices, false)); // Pin them - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); absl::Status status = store.Save(hashes); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for completion bool done = false; @@ -1924,7 +1920,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveMultiWorkerSuccess) { EXPECT_THAT(mock_mgr_1.last_dst_offsets, ElementsAre(0, 1)); auto lookup_res = PeekLookup(store, hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_EQ((*lookup_res)[0].second.host_block_id, 0); @@ -1960,10 +1956,10 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadMultiWorkerSuccess) { ASSERT_TRUE(InsertResident(store, hashes, slices, true)); // Pin them - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); absl::Status status = store.Load(hashes, {2, 3}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for completion bool done = false; @@ -1993,7 +1989,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, LoadMultiWorkerSuccess) { EXPECT_THAT(mock_mgr_1.last_dst_offsets, ElementsAre(2, 3)); auto lookup_res = PeekLookup(store, hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_EQ((*lookup_res)[0].second.host_block_id, 0); @@ -2025,11 +2021,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, SaveWriteThrough) { // 4. Insert them as HBM blocks locally and pin them ASSERT_TRUE(InsertResident(store, hashes, slices, false)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); // 5. Call Save on the store absl::Status status = store.Save(hashes); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // 6. Poll for completion bool done = false; @@ -2138,11 +2134,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, RaidenBlockId(rid, -1, i, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, batch, slices, false)); // Lookup takes the pin that Save consumes. - ASSERT_TRUE(store.Lookup(batch).ok()); + ABSL_ASSERT_OK(store.Lookup(batch)); // The pool has not drained, which is the failure this bound prevents. absl::Status status = store.Save(batch); - ASSERT_TRUE(status.ok()) + ABSL_ASSERT_OK(status) << "save " << i << " failed: " << status.ToString(); ASSERT_FALSE(absl::IsResourceExhausted(status)); saved.push_back(hash); @@ -2222,9 +2218,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictByHashesErasesEitherStatus) { RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); - auto alloc_or = controller->AllocateBlockIds(2); - ASSERT_TRUE(alloc_or.ok()); - std::vector host_block_ids = *alloc_or; + TF_ASSERT_OK_AND_ASSIGN(std::vector host_block_ids, + controller->AllocateBlockIds(2)); ASSERT_EQ(host_block_ids.size(), 2); RaidenId rid{"test_job", "0", "test_cache", 0}; @@ -2235,10 +2230,9 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictByHashesErasesEitherStatus) { auto channel = grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); global_registry::GlobalRegistryClient registry_client(channel); - ASSERT_TRUE(registry_client - .Register({{"hash_1", rid, host_block_ids[0]}, - {"hash_2", rid, host_block_ids[1]}}) - .ok()); + ABSL_ASSERT_OK( + registry_client.Register({{"hash_1", rid, host_block_ids[0]}, + {"hash_2", rid, host_block_ids[1]}})); std::vector hashes = {"hash_1", "hash_2"}; std::vector slices = { @@ -2258,7 +2252,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictByHashesErasesEitherStatus) { // is untouched. EXPECT_EQ(PeekLookup(store, {"hash_1", "hash_2"})->size(), 0); auto lookup_res = PeekLookup(store, {"hash_2"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].first, "hash_2"); EXPECT_EQ((*lookup_res)[0].second.status, c.status); @@ -2307,8 +2301,9 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictKeepsASkippedHashRegistered) { auto channel = grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()); global_registry::GlobalRegistryClient registry_client(channel); - ASSERT_OK(registry_client.Register({{"pinned_1", rid, host_block_ids[0]}, - {"free_1", rid, host_block_ids[1]}})); + ABSL_ASSERT_OK( + registry_client.Register({{"pinned_1", rid, host_block_ids[0]}, + {"free_1", rid, host_block_ids[1]}})); std::vector hashes = {"pinned_1", "free_1"}; std::vector slices = { @@ -2317,7 +2312,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictKeepsASkippedHashRegistered) { ASSERT_TRUE(InsertResident(store, hashes, slices, true)); // The pin is what makes "pinned_1" unevictable. - ASSERT_OK(store.Lookup({"pinned_1"})); + ABSL_ASSERT_OK(store.Lookup({"pinned_1"})); ASSERT_EQ(store.GetPinCount("pinned_1"), 1); // Ask for both. Only the unpinned one can go. @@ -2371,8 +2366,9 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictOnSave) { controller_ptr->AllocateBlockIds(2)); ASSERT_EQ(host_block_ids.size(), 2); - ASSERT_OK(registry_client.Register({{"block_A", rid, host_block_ids[0]}, - {"block_B", rid, host_block_ids[1]}})); + ABSL_ASSERT_OK( + registry_client.Register({{"block_A", rid, host_block_ids[0]}, + {"block_B", rid, host_block_ids[1]}})); std::vector hashes = {"block_A", "block_B"}; std::vector slices = { @@ -2384,11 +2380,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, EvictOnSave) { std::vector slices_C = { RaidenBlockId(rid, -1, 0, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, hashes_C, slices_C, false)); - ASSERT_OK(store.Lookup(hashes_C)); + ABSL_ASSERT_OK(store.Lookup(hashes_C)); EXPECT_EQ(controller_ptr->block_manager()->num_free_blocks(), 0); - ASSERT_OK(store.Save(hashes_C)); + ABSL_ASSERT_OK(store.Save(hashes_C)); bool done = false; while (!done) { @@ -2454,8 +2450,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ProactiveEvictionWithCandidates) { ASSERT_TRUE(InsertResident(store, hashes, slices, false)); // 2. Save A and B (allocates host blocks for both) - ASSERT_TRUE(store.Lookup(hashes).ok()); - ASSERT_TRUE(store.Save(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); + ABSL_ASSERT_OK(store.Save(hashes)); // Poll for Save completion bool save_done = false; @@ -2476,7 +2472,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ProactiveEvictionWithCandidates) { // through the pinned list and back, which reorders the very LRU under test. { auto lookup_res = PeekLookup(store, {"hash_B", "hash_A"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 2); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_EQ((*lookup_res)[1].second.status, BlockStatus::HOST_AND_HBM); @@ -2510,8 +2506,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ProactiveEvictionWithCandidates) { // Controller free host blocks: 0 (used by A and B). // It should pick candidate B for eviction and deallocate its host block. // A (candidate HOST_AND_HBM) should not be affected. - ASSERT_TRUE(store.Lookup(hash_D).ok()); - ASSERT_TRUE(store.Save(hash_D).ok()); + ABSL_ASSERT_OK(store.Lookup(hash_D)); + ABSL_ASSERT_OK(store.Save(hash_D)); // Poll for Save completion save_done = false; @@ -2533,12 +2529,12 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ProactiveEvictionWithCandidates) { // - D should be HOST_AND_HBM { auto lookup_res = PeekLookup(store, {"hash_B"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 0); } { auto lookup_res = PeekLookup(store, {"hash_D"}); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); ASSERT_EQ(lookup_res->size(), 1); EXPECT_EQ((*lookup_res)[0].second.status, BlockStatus::HOST_AND_HBM); EXPECT_NE((*lookup_res)[0].second.host_block_id, -1); @@ -2563,11 +2559,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { 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()); + TF_ASSERT_OK_AND_ASSIGN( + auto src_backend_raw, + HostOffloadBackend::Create(src_config, controller.get())); auto src_backend = - std::dynamic_pointer_cast(*src_backend_or); + std::dynamic_pointer_cast(src_backend_raw); ASSERT_NE(src_backend, nullptr); std::vector src_slices = { @@ -2576,16 +2572,16 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { src_backend->Insert({"hash_0"}, src_slices, /*on_host=*/true); auto src_store_server = KVCacheStoreServer::Create(); - ASSERT_OK(src_store_server->StartServer(src_backend.get(), controller.get(), - "127.0.0.1")); + ABSL_ASSERT_OK(src_store_server->StartServer(src_backend.get(), + 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(), - controller->controller_address())); + ABSL_ASSERT_OK(registry_client->RegisterStore( + src_raiden_id, src_store_server->GetServerAddress(), + controller->controller_address())); KVCacheStore store(10, std::move(controller), registry_address, rid, std::nullopt, /*store_server_ip=*/"127.0.0.1"); @@ -2596,7 +2592,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { const std::vector device_blocks = {7}; absl::Status status = store.ReadRemote(hashes, slices, device_blocks); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for completion bool done = false; @@ -2616,14 +2612,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteSuccess) { // 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. auto lookup_res = store.Lookup(hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_TRUE(lookup_res->empty()); // ...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 registry_lookup = registry_client->Lookup(hashes); - ASSERT_TRUE(registry_lookup.ok()); + ABSL_ASSERT_OK(registry_lookup); EXPECT_TRUE(registry_lookup->empty()) << "read_remote must not advertise the read block to the registry"; } @@ -2650,11 +2646,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteReturnsStagingOnSuccess) { 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()); + TF_ASSERT_OK_AND_ASSIGN( + auto src_backend_raw, + HostOffloadBackend::Create(src_config, dst_controller.get())); auto src_backend = - std::dynamic_pointer_cast(*src_backend_or); + std::dynamic_pointer_cast(src_backend_raw); ASSERT_NE(src_backend, nullptr); std::vector src_slices = { @@ -2667,14 +2663,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteReturnsStagingOnSuccess) { } auto src_store_server = KVCacheStoreServer::Create(); - ASSERT_OK(src_store_server->StartServer(src_backend.get(), - dst_controller.get(), "127.0.0.1")); + ABSL_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( + ABSL_ASSERT_OK(registry_client->RegisterStore( src_raiden_id, src_store_server->GetServerAddress(), dst_controller->controller_address())); @@ -2689,7 +2685,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteReturnsStagingOnSuccess) { std::vector hashes = {absl::StrCat("hash_", round, "_a"), absl::StrCat("hash_", round, "_b")}; absl::Status status = store.ReadRemote(hashes, slices, {7, 8}); - ASSERT_TRUE(status.ok()) + ABSL_ASSERT_OK(status) << "round " << round << " failed to launch: " << status.message() << " -- staging blocks from an earlier round were not reclaimed"; @@ -2724,7 +2720,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteWithoutRegistryFails) { RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; absl::Status status = store.ReadRemote(hashes, slices, {7}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); bool failed = false; for (int attempt = 0; attempt < 100; ++attempt) { @@ -2746,8 +2742,8 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteWithoutRegistryFails) { TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemotePeerWithoutControllerAddressFails) { RaidenId src_raiden_id{"src_job", "0", "src_data", 0}; - ASSERT_OK(PublishPeerController(registry_address_, src_raiden_id, - /*controller_address=*/"")); + ABSL_ASSERT_OK(PublishPeerController(registry_address_, src_raiden_id, + /*controller_address=*/"")); auto dst_controller = MakeController(); RegisterAndInitWorker(*dst_controller, "worker_0", @@ -2762,7 +2758,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; absl::Status status = store.ReadRemote(hashes, slices, {7}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); bool failed = false; for (int attempt = 0; attempt < 100; ++attempt) { @@ -2810,10 +2806,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, 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()); + TF_ASSERT_OK_AND_ASSIGN( + auto src_backend_raw, + HostOffloadBackend::Create(src_config, controller_ptr)); auto src_backend = - std::dynamic_pointer_cast(*src_backend_or); + std::dynamic_pointer_cast(src_backend_raw); ASSERT_NE(src_backend, nullptr); std::vector src_slices = { @@ -2825,13 +2822,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, /*on_host=*/true); auto old_src = KVCacheStoreServer::Create(); - ASSERT_OK( + ABSL_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())); + ABSL_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, @@ -2844,7 +2842,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, std::vector hashes = {hash}; std::vector slices = { RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; - EXPECT_TRUE(store.ReadRemote(hashes, slices, {7}).ok()); + ABSL_EXPECT_OK(store.ReadRemote(hashes, slices, {7})); for (int attempt = 0; attempt < 300; ++attempt) { auto [done, failed, pending] = store.PollRemoteReadStatus(); if (!done.empty()) return true; @@ -2866,10 +2864,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, // The peer restarts: old store server is shut down and new server starts. old_src->Shutdown(); auto new_src = KVCacheStoreServer::Create(); - ASSERT_OK( + ABSL_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())); + ABSL_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. @@ -2907,15 +2906,15 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteFailure) { src_raiden_id.data_name = "src_data"; src_raiden_id.data_replica_idx = 0; - ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, - src_controller_server->server_address)); + ABSL_ASSERT_OK(PublishPeerController(registry_address, src_raiden_id, + src_controller_server->server_address)); 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(); + ABSL_ASSERT_OK(status); }; register_src_worker("worker_0", "src_worker_0_addr", "src_worker_0_transfer"); @@ -2949,7 +2948,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteFailure) { dst_transfer_mock_->fail_transfers = true; absl::Status status = store.ReadRemote(hashes, slices, {7}); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for failure bool failed = false; @@ -2973,14 +2972,14 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteFailure) { // remove. { auto lookup_res = store.Lookup(hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_TRUE(lookup_res->empty()); } // The staging blocks went back to the pool, so the read can be retried. If // the failure path leaked them this second launch would be the one to fail. dst_transfer_mock_->fail_transfers = false; - EXPECT_TRUE(store.ReadRemote(hashes, slices, {7}).ok()) + ABSL_EXPECT_OK(store.ReadRemote(hashes, slices, {7})) << "a failed read must return its staging blocks"; registry_server->Shutdown(); @@ -3002,23 +3001,23 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, 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()); + TF_ASSERT_OK_AND_ASSIGN( + auto src_backend_raw, + HostOffloadBackend::Create(src_config, controller.get())); auto src_backend = - std::dynamic_pointer_cast(*src_backend_or); + std::dynamic_pointer_cast(src_backend_raw); 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")); + ABSL_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())); + ABSL_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(2, std::move(controller), registry_address_, rid, @@ -3028,7 +3027,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, std::vector slices = { RaidenBlockId(src_raiden_id, 42, BlockStatus::REMOTE)}; - ASSERT_TRUE(store.ReadRemote(hashes, slices, {7}).ok()); + ABSL_ASSERT_OK(store.ReadRemote(hashes, slices, {7})); bool failed = false; for (int attempt = 0; attempt < 100; ++attempt) { @@ -3045,7 +3044,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ASSERT_TRUE(failed); // Nothing was recorded locally, on this path as on every other. auto lookup_res = PeekLookup(store, hashes); - ASSERT_TRUE(lookup_res.ok()); + ABSL_ASSERT_OK(lookup_res); EXPECT_TRUE(lookup_res->empty()); } @@ -3090,15 +3089,15 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteDuplicateFails) { src_raiden_id.data_name = "src_data"; src_raiden_id.data_replica_idx = 0; - ASSERT_OK(PublishPeerController(registry_address_, src_raiden_id, - src_controller_server->server_address)); + ABSL_ASSERT_OK(PublishPeerController(registry_address_, src_raiden_id, + src_controller_server->server_address)); 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(); + ABSL_ASSERT_OK(status); }; register_src_worker("worker_0", "src_worker_0_addr", "src_worker_0_transfer"); @@ -3130,7 +3129,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteDuplicateFails) { // First call succeeds absl::Status status1 = store.ReadRemote(hashes, slices, {7}); - ASSERT_TRUE(status1.ok()) << status1.message(); + ABSL_ASSERT_OK(status1); // Second call fails with FailedPreconditionError absl::Status status2 = store.ReadRemote(hashes, slices, {8}); @@ -3156,17 +3155,17 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteMultipleSources) { 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()); + TF_ASSERT_OK_AND_ASSIGN( + auto src_backend_raw_1, + HostOffloadBackend::Create(src_config_1, controller.get())); auto src_backend_1 = - std::dynamic_pointer_cast(*src_backend_or_1); + std::dynamic_pointer_cast(src_backend_raw_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")); + ABSL_ASSERT_OK(src_server_1->StartServer(src_backend_1.get(), + controller.get(), "127.0.0.1")); // Source 2 BackendConfig src_config_2; @@ -3174,27 +3173,27 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, ReadRemoteMultipleSources) { 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()); + TF_ASSERT_OK_AND_ASSIGN( + auto src_backend_raw_2, + HostOffloadBackend::Create(src_config_2, controller.get())); auto src_backend_2 = - std::dynamic_pointer_cast(*src_backend_or_2); + std::dynamic_pointer_cast(src_backend_raw_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")); + ABSL_ASSERT_OK(src_server_2->StartServer(src_backend_2.get(), + controller.get(), "127.0.0.1")); 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())); + ABSL_ASSERT_OK(client.RegisterStore(src_raiden_id_1, + src_server_1->GetServerAddress(), + controller->controller_address())); + ABSL_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(controller), registry_address, rid, @@ -3223,11 +3222,11 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, RegisterAndInitWorker(*controller, "worker_0", test_server_->server_address); MetadataRegion region(10); - auto metadata_or = KVCacheMetadata::Format(region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 10)); RaidenId rid{"test_job", "0", "test_cache", 0}; - KVCacheStore store(10, std::move(controller), "", rid, *metadata_or, + KVCacheStore store(10, std::move(controller), "", rid, metadata, /*store_server_ip=*/"127.0.0.1"); std::vector hashes = {"hash_1", "hash_2"}; @@ -3236,16 +3235,16 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, RaidenBlockId(rid, -1, 1, BlockStatus::HBM)}; ASSERT_TRUE(InsertResident(store, hashes, slices, false)); - ASSERT_TRUE(store.Lookup(hashes).ok()); + ABSL_ASSERT_OK(store.Lookup(hashes)); // Insert has already called SetMetadataEntry for both slices, but their HBM // status fails its data-lives-in-local-host-memory filter: the data exists // only in HBM at this point, so the LRU registration alone must leave the // table empty. - EXPECT_THAT(metadata_or->ValidEntries(), ::testing::IsEmpty()); + EXPECT_THAT(metadata.ValidEntries(), ::testing::IsEmpty()); absl::Status status = store.Save(hashes); - ASSERT_TRUE(status.ok()) << status.message(); + ABSL_ASSERT_OK(status); // Poll for completion bool done = false; @@ -3265,7 +3264,7 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, // Save completion lands the data on host blocks 0 and 1, which is when the // bindings enter the table. - EXPECT_THAT(metadata_or->ValidEntries(), + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(::testing::FieldsAre(0, "hash_1", 0), ::testing::FieldsAre(1, "hash_2", 1))); } @@ -3274,26 +3273,25 @@ TEST_F(KVCacheStoreEmbeddedControllerTest, TEST(KVCacheStoreTest, RecoverFromLocalManifestRebuildsLruCache) { RaidenId rid{"manifest_job", "0", "kv_cache", 0}; MetadataRegion region(10); - auto metadata_or = KVCacheMetadata::Format(region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 10)); // Table left behind by the previous incarnation of this store. - ASSERT_TRUE(metadata_or->Set(5, "hash_b", 3).ok()); - ASSERT_TRUE(metadata_or->Set(7, "hash_a", 4).ok()); - ASSERT_TRUE(metadata_or->Set(9, "hash_c", 8).ok()); + ABSL_ASSERT_OK(metadata.Set(5, "hash_b", 3)); + ABSL_ASSERT_OK(metadata.Set(7, "hash_a", 4)); + ABSL_ASSERT_OK(metadata.Set(9, "hash_c", 8)); auto controller = MakeRecoveryController(rid, 10); auto* controller_ptr = controller.get(); KVCacheStore store(10, std::move(controller), - /*global_registry_address=*/"", rid, *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); - auto recovered_or = store.RecoverFromLocalManifest(); - ASSERT_TRUE(recovered_or.ok()) << recovered_or.status().ToString(); - EXPECT_EQ(*recovered_or, 3); + TF_ASSERT_OK_AND_ASSIGN(auto recovered, store.RecoverFromLocalManifest()); + EXPECT_EQ(recovered, 3); auto lookup = PeekLookup(store, {"hash_a", "hash_b", "hash_c"}); - ASSERT_TRUE(lookup.ok()); + ABSL_ASSERT_OK(lookup); ASSERT_EQ(lookup->size(), 3); EXPECT_EQ((*lookup)[0].second.status, BlockStatus::HOST); EXPECT_EQ((*lookup)[0].second.host_block_id, 7); @@ -3310,7 +3308,7 @@ TEST(KVCacheStoreTest, RecoverFromLocalManifestRebuildsLruCache) { // insert is stamped 9, not 0. ASSERT_TRUE( InsertResident(store, {"hash_d"}, {RaidenBlockId(rid, 0, BlockStatus::HOST)}, true)); - EXPECT_THAT(metadata_or->ValidEntries(), + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(::testing::FieldsAre(0, "hash_d", 9), ::testing::FieldsAre(5, "hash_b", 3), ::testing::FieldsAre(7, "hash_a", 4), @@ -3320,47 +3318,45 @@ TEST(KVCacheStoreTest, RecoverFromLocalManifestRebuildsLruCache) { TEST(KVCacheStoreTest, RecoverFromLocalManifestRebuildsLruOrder) { RaidenId rid{"manifest_job_order", "0", "kv_cache", 0}; MetadataRegion region(10); - auto metadata_or = KVCacheMetadata::Format(region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 10)); - ASSERT_TRUE(metadata_or->Set(5, "hash_b", 3).ok()); - ASSERT_TRUE(metadata_or->Set(7, "hash_a", 4).ok()); - ASSERT_TRUE(metadata_or->Set(9, "hash_c", 8).ok()); + ABSL_ASSERT_OK(metadata.Set(5, "hash_b", 3)); + ABSL_ASSERT_OK(metadata.Set(7, "hash_a", 4)); + ABSL_ASSERT_OK(metadata.Set(9, "hash_c", 8)); // The table also records eviction candidates, so it may hold more entries // than the LRU cache capacity. With capacity 2 the oldest entry overflows // into a candidate again, keeping its block and metadata entry. KVCacheStore store(2, MakeRecoveryController(rid, 10), - /*global_registry_address=*/"", rid, *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); - auto recovered_or = store.RecoverFromLocalManifest(); - ASSERT_TRUE(recovered_or.ok()) << recovered_or.status().ToString(); - EXPECT_EQ(*recovered_or, 3); + TF_ASSERT_OK_AND_ASSIGN(auto recovered, store.RecoverFromLocalManifest()); + EXPECT_EQ(recovered, 3); EXPECT_THAT(KVCacheStoreTest::GetEvictCandidateKeys(store), ElementsAre("hash_b")); - EXPECT_EQ(metadata_or->ValidEntries().size(), 3); + EXPECT_EQ(metadata.ValidEntries().size(), 3); } TEST(KVCacheStoreTest, RecoverFromLocalManifestKeepsNewestDuplicate) { RaidenId rid{"manifest_job_dup", "0", "kv_cache", 0}; MetadataRegion region(10); - auto metadata_or = KVCacheMetadata::Format(region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 10)); - ASSERT_TRUE(metadata_or->Set(2, "dup_hash", 1).ok()); - ASSERT_TRUE(metadata_or->Set(4, "other", 3).ok()); - ASSERT_TRUE(metadata_or->Set(6, "dup_hash", 5).ok()); + ABSL_ASSERT_OK(metadata.Set(2, "dup_hash", 1)); + ABSL_ASSERT_OK(metadata.Set(4, "other", 3)); + ABSL_ASSERT_OK(metadata.Set(6, "dup_hash", 5)); auto controller = MakeRecoveryController(rid, 10); auto* controller_ptr = controller.get(); KVCacheStore store(10, std::move(controller), - /*global_registry_address=*/"", rid, *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); - auto recovered_or = store.RecoverFromLocalManifest(); - ASSERT_TRUE(recovered_or.ok()) << recovered_or.status().ToString(); - EXPECT_EQ(*recovered_or, 2); + TF_ASSERT_OK_AND_ASSIGN(auto recovered, store.RecoverFromLocalManifest()); + EXPECT_EQ(recovered, 2); // The newest binding wins; the stale block is neither tracked nor // allocated, and its entry is cleared from the table. @@ -3369,7 +3365,7 @@ TEST(KVCacheStoreTest, RecoverFromLocalManifestKeepsNewestDuplicate) { EXPECT_EQ((*lookup)[0].second.host_block_id, 6); EXPECT_TRUE(controller_ptr->block_manager()->IsAllocated(6)); EXPECT_FALSE(controller_ptr->block_manager()->IsAllocated(2)); - EXPECT_THAT(metadata_or->ValidEntries(), + EXPECT_THAT(metadata.ValidEntries(), ElementsAre(::testing::FieldsAre(4, "other", 3), ::testing::FieldsAre(6, "dup_hash", 5))); } @@ -3381,29 +3377,28 @@ TEST(KVCacheStoreTest, RecoverFromLocalManifestKeepsNewestDuplicate) { TEST(KVCacheStoreTest, RecoverFromLocalManifestFailsOnAllocatorConflict) { RaidenId rid{"manifest_job_conflict", "0", "kv_cache", 0}; MetadataRegion region(10); - auto metadata_or = KVCacheMetadata::Format(region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); - ASSERT_TRUE(metadata_or->Set(0, "rh1", 0).ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 10)); + ABSL_ASSERT_OK(metadata.Set(0, "rh1", 0)); auto controller = MakeRecoveryController(rid, 10); // Block 0 is already taken locally before recovery runs. - ASSERT_TRUE(controller->AllocateBlockIds(1).ok()); + ABSL_ASSERT_OK(controller->AllocateBlockIds(1)); KVCacheStore store(10, std::move(controller), - /*global_registry_address=*/"", rid, *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); - auto recovered_or = store.RecoverFromLocalManifest(); - EXPECT_EQ(recovered_or.status().code(), - absl::StatusCode::kFailedPrecondition); + auto recovered = store.RecoverFromLocalManifest(); + EXPECT_EQ(recovered.status().code(), absl::StatusCode::kFailedPrecondition); EXPECT_EQ(store.Lookup({"rh1"})->size(), 0); - EXPECT_EQ(metadata_or->ValidEntries().size(), 1); + EXPECT_EQ(metadata.ValidEntries().size(), 1); } TEST(KVCacheStoreTest, RecoverFromLocalManifestPreconditions) { RaidenId rid{"manifest_job_pre", "0", "kv_cache", 0}; MetadataRegion region(10); - auto metadata_or = KVCacheMetadata::Format(region.span(), 10); - ASSERT_TRUE(metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto metadata, + KVCacheMetadata::Format(region.span(), 10)); // A controller-less store is unrepresentable under the construction rules, // so the old no-controller sub-case is gone. @@ -3418,8 +3413,7 @@ TEST(KVCacheStoreTest, RecoverFromLocalManifestPreconditions) { // Non-empty LRU cache. KVCacheStore store_non_empty(10, MakeRecoveryController(rid, 10), - /*global_registry_address=*/"", rid, - *metadata_or, + /*global_registry_address=*/"", rid, metadata, /*store_server_ip=*/"127.0.0.1"); ASSERT_TRUE( InsertResident( @@ -3429,15 +3423,14 @@ TEST(KVCacheStoreTest, RecoverFromLocalManifestPreconditions) { // Empty table: recovery succeeds with zero blocks. MetadataRegion empty_region(10); - auto empty_metadata_or = KVCacheMetadata::Format(empty_region.span(), 10); - ASSERT_TRUE(empty_metadata_or.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto empty_metadata, + KVCacheMetadata::Format(empty_region.span(), 10)); KVCacheStore store_empty(10, MakeRecoveryController(rid, 10), - /*global_registry_address=*/"", rid, - *empty_metadata_or, + /*global_registry_address=*/"", rid, empty_metadata, /*store_server_ip=*/"127.0.0.1"); - auto recovered_or = store_empty.RecoverFromLocalManifest(); - ASSERT_TRUE(recovered_or.ok()) << recovered_or.status().ToString(); - EXPECT_EQ(*recovered_or, 0); + TF_ASSERT_OK_AND_ASSIGN(auto recovered, + store_empty.RecoverFromLocalManifest()); + EXPECT_EQ(recovered, 0); } TEST(KVCacheStoreTest, MultiBackendPriorityLookupChain) { @@ -3463,8 +3456,8 @@ TEST(KVCacheStoreTest, MultiBackendPriorityLookupChain) { auto lookup_res = store.Lookup({"h1", "h2", "h3", "h4"}, /*enable_global=*/true); - ASSERT_TRUE(lookup_res.ok()); - ASSERT_EQ(lookup_res->size(), 4); + ABSL_ASSERT_OK(lookup_res); + EXPECT_EQ(lookup_res->size(), 4); EXPECT_EQ((*lookup_res)[0].first, "h1"); EXPECT_EQ((*lookup_res)[0].second.host_block_id, 1); EXPECT_EQ((*lookup_res)[1].first, "h2"); @@ -3491,14 +3484,14 @@ TEST(KVCacheStoreTest, MultiBackendLocalLookupWhenGlobalDisabled) { // enable_global = false => searches all local backends (b1 and b2) auto gated_res = store.Lookup({"h1", "h2"}, /*enable_global=*/false); - ASSERT_TRUE(gated_res.ok()); + ABSL_ASSERT_OK(gated_res); ASSERT_EQ(gated_res->size(), 2); EXPECT_EQ((*gated_res)[0].first, "h1"); EXPECT_EQ((*gated_res)[1].first, "h2"); // enable_global = true => queries all backends as well auto ungated_res = store.Lookup({"h1", "h2"}, /*enable_global=*/true); - ASSERT_TRUE(ungated_res.ok()); + ABSL_ASSERT_OK(ungated_res); ASSERT_EQ(ungated_res->size(), 2); EXPECT_EQ((*ungated_res)[0].first, "h1"); EXPECT_EQ((*ungated_res)[1].first, "h2"); @@ -3588,12 +3581,11 @@ TEST_F(StoreDiscoveryTest, PublishesStoreAddressToTheRegistry) { EXPECT_EQ(store.store_server_address(), absl::StrCat("127.0.0.1:", store.store_server()->GetGrpcPort())); - auto resolved = client_->ResolveStore(rid); - ASSERT_TRUE(resolved.ok()) << resolved.status().ToString(); - EXPECT_EQ(resolved->store_server_address(), store.store_server_address()); + TF_ASSERT_OK_AND_ASSIGN(auto resolved, client_->ResolveStore(rid)); + EXPECT_EQ(resolved.store_server_address(), store.store_server_address()); // The controller address rides along: it is what a peer dials to acquire a // read lease against this store. - EXPECT_EQ(resolved->controller_address(), store.raiden_controller_address()); + EXPECT_EQ(resolved.controller_address(), store.raiden_controller_address()); } // store_server_ip is bind-and-advertise, so the published address is @@ -3605,7 +3597,7 @@ TEST_F(StoreDiscoveryTest, PublishedAddressIsConnectable) { /*store_server_ip=*/"127.0.0.1"); auto resolved = client_->ResolveStore(rid); - ASSERT_TRUE(resolved.ok()) << resolved.status().ToString(); + ABSL_ASSERT_OK(resolved); auto peer_channel = grpc::CreateChannel(resolved->store_server_address(), grpc::InsecureChannelCredentials()); @@ -3620,7 +3612,7 @@ TEST_F(StoreDiscoveryTest, DestructorUnpublishes) { KVCacheStore store(/*capacity=*/16, registry_address_, rid, /*num_shards=*/1, /*shard_size_bytes=*/512, /*store_server_ip=*/"127.0.0.1"); - ASSERT_TRUE(client_->ResolveStore(rid).ok()); + ABSL_ASSERT_OK(client_->ResolveStore(rid)); } EXPECT_TRUE(absl::IsNotFound(client_->ResolveStore(rid).status())); } @@ -3643,7 +3635,7 @@ TEST_F(StoreDiscoveryTest, RestartReplacesPublishedAddress) { /*num_shards=*/1, /*shard_size_bytes=*/512, /*store_server_ip=*/"127.0.0.1"); auto resolved = client_->ResolveStore(rid); - ASSERT_TRUE(resolved.ok()) << resolved.status().ToString(); + ABSL_ASSERT_OK(resolved); EXPECT_EQ(resolved->store_server_address(), restarted.store_server_address()); EXPECT_NE(resolved->store_server_address(), first_address); } @@ -3669,12 +3661,13 @@ TEST_F(StoreDiscoveryTest, AdoptsAndPublishesTheBackendsServer) { pooling_config.raiden_id = rid; const BackendConfig configs[] = {host_config, pooling_config}; - auto store_or = KVCacheStore::Create( - absl::MakeConstSpan(configs), /*capacity=*/16, registry_address_, rid, - /*num_shards=*/1, /*shard_size_bytes=*/512, - /*store_server_ip=*/"127.0.0.1"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - auto& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN( + auto store_ptr, + KVCacheStore::Create(absl::MakeConstSpan(configs), /*capacity=*/16, + registry_address_, rid, + /*num_shards=*/1, /*shard_size_bytes=*/512, + /*store_server_ip=*/"127.0.0.1")); + auto& store = *store_ptr; auto* pooling = dynamic_cast(store.backends()[1].get()); ASSERT_NE(pooling, nullptr); @@ -3685,7 +3678,7 @@ TEST_F(StoreDiscoveryTest, AdoptsAndPublishesTheBackendsServer) { // Published under the supplied ip, not the backend's hardcoded wildcard. EXPECT_THAT(store.store_server_address(), StartsWith("127.0.0.1:")); auto resolved = client_->ResolveStore(rid); - ASSERT_TRUE(resolved.ok()) << resolved.status().ToString(); + ABSL_ASSERT_OK(resolved); EXPECT_EQ(resolved->store_server_address(), store.store_server_address()); } @@ -3703,7 +3696,7 @@ TEST_F(StoreDiscoveryTest, AdoptsTheBackendsServerRatherThanOwningASecond) { auto bootstrap_controller = MakeRecoveryController(rid, /*num_blocks=*/16); auto backend = std::make_shared( /*capacity=*/16, std::nullopt, rid, bootstrap_controller.get()); - ASSERT_TRUE(backend->StartServer("127.0.0.1").ok()); + ABSL_ASSERT_OK(backend->StartServer("127.0.0.1")); ASSERT_NE(backend->store_server(), nullptr); KVCacheStore store(std::vector>{backend}, @@ -3747,38 +3740,35 @@ TEST_F(StoreDiscoveryTest, StoreMonitorHeartbeatsTheRegistration) { config.monitor_config.enable = true; config.monitor_config.heartbeat_period = absl::Milliseconds(300); - auto store_or = KVCacheStore::Create(config, /*capacity=*/16, - registry_address_, rid, - /*num_shards=*/1, - /*shard_size_bytes=*/512, - /*store_server_ip=*/"127.0.0.1"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); + TF_ASSERT_OK_AND_ASSIGN( + auto store_ptr, + KVCacheStore::Create(config, /*capacity=*/16, registry_address_, rid, + /*num_shards=*/1, /*shard_size_bytes=*/512, + /*store_server_ip=*/"127.0.0.1")); // The registration's TTL is three heartbeat periods (900ms here); well // past it, only the heartbeats keep the store resolvable. absl::SleepFor(absl::Seconds(2)); auto resolved = client_->ResolveStore(rid); - ASSERT_TRUE(resolved.ok()) << resolved.status().ToString(); + ABSL_ASSERT_OK(resolved); EXPECT_EQ(resolved->store_server_address(), - (*store_or)->store_server_address()); + store_ptr->store_server_address()); // A tier-0 caller in the same group is offered this store: the // kv_pool_group and evict_tier from the BackendConfig were published. RaidenId caller{"disco_caller", "0", "kv_cache", 0}; - ASSERT_TRUE(client_ - ->RegisterStore(caller, "10.0.0.7:1111", - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration(), "groupA", - /*evict_tier=*/0) - .ok()); + ABSL_ASSERT_OK(client_->RegisterStore(caller, "10.0.0.7:1111", + /*controller_address=*/"", + /*ttl=*/absl::ZeroDuration(), "groupA", + /*evict_tier=*/0)); auto targets = client_->GetPlacementTargets(caller, /*max_targets=*/8); - ASSERT_TRUE(targets.ok()) << targets.status().ToString(); + ABSL_ASSERT_OK(targets); ASSERT_EQ(targets->size(), 1); EXPECT_EQ((*targets)[0].raiden_id().job_name(), "disco_job_monitored"); // Destruction stops the monitor before unpublishing, so no late heartbeat // re-registers the store after this. - store_or->reset(); + store_ptr.reset(); EXPECT_TRUE(absl::IsNotFound(client_->ResolveStore(rid).status())); } @@ -3797,13 +3787,12 @@ TEST_F(StoreDiscoveryTest, LapsedRegistrationRepublishesBlockEntries) { config.monitor_config.enable = true; config.monitor_config.heartbeat_period = absl::Milliseconds(300); - auto store_or = KVCacheStore::Create(config, /*capacity=*/16, - registry_address_, rid, - /*num_shards=*/1, - /*shard_size_bytes=*/512, - /*store_server_ip=*/"127.0.0.1"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN( + auto store_ptr, + KVCacheStore::Create(config, /*capacity=*/16, registry_address_, rid, + /*num_shards=*/1, /*shard_size_bytes=*/512, + /*store_server_ip=*/"127.0.0.1")); + KVCacheStore& store = *store_ptr; // Host-resident blocks with no registry entries: exactly the state the // dead-store cascade leaves behind (and what Insert alone produces -- @@ -3812,13 +3801,13 @@ TEST_F(StoreDiscoveryTest, LapsedRegistrationRepublishesBlockEntries) { std::vector slices = { RaidenBlockId(rid, /*host_id=*/5, BlockStatus::HOST), RaidenBlockId(rid, /*host_id=*/6, BlockStatus::HOST)}; - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); auto looked = client_->Lookup(hashes); - ASSERT_TRUE(looked.ok()); + ABSL_ASSERT_OK(looked); ASSERT_EQ(looked->size(), 0); // Drop the store registration; the heartbeat has to discover the lapse. - ASSERT_TRUE(client_->UnregisterStore(rid).ok()); + ABSL_ASSERT_OK(client_->UnregisterStore(rid)); // The next heartbeat gets NotFound, re-registers, and republishes. const absl::Time deadline = absl::Now() + absl::Seconds(5); @@ -3834,7 +3823,7 @@ TEST_F(StoreDiscoveryTest, LapsedRegistrationRepublishesBlockEntries) { absl::SleepFor(absl::Milliseconds(100)); } EXPECT_EQ(republished, hashes.size()); - EXPECT_TRUE(client_->ResolveStore(rid).ok()); + ABSL_EXPECT_OK(client_->ResolveStore(rid)); } // The flag promises heartbeats; without a registry there is no registration @@ -3847,13 +3836,13 @@ TEST_F(StoreDiscoveryTest, StoreMonitorWithoutARegistryIsAnError) { config.raiden_id = rid; config.monitor_config.enable = true; - auto store_or = KVCacheStore::Create(config, /*capacity=*/16, - /*global_registry_address=*/"", rid, - /*num_shards=*/1, - /*shard_size_bytes=*/512, - /*store_server_ip=*/"127.0.0.1"); - EXPECT_TRUE(absl::IsFailedPrecondition(store_or.status())) - << store_or.status().ToString(); + auto store_result = KVCacheStore::Create(config, /*capacity=*/16, + /*global_registry_address=*/"", rid, + /*num_shards=*/1, + /*shard_size_bytes=*/512, + /*store_server_ip=*/"127.0.0.1"); + EXPECT_TRUE(absl::IsFailedPrecondition(store_result.status())) + << store_result.status().ToString(); } // A tier-0 backend whose Lookup parks, so a peer's Fetch can be held inside the @@ -3984,17 +3973,17 @@ TEST_F(StoreDiscoveryTest, CapacityConstructedStoreJoinsTheGlobalTier) { // The backend holds a real registry client, so what it publishes reaches // the global tier. Inserting does not publish; a completed save does, // through this path -- so drive it directly. - ASSERT_OK(backend->RegisterBlocksAsync({"tiered_hash"}, {3}).Await()); + ABSL_ASSERT_OK(backend->RegisterBlocksAsync({"tiered_hash"}, {3}).Await()); auto looked_up = client_->Lookup({"tiered_hash"}); - ASSERT_TRUE(looked_up.ok()) << looked_up.status().ToString(); + ABSL_ASSERT_OK(looked_up); ASSERT_EQ(looked_up->size(), 1); EXPECT_EQ((*looked_up)[0].block_id(), 3); // A local miss consults tier 1 and comes back with the owning peer. RaidenId peer{"some_peer", "0", "kv_cache", 0}; - ASSERT_TRUE(client_->Register({{"peer_hash", peer, 9}}).ok()); + ABSL_ASSERT_OK(client_->Register({{"peer_hash", peer, 9}})); auto result = backend->Lookup({"peer_hash"}); - ASSERT_TRUE(result.ok()) << result.status().ToString(); + ABSL_ASSERT_OK(result); ASSERT_EQ(result->size(), 1); EXPECT_EQ((*result)[0].second.raiden_id, peer); EXPECT_EQ((*result)[0].second.status, BlockStatus::REMOTE); @@ -4025,7 +4014,7 @@ TEST_F(StoreDiscoveryTest, DestructorShutsDownABackendStartedServer) { backend_ref = store.backend(); auto* backend = dynamic_cast(backend_ref.get()); ASSERT_NE(backend, nullptr); - ASSERT_TRUE(backend->StartServer("127.0.0.1").ok()); + ABSL_ASSERT_OK(backend->StartServer("127.0.0.1")); ASSERT_NE(backend->store_server(), nullptr); ASSERT_GT(backend->store_server()->GetGrpcPort(), 0); } @@ -4045,7 +4034,7 @@ TEST_F(StoreDiscoveryTest, DestructorSweepSkipsTheAdoptedServer) { auto bootstrap_controller = MakeRecoveryController(rid, /*num_blocks=*/16); auto backend = std::make_shared( /*capacity=*/16, std::nullopt, rid, bootstrap_controller.get()); - ASSERT_TRUE(backend->StartServer("127.0.0.1").ok()); + ABSL_ASSERT_OK(backend->StartServer("127.0.0.1")); ASSERT_GT(backend->store_server()->GetGrpcPort(), 0); { @@ -4081,10 +4070,8 @@ TEST_F(StoreDiscoveryTest, FailedLoadDropsTheCachedPeerClient) { ASSERT_NE(backend, nullptr); // Port 1 is reserved and never listening, so this connect always fails. - ASSERT_TRUE(client_ - ->RegisterStore(peer_rid, "127.0.0.1:1", - /*controller_address=*/"") - .ok()); + ABSL_ASSERT_OK(client_->RegisterStore(peer_rid, "127.0.0.1:1", + /*controller_address=*/"")); // Load requires a destination for every hash. The device block is never written // as the loads below fail while resolving or querying the peer. auto first = backend->Load(peer_rid, {"h"}, {0}).Await(); @@ -4448,10 +4435,9 @@ class RemoteWriteSourceTest : public StoreDiscoveryTest { ::tpu_raiden::core::controller::RaidenControllerClient client( store.raiden_controller_address()); - auto status = + ABSL_ASSERT_OK( client.RegisterWorker("worker_0", worker_server_->server_address, - {{worker_server_->server_address, {}}}); - ASSERT_TRUE(status.ok()) << status.message(); + {{worker_server_->server_address, {}}})); } // Puts `hashes` in `store` as host-resident, which is the precondition for @@ -4465,7 +4451,7 @@ class RemoteWriteSourceTest : public StoreDiscoveryTest { } // Insert(), not InsertResident(): these cases are about offering blocks to // a peer, and a remote save requires -- and consumes -- the caller's pin. - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); } // Reads PollSaveStatus() until the write leaves the pending set. @@ -4496,12 +4482,10 @@ class RemoteWriteSourceTest : public StoreDiscoveryTest { builder.RegisterService(&fake_destination_); fake_destination_server_ = builder.BuildAndStart(); ASSERT_NE(fake_destination_server_, nullptr); - ASSERT_TRUE(client_ - ->RegisterStore(dst, "127.0.0.1:" + std::to_string(port), - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration(), kv_pool_group, - evict_tier) - .ok()); + ABSL_ASSERT_OK(client_->RegisterStore( + dst, "127.0.0.1:" + std::to_string(port), + /*controller_address=*/"", + /*ttl=*/absl::ZeroDuration(), kv_pool_group, evict_tier)); } void TearDown() override { @@ -4580,7 +4564,7 @@ TEST_F(RemoteWriteSourceTest, RefusesToOfferAnHbmOnlyBlock) { std::vector slices = { RaidenBlockId(src, /*host_block_id=*/-1, /*device_block_id=*/0, BlockStatus::HBM)}; - ASSERT_TRUE(store->Insert({"a"}, slices, /*on_host=*/false).ok()); + ABSL_ASSERT_OK(store->Insert({"a"}, slices, /*on_host=*/false)); auto status = store->Save({"a"}, RaidenId{"rw_dst_hbm", "0", "kv", 0}); EXPECT_TRUE(absl::IsFailedPrecondition(status)) << status.ToString(); @@ -4610,7 +4594,7 @@ TEST_F(RemoteWriteSourceTest, AllExistSettlesDoneWithoutATransfer) { {"a", "b"}, {RaidenBlockId(dst, 5, BlockStatus::HOST), RaidenBlockId(dst, 6, BlockStatus::HOST)})); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); auto [done, failed, pending, existing, unregistered] = src_store->PollSaveStatus(); @@ -4632,7 +4616,7 @@ TEST_F(RemoteWriteSourceTest, PartialExistIsReportedAndNotRetried) { ASSERT_TRUE(dst_store->backend()->InsertAllOrNothing( {"a"}, {RaidenBlockId(dst, 5, BlockStatus::HOST)})); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); auto [done, failed, pending, existing, unregistered] = src_store->PollSaveStatus(); @@ -4656,7 +4640,7 @@ TEST_F(RemoteWriteSourceTest, AnAcceptedOfferSettlesOnATerminalVerdict) { RegisterWorker(*src_store); Populate(*src_store, src, {"a", "b"}); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); EXPECT_TRUE(done.empty()); @@ -4664,7 +4648,7 @@ TEST_F(RemoteWriteSourceTest, AnAcceptedOfferSettlesOnATerminalVerdict) { // The internal pin is gone and the hashes are no longer marked as writing, // so the same blocks can be offered again. - EXPECT_TRUE(src_store->Save({"a", "b"}, dst).ok()) + ABSL_EXPECT_OK(src_store->Save({"a", "b"}, dst)) << "the first operation never released its claim on these hashes"; AwaitWriteSettled(*src_store); } @@ -4681,7 +4665,7 @@ TEST_F(RemoteWriteSourceTest, RefusesASecondConcurrentOfferOfTheSameHash) { verdict.set_state(proto::PollWriteRemoteResponse::PENDING); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); auto second = src_store->Save({"a"}, dst); EXPECT_TRUE(absl::IsFailedPrecondition(second)) << second.ToString(); } @@ -4703,7 +4687,7 @@ TEST_F(RemoteWriteSourceTest, StoredUnregisteredIsFailedAndNamesTheBlocks) { verdict.add_unregistered_hashes("b"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); EXPECT_TRUE(done.empty()) @@ -4715,7 +4699,7 @@ TEST_F(RemoteWriteSourceTest, StoredUnregisteredIsFailedAndNamesTheBlocks) { // The internal pin is released either way, so the caller can act on the // list -- including by offering the same blocks somewhere else. - EXPECT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_EXPECT_OK(src_store->Save({"a", "b"}, dst)); AwaitWriteSettled(*src_store); } @@ -4734,7 +4718,7 @@ TEST_F(RemoteWriteSourceTest, CommittedIsReportedAsDone) { verdict.add_committed_hashes("b"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); EXPECT_THAT(done, ::testing::UnorderedElementsAre("a", "b")); @@ -4766,9 +4750,9 @@ TEST_F(RemoteWriteSourceTest, ARefusalKeepsTheStoreClient) { // Unregistering the peer makes re-resolution impossible, so the second // offer can only reach the destination through the client kept from the // first one. - ASSERT_TRUE(client_->UnregisterStore(dst).ok()); + ABSL_ASSERT_OK(client_->UnregisterStore(dst)); fake_destination_.SetWriteRemoteStatus(::grpc::Status::OK); - EXPECT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_EXPECT_OK(src_store->Save({"a"}, dst)); EXPECT_EQ(fake_destination_.write_calls(), 2); AwaitWriteSettled(*src_store); } @@ -4790,7 +4774,7 @@ TEST_F(RemoteWriteSourceTest, ATransportErrorInvalidatesTheStoreClient) { fake_destination_.SetPollResponse(verdict); // A first offer establishes the store client for dst. - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); AwaitWriteSettled(*src_store); fake_destination_server_->Shutdown(); @@ -4800,9 +4784,9 @@ TEST_F(RemoteWriteSourceTest, ATransportErrorInvalidatesTheStoreClient) { // The destination comes back under the same identity on a fresh port. A // client kept across the transport error would still dial the dead // one; this offer succeeds only by re-resolving. - ASSERT_TRUE(client_->UnregisterStore(dst).ok()); + ABSL_ASSERT_OK(client_->UnregisterStore(dst)); StartFakeDestination(dst); - EXPECT_TRUE(src_store->Save({"b"}, dst).ok()); + ABSL_EXPECT_OK(src_store->Save({"b"}, dst)); AwaitWriteSettled(*src_store); } @@ -4820,7 +4804,7 @@ TEST_F(RemoteWriteSourceTest, UnknownIsReportedAsAPlainFailure) { ::tpu_raiden::kv_cache::proto::PollWriteRemoteResponse::UNKNOWN); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); EXPECT_TRUE(done.empty()); @@ -4843,7 +4827,7 @@ TEST_F(RemoteWriteSourceTest, TheOfferAsksForLessThanTheSourceWillHold) { verdict.add_committed_hashes("a"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); AwaitWriteSettled(*src_store); EXPECT_EQ(fake_destination_.write_calls(), 1); @@ -4866,7 +4850,7 @@ TEST_F(RemoteWriteSourceTest, DestinationCommitsSettlesWithZeroPollCalls) { verdict.add_committed_hashes("a"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); EXPECT_THAT(done, ::testing::UnorderedElementsAre("a")); @@ -4893,7 +4877,7 @@ TEST_F(RemoteWriteSourceTest, verdict.add_unregistered_hashes("y"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"x", "y"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"x", "y"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); EXPECT_TRUE(done.empty()); @@ -4915,7 +4899,7 @@ TEST_F(RemoteWriteSourceTest, StreamBrokenMidFlightRecoversViaWaitingPoll) { pending_verdict.set_state(proto::PollWriteRemoteResponse::PENDING); fake_destination_.SetPollResponse(pending_verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); // The offer was accepted and is in flight. Internal pin is held. EXPECT_EQ(src_store->GetPinCount("a"), 2); @@ -4954,7 +4938,7 @@ TEST_F(RemoteWriteSourceTest, pending_verdict.set_state(proto::PollWriteRemoteResponse::PENDING); fake_destination_.SetPollResponse(pending_verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); EXPECT_EQ(src_store->InFlightRemoteWritesCountForTesting(), 1); // Destroy the store while stream is in flight: destructor takes and cancels. @@ -4985,7 +4969,7 @@ TEST_F(RemoteWriteSourceTest, ConcurrentPollsSettleAnOfferExactlyOnce) { verdict.add_committed_hashes("b"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); // Each thread drains PollSaveStatus and pools what it sees. The invariant: // every hash appears exactly once across all of them. @@ -5045,7 +5029,7 @@ TEST_F(RemoteWriteSourceTest, ASuccessfulRemoteSaveConsumesTheCallerPin) { verdict.add_committed_hashes("a"); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); ASSERT_THAT(done, ::testing::ElementsAre("a")); @@ -5066,7 +5050,7 @@ TEST_F(RemoteWriteSourceTest, AFailedRemoteSaveKeepsTheCallerPin) { verdict.set_state(proto::PollWriteRemoteResponse::FAILED); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); auto [done, failed, existing, unregistered] = AwaitWriteSettled(*src_store); ASSERT_THAT(failed, ::testing::ElementsAre("a")); @@ -5085,7 +5069,7 @@ TEST_F(RemoteWriteSourceTest, AnAllExistRemoteSaveConsumesTheCallerPinToo) { fake_destination_.SetWriteExistState(proto::WRITE_ALL_EXIST); ASSERT_EQ(src_store->GetPinCount("a"), 1); - ASSERT_TRUE(src_store->Save({"a", "b"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a", "b"}, dst)); // Settled inside the call: no poll was needed to get here. EXPECT_EQ(src_store->GetPinCount("a"), 0); @@ -5111,7 +5095,7 @@ TEST_F(RemoteWriteSourceTest, DestroyingAStoreMidOfferReleasesItsInternalPin) { verdict.set_state(proto::PollWriteRemoteResponse::PENDING); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); // Populate's pin plus the offer's internal one. EXPECT_EQ(src_store->GetPinCount("a"), 2); @@ -5136,7 +5120,7 @@ TEST_F(RemoteWriteSourceTest, DestroyingAStoreMidOfferDoesNotWaitForTheHold) { verdict.set_state(proto::PollWriteRemoteResponse::PENDING); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); { auto [done, failed, pending, existing, unregistered] = src_store->PollSaveStatus(); @@ -5180,7 +5164,7 @@ TEST_F(RemoteWriteSourceTest, AckDeadlineReleasesThePinAndTheRecord) { verdict.set_state(proto::PollWriteRemoteResponse::COMMITTED); verdict.add_committed_hashes("a"); fake_destination_.SetPollResponse(verdict); - EXPECT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_EXPECT_OK(src_store->Save({"a"}, dst)); AwaitWriteSettled(*src_store); } @@ -5217,7 +5201,7 @@ TEST_F(RemoteWriteSourceTest, ALostAnswerUndoesTheOfferAndAllowsARetry) { verdict.set_state(proto::PollWriteRemoteResponse::COMMITTED); verdict.add_committed_hashes("a"); fake_destination_.SetPollResponse(verdict); - EXPECT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_EXPECT_OK(src_store->Save({"a"}, dst)); AwaitWriteSettled(*src_store); } @@ -5262,7 +5246,7 @@ TEST_F(RemoteWriteSourceTest, AGrantPastTheHoldKeepsTheBlocksForTheGrant) { // its own deadline with the operation still live on the destination. fake_destination_.SetGrantedDeadlineMs(20000); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); EXPECT_EQ(src_store->GetPinCount("a"), 2); // The HOLD runs out. What must NOT happen is the operation being settled @@ -5299,7 +5283,7 @@ TEST_F(RemoteWriteSourceTest, AHoldThatRunsOutReleasesThePinAndFailsTheBatch) { StartFakeDestination(dst); // Accepted, and then nothing: no result, and no grant to outlive the hold. - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); EXPECT_EQ(src_store->GetPinCount("a"), 2); auto [done, failed, existing, unregistered] = @@ -5326,7 +5310,7 @@ TEST_F(RemoteWriteSourceTest, TeardownEndsTheCallsItAbandons) { // Accepted, and then silent: the destination sits on the stream, so within // the two seconds observed below, the call ending can only be the // cancellation -- the six-second hold has not run out yet. - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); ASSERT_EQ(fake_destination_.open_streams(), 1); src_store.reset(); @@ -5351,7 +5335,7 @@ TEST_F(RemoteWriteSourceTest, TeardownDoesNotWaitForARecoveryAskToBeAnswered) { // Accepted, then the stream drops -- and the destination then sits on the // recovery ask rather than answering it (the fake's poll parks while its // response is unset). - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); fake_destination_.BreakActiveStreams( ::grpc::Status(::grpc::StatusCode::UNAVAILABLE, "stream lost")); for (int i = 0; i < 500 && fake_destination_.poll_calls() == 0; ++i) { @@ -5397,7 +5381,7 @@ TEST_F(RemoteWriteSourceTest, TakeRemoteWriteSettlesExactlyOnce) { verdict.set_state(proto::PollWriteRemoteResponse::PENDING); fake_destination_.SetPollResponse(verdict); - ASSERT_TRUE(src_store->Save({"a"}, dst).ok()); + ABSL_ASSERT_OK(src_store->Save({"a"}, dst)); ASSERT_EQ(src_store->InFlightRemoteWritesCountForTesting(), 1); constexpr int kNumThreads = 10; @@ -5482,14 +5466,14 @@ class EvictSweepTest : public RemoteWriteSourceTest { // eviction raises it back. void PopulateCold(KVCacheStore& store, const RaidenId& id, const std::vector& hashes) { - auto ids_or = store.raiden_controller()->AllocateBlockIds(hashes.size()); - ASSERT_TRUE(ids_or.ok()) << ids_or.status().ToString(); + TF_ASSERT_OK_AND_ASSIGN( + auto ids, store.raiden_controller()->AllocateBlockIds(hashes.size())); std::vector slices; for (size_t i = 0; i < hashes.size(); ++i) { - slices.push_back(RaidenBlockId(id, (*ids_or)[i], BlockStatus::HOST)); + slices.push_back(RaidenBlockId(id, ids[i], BlockStatus::HOST)); } // Insert() pins; the sweep only takes unpinned blocks, so release. - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); store.Release(hashes); } @@ -5497,13 +5481,13 @@ class EvictSweepTest : public RemoteWriteSourceTest { // requires the caller's pin and consumes it on success. void PopulatePinned(KVCacheStore& store, const RaidenId& id, const std::vector& hashes) { - auto ids_or = store.raiden_controller()->AllocateBlockIds(hashes.size()); - ASSERT_TRUE(ids_or.ok()) << ids_or.status().ToString(); + TF_ASSERT_OK_AND_ASSIGN( + auto ids, store.raiden_controller()->AllocateBlockIds(hashes.size())); std::vector slices; for (size_t i = 0; i < hashes.size(); ++i) { - slices.push_back(RaidenBlockId(id, (*ids_or)[i], BlockStatus::HOST)); + slices.push_back(RaidenBlockId(id, ids[i], BlockStatus::HOST)); } - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); } // The sweep runs on the monitor's thread; wait until it has raised the free @@ -5543,9 +5527,9 @@ TEST_F(EvictSweepTest, DemotesColdBlocksToAPlacementTarget) { RaidenId dst{"sweep_dst", "0", "kv", 0}; // Large enough that reaching the high watermark takes several batches of // the built-in batch cap (128). - auto store_or = MakeSweepStore(src, "sweepgroup", /*capacity=*/600); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN(auto store_ptr, + MakeSweepStore(src, "sweepgroup", /*capacity=*/600)); + KVCacheStore& store = *store_ptr; ASSERT_EQ(store.raiden_controller()->block_manager()->total_blocks(), 600); StartFakeDestination(dst, "sweepgroup", /*evict_tier=*/1); @@ -5574,21 +5558,19 @@ TEST_F(EvictSweepTest, SkipsAnUnreachableTargetForTheNextOne) { RaidenId src{"sweep_src_skip", "0", "kv", 0}; RaidenId dead{"sweep_dead", "0", "kv", 0}; RaidenId dst{"sweep_dst_skip", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "skipgroup"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN(auto store_ptr, MakeSweepStore(src, "skipgroup")); + KVCacheStore& store = *store_ptr; // A dead peer, ranked first: registered at tier 1 with the most reported // free blocks, but nothing listens on its address. - ASSERT_TRUE(client_ - ->RegisterStore(dead, "127.0.0.1:1", - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration(), "skipgroup", - /*evict_tier=*/1) - .ok()); + ABSL_ASSERT_OK(client_->RegisterStore(dead, "127.0.0.1:1", + /*controller_address=*/"", + /*ttl=*/absl::ZeroDuration(), + "skipgroup", + /*evict_tier=*/1)); global_registry::StoreStatus roomy; roomy.set_free_blocks(1000); - ASSERT_TRUE(client_->Heartbeat(dead, roomy).ok()); + ABSL_ASSERT_OK(client_->Heartbeat(dead, roomy)); StartFakeDestination(dst, "skipgroup", /*evict_tier=*/1); proto::PollWriteRemoteResponse verdict; verdict.set_state(proto::PollWriteRemoteResponse::COMMITTED); @@ -5608,9 +5590,8 @@ TEST_F(EvictSweepTest, SkipsAnUnreachableTargetForTheNextOne) { TEST_F(EvictSweepTest, ATransferFailureFallsBackToLocalDrop) { RaidenId src{"sweep_src_fail", "0", "kv", 0}; RaidenId dst{"sweep_dst_fail", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "failgroup"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN(auto store_ptr, MakeSweepStore(src, "failgroup")); + KVCacheStore& store = *store_ptr; StartFakeDestination(dst, "failgroup", /*evict_tier=*/1); proto::PollWriteRemoteResponse verdict; @@ -5631,9 +5612,9 @@ TEST_F(EvictSweepTest, ATransferFailureFallsBackToLocalDrop) { // be forced into later, done early. TEST_F(EvictSweepTest, DropsLocallyWhenThereAreNoTargets) { RaidenId src{"sweep_src_bottom", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "sweepgroup_bottom"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN(auto store_ptr, + MakeSweepStore(src, "sweepgroup_bottom")); + KVCacheStore& store = *store_ptr; PopulateCold(store, src, {"a", "b", "c", "d", "e", "f"}); @@ -5657,9 +5638,8 @@ TEST_F(EvictSweepTest, DropsLocallyWhenThereAreNoTargets) { TEST_F(EvictSweepTest, AnApplicationPollNeitherSeesNorStealsSweepVerdicts) { RaidenId src{"sweep_src_owner", "0", "kv", 0}; RaidenId dst{"sweep_dst_owner", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "ownergroup"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN(auto store_ptr, MakeSweepStore(src, "ownergroup")); + KVCacheStore& store = *store_ptr; StartFakeDestination(dst, "ownergroup", /*evict_tier=*/1); // Held in flight: the destination keeps saying "still working", so the @@ -5723,9 +5703,9 @@ TEST_F(EvictSweepTest, AnApplicationPollNeitherSeesNorStealsSweepVerdicts) { TEST_F(EvictSweepTest, TheSweepDoesNotSwallowAnApplicationVerdict) { RaidenId src{"sweep_src_appverdict", "0", "kv", 0}; RaidenId dst{"sweep_dst_appverdict", "0", "kv", 0}; - auto store_or = MakeSweepStore(src, "appverdictgroup"); - ASSERT_TRUE(store_or.ok()) << store_or.status().ToString(); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN(auto store_ptr, + MakeSweepStore(src, "appverdictgroup")); + KVCacheStore& store = *store_ptr; StartFakeDestination(dst, "appverdictgroup", /*evict_tier=*/1); // ALL_EXIST settles an offer synchronously inside Save, so the @@ -5736,7 +5716,7 @@ TEST_F(EvictSweepTest, TheSweepDoesNotSwallowAnApplicationVerdict) { // Two blocks, still under the low watermark for free blocks (6 of 8 free), // so saving them does not itself start a pressure episode. PopulatePinned(store, src, {"app0", "app1"}); - ASSERT_TRUE(store.Save({"app0", "app1"}, dst).ok()); + ABSL_ASSERT_OK(store.Save({"app0", "app1"}, dst)); // Now push the store under the watermark and let the sweep run a batch, // draining as it goes. @@ -5796,33 +5776,35 @@ BackendConfig MakeHostBackendConfig() { } TEST(KVCacheStoreConstructionRulesTest, CreateRejectsEmptyStoreServerIp) { - auto store_or = KVCacheStore::Create(MakeHostBackendConfig(), /*capacity=*/4, - /*global_registry_address=*/"", - RaidenId{}, /*num_shards=*/1, - /*shard_size_bytes=*/512, - /*store_server_ip=*/""); - EXPECT_TRUE(absl::IsInvalidArgument(store_or.status())) << store_or.status(); + auto store_result = + KVCacheStore::Create(MakeHostBackendConfig(), /*capacity=*/4, + /*global_registry_address=*/"", RaidenId{}, + /*num_shards=*/1, /*shard_size_bytes=*/512, + /*store_server_ip=*/""); + EXPECT_TRUE(absl::IsInvalidArgument(store_result.status())) + << store_result.status(); } TEST(KVCacheStoreConstructionRulesTest, CreateRejectsWildcardStoreServerIp) { for (const char* wildcard : {"[::]", "0.0.0.0", "::"}) { - auto store_or = KVCacheStore::Create( + auto store_result = KVCacheStore::Create( MakeHostBackendConfig(), /*capacity=*/4, /*global_registry_address=*/"", RaidenId{}, /*num_shards=*/1, /*shard_size_bytes=*/512, /*store_server_ip=*/wildcard); - EXPECT_TRUE(absl::IsInvalidArgument(store_or.status())) - << "wildcard \"" << wildcard << "\": " << store_or.status(); + EXPECT_TRUE(absl::IsInvalidArgument(store_result.status())) + << "wildcard \"" << wildcard << "\": " << store_result.status(); } } TEST(KVCacheStoreConstructionRulesTest, CreateRejectsZeroShards) { - auto store_or = KVCacheStore::Create(MakeHostBackendConfig(), /*capacity=*/4, - /*global_registry_address=*/"", - RaidenId{}, /*num_shards=*/0, - /*shard_size_bytes=*/512, - /*store_server_ip=*/"127.0.0.1"); - EXPECT_TRUE(absl::IsInvalidArgument(store_or.status())) << store_or.status(); + auto store_result = + KVCacheStore::Create(MakeHostBackendConfig(), /*capacity=*/4, + /*global_registry_address=*/"", RaidenId{}, + /*num_shards=*/0, /*shard_size_bytes=*/512, + /*store_server_ip=*/"127.0.0.1"); + EXPECT_TRUE(absl::IsInvalidArgument(store_result.status())) + << store_result.status(); } TEST(KVCacheStoreConstructionRulesDeathTest, CapacityCtorDiesOnEmptyIp) { @@ -5866,12 +5848,12 @@ TEST(KVCacheStoreConstructionRulesTest, CreateFailsWhenRegistryPublishFails) { // A reachable registry that genuinely rejects RegisterStore (empty // job_name) -- registered, valid construction args, but a real runtime // publish failure. - auto store_or = KVCacheStore::Create( + auto store_result = KVCacheStore::Create( MakeHostBackendConfig(), /*capacity=*/4, server_address, RaidenId{}, /*num_shards=*/1, /*shard_size_bytes=*/512, /*store_server_ip=*/"127.0.0.1"); - EXPECT_FALSE(store_or.ok()) << "expected RegisterStore's rejection to " - "surface as a Create() failure"; + EXPECT_FALSE(store_result.ok()) << "expected RegisterStore's rejection to " + "surface as a Create() failure"; server->Shutdown(); } @@ -5925,14 +5907,14 @@ TEST(KVCacheStoreTest, LookupAndPinWorkflow) { std::vector hashes = {"h1", "h2"}; std::vector slices = {RaidenBlockId(id, 1, BlockStatus::HOST), RaidenBlockId(id, 2, BlockStatus::HOST)}; - ASSERT_TRUE(store.Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(store.Insert(hashes, slices, /*on_host=*/true)); // Insert pins what it takes. This case is about the pin LOOKUP grants, so // hand the insert's back first and count from zero. store.Release(hashes); - auto res = store.Lookup(hashes, LookupOptions{.pin_found = true}); - ASSERT_TRUE(res.ok()); - EXPECT_EQ(res->size(), 2); + TF_ASSERT_OK_AND_ASSIGN( + auto res, store.Lookup(hashes, LookupOptions{.pin_found = true})); + EXPECT_EQ(res.size(), 2); EXPECT_EQ(store.GetPinCount("h1"), 1); EXPECT_EQ(store.GetPinCount("h2"), 1); @@ -5952,7 +5934,7 @@ TEST(KVCacheStoreTest, LookupAndPinWorkflow) { EXPECT_EQ(store.GetPinCount("h2"), 0); // Now inserting h3 succeeds, evicting unpinned h2 (tail of sequence). - EXPECT_TRUE(store.Insert(new_hash, new_slice, /*on_host=*/true).ok()); + ABSL_EXPECT_OK(store.Insert(new_hash, new_slice, /*on_host=*/true)); EXPECT_TRUE(PeekLookup(store, {"h2"})->empty()); EXPECT_EQ(PeekLookup(store, {"h1"})->size(), 1); } @@ -5969,14 +5951,14 @@ TEST(KVCacheStoreTest, LookupPinFoundFalseObservesWithoutPinning) { ASSERT_TRUE(InsertResident(store, {"h1"}, {RaidenBlockId(id, 1, BlockStatus::HOST)}, true)); - auto res = - store.Lookup({"h1"}, /*enable_global=*/false, /*pin_found=*/false); - ASSERT_TRUE(res.ok()); - EXPECT_EQ(res->size(), 1); + TF_ASSERT_OK_AND_ASSIGN(auto res, + store.Lookup({"h1"}, /*enable_global=*/false, + /*pin_found=*/false)); + EXPECT_EQ(res.size(), 1); EXPECT_EQ(store.GetPinCount("h1"), 0); // The default still pins. - ASSERT_TRUE(store.Lookup({"h1"}).ok()); + ABSL_ASSERT_OK(store.Lookup({"h1"})); EXPECT_EQ(store.GetPinCount("h1"), 1); store.Release({"h1"}); } @@ -6016,11 +5998,12 @@ TEST(KVCacheStoreTest, LookupAndPinCapacityTruncation) { EXPECT_EQ(store.capacity(), 2); - auto res = store.Lookup({"h1", "h2", "h3"}, LookupOptions{.pin_found = true}); - ASSERT_TRUE(res.ok()); - ASSERT_EQ(res->size(), 2); - EXPECT_EQ((*res)[0].first, "h1"); - EXPECT_EQ((*res)[1].first, "h2"); + TF_ASSERT_OK_AND_ASSIGN( + auto res, + store.Lookup({"h1", "h2", "h3"}, LookupOptions{.pin_found = true})); + ASSERT_EQ(res.size(), 2); + EXPECT_EQ(res[0].first, "h1"); + EXPECT_EQ(res[1].first, "h2"); EXPECT_EQ(b1->GetPinCount("h1"), 1); EXPECT_EQ(b1->GetPinCount("h2"), 1); @@ -6046,19 +6029,20 @@ TEST(KVCacheStoreTest, RegisterKVTransferSpecFromWorkersPublishesToRegistry) { BackendConfig config; config.type = "HostOffloadBackend"; - auto store_or = KVCacheStore::Create( - config, /*capacity=*/4, registry_address, RaidenId{"job", "0", "data", 0}, - /*num_shards=*/2, /*shard_size_bytes=*/512, - /*store_server_ip=*/"127.0.0.1"); - ASSERT_OK(store_or.status()); - KVCacheStore& store = **store_or; + TF_ASSERT_OK_AND_ASSIGN( + auto store_ptr, + KVCacheStore::Create(config, /*capacity=*/4, registry_address, + RaidenId{"job", "0", "data", 0}, + /*num_shards=*/2, /*shard_size_bytes=*/512, + /*store_server_ip=*/"127.0.0.1")); + KVCacheStore& store = *store_ptr; // A live WorkerService to register: registration probes the worker's // endpoint with an empty CreateBuffers RPC and rejects an unreachable one. auto worker_server = ::tpu_raiden::controller::CreateTestWorkerServer(); core::controller::RaidenControllerClient controller_client( store.raiden_controller_address()); - ASSERT_OK(controller_client.RegisterWorker( + ABSL_ASSERT_OK(controller_client.RegisterWorker( "worker_0", worker_server->server_address, {::tpu_raiden::RaidenTransferEndpoint{worker_server->server_address, {}}}, /*node_id=*/0, /*block_array_bytes=*/{4096, 512}, @@ -6066,21 +6050,20 @@ TEST(KVCacheStoreTest, RegisterKVTransferSpecFromWorkersPublishesToRegistry) { auto* backend = dynamic_cast(store.backend().get()); ASSERT_NE(backend, nullptr); - ASSERT_OK(backend->RegisterKVTransferSpecFromWorkers()); + ABSL_ASSERT_OK(backend->RegisterKVTransferSpecFromWorkers()); // No kv_pool_group was configured, so the publish fell back to the // store's raiden_id.job_name as the group. global_registry::GlobalRegistryClient registry_client(grpc::CreateChannel( registry_address, grpc::InsecureChannelCredentials())); - auto spec_or = registry_client.GetKVTransferSpec("job"); - ASSERT_OK(spec_or.status()); - ASSERT_EQ(spec_or->block_arrays_size(), 2); - EXPECT_EQ(spec_or->block_arrays(0).block_bytes(), 4096); - EXPECT_EQ(spec_or->block_arrays(1).block_bytes(), 512); - EXPECT_EQ(spec_or->num_kv_shards(), 2); - EXPECT_EQ(spec_or->num_workers(), 1); - - store_or->reset(); + TF_ASSERT_OK_AND_ASSIGN(auto spec, registry_client.GetKVTransferSpec("job")); + ASSERT_EQ(spec.block_arrays_size(), 2); + EXPECT_EQ(spec.block_arrays(0).block_bytes(), 4096); + EXPECT_EQ(spec.block_arrays(1).block_bytes(), 512); + EXPECT_EQ(spec.num_kv_shards(), 2); + EXPECT_EQ(spec.num_workers(), 1); + + store_ptr.reset(); registry_server->Shutdown(); } @@ -6126,33 +6109,33 @@ std::unique_ptr MakeStoreInterleaveFixture( TEST(KVCacheStoreTest, LookupInterleavesLocalAndRemoteThroughTheStore) { auto f = MakeStoreInterleaveFixture(); - auto res = f->store->Lookup({"r1", "l1", "r2", "l2", "nowhere"}, - /*enable_global=*/true); - ASSERT_TRUE(res.ok()); - ASSERT_EQ(res->size(), 4); + TF_ASSERT_OK_AND_ASSIGN(auto res, + f->store->Lookup({"r1", "l1", "r2", "l2", "nowhere"}, + /*enable_global=*/true)); + ASSERT_EQ(res.size(), 4); - EXPECT_EQ((*res)[0].second.status, BlockStatus::REMOTE); - EXPECT_EQ((*res)[0].second.raiden_id, f->peer_id); - EXPECT_EQ((*res)[0].second.host_block_id, 42); + EXPECT_EQ(res[0].second.status, BlockStatus::REMOTE); + EXPECT_EQ(res[0].second.raiden_id, f->peer_id); + EXPECT_EQ(res[0].second.host_block_id, 42); - EXPECT_EQ((*res)[1].second.status, BlockStatus::HOST); - EXPECT_EQ((*res)[1].second.raiden_id, f->store_id); - EXPECT_EQ((*res)[1].second.host_block_id, 11); + EXPECT_EQ(res[1].second.status, BlockStatus::HOST); + EXPECT_EQ(res[1].second.raiden_id, f->store_id); + EXPECT_EQ(res[1].second.host_block_id, 11); - EXPECT_EQ((*res)[2].second.status, BlockStatus::REMOTE); - EXPECT_EQ((*res)[2].second.host_block_id, 43); + EXPECT_EQ(res[2].second.status, BlockStatus::REMOTE); + EXPECT_EQ(res[2].second.host_block_id, 43); - EXPECT_EQ((*res)[3].second.status, BlockStatus::HOST); - EXPECT_EQ((*res)[3].second.host_block_id, 12); + EXPECT_EQ(res[3].second.status, BlockStatus::HOST); + EXPECT_EQ(res[3].second.host_block_id, 12); } TEST(KVCacheStoreTest, LookupInterleavedWithPinFoundPinsOnlyLocalEntries) { auto f = MakeStoreInterleaveFixture(); - auto res = f->store->Lookup({"r1", "l1", "r2", "l2", "nowhere"}, - LookupOptions{.pin_found = true}); - ASSERT_TRUE(res.ok()); - ASSERT_EQ(res->size(), 4); + TF_ASSERT_OK_AND_ASSIGN(auto res, + f->store->Lookup({"r1", "l1", "r2", "l2", "nowhere"}, + LookupOptions{.pin_found = true})); + ASSERT_EQ(res.size(), 4); EXPECT_EQ(f->store->GetPinCount("l1"), 1); EXPECT_EQ(f->store->GetPinCount("l2"), 1); // A registry descriptor has no local entry behind it, so there is nothing to @@ -6170,30 +6153,31 @@ TEST(KVCacheStoreTest, LookupInterleavedDisabledThroughTheStore) { // The sweep stops at "r1", so "l1" can only be looked for in the registry, // which no longer answers for it. - auto legacy = f->store->Lookup( - {"r1", "l1"}, LookupOptions{.enable_interleaved_lookup = false}); - ASSERT_TRUE(legacy.ok()); - ASSERT_EQ(legacy->size(), 1); - EXPECT_EQ((*legacy)[0].first, "r1"); + TF_ASSERT_OK_AND_ASSIGN( + auto legacy, + f->store->Lookup({"r1", "l1"}, + LookupOptions{.enable_interleaved_lookup = false})); + ASSERT_EQ(legacy.size(), 1); + EXPECT_EQ(legacy[0].first, "r1"); - auto interleaved = f->store->Lookup({"r1", "l1"}, /*enable_global=*/true); - ASSERT_TRUE(interleaved.ok()); - ASSERT_EQ(interleaved->size(), 2); - EXPECT_EQ((*interleaved)[1].first, "l1"); - EXPECT_EQ((*interleaved)[1].second.status, BlockStatus::HOST); + TF_ASSERT_OK_AND_ASSIGN( + auto interleaved, f->store->Lookup({"r1", "l1"}, /*enable_global=*/true)); + ASSERT_EQ(interleaved.size(), 2); + EXPECT_EQ(interleaved[1].first, "l1"); + EXPECT_EQ(interleaved[1].second.status, BlockStatus::HOST); } TEST(KVCacheStoreTest, LookupInterleavedWithoutGlobalStopsAtTheLocalMiss) { auto f = MakeStoreInterleaveFixture(); for (bool interleaved : {true, false}) { - auto res = f->store->Lookup( - {"l1", "r1", "l2"}, - LookupOptions{.enable_global = false, - .enable_interleaved_lookup = interleaved}); - ASSERT_TRUE(res.ok()); - ASSERT_EQ(res->size(), 1) << "interleaved=" << interleaved; - EXPECT_EQ((*res)[0].first, "l1"); + TF_ASSERT_OK_AND_ASSIGN( + auto res, f->store->Lookup( + {"l1", "r1", "l2"}, + LookupOptions{.enable_global = false, + .enable_interleaved_lookup = interleaved})); + ASSERT_EQ(res.size(), 1) << "interleaved=" << interleaved; + EXPECT_EQ(res[0].first, "l1"); } } @@ -6203,12 +6187,12 @@ TEST(KVCacheStoreTest, LookupInterleavedTruncatesToCapacityAndUnwindsPins) { // for a registry descriptor it never pinned, which must stay harmless. auto f = MakeStoreInterleaveFixture(/*capacity=*/2); - auto res = f->store->Lookup({"r1", "l1", "r2", "l2"}, - LookupOptions{.pin_found = true}); - ASSERT_TRUE(res.ok()); - ASSERT_EQ(res->size(), 2); - EXPECT_EQ((*res)[0].first, "r1"); - EXPECT_EQ((*res)[1].first, "l1"); + TF_ASSERT_OK_AND_ASSIGN(auto res, + f->store->Lookup({"r1", "l1", "r2", "l2"}, + LookupOptions{.pin_found = true})); + ASSERT_EQ(res.size(), 2); + EXPECT_EQ(res[0].first, "r1"); + EXPECT_EQ(res[1].first, "l1"); EXPECT_EQ(f->store->GetPinCount("l1"), 1); EXPECT_EQ(f->store->GetPinCount("l2"), 0); diff --git a/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc b/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc index faffc1b0..f0adefd8 100644 --- a/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc +++ b/tpu_sync/kv_cache/kv_cache_store_wrapper_test.cc @@ -26,9 +26,11 @@ #include #include +#include "absl/status/status_matchers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/common/raiden_id.h" #include "tpu_sync/kv_cache/global_registry/test_util.h" #include "tpu_sync/kv_cache/kv_cache_metadata_shm.h" @@ -128,7 +130,7 @@ class KVCacheStoreWrapperTest : public ::testing::Test { for (int i = 0; i < static_cast(hashes.size()); ++i) { slices.push_back(RaidenBlockId(rid, i, BlockStatus::HOST)); } - ASSERT_TRUE(wrapper->Insert(hashes, slices, /*on_host=*/true).ok()); + ABSL_ASSERT_OK(wrapper->Insert(hashes, slices, /*on_host=*/true)); // Insert pins what it takes; these cases only want the blocks resident, // and a held pin would make them unevictable. wrapper->Release(hashes); @@ -158,9 +160,9 @@ TEST_F(KVCacheStoreWrapperTest, ColdStartCreatesMetadataTable) { auto wrapper = MakeWrapper(/*capacity=*/4, /*num_shards=*/1); EXPECT_TRUE( MetadataSegmentExists(absl::StrCat("_metadata", kIdentitySuffix))); - auto lookup_or = (*wrapper)->Lookup({"host_1"}, LookupOptions{}); - ASSERT_TRUE(lookup_or.ok()); - EXPECT_THAT(*lookup_or, IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto lookup, + (*wrapper)->Lookup({"host_1"}, LookupOptions{})); + EXPECT_THAT(lookup, IsEmpty()); } TEST_F(KVCacheStoreWrapperTest, ServerNameSuffixesMetadataSegment) { @@ -221,15 +223,15 @@ TEST_F(KVCacheStoreWrapperTest, RecoversHostBlocksAfterRestart) { wrapper.reset(); wrapper = MakeWrapper(/*capacity=*/4, /*num_shards=*/1); - auto lookup_or = (*wrapper)->Lookup({"host_1", "host_2"}, LookupOptions{}); - ASSERT_TRUE(lookup_or.ok()); - ASSERT_EQ(lookup_or->size(), 2); - EXPECT_EQ((*lookup_or)[0].first, "host_1"); - EXPECT_EQ((*lookup_or)[0].second.status, BlockStatus::HOST); - EXPECT_EQ((*lookup_or)[0].second.host_block_id, 0); - EXPECT_EQ((*lookup_or)[1].first, "host_2"); - EXPECT_EQ((*lookup_or)[1].second.status, BlockStatus::HOST); - EXPECT_EQ((*lookup_or)[1].second.host_block_id, 1); + TF_ASSERT_OK_AND_ASSIGN( + auto lookup, (*wrapper)->Lookup({"host_1", "host_2"}, LookupOptions{})); + ASSERT_EQ(lookup.size(), 2); + EXPECT_EQ(lookup[0].first, "host_1"); + EXPECT_EQ(lookup[0].second.status, BlockStatus::HOST); + EXPECT_EQ(lookup[0].second.host_block_id, 0); + EXPECT_EQ(lookup[1].first, "host_2"); + EXPECT_EQ(lookup[1].second.status, BlockStatus::HOST); + EXPECT_EQ(lookup[1].second.host_block_id, 1); } // Exercises the env contract StoreMonitorConfigFromEnv documents; the @@ -324,9 +326,9 @@ TEST_F(KVCacheStoreWrapperTest, ModelUidMismatchColdStarts) { setenv("RAIDEN_SHM_MODEL_UID", "model_b", /*overwrite=*/1); wrapper = MakeWrapper(/*capacity=*/4, /*num_shards=*/1); - auto lookup_or = (*wrapper)->Lookup({"host_1"}, LookupOptions{}); - ASSERT_TRUE(lookup_or.ok()); - EXPECT_THAT(*lookup_or, IsEmpty()); + TF_ASSERT_OK_AND_ASSIGN(auto lookup, + (*wrapper)->Lookup({"host_1"}, LookupOptions{})); + EXPECT_THAT(lookup, IsEmpty()); } } // namespace diff --git a/tpu_sync/kv_cache/logical_block_manager_test.cc b/tpu_sync/kv_cache/logical_block_manager_test.cc index 5753eeda..a62989b6 100644 --- a/tpu_sync/kv_cache/logical_block_manager_test.cc +++ b/tpu_sync/kv_cache/logical_block_manager_test.cc @@ -20,12 +20,15 @@ #include #include #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/status/statusor.h" +#include "xla/tsl/platform/statusor.h" namespace tpu_raiden { namespace kv_cache { namespace { +using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::Eq; using ::testing::Optional; @@ -45,9 +48,8 @@ TEST(LogicalBlockManagerTest, InitialState) { TEST(LogicalBlockManagerTest, BasicAllocation) { LogicalBlockManager manager(5); - auto blocks_or = manager.Allocate(3, /*lock=*/false); - ASSERT_TRUE(blocks_or.ok()); - EXPECT_THAT(*blocks_or, ElementsAre(0, 1, 2)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks, manager.Allocate(3, /*lock=*/false)); + EXPECT_THAT(blocks, ElementsAre(0, 1, 2)); EXPECT_EQ(manager.num_free_blocks(), 2); EXPECT_EQ(manager.num_allocated_blocks(), 3); @@ -61,9 +63,8 @@ TEST(LogicalBlockManagerTest, BasicAllocation) { TEST(LogicalBlockManagerTest, AllocationWithLocking) { LogicalBlockManager manager(5); - auto blocks_or = manager.Allocate(2, /*lock=*/true); - ASSERT_TRUE(blocks_or.ok()); - EXPECT_THAT(*blocks_or, ElementsAre(0, 1)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks, manager.Allocate(2, /*lock=*/true)); + EXPECT_THAT(blocks, ElementsAre(0, 1)); EXPECT_EQ(manager.num_locked_blocks(), 2); EXPECT_TRUE(manager.IsLocked(0)); @@ -74,72 +75,64 @@ TEST(LogicalBlockManagerTest, LruEvictionOrder) { LogicalBlockManager manager(4); // Allocate 2 blocks to entity 10 (unlocked). - auto blocks1 = manager.Allocate(2); - ASSERT_TRUE(blocks1.ok()); - EXPECT_THAT(*blocks1, ElementsAre(0, 1)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks1, manager.Allocate(2)); + EXPECT_THAT(blocks1, ElementsAre(0, 1)); // Allocate 2 blocks to entity 20 (unlocked). - auto blocks2 = manager.Allocate(2); - ASSERT_TRUE(blocks2.ok()); - EXPECT_THAT(*blocks2, ElementsAre(2, 3)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks2, manager.Allocate(2)); + EXPECT_THAT(blocks2, ElementsAre(2, 3)); EXPECT_EQ(manager.num_free_blocks(), 0); // Requesting 2 blocks for entity 30 should evict entity 10's blocks // because they were allocated earlier (LRU). - auto blocks3 = manager.Allocate(2); - ASSERT_TRUE(blocks3.ok()); - EXPECT_THAT(*blocks3, ElementsAre(0, 1)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks3, manager.Allocate(2)); + EXPECT_THAT(blocks3, ElementsAre(0, 1)); } TEST(LogicalBlockManagerTest, AccessUpdatesLruOrder) { LogicalBlockManager manager(4); - auto blocks1 = manager.Allocate(2); - ASSERT_TRUE(blocks1.ok()); - auto blocks2 = manager.Allocate(2); - ASSERT_TRUE(blocks2.ok()); + TF_ASSERT_OK_AND_ASSIGN(auto blocks1, manager.Allocate(2)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks2, manager.Allocate(2)); // Access entity 10's blocks, making entity 20's blocks the least recently // used. - EXPECT_TRUE(manager.AccessBlock(0).ok()); - EXPECT_TRUE(manager.AccessBlock(1).ok()); + ABSL_EXPECT_OK(manager.AccessBlock(0)); + ABSL_EXPECT_OK(manager.AccessBlock(1)); // Allocate 2 blocks for entity 30. Should evict entity 20's blocks (2 and 3). - auto blocks3 = manager.Allocate(2); - ASSERT_TRUE(blocks3.ok()); - EXPECT_THAT(*blocks3, ElementsAre(2, 3)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks3, manager.Allocate(2)); + EXPECT_THAT(blocks3, ElementsAre(2, 3)); } TEST(LogicalBlockManagerTest, LockedBlocksPreventEviction) { LogicalBlockManager manager(4); // Allocate 2 locked blocks to entity 10. - ASSERT_TRUE(manager.Allocate(2, /*lock=*/true).ok()); + ABSL_ASSERT_OK(manager.Allocate(2, /*lock=*/true)); // Allocate 2 unlocked blocks. - ASSERT_TRUE(manager.Allocate(2, /*lock=*/false).ok()); + ABSL_ASSERT_OK(manager.Allocate(2, /*lock=*/false)); // Requesting 3 blocks should fail since only 2 blocks are evictable. - auto failed_or = manager.Allocate(3); - EXPECT_FALSE(failed_or.ok()); - EXPECT_TRUE(absl::IsResourceExhausted(failed_or.status())); + EXPECT_THAT(manager.Allocate(3), + StatusIs(absl::StatusCode::kResourceExhausted)); } TEST(LogicalBlockManagerTest, UnlockAllowsEviction) { LogicalBlockManager manager(4); - ASSERT_TRUE(manager.Allocate(2, /*lock=*/true).ok()); - ASSERT_TRUE(manager.Allocate(2, /*lock=*/false).ok()); + ABSL_ASSERT_OK(manager.Allocate(2, /*lock=*/true)); + ABSL_ASSERT_OK(manager.Allocate(2, /*lock=*/false)); // Unlock entity 10's blocks. std::vector to_unlock = {0, 1}; - EXPECT_TRUE(manager.Unlock(to_unlock).ok()); + ABSL_EXPECT_OK(manager.Unlock(to_unlock)); EXPECT_EQ(manager.num_locked_blocks(), 0); // Now requesting 3 blocks succeeds. - auto blocks_or = manager.Allocate(3); - EXPECT_TRUE(blocks_or.ok()); - EXPECT_EQ(blocks_or->size(), 3); + TF_ASSERT_OK_AND_ASSIGN(auto blocks, manager.Allocate(3)); + EXPECT_EQ(blocks.size(), 3); } TEST(LogicalBlockManagerTest, InvalidArguments) { @@ -160,7 +153,7 @@ TEST(LogicalBlockManagerTest, InvalidArguments) { TEST(LogicalBlockManagerTest, AllocateTargetMarksBlocksAllocatedAndLocked) { LogicalBlockManager manager(5); - ASSERT_TRUE(manager.AllocateTarget({1, 3}).ok()); + ABSL_ASSERT_OK(manager.AllocateTarget({1, 3})); EXPECT_TRUE(manager.IsAllocated(1)); EXPECT_TRUE(manager.IsLocked(1)); EXPECT_TRUE(manager.IsAllocated(3)); @@ -168,26 +161,25 @@ TEST(LogicalBlockManagerTest, AllocateTargetMarksBlocksAllocatedAndLocked) { EXPECT_EQ(manager.num_free_blocks(), 3); // Target-allocated blocks are never handed out by subsequent allocations. - auto blocks_or = manager.Allocate(3, /*lock=*/true); - ASSERT_TRUE(blocks_or.ok()); - EXPECT_THAT(*blocks_or, ElementsAre(0, 2, 4)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks, manager.Allocate(3, /*lock=*/true)); + EXPECT_THAT(blocks, ElementsAre(0, 2, 4)); // Everything is locked now: further allocation must fail. EXPECT_FALSE(manager.Allocate(1).ok()); } TEST(LogicalBlockManagerTest, AllocateTargetValidatesAtomically) { LogicalBlockManager manager(5); - ASSERT_TRUE(manager.Allocate(1).ok()); // Block 0 becomes allocated. + ABSL_ASSERT_OK(manager.Allocate(1)); // Block 0 becomes allocated. // Out of range. - EXPECT_EQ(manager.AllocateTarget({1, 5}).code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(manager.AllocateTarget({1, 5}), + StatusIs(absl::StatusCode::kInvalidArgument)); // Already allocated (even though unlocked). - EXPECT_EQ(manager.AllocateTarget({1, 0}).code(), - absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(manager.AllocateTarget({1, 0}), + StatusIs(absl::StatusCode::kFailedPrecondition)); // Duplicate ID within the batch. - EXPECT_EQ(manager.AllocateTarget({2, 2}).code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(manager.AllocateTarget({2, 2}), + StatusIs(absl::StatusCode::kInvalidArgument)); // Failed calls must not have modified any state. EXPECT_FALSE(manager.IsAllocated(1)); @@ -197,12 +189,12 @@ TEST(LogicalBlockManagerTest, AllocateTargetValidatesAtomically) { TEST(LogicalBlockManagerTest, DeallocateReturnsBlocksToFreePool) { LogicalBlockManager manager(4); - ASSERT_TRUE(manager.Allocate(2, /*lock=*/true).ok()); - ASSERT_TRUE(manager.Allocate(1, /*lock=*/false).ok()); + ABSL_ASSERT_OK(manager.Allocate(2, /*lock=*/true)); + ABSL_ASSERT_OK(manager.Allocate(1, /*lock=*/false)); EXPECT_EQ(manager.num_free_blocks(), 1); // Deallocation works on locked and unlocked blocks alike. - ASSERT_TRUE(manager.Deallocate({0, 2}).ok()); + ABSL_ASSERT_OK(manager.Deallocate({0, 2})); EXPECT_EQ(manager.num_free_blocks(), 3); EXPECT_EQ(manager.num_allocated_blocks(), 1); EXPECT_EQ(manager.num_locked_blocks(), 1); @@ -211,42 +203,40 @@ TEST(LogicalBlockManagerTest, DeallocateReturnsBlocksToFreePool) { EXPECT_FALSE(manager.IsAllocated(2)); // Deallocated blocks are free again for both allocation paths. - ASSERT_TRUE(manager.AllocateTarget({0}).ok()); - auto blocks_or = manager.Allocate(1); - ASSERT_TRUE(blocks_or.ok()); - EXPECT_THAT(*blocks_or, ElementsAre(2)); + ABSL_ASSERT_OK(manager.AllocateTarget({0})); + TF_ASSERT_OK_AND_ASSIGN(auto blocks, manager.Allocate(1)); + EXPECT_THAT(blocks, ElementsAre(2)); } TEST(LogicalBlockManagerTest, DeallocateValidatesAtomically) { LogicalBlockManager manager(3); - ASSERT_TRUE(manager.Allocate(1, /*lock=*/true).ok()); // Block 0. + ABSL_ASSERT_OK(manager.Allocate(1, /*lock=*/true)); // Block 0. // Out of range. - EXPECT_EQ(manager.Deallocate({0, 3}).code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(manager.Deallocate({0, 3}), + StatusIs(absl::StatusCode::kInvalidArgument)); // Not allocated. - EXPECT_EQ(manager.Deallocate({0, 1}).code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(manager.Deallocate({0, 1}), + StatusIs(absl::StatusCode::kInvalidArgument)); // Failed calls must not have modified any state. EXPECT_TRUE(manager.IsAllocated(0)); EXPECT_TRUE(manager.IsLocked(0)); - ASSERT_TRUE(manager.Deallocate({0}).ok()); + ABSL_ASSERT_OK(manager.Deallocate({0})); // Double deallocation fails. - EXPECT_EQ(manager.Deallocate({0}).code(), - absl::StatusCode::kInvalidArgument); + EXPECT_THAT(manager.Deallocate({0}), + StatusIs(absl::StatusCode::kInvalidArgument)); } TEST(LogicalBlockManagerTest, TargetAllocatedBlocksReusableAfterUnlock) { LogicalBlockManager manager(3); - ASSERT_TRUE(manager.AllocateTarget({0, 1, 2}).ok()); - ASSERT_TRUE(manager.Unlock({1}).ok()); + ABSL_ASSERT_OK(manager.AllocateTarget({0, 1, 2})); + ABSL_ASSERT_OK(manager.Unlock({1})); // The unlocked target-allocated block is evictable and gets reused. - auto blocks_or = manager.Allocate(1); - ASSERT_TRUE(blocks_or.ok()); - EXPECT_THAT(*blocks_or, ElementsAre(1)); + TF_ASSERT_OK_AND_ASSIGN(auto blocks, manager.Allocate(1)); + EXPECT_THAT(blocks, ElementsAre(1)); } } // namespace diff --git a/tpu_sync/kv_cache/pool_layout.cc b/tpu_sync/kv_cache/pool_layout.cc index 1794ff8e..781aef6b 100644 --- a/tpu_sync/kv_cache/pool_layout.cc +++ b/tpu_sync/kv_cache/pool_layout.cc @@ -23,6 +23,7 @@ #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/types/span.h" @@ -159,28 +160,29 @@ absl::Status RegionSpec::Validate(int64_t slot_bytes) const { return absl::InvalidArgumentError( absl::StrCat("region ", name, " stride_bytes must be positive")); } - auto packed_bytes = + ABSL_ASSIGN_OR_RETURN( + int64_t packed_bytes, CheckedMul(units_per_stride, unit_bytes, - absl::StrCat("region ", name, " packed unit bytes")); - if (!packed_bytes.ok()) return packed_bytes.status(); - if (num_units > 0 && stride_bytes < *packed_bytes) { + absl::StrCat("region ", name, " packed unit bytes"))); + if (num_units > 0 && stride_bytes < packed_bytes) { return absl::InvalidArgumentError(absl::StrCat( "region ", name, " stride_bytes is smaller than packed units")); } - auto stride_extent = + ABSL_ASSIGN_OR_RETURN( + int64_t stride_extent, CheckedMul(num_units > 0 ? num_units - 1 : 0, stride_bytes, - absl::StrCat("region ", name, " stride extent")); - if (!stride_extent.ok()) return stride_extent.status(); - auto extent_start = - CheckedAdd(offset_bytes, *stride_extent, - absl::StrCat("region ", name, " extent start")); - if (!extent_start.ok()) return extent_start.status(); - auto extent_end = CheckedAdd(*extent_start, *packed_bytes, - absl::StrCat("region ", name, " extent end")); - if (!extent_end.ok()) return extent_end.status(); - if (*extent_end > slot_bytes) { + absl::StrCat("region ", name, " stride extent"))); + ABSL_ASSIGN_OR_RETURN( + int64_t extent_start, + CheckedAdd(offset_bytes, stride_extent, + absl::StrCat("region ", name, " extent start"))); + ABSL_ASSIGN_OR_RETURN( + int64_t extent_end, + CheckedAdd(extent_start, packed_bytes, + absl::StrCat("region ", name, " extent end"))); + if (extent_end > slot_bytes) { return absl::InvalidArgumentError( - absl::StrCat("region ", name, " exceeds slot bytes: end=", *extent_end, + absl::StrCat("region ", name, " exceeds slot bytes: end=", extent_end, " slot=", slot_bytes)); } return absl::OkStatus(); @@ -280,59 +282,54 @@ absl::StatusOr> ComputePoolBlockCopyExtents( block_id, " out of range [0, ", pool.num_blocks, ")")); } - auto block_delta_or = + ABSL_ASSIGN_OR_RETURN( + const int64_t block_delta, CheckedMul(block_id, pool.block_stride_bytes, - absl::StrCat("pool ", pool.tag, " block offset")); - if (!block_delta_or.ok()) return block_delta_or.status(); - const int64_t block_delta = *block_delta_or; - auto block_base_or = + absl::StrCat("pool ", pool.tag, " block offset"))); + ABSL_ASSIGN_OR_RETURN( + const int64_t block_base, CheckedAdd(pool.base_offset_bytes, block_delta, - absl::StrCat("pool ", pool.tag, " block base")); - if (!block_base_or.ok()) return block_base_or.status(); - const int64_t block_base = *block_base_or; + absl::StrCat("pool ", pool.tag, " block base"))); for (const RegionSpec& region : pool.regions) { - auto packed_bytes_or = + ABSL_ASSIGN_OR_RETURN( + const int64_t packed_bytes, CheckedMul(region.unit_bytes, region.units_per_stride, absl::StrCat("pool ", pool.tag, " region ", region.name, - " packed bytes")); - if (!packed_bytes_or.ok()) return packed_bytes_or.status(); - const int64_t packed_bytes = *packed_bytes_or; + " packed bytes"))); if (region.num_units == 0 || packed_bytes == 0) continue; // A tightly packed strided region is one contiguous DMA extent. Sparse // regions retain one extent per unit so padding is never staged. if (region.stride_bytes == packed_bytes) { - auto region_bytes_or = + ABSL_ASSIGN_OR_RETURN( + const int64_t region_bytes, CheckedMul(region.num_units, packed_bytes, absl::StrCat("pool ", pool.tag, " region ", region.name, - " live bytes")); - if (!region_bytes_or.ok()) return region_bytes_or.status(); - const int64_t region_bytes = *region_bytes_or; - auto offset_or = CheckedAdd(block_base, region.offset_bytes, - absl::StrCat("pool ", pool.tag, " region ", - region.name, " offset")); - if (!offset_or.ok()) return offset_or.status(); - const int64_t offset = *offset_or; + " live bytes"))); + ABSL_ASSIGN_OR_RETURN( + const int64_t offset, + CheckedAdd(block_base, region.offset_bytes, + absl::StrCat("pool ", pool.tag, " region ", region.name, + " offset"))); extents.push_back({.offset_bytes = offset, .size_bytes = region_bytes}); continue; } for (int64_t unit = 0; unit < region.num_units; ++unit) { - auto unit_delta_or = + ABSL_ASSIGN_OR_RETURN( + const int64_t unit_delta, CheckedMul(unit, region.stride_bytes, absl::StrCat("pool ", pool.tag, " region ", region.name, - " unit offset")); - if (!unit_delta_or.ok()) return unit_delta_or.status(); - const int64_t unit_delta = *unit_delta_or; - auto region_base_or = CheckedAdd( - block_base, region.offset_bytes, - absl::StrCat("pool ", pool.tag, " region ", region.name, " base")); - if (!region_base_or.ok()) return region_base_or.status(); - const int64_t region_base = *region_base_or; - auto offset_or = CheckedAdd(region_base, unit_delta, - absl::StrCat("pool ", pool.tag, " region ", - region.name, " unit address")); - if (!offset_or.ok()) return offset_or.status(); - const int64_t offset = *offset_or; + " unit offset"))); + ABSL_ASSIGN_OR_RETURN( + const int64_t region_base, + CheckedAdd(block_base, region.offset_bytes, + absl::StrCat("pool ", pool.tag, " region ", region.name, + " base"))); + ABSL_ASSIGN_OR_RETURN( + const int64_t offset, + CheckedAdd(region_base, unit_delta, + absl::StrCat("pool ", pool.tag, " region ", region.name, + " unit address"))); extents.push_back({.offset_bytes = offset, .size_bytes = packed_bytes}); } } @@ -348,21 +345,20 @@ absl::StatusOr> ComputePoolBlockCopyExtents( std::vector merged; merged.reserve(extents.size()); for (const PoolBlockCopyExtent& extent : extents) { - auto extent_end_or = + ABSL_ASSIGN_OR_RETURN( + const int64_t extent_end, CheckedAdd(extent.offset_bytes, extent.size_bytes, - absl::StrCat("pool ", pool.tag, " copy extent end")); - if (!extent_end_or.ok()) return extent_end_or.status(); - const int64_t extent_end = *extent_end_or; + absl::StrCat("pool ", pool.tag, " copy extent end"))); if (merged.empty()) { merged.push_back(extent); continue; } PoolBlockCopyExtent& previous = merged.back(); - auto previous_end_or = CheckedAdd( - previous.offset_bytes, previous.size_bytes, - absl::StrCat("pool ", pool.tag, " previous copy extent end")); - if (!previous_end_or.ok()) return previous_end_or.status(); - const int64_t previous_end = *previous_end_or; + ABSL_ASSIGN_OR_RETURN( + const int64_t previous_end, + CheckedAdd( + previous.offset_bytes, previous.size_bytes, + absl::StrCat("pool ", pool.tag, " previous copy extent end"))); if (extent.offset_bytes > previous_end) { merged.push_back(extent); continue; @@ -543,9 +539,8 @@ absl::StatusOr PoolSpecFromProto( pool.num_blocks = proto.num_blocks(); pool.regions.reserve(proto.regions_size()); for (const auto& region_proto : proto.regions()) { - absl::StatusOr region = RegionSpecFromProto(region_proto); - if (!region.ok()) return region.status(); - pool.regions.push_back(*region); + ABSL_ASSIGN_OR_RETURN(RegionSpec region, RegionSpecFromProto(region_proto)); + pool.regions.push_back(std::move(region)); } pool.dtype_tag = proto.dtype_tag(); return pool; @@ -583,13 +578,13 @@ absl::StatusOr> TranslateLiveCopy( int64_t src_cursor = src_offset; int64_t dst_cursor = dst_offset; while (remaining > 0) { - auto src_loc = locate(src_segments, src_cursor); - if (!src_loc.ok()) return src_loc.status(); - auto dst_loc = locate(dst_segments, dst_cursor); - if (!dst_loc.ok()) return dst_loc.status(); + ABSL_ASSIGN_OR_RETURN(const auto& src_loc, + locate(src_segments, src_cursor)); + ABSL_ASSIGN_OR_RETURN(const auto& dst_loc, + locate(dst_segments, dst_cursor)); const int64_t chunk_size = - std::min({remaining, src_loc->second, dst_loc->second}); - chunks.push_back(LiveCopyChunk{src_loc->first, dst_loc->first, chunk_size}); + std::min({remaining, src_loc.second, dst_loc.second}); + chunks.push_back(LiveCopyChunk{src_loc.first, dst_loc.first, chunk_size}); if (static_cast(chunks.size()) > kMaxLiveSegments) { return absl::InvalidArgumentError( "Compact-live copy exceeds the segment expansion bound"); diff --git a/tpu_sync/kv_cache/pool_layout_test.cc b/tpu_sync/kv_cache/pool_layout_test.cc index d0394384..1f9859dd 100644 --- a/tpu_sync/kv_cache/pool_layout_test.cc +++ b/tpu_sync/kv_cache/pool_layout_test.cc @@ -23,6 +23,7 @@ #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/rpc/raiden_service.pb.h" namespace tpu_raiden { @@ -125,15 +126,14 @@ TEST(PoolSpecTest, AdmitsGdnSsmWhenLastLiveByteFitsSharedRawStorage) { TEST(PoolSpecTest, ProtoRoundTrip) { PoolSpec pool = ValidPool(); tpu_sync::rpc::PoolSpecProto proto = ToProto(pool); - auto roundtrip = PoolSpecFromProto(proto); - ABSL_ASSERT_OK(roundtrip); - EXPECT_EQ(roundtrip->tag, pool.tag); - EXPECT_EQ(roundtrip->storage_index, pool.storage_index); - EXPECT_EQ(roundtrip->base_offset_bytes, pool.base_offset_bytes); - EXPECT_EQ(roundtrip->block_stride_bytes, pool.block_stride_bytes); - EXPECT_EQ(roundtrip->num_blocks, pool.num_blocks); - ASSERT_EQ(roundtrip->regions.size(), pool.regions.size()); - EXPECT_EQ(roundtrip->regions[0].name, pool.regions[0].name); + TF_ASSERT_OK_AND_ASSIGN(auto roundtrip, PoolSpecFromProto(proto)); + EXPECT_EQ(roundtrip.tag, pool.tag); + EXPECT_EQ(roundtrip.storage_index, pool.storage_index); + EXPECT_EQ(roundtrip.base_offset_bytes, pool.base_offset_bytes); + EXPECT_EQ(roundtrip.block_stride_bytes, pool.block_stride_bytes); + EXPECT_EQ(roundtrip.num_blocks, pool.num_blocks); + ASSERT_EQ(roundtrip.regions.size(), pool.regions.size()); + EXPECT_EQ(roundtrip.regions[0].name, pool.regions[0].name); } TEST(RegionsCoverRangeTest, Basic) { @@ -171,9 +171,8 @@ TEST(ComputePoolBlockCopyExtentsTest, Basic) { .num_units = 1, .units_per_stride = 1, }}; - auto extents_or = ComputePoolBlockCopyExtents(pool, {0, 1, 3, 4, 5}); - ABSL_ASSERT_OK(extents_or); - const auto& extents = *extents_or; + TF_ASSERT_OK_AND_ASSIGN(auto extents, + ComputePoolBlockCopyExtents(pool, {0, 1, 3, 4, 5})); ASSERT_EQ(extents.size(), 2); EXPECT_EQ(extents[0].offset_bytes, 0); EXPECT_EQ(extents[0].size_bytes, 2000); // block 0 and 1 coalesced @@ -196,17 +195,17 @@ TEST(ComputePoolBlockCopyExtentsTest, OmitsInterRegionAndTrailingPadding) { }}; ABSL_ASSERT_OK(pool.Validate(/*storage_bytes=*/184)); - auto extents_or = ComputePoolBlockCopyExtents(pool, {0, 1}); - ABSL_ASSERT_OK(extents_or); - ASSERT_EQ(extents_or->size(), 4); - EXPECT_EQ((*extents_or)[0].offset_bytes, 8); - EXPECT_EQ((*extents_or)[0].size_bytes, 16); - EXPECT_EQ((*extents_or)[1].offset_bytes, 40); - EXPECT_EQ((*extents_or)[1].size_bytes, 16); - EXPECT_EQ((*extents_or)[2].offset_bytes, 136); - EXPECT_EQ((*extents_or)[2].size_bytes, 16); - EXPECT_EQ((*extents_or)[3].offset_bytes, 168); - EXPECT_EQ((*extents_or)[3].size_bytes, 16); + TF_ASSERT_OK_AND_ASSIGN(auto extents, + ComputePoolBlockCopyExtents(pool, {0, 1})); + ASSERT_EQ(extents.size(), 4); + EXPECT_EQ(extents[0].offset_bytes, 8); + EXPECT_EQ(extents[0].size_bytes, 16); + EXPECT_EQ(extents[1].offset_bytes, 40); + EXPECT_EQ(extents[1].size_bytes, 16); + EXPECT_EQ(extents[2].offset_bytes, 136); + EXPECT_EQ(extents[2].size_bytes, 16); + EXPECT_EQ(extents[3].offset_bytes, 168); + EXPECT_EQ(extents[3].size_bytes, 16); } TEST(ComputePoolBlockCopyExtentsTest, @@ -231,24 +230,21 @@ TEST(ComputePoolBlockCopyExtentsTest, } ABSL_ASSERT_OK(pool.Validate(/*storage_bytes=*/1884160)); - auto extents_or = ComputePoolBlockCopyExtents(pool, {0}); - ABSL_ASSERT_OK(extents_or); - ASSERT_EQ(extents_or->size(), 4); + TF_ASSERT_OK_AND_ASSIGN(auto extents, ComputePoolBlockCopyExtents(pool, {0})); + ASSERT_EQ(extents.size(), 4); for (int64_t head_group = 0; head_group < 4; ++head_group) { - EXPECT_EQ((*extents_or)[head_group].offset_bytes, head_group * 540672); - EXPECT_EQ((*extents_or)[head_group].size_bytes, 262144); + EXPECT_EQ(extents[head_group].offset_bytes, head_group * 540672); + EXPECT_EQ(extents[head_group].size_bytes, 262144); } } TEST(PoolLiveSegmentsTest, ExpandsStridedRegionsToCompactRuns) { - absl::StatusOr> segments = - ExpandPoolLiveSegments(ValidPool()); - ASSERT_TRUE(segments.ok()) << segments.status().ToString(); - ASSERT_EQ(segments->size(), 5u); + TF_ASSERT_OK_AND_ASSIGN(auto segments, ExpandPoolLiveSegments(ValidPool())); + ASSERT_EQ(segments.size(), 5u); for (int64_t unit = 0; unit < 5; ++unit) { - EXPECT_EQ((*segments)[unit].logical_offset, unit * 20); - EXPECT_EQ((*segments)[unit].physical_offset, unit * 100); - EXPECT_EQ((*segments)[unit].size, 20); + EXPECT_EQ(segments[unit].logical_offset, unit * 20); + EXPECT_EQ(segments[unit].physical_offset, unit * 100); + EXPECT_EQ(segments[unit].size, 20); } } @@ -272,13 +268,11 @@ TEST(PoolLiveSegmentsTest, CoalescesAbuttingRegionsAndRejectsOverlap) { .units_per_stride = 1, }, }; - absl::StatusOr> segments = - ExpandPoolLiveSegments(pool); - ASSERT_TRUE(segments.ok()) << segments.status().ToString(); - ASSERT_EQ(segments->size(), 1u); - EXPECT_EQ((*segments)[0].logical_offset, 0); - EXPECT_EQ((*segments)[0].physical_offset, 0); - EXPECT_EQ((*segments)[0].size, 50); + TF_ASSERT_OK_AND_ASSIGN(auto segments, ExpandPoolLiveSegments(pool)); + ASSERT_EQ(segments.size(), 1u); + EXPECT_EQ(segments[0].logical_offset, 0); + EXPECT_EQ(segments[0].physical_offset, 0); + EXPECT_EQ(segments[0].size, 50); pool.regions[1].offset_bytes = 10; EXPECT_THAT(ExpandPoolLiveSegments(pool).status(), @@ -287,19 +281,16 @@ TEST(PoolLiveSegmentsTest, CoalescesAbuttingRegionsAndRejectsOverlap) { } TEST(PhysicalLiveRangeToLogicalTest, MapsWithinRunsAndRejectsPaddingCross) { - absl::StatusOr> segments = - ExpandPoolLiveSegments(ValidPool()); - ASSERT_TRUE(segments.ok()) << segments.status().ToString(); + TF_ASSERT_OK_AND_ASSIGN(auto segments, ExpandPoolLiveSegments(ValidPool())); - absl::StatusOr> range = - PhysicalLiveRangeToLogical(*segments, /*physical_offset=*/105, - /*size=*/10); - ASSERT_TRUE(range.ok()) << range.status().ToString(); - EXPECT_EQ(range->first, 25); - EXPECT_EQ(range->second, 35); + TF_ASSERT_OK_AND_ASSIGN( + auto range, PhysicalLiveRangeToLogical(segments, /*physical_offset=*/105, + /*size=*/10)); + EXPECT_EQ(range.first, 25); + EXPECT_EQ(range.second, 35); EXPECT_THAT( - PhysicalLiveRangeToLogical(*segments, /*physical_offset=*/15, /*size=*/10) + PhysicalLiveRangeToLogical(segments, /*physical_offset=*/15, /*size=*/10) .status(), StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("crosses padding"))); diff --git a/tpu_sync/kv_cache/pool_staging_test.cc b/tpu_sync/kv_cache/pool_staging_test.cc index 873ea408..a0d490aa 100644 --- a/tpu_sync/kv_cache/pool_staging_test.cc +++ b/tpu_sync/kv_cache/pool_staging_test.cc @@ -23,9 +23,12 @@ #include #include +#include #include #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/time/time.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/kv_cache/kv_cache_manager_base.h" #include "tpu_sync/kv_cache/pool_layout.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -35,6 +38,8 @@ namespace tpu_raiden { namespace kv_cache { namespace { +using ::absl_testing::StatusIs; + class StagingTestManager : public KVCacheManagerBase { public: StagingTestManager(size_t num_layers, size_t num_shards, @@ -118,7 +123,7 @@ TEST(PoolStagingTest, BoundedArenaLeasesAndAddressing) { manager.RegisterPools({DensePool("fa", 0, 0, kStride, kNumBlocks, /*staging_blocks_per_request=*/2)}, /*staging_leases=*/2); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_TRUE(manager.PoolStorageStagingBounded(0)); const size_t host_size = manager.GetHostSize(/*layer_idx=*/0, 0); EXPECT_GE(host_size, static_cast(4 * kStride)); @@ -134,8 +139,8 @@ TEST(PoolStagingTest, BoundedArenaLeasesAndAddressing) { static_cast(4 * kStride)); // No standing host residency on a bounded storage. - EXPECT_EQ(manager.GetPoolBlockRef(0, 0, 5).status().code(), - absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(manager.GetPoolBlockRef(0, 0, 5), + StatusIs(absl::StatusCode::kFailedPrecondition)); EXPECT_EQ(manager.GetBlockHostPointer(/*layer_idx=*/0, 0, /*block_id=*/5), nullptr); @@ -143,11 +148,11 @@ TEST(PoolStagingTest, BoundedArenaLeasesAndAddressing) { status = manager.AcquirePoolStagingLease(/*uuid=*/7, /*storage_idx=*/0, std::vector{5, 9}, absl::Milliseconds(50)); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); // Re-acquiring the same ids is a no-op; adding one more takes slot 2. status = manager.AcquirePoolStagingLease(7, 0, std::vector{9, 11}, absl::Milliseconds(50)); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_EQ(manager.PoolStagingSummary()[0].free_slots, 1); // Receiver-side chunk resolution lands dst block 9 at its slot, not at @@ -164,7 +169,7 @@ TEST(PoolStagingTest, BoundedArenaLeasesAndAddressing) { entry->set_src_offset_bytes(0); entry->set_size_bytes(16); status = manager.RegisterActivePlan(7, request, /*is_sender=*/false); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); uint8_t* host_base = manager.GetHostPointer(/*layer_idx=*/0, 0); std::vector chunks = manager.GetBlockChunks( /*layer_idx=*/0, /*shard_idx=*/0, std::vector{9}, @@ -181,21 +186,20 @@ TEST(PoolStagingTest, BoundedArenaLeasesAndAddressing) { // Another transfer needing more slots than are free waits, then fails. status = manager.AcquirePoolStagingLease( /*uuid=*/8, 0, std::vector{1, 2}, absl::Milliseconds(20)); - EXPECT_EQ(status.code(), absl::StatusCode::kResourceExhausted) - << status.ToString(); + EXPECT_THAT(status, StatusIs(absl::StatusCode::kResourceExhausted)); // More blocks than the whole arena is rejected outright. status = manager.AcquirePoolStagingLease(/*uuid=*/9, 0, std::vector{1, 2, 3, 4, 6}, absl::Milliseconds(20)); - EXPECT_EQ(status.code(), absl::StatusCode::kResourceExhausted); + EXPECT_THAT(status, StatusIs(absl::StatusCode::kResourceExhausted)); // Releasing uuid 7 returns its three slots; uuid 8 now fits. - ASSERT_TRUE(manager.UnregisterActivePlan(7).ok()); + ABSL_ASSERT_OK(manager.UnregisterActivePlan(7)); manager.ReleasePoolStagingLeases(7); EXPECT_EQ(manager.PoolStagingSummary()[0].free_slots, 4); status = manager.AcquirePoolStagingLease(8, 0, std::vector{1, 2}, absl::Milliseconds(20)); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_EQ(manager.PoolStagingSummary()[0].free_slots, 2); manager.ReleasePoolStagingLeases(8); // Releasing an unknown uuid is a no-op. @@ -217,18 +221,15 @@ TEST(PoolStagingTest, FallsBackToFullMirrorWithoutHintsOrLeases) { variant == 0 ? 0 : 2); absl::Status status = manager.RegisterPools({pool}, /*staging_leases=*/variant == 0 ? 2 : 0); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_FALSE(manager.PoolStorageStagingBounded(0)); EXPECT_GE(manager.GetHostSize(0, 0), static_cast(kStride * kNumBlocks)); - auto ref = manager.GetPoolBlockRef(0, 0, 9); - ASSERT_TRUE(ref.ok()) << ref.status().ToString(); - EXPECT_EQ(ref->ptr, manager.GetHostPointer(0, 0) + 9 * kStride); + TF_ASSERT_OK_AND_ASSIGN(auto ref, manager.GetPoolBlockRef(0, 0, 9)); + EXPECT_EQ(ref.ptr, manager.GetHostPointer(0, 0) + 9 * kStride); // Leases are no-ops on unbounded storages. - EXPECT_TRUE(manager - .AcquirePoolStagingLease(1, 0, std::vector{9}, - absl::Milliseconds(1)) - .ok()); + ABSL_EXPECT_OK(manager.AcquirePoolStagingLease( + 1, 0, std::vector{9}, absl::Milliseconds(1))); EXPECT_FALSE(manager.PoolStagingSummary()[0].bounded); } } @@ -245,7 +246,7 @@ TEST(PoolStagingTest, SmallPoolIdentitySharedStorageAndStrideMismatch) { absl::Status status = small.RegisterPools({DensePool("fa", 0, 0, kStride, /*num_blocks=*/4, 2)}, /*staging_leases=*/2); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_FALSE(small.PoolStorageStagingBounded(0)); StagingTestManager shared(/*num_layers=*/1, /*num_shards=*/1, @@ -255,14 +256,11 @@ TEST(PoolStagingTest, SmallPoolIdentitySharedStorageAndStrideMismatch) { {StridedPool("gdn.conv", 0, /*base_offset=*/0, kStride, 32, 1), StridedPool("gdn.ssm", 0, /*base_offset=*/16, kStride, 32, 1)}, /*staging_leases=*/3); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_TRUE(shared.PoolStorageStagingBounded(0)); EXPECT_EQ(shared.PoolStagingSummary()[0].num_slots, 3); - ASSERT_TRUE(shared - .AcquirePoolStagingLease(/*uuid=*/5, 0, - std::vector{20}, - absl::Milliseconds(10)) - .ok()); + ABSL_ASSERT_OK(shared.AcquirePoolStagingLease( + /*uuid=*/5, 0, std::vector{20}, absl::Milliseconds(10))); EXPECT_EQ(shared.PoolStagingSummary()[0].free_slots, 2); // Both pools of the storage address device page 20 through the same slot // (slot 0), each at its own base offset. @@ -279,7 +277,7 @@ TEST(PoolStagingTest, SmallPoolIdentitySharedStorageAndStrideMismatch) { entry->set_size_bytes(16); entry->set_layer_idx(pool_idx); } - ASSERT_TRUE(shared.RegisterActivePlan(5, request, /*is_sender=*/false).ok()); + ABSL_ASSERT_OK(shared.RegisterActivePlan(5, request, /*is_sender=*/false)); uint8_t* base = shared.GetHostPointer(0, 0); auto conv_chunks = shared.GetBlockChunks(/*layer_idx=*/0, 0, std::vector{20}, 16, 5, @@ -298,7 +296,7 @@ TEST(PoolStagingTest, SmallPoolIdentitySharedStorageAndStrideMismatch) { status = mixed.RegisterPools({DensePool("a", 0, 0, kStride, 32, 1), DensePool("b", 0, 0, kStride / 2, 64, 1)}, /*staging_leases=*/3); - ASSERT_TRUE(status.ok()) << status.ToString(); + ABSL_ASSERT_OK(status); EXPECT_FALSE(mixed.PoolStorageStagingBounded(0)); } diff --git a/tpu_sync/kv_cache/store_monitor_test.cc b/tpu_sync/kv_cache/store_monitor_test.cc index d2b58330..dfa9209e 100644 --- a/tpu_sync/kv_cache/store_monitor_test.cc +++ b/tpu_sync/kv_cache/store_monitor_test.cc @@ -22,10 +22,12 @@ #include #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/time/clock.h" #include "absl/time/time.h" +#include "xla/tsl/platform/statusor.h" #include "tpu_sync/common/raiden_id.h" #include "tpu_sync/kv_cache/global_registry/global_registry.pb.h" #include "tpu_sync/kv_cache/global_registry/global_registry_client.h" @@ -61,12 +63,10 @@ class StoreMonitorTest : public ::testing::Test { TEST_F(StoreMonitorTest, HeartbeatsKeepTheRegistrationAlive) { RaidenId id = {"monitored", "r0", "dataS", 0}; - ASSERT_TRUE(client_ - ->RegisterStore(id, "10.0.0.7:1111", - /*controller_address=*/"", - /*ttl=*/absl::Seconds(2), "groupA", - /*evict_tier=*/1) - .ok()); + ABSL_ASSERT_OK(client_->RegisterStore(id, "10.0.0.7:1111", + /*controller_address=*/"", + /*ttl=*/absl::Seconds(2), "groupA", + /*evict_tier=*/1)); std::atomic reregister_calls{0}; StoreMonitor monitor( @@ -80,7 +80,7 @@ TEST_F(StoreMonitorTest, HeartbeatsKeepTheRegistrationAlive) { // Well past the original 2s TTL: only the heartbeats keep it alive. absl::SleepFor(absl::Seconds(3)); - EXPECT_TRUE(client_->ResolveStore(id).ok()); + ABSL_EXPECT_OK(client_->ResolveStore(id)); // Alive the whole time, so the heartbeats never needed the re-register // fallback. EXPECT_EQ(reregister_calls.load(), 0); @@ -112,18 +112,16 @@ TEST_F(StoreMonitorTest, ReregistersWhenTheRegistrationLapses) { absl::SleepFor(absl::Seconds(1)); EXPECT_GE(reregister_calls.load(), 1); - EXPECT_TRUE(client_->ResolveStore(id).ok()); + ABSL_EXPECT_OK(client_->ResolveStore(id)); } TEST_F(StoreMonitorTest, ReportedStatusFeedsThePlacementRanking) { auto register_store = [&](absl::string_view job, int32_t tier) { RaidenId id = {std::string(job), "r0", "dataS", 0}; - EXPECT_TRUE(client_ - ->RegisterStore(id, absl::StrCat(job, ":1111"), - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration(), "groupA", - tier) - .ok()); + ABSL_EXPECT_OK(client_->RegisterStore(id, absl::StrCat(job, ":1111"), + /*controller_address=*/"", + /*ttl=*/absl::ZeroDuration(), + "groupA", tier)); return id; }; RaidenId caller = register_store("caller", 0); @@ -132,7 +130,7 @@ TEST_F(StoreMonitorTest, ReportedStatusFeedsThePlacementRanking) { StoreStatus crowded_status; crowded_status.set_free_blocks(10); - ASSERT_TRUE(client_->Heartbeat(crowded, crowded_status).ok()); + ABSL_ASSERT_OK(client_->Heartbeat(crowded, crowded_status)); StoreMonitor monitor( StoreMonitor::Options{.heartbeat_period = absl::Milliseconds(300)}, @@ -141,11 +139,11 @@ TEST_F(StoreMonitorTest, ReportedStatusFeedsThePlacementRanking) { monitor.Start(); absl::SleepFor(absl::Seconds(1)); // At least one heartbeat. - auto targets = client_->GetPlacementTargets(caller, /*max_targets=*/8); - ASSERT_TRUE(targets.ok()); - ASSERT_EQ(targets->size(), 2); - EXPECT_EQ((*targets)[0].raiden_id().job_name(), "roomy"); - EXPECT_EQ((*targets)[1].raiden_id().job_name(), "crowded"); + TF_ASSERT_OK_AND_ASSIGN( + auto targets, client_->GetPlacementTargets(caller, /*max_targets=*/8)); + ASSERT_EQ(targets.size(), 2); + EXPECT_EQ(targets[0].raiden_id().job_name(), "roomy"); + EXPECT_EQ(targets[1].raiden_id().job_name(), "crowded"); } TEST_F(StoreMonitorTest, SweepRunsOnItsPeriod) { @@ -249,11 +247,9 @@ TEST_F(StoreMonitorTest, ASuccessfulReregisterTriggersTheRepublish) { TEST_F(StoreMonitorTest, ARequestStartsARepublishWithoutALapse) { RaidenId id = {"monitored", "r0", "dataS", 0}; - ASSERT_TRUE(client_ - ->RegisterStore(id, "10.0.0.7:1111", - /*controller_address=*/"", - /*ttl=*/absl::ZeroDuration()) - .ok()); + ABSL_ASSERT_OK(client_->RegisterStore(id, "10.0.0.7:1111", + /*controller_address=*/"", + /*ttl=*/absl::ZeroDuration())); std::atomic republish_calls{0}; StoreMonitor monitor( StoreMonitor::Options{.heartbeat_period = absl::Hours(1)}, client_, id,