From 4350cf566ed7fea4c0c326ccc9088df1761e68d5 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:47:54 +0000 Subject: [PATCH 1/3] [TRTLLM-17105][test] Bounce gtests: do not start the NIXL listen thread The bounce gtests wire agents with loadRemoteAgent(AgentDesc) and never use the NIXL listen thread. With it on, NixlTransferAgent locks /tmp/trtllm_nixl_port.lock to pick a port. On a shared machine that file may belong to another user (root is locked out too under fs.protected_regular=2), so the agent fails to build and every transport test SKIPs with "NIXL agent/backend unavailable", hiding regressions. Build every bounce test agent with useListenThread=false, through one shared bounce_test::makeAgentConfig that makeNode and the agent E2E test both use. Production is unchanged: the legacy C++ transceiver still uses the listen thread, and the Python transceiver already passes use_listen_thread=False. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../executor/bounce/bounceAgentE2ETest.cpp | 5 +++-- .../unit_tests/executor/bounce/bounceTestNixlNode.h | 13 ++++++++++++- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp index 7fa215f170e6..58ce2fd966a2 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp @@ -41,6 +41,7 @@ namespace kvc = tensorrt_llm::executor::kv_cache; using bounce_test::freeXferBufs; using bounce_test::hasCuda; +using bounce_test::makeAgentConfig; using bounce_test::makeXferBufs; using bounce_test::verifyXferBufs; using bounce_test::XferBufs; @@ -108,7 +109,7 @@ int runConcurrentFlows( kvc::BaseAgentConfig makeBounceConfig( std::string name, std::size_t arenaSizeMb = 2, char const* granularityBytes = "256") { - kvc::BaseAgentConfig cfg{std::move(name), true, false, true}; + auto cfg = makeAgentConfig(std::move(name)); cfg.agentBufferSizeMb = arenaSizeMb; if (arenaSizeMb > 0) { @@ -212,7 +213,7 @@ TEST(BounceAgentE2E, EffectiveChunkCapFallsBackToStandardNixl) // cap. The descriptor below sits between the two limits. auto makeCapConfig = [](char const* name) { - kvc::BaseAgentConfig cfg{name, true, false, true}; + auto cfg = makeAgentConfig(name); cfg.agentBufferSizeMb = 3; cfg.bounceParams = { {"min_descriptor_count", "1"}, diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h b/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h index 2f4dc8c74a53..e5153f8b4cf0 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h +++ b/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h @@ -41,6 +41,7 @@ #include #include #include +#include #include namespace bounce_test @@ -165,6 +166,16 @@ struct Node } }; +// Config for every agent the bounce tests create: no NIXL listen thread, because the tests exchange +// AgentDescs directly. With the listen thread on, the agent locks /tmp/trtllm_nixl_port.lock to pick a +// port; on a shared machine that file may belong to another user, so the agent fails to build and +// every transport test silently SKIPs. +inline kvc::BaseAgentConfig makeAgentConfig(std::string name) +{ + return kvc::BaseAgentConfig{ + std::move(name), /*useProgThread=*/true, /*multiThread=*/false, /*useListenThread=*/false}; +} + // Build one full bounce node (agent+arena+exec+channel+transport). Returns nullptr if the NIXL // agent/backend can't init or the arena can't be registered (caller GTEST_SKIPs). `cfg` supplies // arenaSizeBytes/arenaAllocationGranularityBytes/maxInflightChunksPerRequest etc., so callers @@ -177,7 +188,7 @@ inline std::unique_ptr makeNode(std::string const& name, b::BounceConfig c n->name = name; try { - kvc::BaseAgentConfig c{name, /*useProgThread=*/true, /*multiThread=*/false, /*useListenThread=*/true}; + auto const c = makeAgentConfig(name); n->agent = makeAgent ? makeAgent(c) : std::make_unique(c); } catch (std::exception const&) From e4b2ca1d3a503c995a84369d3d0a5be924a29e49 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:59:05 +0000 Subject: [PATCH 2/3] [TRTLLM-17105][fix] Report a bounce failure only after the receiver has answered every sent chunk Only the receiver's scatter of a received DATA writes the destination, and the receiver answers (ACK or NACK) after that scatter's last write. A failed request used to resolve kFAILURE at once, even while chunks whose DATA was already sent were queued or running on the receiver, so a caller that reused the destination could have it overwritten. Contract: BounceResult gains `quiesced`. A failed request whose sent chunks are not all answered now waits for those answers: it posts, gathers and consumes credits no more, releases its staging and sends its cancel as before, and resolves FAILURE with quiesced=true once the last answer arrives. Answers match by (peer, rid, chunk, region), and failRequest is idempotent, so a later forgetPeer keeps the first reason and the answer deadline. The receiver drains the exec stream before NACKing a failed launch, so a NACK never precedes a write. In-doubt: with a sent chunk still unanswered, a receiver silent for requestTimeoutMs after the failure or its last answer (the dead-peer assumption the lease and quarantine already make), a kNoProgressTimeout, or shutdown resolves the request with quiesced=false and a WARNING. BounceTransferStatus moves to a header-only file; isCompleted() is true on SUCCESS and on a quiesced FAILURE, and the status string marks an in-doubt one. The bounce reactor gtests get a 300 s per-test ctest TIMEOUT, so a hang fails fast under its own name instead of consuming the 2700 s stage budget. The success path is unchanged: no new messages, locks, clock reads or allocations. No wire or Python change. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../tensorrt_llm/executor/transferAgent.h | 8 +- .../nixl_utils/bounce/BounceConfig.h | 4 +- .../nixl_utils/bounce/BounceTransferStatus.h | 103 ++ .../nixl_utils/bounce/BounceTransport.cpp | 234 +++- .../nixl_utils/bounce/BounceTransport.h | 73 +- .../nixl_utils/transferAgent.cpp | 57 +- .../unit_tests/executor/bounce/CMakeLists.txt | 15 +- .../executor/bounce/bounceAgentE2ETest.cpp | 40 + .../bounce/bounceTransportFailureTest.cpp | 1024 ++++++++++++++++- 9 files changed, 1415 insertions(+), 143 deletions(-) create mode 100644 cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 247b7c0c5c06..73f672e26507 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -374,9 +374,11 @@ class TransferStatus { public: virtual ~TransferStatus() = default; - /// NIXL and bounce statuses return true only on success; backends that cannot distinguish - /// failure from completion (mooncake) return true on any terminal state — use - /// wait()/getLastStatus for the outcome. + /// NIXL statuses return true only on success. Bounce statuses return true on success and on a + /// failure proven never to write the destination again (a source read by an orphaned local + /// gather may still be in flight); a failure without that proof stays false. Backends that + /// cannot distinguish failure from completion (mooncake) return true on any terminal state — + /// use wait() for the outcome. [[nodiscard]] virtual bool isCompleted() const = 0; virtual TransferState wait(int64_t timeout_ms = -1) const = 0; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h index f5acdefb0abe..bd5acabd628e 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h @@ -166,7 +166,9 @@ struct BounceConfig int requestTimeoutMs{30000}; // request_timeout_ms; must be > 0 — the whole failure model // (abandoned-flow resolution, receiver lease, quarantine) hangs off // this timer, so applyParam rejects 0 (negatives already fail the - // strict unsigned parse) and keeps the current value. + // strict unsigned parse) and keeps the current value. It also bounds + // how long a failed request waits for the receiver's answers (after + // the failure, and again after each answer) before resolving in-doubt. // Receiver-side lease on granted regions — DERIVED in deriveDependentTimeouts() as // 2 x requestTimeoutMs, not an independent knob (tests may still set the field directly). A // dead sender emits neither DATA nor a cancel, which is unobservable through the protocol diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h new file mode 100644 index 000000000000..3724e081ca20 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h" +#include "tensorrt_llm/executor/transferAgent.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +/// TransferStatus over the bounce transport's completion future (the value Python's wait() reads). +/// Header-only so its contract can be unit-tested over hand-built futures. +class BounceTransferStatus final : public TransferStatus +{ +public: + explicit BounceTransferStatus(std::shared_future fut) + : mFut(std::move(fut)) + { + } + + [[nodiscard]] bool isCompleted() const override + { + // Terminal AND the destination is quiesced (BounceResult::quiesced): kSUCCESS, or a kFAILURE the + // receiver has fully answered; an in-doubt kFAILURE stays false. Callers read the outcome from + // wait()/getLastStatusStr. get() cannot throw: every erase/shutdown path resolves the promise + // with set_value first, so broken_promise is unreachable. + if (!isReady()) + { + return false; + } + auto const& result = mFut.get(); + return result.state == TransferState::kSUCCESS || result.quiesced; + } + + [[nodiscard]] TransferState wait(int64_t timeoutMs) const override + { + if (!mFut.valid()) + { + return TransferState::kFAILURE; + } + if (timeoutMs < 0) + { + return mFut.get().state; + } + if (mFut.wait_for(std::chrono::milliseconds(timeoutMs)) == std::future_status::ready) + { + return mFut.get().state; + } + return TransferState::kIN_PROGRESS; + } + + /// Failure cause recorded by the bounce transport (empty while in flight / on success). An + /// in-doubt failure says so, since only then may the destination still be written. + [[nodiscard]] std::string getLastStatusStr() const override + { + if (!isReady()) + { + return {}; + } + auto const& result = mFut.get(); + if (result.reason == BounceFailReason::kNone) + { + return {}; + } + std::string detail = toString(result.reason); + if (!result.quiesced) + { + detail += "; destination not proven quiescent"; + } + return detail; + } + +private: + [[nodiscard]] bool isReady() const + { + return mFut.valid() && mFut.wait_for(std::chrono::seconds(0)) == std::future_status::ready; + } + + std::shared_future mFut; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp index 5657b249019a..3e0f68085c09 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp @@ -684,8 +684,8 @@ void BounceReceiver::scatterWorkerLoop() // Scatter into the final dst, then wait so the data is at dst before we ACK. A bad source, a // failed launch, OR a stream error must NOT produce an ACK — an ACK tells the sender its KV // data landed, so a false ACK here is silent corruption. On error the worker sends a NACK - // instead (the sender fails immediately with kPeerNack) and the done queue still releases - // the region. + // instead (the sender fails the request with kPeerNack without waiting out requestTimeoutMs) + // and the done queue still releases the region. cudaError_t launchErr = cudaErrorInvalidValue; { // Prep leg: bounds-check + plan-array build + kernel launch (host-side cost). @@ -697,7 +697,7 @@ void BounceReceiver::scatterWorkerLoop() // flow's granted region), then fill the pinned plan buffers DIRECTLY (no intermediate // vectors). Piece counts come from the peer's DATA message — a peer built with a larger // maxChunkSizeBytes can carry more pieces than our pinned/scratch hold; reject rather than - // overflow (no launch -> NACK, the sender fails immediately; never a false ACK). + // overflow (no launch -> NACK, the sender fails the request; never a false ACK). std::size_t const maxEntries = maxPlanEntries(ctx); std::uint64_t rawPieces = 0; for (std::uint32_t i = 0; i < n; ++i) @@ -771,6 +771,15 @@ void BounceReceiver::scatterWorkerLoop() bool const ok = srcInBounds && launchErr == cudaSuccess && syncErr == cudaSuccess; if (!ok) { + if (launchErr != cudaSuccess) + { + // A failed launch can still have enqueued the scatter: launchBatchedCopy reports + // cudaGetLastError(), which may be a stale error from an earlier call on this thread + // (e.g. the warn-only cudaSetDevice above). The NACK below tells the sender this + // chunk will not write dst again, so drain the stream before sending it. (Also runs + // when nothing was launched — launchErr keeps its initial value — on an idle stream.) + (void) cudaStreamSynchronize(ctx->stream); + } // Clear any non-sticky error so the next launch does not inherit it; a sticky // (context-level) error cannot be cleared and will surface on the next CUDA call. (void) cudaGetLastError(); @@ -789,8 +798,8 @@ void BounceReceiver::scatterWorkerLoop() // the sender's ackWait critical path. Region bookkeeping (scheduler free / re-grant) still // goes through drainScatterDone on the IO thread. A failed scatter sends NO ACK — a false // ACK would tell the sender corrupt/absent data landed. Instead the worker sends a NACK - // (also covering the plan-overflow rejection above) so the sender fails now rather than at - // requestTimeoutMs. + // (also covering the plan-overflow rejection above) so the sender fails the request without + // waiting out requestTimeoutMs. if (ok) { BounceNvtxScope ackScope( @@ -851,13 +860,14 @@ std::shared_future BounceSender::submit( // silently absorbed as a slow success. Whether to retry is the caller's decision. TLLM_LOG_WARNING( "BounceTransport(%s): plan build for peer %s rejected: %s", mCtx.selfName.c_str(), peer.c_str(), e.what()); - promise->set_value({TransferState::kFAILURE, BounceFailReason::kPlanRejected}); + // Nothing was sent, so nothing can write the destination. + promise->set_value({TransferState::kFAILURE, BounceFailReason::kPlanRejected, /*quiesced=*/true}); return fut; } auto const numChunks = static_cast(plan.numChunks()); if (numChunks == 0) { - promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); + promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone, /*quiesced=*/true}); return fut; } @@ -931,6 +941,10 @@ void BounceSender::onGrant(std::string const& peer, BounceMsgHeader const& h, st mCtx.selfName.c_str(), peer.c_str(), req.peer.c_str(), static_cast(h.requestId)); return; } + if (req.failure) + { + return; // awaiting answers: nothing will be written; its cancel reclaims what this GRANT holds + } bounceRangeEnd(req.nvtxGrantWait); // first GRANT ends the credit-wait span (no-op on later GRANTs) bounceRangeEnd(req.nvtxCreditStarved); // a GRANT ends the current starvation period (pump may reopen one) for (auto const& credit : credits) @@ -1021,6 +1035,10 @@ void BounceSender::attachCredits(std::uint64_t rid, Request& req) void BounceSender::pumpRequest(std::uint64_t rid, Request& req) { + if (req.failure) + { + return; // awaiting answers: posted and pendingCredits stay empty + } if (req.abandonReason != BounceFailReason::kNone) { // Abandoned (GRANT mispair / plan overflow): stop attaching credits and launching gathers. @@ -1361,8 +1379,8 @@ void BounceSender::onNack(std::string const& peer, BounceMsgHeader const& h) auto it = mRequests.find(h.requestId); if (it == mRequests.end()) { - // Expected late arrival: the request already failed / was cancelled (e.g. a scatter NACK - // racing the timeout or a second chunk's NACK after the first one failed it). + // Expected late arrival: the request already resolved (e.g. an answer arriving after it + // resolved in-doubt on requestTimeoutMs, or a NACK racing kNoProgressTimeout). TLLM_LOG_DEBUG("BounceTransport(%s): dropping NACK for unknown rid=%llu peer=%s chunk=%u", mCtx.selfName.c_str(), static_cast(h.requestId), peer.c_str(), h.chunkIdx); return; @@ -1375,8 +1393,18 @@ void BounceSender::onNack(std::string const& peer, BounceMsgHeader const& h) h.chunkIdx); return; } - // Same terminal path as a failed RDMA write (pollSenderHandles): failRequest resolves the future, - // recycles/defers the staging regions and sends the cancel that reclaims the receiver's regions. + if (req.failure) + { + onAnswerAfterFailure(req, h); + return; + } + // The NACK answers its own chunk (the receiver never scattered it, or synced before NACKing), so + // recycle that chunk now: failRequest must not wait for an answer that already arrived. A + // request-level NACK (rejected WANT) arrives before any chunk was sent and matches none. + releaseAnsweredChunk(req, h); + // Same terminal path as a failed RDMA write (pollSenderHandles): failRequest recycles/defers the + // staging regions, sends the cancel that reclaims the receiver's regions and resolves the future + // once the receiver has answered every other sent chunk. failRequest(h.requestId, req, BounceFailReason::kPeerNack); } @@ -1401,22 +1429,14 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) h.chunkIdx); return; } - bool found = false; - for (auto pit = req.posted.begin(); pit != req.posted.end(); ++pit) + if (req.failure) { - // Only a Sent chunk has finished reading its local staging region and may recycle it on ACK. - if (pit->chunkIdx == h.chunkIdx && pit->remoteHandle == h.regionHandle && pit->state == PostState::Sent) - { - bounceRangeEnd(pit->nvtxAckWait); - // Return the gather-staging region to the shared arena; re-schedule may hand the freed - // bytes to a waiting remote flow. - mCtx.sendGrants(mCtx.scheduler.releaseLocal(pit->localOffset)); - req.posted.erase(pit); - found = true; - break; - } + // A failed request resolves FAILURE even if every chunk ends up ACKed (e.g. forgetPeer after + // all were Sent): never take the SUCCESS path below. + onAnswerAfterFailure(req, h); + return; } - if (!found) + if (!releaseAnsweredChunk(req, h)) { TLLM_LOG_WARNING("BounceTransport(%s): dropping stale/invalid ACK peer=%s rid=%llu chunk=%u region=%llu", mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, @@ -1429,23 +1449,54 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk was ACKed if (req.acked >= req.numChunks) { - bounceRangeEnd(req.nvtxGrantWait); - bounceRangeEnd(req.nvtxCreditStarved); - bounceRangeEnd(req.nvtxArenaStarved); - bounceRangeEnd(req.nvtxReq); - try - { - req.promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); - } - catch (...) - { - // set_value throws std::future_error ONLY if the promise is already satisfied — a benign - // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. - } + resolve(req, {TransferState::kSUCCESS, BounceFailReason::kNone, /*quiesced=*/true}); mRequests.erase(it); } } +bool BounceSender::releaseAnsweredChunk(Request& req, BounceMsgHeader const& h) +{ + // Only a Sent chunk can be answered: its DATA is out, and its write has finished reading the local + // staging region, so the answer may recycle that region. + auto const answered = std::find_if(req.posted.begin(), req.posted.end(), + [&h](Posted const& p) + { return p.chunkIdx == h.chunkIdx && p.remoteHandle == h.regionHandle && p.state == PostState::Sent; }); + if (answered == req.posted.end()) + { + return false; + } + bounceRangeEnd(answered->nvtxAckWait); + // Return the gather-staging region to the shared arena; re-schedule may hand the freed bytes to a + // waiting remote flow. + mCtx.sendGrants(mCtx.scheduler.releaseLocal(answered->localOffset)); + req.posted.erase(answered); + return true; +} + +void BounceSender::onAnswerAfterFailure(Request& req, BounceMsgHeader const& h) +{ + auto const answered = std::find_if(req.awaiting.begin(), req.awaiting.end(), + [&h](AnswerKey const& k) { return k.chunkIdx == h.chunkIdx && k.regionHandle == h.regionHandle; }); + if (answered == req.awaiting.end()) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): dropping stale/invalid answer for failed request peer=%s rid=%llu chunk=%u " + "region=%llu", + mCtx.selfName.c_str(), req.peer.c_str(), static_cast(h.requestId), h.chunkIdx, + static_cast(h.regionHandle)); + return; + } + // Its staging region was already released when the request failed; only the bookkeeping is left. + req.awaiting.erase(answered); + // A live receiver gets a fresh requestTimeoutMs for each remaining answer. + req.lastProgress = std::chrono::steady_clock::now(); + if (req.awaiting.empty()) + { + resolveFailure(h.requestId, req); + mRequests.erase(h.requestId); + } +} + void BounceSender::checkTimeouts() { if (mCtx.cfg.requestTimeoutMs <= 0) @@ -1470,12 +1521,21 @@ void BounceSender::checkTimeouts() for (auto rid : stuck) { auto it = mRequests.find(rid); - if (it != mRequests.end()) + if (it == mRequests.end()) { - // Peer never granted or stopped making progress (unreachable / not bounce-ready / - // congested). Fail the request so wait() returns FAILURE instead of hanging. - failRequest(rid, it->second, BounceFailReason::kNoProgressTimeout); + continue; } + Request& req = it->second; + if (req.failure) + { + // No answer for requestTimeoutMs while chunks are still awaited: in-doubt (see failRequest). + resolveFailure(rid, req); + mRequests.erase(it); + continue; + } + // Peer never granted or stopped making progress (unreachable / not bounce-ready / + // congested). Fail the request so wait() returns FAILURE instead of hanging. + failRequest(rid, req, BounceFailReason::kNoProgressTimeout); } } } @@ -1599,16 +1659,20 @@ bool BounceSender::drainOrphanGather() void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason reason) { - // An abandoned flow (GRANT mispair / plan overflow) reaches here through the timeout path; - // report the specific abandon cause, not the generic timeout. - if (req.abandonReason != BounceFailReason::kNone) + if (req.failure) { - reason = req.abandonReason; + // Already awaiting answers (forget() reaches here again): keep the first reason and the answer + // deadline; staging and the cancel were handled at the first failure. + return; } + // An abandoned flow (GRANT mispair / plan overflow) reaches here through the timeout path; + // report the specific abandon cause, not the generic timeout. + BounceFailReason const reported = req.abandonReason != BounceFailReason::kNone ? req.abandonReason : reason; + req.failure = reported; // The one log line every sender-side failure passes through (timeout, peer drop, write/gather // failure) — keep it WARNING and keep the progress context. TLLM_LOG_WARNING("BounceTransport(%s): request FAILED rid=%llu peer=%s reason=\"%s\" chunks acked=%u posted=%u/%u", - mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(reason), req.acked, + mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(reported), req.acked, req.nextPost, req.numChunks); // Release in-flight transfer handles and return each gather-staging region to the shared arena — // but only once nothing is still touching the region's memory, else recycling races a live DMA. @@ -1622,7 +1686,8 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason // cudaStreamQuery(ctx->stream) reports the stream idle. // - Gathered: gather done (event observed), ctx already returned -> recycle now. // - Sent: write landed (poll==kDone), xfer already released in pollSenderHandles, NIC done - // reading -> recycle now. + // reading -> recycle now. Its DATA is out, though: the receiver may still scatter it into the + // destination, so remember it as awaiting the receiver's answer. bool deferredWrite = false; for (auto& p : req.posted) { @@ -1644,12 +1709,21 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason p.ctx = nullptr; continue; // do NOT release ctx or recycle the region yet } + if (p.state == PostState::Sent) + { + req.awaiting.push_back(AnswerKey{p.chunkIdx, p.remoteHandle}); + } mCtx.sendGrants(mCtx.scheduler.releaseLocal(p.localOffset)); } + // Every chunk is now recycled, parked as an orphan or awaiting its answer, so the drain loops have + // nothing left to gather, post or poll for this request; its parked credits die with the cancel. + req.posted.clear(); + req.pendingCredits.clear(); // Retract the credit request so the receiver stops holding/granting for it. If any RDMA write is // still in flight it is landing on the receiver's region; sending the cancel NOW would let the // receiver reclaim+re-grant that region under the write -> corruption. Defer the cancel until the // flow's writes drain (drainOrphanLocal sends it). With no in-flight write, send it immediately. + // The cancel does not stop the receiver's queued scatters: they still run and answer. if (deferredWrite) { mPendingCancel[rid] = req.peer; @@ -1658,25 +1732,63 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason { mCtx.channel->sendTo(req.peer, encodeCancel(rid, mCtx.channel->localEndpoint())); } + // Resolve at once when nothing is awaited; on kNoProgressTimeout (checkTimeouts, abandoned flows + // included), which already saw no progress for requestTimeoutMs, so an unanswered chunk is reported + // in-doubt rather than waited for again (its DATA may be recent); and when requestTimeoutMs <= 0 + // (direct test configs only), which leaves no bound to wait under. This tests the caller's `reason`, + // not `reported`: an abandoned flow reports kProtocolError but arrives here from checkTimeouts. + bool const waitForAnswers + = !req.awaiting.empty() && reason != BounceFailReason::kNoProgressTimeout && mCtx.cfg.requestTimeoutMs > 0; + if (waitForAnswers) + { + // Resolving now would let the caller reuse the destination while the receiver may still + // scatter into it. lastProgress restarts here and at each answer; checkTimeouts resolves + // in-doubt after requestTimeoutMs without one. The pipeline is over: only nvtxReq stays open. + bounceRangeEnd(req.nvtxGrantWait); + bounceRangeEnd(req.nvtxCreditStarved); + bounceRangeEnd(req.nvtxArenaStarved); + req.lastProgress = std::chrono::steady_clock::now(); + return; + } + resolveFailure(rid, req); + mRequests.erase(rid); +} + +void BounceSender::resolveFailure(std::uint64_t rid, Request& req) +{ + // A sent chunk the receiver has not answered may still be scattered into the destination. + bool const quiesced = req.awaiting.empty(); + if (!quiesced) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): request rid=%llu peer=%s resolved FAILURE (\"%s\") with %zu sent chunk(s) " + "unanswered -> destination not proven quiescent (the receiver may still write it)", + mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(*req.failure), + req.awaiting.size()); + } + resolve(req, {TransferState::kFAILURE, *req.failure, quiesced}); +} + +void BounceSender::resolve(Request& req, BounceResult const& result) +{ bounceRangeEnd(req.nvtxGrantWait); bounceRangeEnd(req.nvtxCreditStarved); bounceRangeEnd(req.nvtxArenaStarved); bounceRangeEnd(req.nvtxReq); try { - req.promise->set_value({TransferState::kFAILURE, reason}); + req.promise->set_value(result); } catch (...) { // set_value throws std::future_error ONLY if the promise is already satisfied — a benign // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. } - mRequests.erase(rid); } void BounceSender::forget(std::string const& peer) { - // Fail any in-flight request targeting the gone peer so its wait() returns. + // Fail every request to the gone peer (failRequest leaves one already awaiting answers alone). std::lock_guard lk(mReqMu); std::vector toFail; for (auto const& [rid, req] : mRequests) @@ -1775,20 +1887,18 @@ void BounceSender::failAll() mCtx.exec->release(p.ctx); // still-gathering chunk: return its borrowed exec context p.ctx = nullptr; } + if (p.state == PostState::Sent) + { + req.awaiting.push_back(AnswerKey{p.chunkIdx, p.remoteHandle}); + } } - bounceRangeEnd(req.nvtxGrantWait); - bounceRangeEnd(req.nvtxCreditStarved); - bounceRangeEnd(req.nvtxArenaStarved); - bounceRangeEnd(req.nvtxReq); - try - { - req.promise->set_value({TransferState::kFAILURE, BounceFailReason::kShutdown}); - } - catch (...) + if (!req.failure) { - // set_value throws std::future_error ONLY if the promise is already satisfied — a benign - // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. + req.failure = BounceFailReason::kShutdown; // a request already awaiting answers keeps its reason } + // Any sent chunk still unanswered makes the result in-doubt: the caller's device sync drained our + // kernels, not the receiver's scatter. + resolveFailure(rid, req); } mRequests.clear(); // Release the handles of deferred writes (waited on above) without recycling their regions or diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h index 42f12ee0820c..b4fb96af7500 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h @@ -36,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -78,6 +79,7 @@ namespace tensorrt_llm::executor::kv_cache::bounce // availability have been observed. // [S] ACK(rid, chunk, h) -> releaseLocal() that chunk's region; acked++; // all acked -> resolve the request's future SUCCESS. +// (A request awaiting answers only matches them; see failRequest.) // [R] WANT(rid, sizes[]) -> scheduler.onWant(peer:rid, sizes) -> send GRANT(s). // [R] DATA(rid, chunk, h, plan) -> enqueue a scatter job over region `h`. // 2. poll every in-flight write (TransferStatus::wait(0)); on Done send DATA (data has @@ -94,6 +96,9 @@ namespace tensorrt_llm::executor::kv_cache::bounce // FAILURE on transfer error, peer invalidation, shutdown, or request timeout. requestTimeoutMs must // be > 0 (the config layer enforces this): abandoned-flow resolution and the receiver lease both // hang off this timer, so disabling it would let a protocol-error flow hold wait() forever. +// A FAILURE first waits for the receiver's answer to every chunk whose DATA was sent, for at most +// requestTimeoutMs after the failure or after the receiver's last answer, so that +// BounceResult::quiesced can promise the destination is no longer written. // // The agent (data plane) & ControlChannel are injected (not owned) so the same reactor runs // unchanged in tests and production (both over the real NIXL/zmq stack; failure tests subclass @@ -124,6 +129,16 @@ struct BounceResult { TransferState state{TransferState::kFAILURE}; BounceFailReason reason{BounceFailReason::kNone}; + // The receiver will not write the destination again for this request: always on kSUCCESS, and on + // kFAILURE once the receiver has answered (ACK or NACK) every chunk whose DATA was sent, because + // only that chunk's scatter writes the destination and the answer follows it. A failed request + // waits for those answers. It resolves in-doubt (false) when no answer comes within + // requestTimeoutMs after the failure or after the receiver's last answer (assumed dead, as the + // receiver lease and quarantine assume), on kNoProgressTimeout with a sent chunk unanswered (no + // progress for requestTimeoutMs, though the last DATA may be recent), and at shutdown with a sent + // chunk unanswered (see BounceSender::failRequest and failAll). Covers the destination only: a + // gather orphaned by the failure may still be reading the source. + bool quiesced{false}; }; /// Shared dependencies used by both roles. Holds the injected channel/agent/arena/exec (borrowed, @@ -285,6 +300,7 @@ class BounceSender void onAck(std::string const& peer, BounceMsgHeader const& h); /// The receiver reported it cannot complete a chunk/request: fail the request now (kPeerNack) /// instead of waiting out requestTimeoutMs. Unknown / wrong-peer NACKs are dropped like ACKs. + /// For a request that already failed, a NACK is just the answer it is waiting for. void onNack(std::string const& peer, BounceMsgHeader const& h); /// Issue the RDMA write for any chunk whose gather kernel has now completed (poll ctx->event, /// non-blocking). Decouples the IO thread from gather latency on a shared GPU. Returns true if it @@ -300,15 +316,18 @@ class BounceSender bool drainOrphanGather(); /// Retry parked credits for every request (called on the IO loop after ACKs free regions). void drainPendingPosts(); - /// Fail requests that have made no progress within requestTimeoutMs (e.g. peer never granted). + /// Fail requests that have made no progress within requestTimeoutMs (e.g. peer never granted), + /// and resolve in-doubt the failed requests whose receiver has sent no answer for that long. void checkTimeouts(); - /// A peer is gone: fail any in-flight request targeting it so its wait() returns. + /// A peer is gone: fail any in-flight request targeting it. Its wait() returns once its sent + /// chunks are answered (answers still arrive on our ROUTER) or the answer wait times out. void forget(std::string const& peer); - /// Shutdown: fail every still-pending request, wait (bounded) for in-flight outbound writes to - /// reach a terminal state so the NIC is done reading the arena before it is freed, then release - /// each write's TransferStatus handle and return the parked gathers' exec contexts. Recycles no - /// regions (all are left to arena destruction) and sends no grants. Called after the IO/workers - /// are joined and local CUDA work has been synced. + /// Shutdown: fail every still-pending request (a request already awaiting answers keeps its reason; + /// any request with an unanswered sent chunk resolves in-doubt), wait (bounded) for in-flight + /// outbound writes to reach a terminal state so the NIC is done reading the arena before it is + /// freed, then release each write's TransferStatus handle and return the parked gathers' exec + /// contexts. Recycles no regions (all are left to arena destruction) and sends no grants. Called + /// after the IO/workers are joined and local CUDA work has been synced. void failAll(); /// True while a local gather region is held or an orphaned write/gather is still in flight @@ -360,6 +379,14 @@ class BounceSender std::uint64_t nvtxAckWait{0}; // DATA sent -> ACK received }; + // A chunk whose DATA was sent, identified the way the receiver's ACK/NACK names it (the rid and + // peer are the owning request's). + struct AnswerKey + { + std::uint32_t chunkIdx{}; + std::uint64_t regionHandle{}; + }; + struct Request { std::string peer; @@ -383,9 +410,14 @@ class BounceSender // Set by the abandon sites (GRANT mispair / plan overflow): the request then dies through // the timeout path, but failRequest reports THIS more specific cause instead. BounceFailReason abandonReason{BounceFailReason::kNone}; - // Last time this request made forward progress (granted+posted a chunk, or got an ACK). - // A request stuck with no progress for requestTimeoutMs (e.g. the peer never GRANTs because - // it is unreachable / not bounce-ready) is failed rather than hanging wait() forever. + // Set once the request has failed. A failed request stays in mRequests only while awaiting + // answers, and then it only collects them: no more posts, gathers or credits. + std::optional failure; + std::vector awaiting; // sent chunks the receiver has not answered yet (once failed) + // Last time this request made forward progress (granted+posted a chunk, or got an ACK; for a + // failed request: the failure, or an answer). A request stuck with no progress for + // requestTimeoutMs (e.g. the peer never GRANTs because it is unreachable / not bounce-ready) + // is failed rather than hanging wait() forever. std::chrono::steady_clock::time_point lastProgress; // Cross-thread NVTX spans (0 == none; ended on failure paths too). std::uint64_t nvtxReq{0}; // submit -> promise resolved @@ -445,10 +477,25 @@ class BounceSender /// (rest parked/retried). /// Caller MUST hold mReqMu. Runs attachCredits() first so credit pairing stays in chunk order. void pumpRequest(std::uint64_t rid, Request& req); - /// Resolve `rid` to kFAILURE, recording `reason` (unless a more specific one was tagged earlier) - /// and logging ONE warning with the request's progress context — the single choke point that - /// keeps every failure mode observable. See the body for the region-recycling rules. + /// Fail `rid` with `reason` (unless a more specific one was tagged earlier), logging ONE warning + /// with the request's progress context — the single choke point that keeps every failure mode + /// observable. See the body for the region-recycling rules and for when it resolves at once. + /// Otherwise its sent chunks move to `awaiting` and it resolves once the receiver has answered them + /// all (quiesced) or has sent no answer for requestTimeoutMs (in-doubt; checkTimeouts). Idempotent: + /// a later failure (forget) keeps the first reason and answer deadline. void failRequest(std::uint64_t rid, Request& req, BounceFailReason reason); + /// Recycle the Sent chunk an ACK/NACK answers (same chunk index and region handle): end its + /// ack-wait span, return its staging region and drop it from `posted`. False if none matches. + bool releaseAnsweredChunk(Request& req, BounceMsgHeader const& h); + /// An ACK or NACK for a request awaiting answers: it only tells us that chunk's scatter will not + /// write the destination again. Once every sent chunk is answered, resolves the request kFAILURE + /// (quiesced) and erases it. + void onAnswerAfterFailure(Request& req, BounceMsgHeader const& h); + /// Resolve a failed request with its `failure` reason: quiesced iff `awaiting` is empty, else + /// in-doubt with a WARNING. The caller erases the request. + void resolveFailure(std::uint64_t rid, Request& req); + /// Deliver the request's one result: end its request-level NVTX spans and set the promise. + static void resolve(Request& req, BounceResult const& result); BounceContext& mCtx; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 61633d141244..6e935fd7b703 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -46,6 +46,7 @@ #ifdef TLLM_BOUNCE_V2 #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.h" @@ -88,60 +89,6 @@ struct NixlBounceState }; } // namespace bounce -namespace -{ -/// TransferStatus over the bounce transport's completion future (the value Python's wait() reads). -class BounceTransferStatus final : public TransferStatus -{ -public: - explicit BounceTransferStatus(std::shared_future fut) - : mFut(std::move(fut)) - { - } - - [[nodiscard]] bool isCompleted() const override - { - // Success only (TransferStatus contract, mirrors NixlTransferStatus): a failed transfer never - // reports completed; callers poll wait()/getLastStatusStr for the outcome. get() cannot throw: - // the promise is only ever resolved with set_value, and every erase/shutdown path resolves it - // first, so broken_promise is unreachable. - return mFut.valid() && mFut.wait_for(std::chrono::seconds(0)) == std::future_status::ready - && mFut.get().state == TransferState::kSUCCESS; - } - - [[nodiscard]] TransferState wait(int64_t timeoutMs) const override - { - if (!mFut.valid()) - { - return TransferState::kFAILURE; - } - if (timeoutMs < 0) - { - return mFut.get().state; - } - if (mFut.wait_for(std::chrono::milliseconds(timeoutMs)) == std::future_status::ready) - { - return mFut.get().state; - } - return TransferState::kIN_PROGRESS; - } - - /// Failure cause recorded by the bounce transport (empty while in flight / on success). - [[nodiscard]] std::string getLastStatusStr() const override - { - if (!mFut.valid() || mFut.wait_for(std::chrono::seconds(0)) != std::future_status::ready) - { - return {}; - } - auto const& result = mFut.get(); - return result.reason == bounce::BounceFailReason::kNone ? std::string{} : bounce::toString(result.reason); - } - -private: - std::shared_future mFut; -}; -} // namespace - void NixlTransferAgent::maybeInitBounce( std::size_t agentBufferSizeMb, std::unordered_map const& bounceParams) { @@ -1049,7 +996,7 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) TLLM_LOG_DEBUG("NixlTransferAgent(%s): bounce path engaged for write to %s (%zu descs)", mName.c_str(), request.getRemoteName().c_str(), request.getSrcDescs().getDescs().size()); auto fut = mBounce->transport->submit(request.getSrcDescs(), request.getDstDescs(), request.getRemoteName()); - return std::make_unique(std::move(fut)); + return std::make_unique(std::move(fut)); } #endif diff --git a/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt b/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt index ceaa55d53818..8b6a6e835527 100644 --- a/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt +++ b/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt @@ -150,11 +150,20 @@ if(NIXL_ROOT) ${NIXL_BUILD_LIBRARY} PkgConfig::TLLM_BOUNCE_ZMQ) # One reactor test target: the bounce sources compiled in, plus the shared - # include/link sets. + # include/link sets. Each test case gets a 300 s ctest timeout (each takes + # seconds), so a test that hangs (e.g. on a closed stream gate) fails under + # its own name within minutes instead of at the stage's time limit. function(add_bounce_reactor_test name src) cmake_parse_arguments(ARG "INSTRUMENTED" "" "" ${ARGN}) - add_tllm_gtest(${name} SOURCES ${src} FULL_STACK LABELS - "unit;full-stack;gpu") + add_tllm_gtest( + ${name} + SOURCES + ${src} + FULL_STACK + TIMEOUT + 300 + LABELS + "unit;full-stack;gpu") if(ARG_INSTRUMENTED) target_sources(${name} PRIVATE ${BOUNCE_TSAN_SRCS}) else() diff --git a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp index 58ce2fd966a2..1e7cb13e2e29 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp @@ -194,6 +194,7 @@ TEST(BounceAgentE2E, SubmitTransferRequestsUsesBounce) ASSERT_NE(status, nullptr); EXPECT_EQ(waitTerminal(status, 30), kvc::TransferState::kSUCCESS) << "bounce transfer via submitTransferRequests did not succeed"; + EXPECT_TRUE(status->isCompleted()) << "a SUCCESS must report the destination quiesced"; EXPECT_TRUE(verifyXferBufs(bufs)); a->shutdown(); @@ -201,6 +202,45 @@ TEST(BounceAgentE2E, SubmitTransferRequestsUsesBounce) freeXferBufs(bufs); } +// A bounce FAILURE that sent no DATA (the peer never grants) is proven not to write the destination, +// so the agent's status reports it completed (isCompleted) while wait() still reports FAILURE. +TEST(BounceAgentE2E, FailureWithNothingSentIsCompleted) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + auto makeShortTimeoutConfig = [](char const* name) + { + auto cfg = makeBounceConfig(name); + cfg.bounceParams["request_timeout_ms"] = "500"; // on both sides: the handshake requires equal values + return cfg; + }; + std::string skipMsg; + auto a = tryMakeAgent(makeShortTimeoutConfig("nsAgentA"), skipMsg); + auto b = a ? tryMakeAgent(makeShortTimeoutConfig("nsAgentB"), skipMsg) : nullptr; + if (!a || !b) + { + GTEST_SKIP() << "NIXL agent/backend unavailable: " << skipMsg; + } + a->loadRemoteAgent("nsAgentB", b->getLocalAgentDesc()); + b->shutdown(); // A keeps B's validated handshake, but nothing will ever GRANT its WANT + + auto bufs = makeXferBufs(/*nDescs=*/24, /*descBytes=*/600, /*seed=*/9); + auto req = makeReq(bufs, "nsAgentB"); + auto status = a->submitTransferRequests(req); + ASSERT_NE(status, nullptr); + EXPECT_EQ(a->getBounceSubmitCount(), 1u) << "the request must take the bounce path"; + EXPECT_EQ(waitTerminal(status, 30), kvc::TransferState::kFAILURE); + EXPECT_TRUE(status->isCompleted()) << "nothing was sent, so the destination is quiescent"; + auto const detail = status->getLastStatusStr(); + EXPECT_NE(detail.find("no GRANT/ACK progress"), std::string::npos) << detail; + EXPECT_EQ(detail.find("not proven quiescent"), std::string::npos) << detail; + + a->shutdown(); + freeXferBufs(bufs); +} + // A non-power-of-two arena can clamp the effective chunk cap below the configured cap. A descriptor // between those limits must use ordinary NIXL instead of entering bounce and failing plan creation. TEST(BounceAgentE2E, EffectiveChunkCapFallsBackToStandardNixl) diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp index 6b90b9dd7da2..fe60c05ab86d 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp @@ -19,9 +19,9 @@ // (SUCCESS or FAILURE) and never hang — including when the peer never grants, the transfer engine // fails, forgetPeer() drops a peer mid-transfer, or shutdown races in-flight requests. Plus a // multi-threaded concurrent-submit test. -// The data plane is real NIXL (via bounceTestNixlNode helpers); the one exception is the -// transfer-failure path, which uses a FakeXferAgent (a NixlTransferAgent subclass overriding the -// postXferRequest primitive) because real NIXL cannot be made to deterministically fail a write. +// The data plane is real NIXL (via bounceTestNixlNode helpers), except where a test must script +// individual writes (complete / stall / fail): those use a FakeXferAgent (a NixlTransferAgent subclass +// overriding postXferRequest), because real NIXL cannot be made to do that. // // The happy-path / concurrency / bidirectional / multi-agent coverage lives in bounceAgentE2ETest, // which drives the SAME pipeline through the production entry point (NixlTransferAgent:: @@ -29,12 +29,24 @@ #include "bounceTestNixlNode.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h" + #include +#include #include #include +#include +#include +#include #include +#include +#include +#include +#include #include +#include #include namespace b = tensorrt_llm::executor::kv_cache::bounce; @@ -42,6 +54,14 @@ namespace kvc = tensorrt_llm::executor::kv_cache; namespace { +// How one scripted write polls (see XferControls::scriptWrite). +enum class WriteScript +{ + kComplete, + kStall, + kFail, +}; + // Test-controlled behavior for every write a FakeXferAgent posts. Shared by the agent and all the // statuses it hands out. struct XferControls @@ -51,6 +71,29 @@ struct XferControls std::atomic releaseSucceeds{true}; // false -> release() fails (handle retained) std::atomic postCount{0}; std::atomic releaseCount{0}; + + // Override how the write with this post index (0 = the agent's first write) polls; may be switched + // while that write is in flight. Unscripted writes follow the switches above. + void scriptWrite(std::uint64_t postIndex, WriteScript script) + { + std::lock_guard lk(mScriptMu); + mScripts[postIndex] = script; + } + + [[nodiscard]] std::optional scriptFor(std::uint64_t postIndex) + { + std::lock_guard lk(mScriptMu); + auto const it = mScripts.find(postIndex); + if (it == mScripts.end()) + { + return std::nullopt; + } + return it->second; + } + +private: + std::mutex mScriptMu; + std::unordered_map mScripts; }; // A TransferStatus driven by XferControls instead of a real backend handle. Unlike @@ -58,8 +101,9 @@ struct XferControls class FakeXferStatus final : public kvc::TransferStatus { public: - explicit FakeXferStatus(std::shared_ptr ctl) + FakeXferStatus(std::shared_ptr ctl, std::uint64_t postIndex) : mCtl(std::move(ctl)) + , mPostIndex(postIndex) { } @@ -70,6 +114,15 @@ class FakeXferStatus final : public kvc::TransferStatus [[nodiscard]] kvc::TransferState wait(int64_t) const override { + if (auto const script = mCtl->scriptFor(mPostIndex)) + { + switch (*script) + { + case WriteScript::kComplete: return kvc::TransferState::kSUCCESS; + case WriteScript::kStall: return kvc::TransferState::kIN_PROGRESS; + case WriteScript::kFail: return kvc::TransferState::kFAILURE; + } + } if (mCtl->failWrites.load(std::memory_order_acquire)) { return kvc::TransferState::kFAILURE; @@ -95,6 +148,7 @@ class FakeXferStatus final : public kvc::TransferStatus private: std::shared_ptr mCtl; + std::uint64_t mPostIndex{}; bool mReleased{false}; }; @@ -113,8 +167,8 @@ class FakeXferAgent final : public kvc::NixlTransferAgent [[nodiscard]] std::unique_ptr postXferRequest(kvc::TransferOp, kvc::TransferDescs const&, kvc::TransferDescs const&, std::string const&, std::optional const&) override { - mCtl->postCount.fetch_add(1, std::memory_order_relaxed); - return std::make_unique(mCtl); + auto const postIndex = mCtl->postCount.fetch_add(1, std::memory_order_acq_rel); + return std::make_unique(mCtl, postIndex); } private: @@ -213,6 +267,359 @@ int countAcks(b::ZmqControlChannel& ch, std::chrono::milliseconds budget, std::u }); return n; } + +// ---- Helpers for the FAILURE-contract tests (last section) ---- + +// Poll `done` every 1 ms until it holds or `budget` elapses; returns whether it held. For state the +// test cannot be notified of (exec-pool and fake-agent counters). +template +bool waitUntil(Pred&& done, std::chrono::milliseconds budget = std::chrono::seconds(5)) +{ + auto const deadline = std::chrono::steady_clock::now() + budget; + while (!done()) + { + if (std::chrono::steady_clock::now() >= deadline) + { + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return true; +} + +// Takes exec contexts away from a node's reactor and gives them back on destruction, so a failed +// ASSERT cannot strand them. Holding every receiver context stops scatters before they launch. +// Holding all but one sender context serializes the gathers, so writes post in chunk order and a +// FakeXferAgent post index names its chunk. +class ExecHold +{ +public: + ExecHold(b::ExecPool& pool, std::size_t leaveFree) + : mPool(pool) + , mLeaveFree(leaveFree) + { + while (mPool.freeCount() > leaveFree) + { + auto* ctx = mPool.tryAcquire(); + if (ctx == nullptr) + { + break; + } + mHeld.push_back(ctx); + } + } + + ~ExecHold() + { + releaseAll(); + } + + ExecHold(ExecHold const&) = delete; + ExecHold& operator=(ExecHold const&) = delete; + + // Whether it holds every context but the `leaveFree` it was asked to leave (i.e. the pool was idle). + [[nodiscard]] bool complete() const + { + return mHeld.size() + mLeaveFree == mPool.size(); + } + + [[nodiscard]] b::ExecCtx* front() const + { + return mHeld.front(); + } + + void releaseFront() + { + mPool.release(mHeld.front()); + mHeld.erase(mHeld.begin()); + } + + void releaseAll() + { + for (auto* ctx : mHeld) + { + mPool.release(ctx); + } + mHeld.clear(); + } + +private: + b::ExecPool& mPool; + std::size_t mLeaveFree; + std::vector mHeld; +}; + +// Launch the batched-copy kernel once (a zero-byte copy) so its module is loaded. Under +// CUDA_MODULE_LOADING=LAZY, the default, the first launch in the process loads the module, which +// needs a device-wide lock and so would wait forever behind a closed StreamGate. Gather and scatter +// both launch this kernel (launchPrepared -> launchBatchedCopy), and the test executable links the +// same copy, so this one launch covers both. +bool loadBatchedCopyKernel() +{ + // One all-zero plan entry (src = dst = 0, size = 0), packed srcs|dsts|sizes like the plan buffers: + // the kernel touches no memory. + constexpr std::size_t kPlanBytes = 3 * sizeof(std::uint64_t); + void* plan = nullptr; + if (cudaSetDevice(0) != cudaSuccess || cudaMalloc(&plan, kPlanBytes) != cudaSuccess) // every node is on device 0 + { + return false; + } + auto const* entry = static_cast(plan); + bool const ok = cudaMemset(plan, 0, kPlanBytes) == cudaSuccess + && b::launchBatchedCopy(/*srcs=*/entry, /*dsts=*/entry + 1, + /*sizes=*/reinterpret_cast(entry + 2), /*n=*/1, /*stream=*/nullptr) + == cudaSuccess + && cudaStreamSynchronize(nullptr) == cudaSuccess; + return cudaFree(plan) == cudaSuccess && ok; +} + +// Blocks a CUDA stream: a host function enqueued on it spins until open(), so everything enqueued +// later on that stream (a scatter kernel) waits behind it. The destructor opens the gate and drains +// the stream, so a failed ASSERT can never leave the stream blocked (the receiver's shutdown would +// hang in cudaDeviceSynchronize). While it is closed, anything that needs a device-wide lock waits +// for it and deadlocks: cudaFree, cudaDeviceSynchronize, and lazy module loading on a kernel's first +// launch. So allocate before closing a gate and free after opening it; the constructor loads the +// bounce kernel before it closes. +class StreamGate +{ +public: + explicit StreamGate(cudaStream_t stream) + : mStream(stream) + { + if (loadBatchedCopyKernel()) // before closing: see the class comment + { + mInstalled = cudaLaunchHostFunc(mStream, &StreamGate::holdUntilOpen, &mOpen) == cudaSuccess; + } + } + + ~StreamGate() + { + open(); + if (mInstalled) + { + (void) cudaStreamSynchronize(mStream); + } + } + + StreamGate(StreamGate const&) = delete; + StreamGate& operator=(StreamGate const&) = delete; + + // Whether the gate was placed on the stream (it then blocks the stream until open()). + [[nodiscard]] bool installed() const + { + return mInstalled; + } + + void open() + { + mOpen.store(true, std::memory_order_release); + } + +private: + static void CUDART_CB holdUntilOpen(void* openFlag) + { + auto const& isOpen = *static_cast const*>(openFlag); + while (!isOpen.load(std::memory_order_acquire)) + { + std::this_thread::yield(); + } + } + + cudaStream_t mStream; + std::atomic mOpen{false}; + bool mInstalled{false}; +}; + +// XferBufs freed on scope exit. cudaFree synchronizes the device, so declare it before any StreamGate. +struct OwnedXferBufs : bounce_test::XferBufs +{ + OwnedXferBufs(std::uint32_t nDescs, std::uint32_t descBytes, std::uint32_t seed) + : bounce_test::XferBufs(bounce_test::makeXferBufs(nDescs, descBytes, seed)) + { + } + + ~OwnedXferBufs() + { + bounce_test::freeXferBufs(*this); + } + + OwnedXferBufs(OwnedXferBufs const&) = delete; + OwnedXferBufs& operator=(OwnedXferBufs const&) = delete; +}; + +// Whether every byte of the device range [base + offset, base + offset + bytes) equals `value`. +bool deviceBytesEqual(void const* base, std::uint64_t offset, std::uint64_t bytes, unsigned char value) +{ + std::vector host(bytes); + if (cudaMemcpy(host.data(), static_cast(base) + offset, bytes, cudaMemcpyDeviceToHost) + != cudaSuccess) + { + return false; + } + return std::all_of(host.begin(), host.end(), [value](unsigned char v) { return v == value; }); +} + +// What a receiver's scatter copies into dst: its arena is prefilled with this byte, because fake +// writes never deliver the source. +constexpr unsigned char kArenaFill = 0xA5; +// What the caller writes into dst once the request has resolved. +constexpr unsigned char kReusedFill = 0x3C; + +// Wait for the sender's (non-cancel) WANT and store its rid in `rid`. +bool waitWant(b::ZmqControlChannel& ch, std::uint64_t& rid, std::chrono::milliseconds budget = std::chrono::seconds(5)) +{ + return pumpChannel(ch, budget, + [&](b::BounceMsgHeader const& h, std::string const& blob) + { + std::vector chunkBytes; + std::string endpoint; + if (static_cast(h.msgType) != b::BounceMsgType::kWANT + || !b::decodeWant(blob, h, chunkBytes, endpoint) || b::isCancelWant(chunkBytes)) + { + return false; + } + rid = h.requestId; + return true; + }); +} + +// Wait until DATA for `chunks` distinct chunks of `rid` has arrived, i.e. the sender marked them Sent. +bool waitData(b::ZmqControlChannel& ch, std::uint64_t rid, std::size_t chunks, + std::chrono::milliseconds budget = std::chrono::seconds(5)) +{ + std::set seen; + return pumpChannel(ch, budget, + [&](b::BounceMsgHeader const& h, std::string const&) + { + if (static_cast(h.msgType) == b::BounceMsgType::kDATA && h.requestId == rid) + { + seen.insert(h.chunkIdx); + } + return seen.size() >= chunks; + }); +} + +// Wait for the sender's cancel (an empty WANT) for `rid`, i.e. proof that it failed the request. +bool waitCancel(b::ZmqControlChannel& ch, std::uint64_t rid, std::chrono::milliseconds budget = std::chrono::seconds(5)) +{ + return pumpChannel(ch, budget, + [&](b::BounceMsgHeader const& h, std::string const& blob) + { + std::vector chunkBytes; + std::string endpoint; + return static_cast(h.msgType) == b::BounceMsgType::kWANT && h.requestId == rid + && b::decodeWant(blob, h, chunkBytes, endpoint) && b::isCancelWant(chunkBytes); + }); +} + +// The region handle a scripted receiver grants to chunk `chunkIdx`. +std::uint64_t fakeRegion(std::uint32_t chunkIdx) +{ + return 100 + chunkIdx; +} + +// A FakeXferAgent sender wired to a scripted receiver: a bare control channel through which the test +// plays the receiver by hand. Writes are fake, so the credits carry fake addresses. +struct ScriptedPeers +{ + // Declared first, so destroyed last: freed only after the sender has shut down. + std::deque requestBufs; + std::shared_ptr ctl{std::make_shared()}; + std::string senderName; + std::string receiverName; + std::uint32_t chunkBytes{}; + std::unique_ptr receiver; + std::unique_ptr sender; // declared last: shut down before the channel goes away + + // GRANT chunks [firstChunk, endChunk) of `rid` in one message; chunk i gets fakeRegion(i). + void grant(std::uint64_t rid, std::uint32_t firstChunk, std::uint32_t endChunk) const + { + std::vector credits; + for (std::uint32_t i = firstChunk; i < endChunk; ++i) + { + credits.push_back({/*addr=*/0, chunkBytes, /*devId=*/0, fakeRegion(i)}); + } + receiver->sendTo(senderName, b::encodeGrant(rid, credits)); + } + + void ack(std::uint64_t rid, std::uint32_t chunkIdx) const + { + receiver->sendTo(senderName, b::encodeAck(rid, chunkIdx, fakeRegion(chunkIdx))); + } + + void nack(std::uint64_t rid, std::uint32_t chunkIdx) const + { + receiver->sendTo(senderName, b::encodeNack(rid, chunkIdx, fakeRegion(chunkIdx))); + } +}; + +// nullptr when the NIXL agent/backend is unavailable (the caller GTEST_SKIPs). +std::unique_ptr makeScriptedPeers(std::string const& name, b::BounceConfig const& c) +{ + auto peers = std::make_unique(); + peers->senderName = name + "Sender"; + peers->receiverName = name + "Receiver"; + peers->chunkBytes = static_cast(c.maxChunkSizeBytes); + peers->receiver = std::make_unique(peers->receiverName); + peers->sender = makeFakeNode(peers->senderName, c, peers->ctl); + if (!peers->sender) + { + return nullptr; + } + EXPECT_TRUE(peers->sender->tx->addPeer(peers->receiverName, peers->receiver->localEndpoint())); + EXPECT_TRUE(peers->receiver->addPeer(peers->senderName, peers->sender->ch->localEndpoint())); + return peers; +} + +// One request to the scripted receiver: `chunks` descriptors of one full chunk each, so chunk i is +// descriptor i. Its buffers live in `peers`. The rid is filled in from the WANT (waitWant). +struct ScriptedRequest +{ + ScriptedRequest(ScriptedPeers& peers, std::uint32_t chunks, std::uint32_t seed) + : chunks(chunks) + , bufs(peers.requestBufs.emplace_back(chunks, peers.chunkBytes, seed)) + , fut(peers.sender->tx->submit(bufs.srcDescs, bufs.dstDescs, peers.receiverName)) + { + } + + std::uint32_t chunks; + OwnedXferBufs const& bufs; + std::shared_future fut; + std::uint64_t rid{}; +}; + +// Learn `req`'s rid from its WANT, grant all its chunks at once and wait until every one is Sent. +bool grantUntilAllSent(ScriptedPeers& peers, ScriptedRequest& req) +{ + if (!waitWant(*peers.receiver, req.rid)) + { + return false; + } + peers.grant(req.rid, /*firstChunk=*/0, /*endChunk=*/req.chunks); + return waitData(*peers.receiver, req.rid, req.chunks); +} + +// FIFO barrier: once this returns true, the sender's reactor has dispatched every message the scripted +// receiver sent before it. A never-granted sentinel request is NACKed and must resolve; the reactor reads +// one peer's messages in order. +bool syncSender(ScriptedPeers& peers) +{ + ScriptedRequest sentinel(peers, /*chunks=*/1, /*seed=*/0); + if (!waitWant(*peers.receiver, sentinel.rid)) + { + return false; + } + peers.nack(sentinel.rid, 0); + return sentinel.fut.wait_for(std::chrono::seconds(5)) == std::future_status::ready; +} + +// A status over a future already resolved with `result`. +b::BounceTransferStatus resolvedStatus(b::BounceResult const& result) +{ + std::promise promise; + promise.set_value(result); + return b::BounceTransferStatus{promise.get_future().share()}; +} } // namespace // A peer that never GRANTs must fail the request on requestTimeoutMs, not hang. The WANT is @@ -236,6 +643,7 @@ TEST(BounceTransportFailure, NoGrantTimesOutNotHang) ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "request hung with no grant"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kNoProgressTimeout); + EXPECT_TRUE(fut.get().quiesced) << "nothing was sent, so nothing can write dst"; t->tx->shutdown(); bounce_test::freeXferBufs(bufs); @@ -263,6 +671,7 @@ TEST(BounceTransportFailure, WriteFailureFailsRequest) ASSERT_EQ(fut.wait_for(std::chrono::seconds(10)), std::future_status::ready) << "write failure hung"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kWriteFailed); + EXPECT_TRUE(fut.get().quiesced) << "no write completed, so no DATA was sent"; EXPECT_GE(ctl->postCount.load(std::memory_order_acquire), 1u); A->tx->shutdown(); @@ -300,6 +709,7 @@ TEST(BounceTransportFailure, ReceiverNackFailsSenderBeforeTimeout) auto const elapsed = std::chrono::steady_clock::now() - start; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kPeerNack); + EXPECT_TRUE(fut.get().quiesced) << "the WANT was rejected, so no DATA was sent"; EXPECT_LT(elapsed, std::chrono::milliseconds(kTimeoutMs / 2)) << "NACK did not short-circuit the timeout"; // The eager gather of the failed chunk holds an ExecCtx + staging region until failRequest (or, @@ -687,6 +1097,7 @@ TEST(BounceTransportFailure, CreditMispairAbandonIsTerminalAndIgnoresLaterGrants ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "abandoned flow never resolved"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kProtocolError); + EXPECT_TRUE(fut.get().quiesced) << "no write was posted, so no DATA was sent"; // The later valid GRANT must not have resurrected the abandoned flow into posting a write. EXPECT_EQ(ctl->postCount.load(std::memory_order_acquire), 0u) << "abandoned flow posted an RDMA write"; @@ -711,6 +1122,7 @@ TEST(BounceTransportFailure, ShutdownFailsInflight) ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "shutdown left request hanging"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); + EXPECT_TRUE(fut.get().quiesced) << "never granted, so no DATA was sent"; bounce_test::freeXferBufs(bufs); } @@ -754,6 +1166,7 @@ TEST(BounceTransportFailure, ShutdownReleaseFailureDoesNotHangOrLoseHandle) ASSERT_EQ(fut.wait_for(std::chrono::seconds(1)), std::future_status::ready); EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); + EXPECT_TRUE(fut.get().quiesced) << "the write never completed, so no DATA was sent"; // failAll attempts the failed release exactly twice (initial + bounded retry), never spins. EXPECT_EQ(ctl->releaseCount.load(std::memory_order_acquire), 2u); bounce_test::freeXferBufs(bufs); @@ -1030,3 +1443,602 @@ TEST(BounceTransportFailure, ConcurrentMultiThreadedSubmit) bounce_test::freeXferBufs(bb); } } + +// ============================================================================ +// A FAILURE resolves only after the receiver has answered every chunk whose DATA was sent +// ---------------------------------------------------------------------------- +// Only the receiver's scatter of a received DATA writes dst, and the receiver answers (ACK or NACK) +// after that scatter's last write. So a failed request waits for those answers before it resolves +// kFAILURE with quiesced == true; a receiver that sends no answer for requestTimeoutMs is assumed dead +// and the request resolves in-doubt (quiesced == false). These tests order events with exec-context +// holds, a stream gate, scripted writes and FIFO barriers, not by sleeping. +// ============================================================================ + +// The status contract the agent hands out (pure: hand-built futures, no NIXL). isCompleted() is true +// exactly when the request is over AND nothing will write the destination again. +TEST(BounceTransferStatus, CompletedOnlyOnceTheDestinationIsQuiesced) +{ + std::promise pending; + b::BounceTransferStatus const inFlight{pending.get_future().share()}; + EXPECT_FALSE(inFlight.isCompleted()); + EXPECT_EQ(inFlight.wait(0), kvc::TransferState::kIN_PROGRESS); + EXPECT_EQ(inFlight.getLastStatusStr(), ""); + + auto const success = resolvedStatus({kvc::TransferState::kSUCCESS, b::BounceFailReason::kNone, /*quiesced=*/true}); + EXPECT_TRUE(success.isCompleted()); + EXPECT_EQ(success.wait(0), kvc::TransferState::kSUCCESS); + EXPECT_EQ(success.getLastStatusStr(), ""); + + auto const quiescedFailure + = resolvedStatus({kvc::TransferState::kFAILURE, b::BounceFailReason::kWriteFailed, /*quiesced=*/true}); + EXPECT_TRUE(quiescedFailure.isCompleted()) << "a destination proven quiescent may be reused"; + EXPECT_EQ(quiescedFailure.wait(0), kvc::TransferState::kFAILURE); + EXPECT_EQ(quiescedFailure.getLastStatusStr(), b::toString(b::BounceFailReason::kWriteFailed)); + + auto const inDoubt + = resolvedStatus({kvc::TransferState::kFAILURE, b::BounceFailReason::kNoProgressTimeout, /*quiesced=*/false}); + EXPECT_FALSE(inDoubt.isCompleted()) << "the receiver may still write the destination"; + EXPECT_EQ(inDoubt.wait(0), kvc::TransferState::kFAILURE); + EXPECT_EQ(inDoubt.getLastStatusStr(), + std::string(b::toString(b::BounceFailReason::kNoProgressTimeout)) + "; destination not proven quiescent"); + + pending.set_value({kvc::TransferState::kSUCCESS, b::BounceFailReason::kNone, /*quiesced=*/true}); + EXPECT_TRUE(inFlight.isCompleted()); +} + +// A request rejected at plan time never sent anything, so its FAILURE is quiesced. +TEST(BounceTransportFailure, PlanRejectionIsQuiesced) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + // One descriptor larger than maxChunkSizeBytes: BounceTransferPlan::build rejects it. + OwnedXferBufs bufs(/*nDescs=*/1, /*descBytes=*/2 * static_cast(c.maxChunkSizeBytes), /*seed=*/81); + auto t = bounce_test::makeNode("prSolo", c, 1024); + if (!t) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "nobody"); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(0)), std::future_status::ready) + << "plan rejection resolves in submit()"; + EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kPlanRejected); + EXPECT_TRUE(fut.get().quiesced); + t->tx->shutdown(); +} + +// A write fails while an earlier chunk's scatter is RUNNING on the receiver. The FAILURE must wait for +// that scatter's answer; resolving first lets the caller reuse dst while the scatter still writes it. +// The scatter is held on the receiver's only scatter worker by a stream gate placed ahead of it on the +// one exec context the worker can get. +TEST(BounceTransportFailure, FailureWaitsForRunningScatter) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + // The timeout never fires here: only the receiver's answers may resolve the failed request. + auto c = cfg(/*timeoutMs=*/30000); + c.scatterWorkerCount = 1; + capRegions(c, c.maxInflightChunksPerRequest); // two chunks, both granted at once + // Chunk i is descriptor i. Declared before the nodes, so it is freed after they shut down. + OwnedXferBufs bufs(/*nDescs=*/2, static_cast(c.maxChunkSizeBytes), /*seed=*/91); + auto ctl = std::make_shared(); + auto sender = makeFakeNode("runScatterSender", c, ctl); + auto receiver = bounce_test::makeNode("runScatterReceiver", c, 1024); + if (!sender || !receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender->tx->addPeer(receiver->name, receiver->ch->localEndpoint())); + ASSERT_TRUE(receiver->tx->addPeer(sender->name, sender->ch->localEndpoint())); + ASSERT_EQ(cudaMemset(receiver->arena->base(), kArenaFill, receiver->arena->bytes()), cudaSuccess); + + ExecHold senderGathers(*sender->exec, /*leaveFree=*/1); // post index == chunk index + ExecHold receiverScatters(*receiver->exec, /*leaveFree=*/0); + ASSERT_TRUE(senderGathers.complete()); + ASSERT_TRUE(receiverScatters.complete()); + StreamGate gate(receiverScatters.front()->stream); + ASSERT_TRUE(gate.installed()); + receiverScatters.releaseFront(); // the worker can get only the gated context + ctl->scriptWrite(0, WriteScript::kComplete); + ctl->scriptWrite(1, WriteScript::kStall); + + auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, receiver->name); + // The worker took the gated context: chunk 0 is Sent and its scatter waits behind the gate. + ASSERT_TRUE(waitUntil([&] { return receiver->exec->freeCount() == 0; })) << "chunk 0 never reached the scatter"; + ASSERT_TRUE(waitUntil([&] { return ctl->postCount.load(std::memory_order_acquire) == 2; })); + ctl->scriptWrite(1, WriteScript::kFail); + // Write 0's handle was released when its DATA went out; write 1's is released by the orphan + // drain that runs only after failRequest. + ASSERT_TRUE(waitUntil([&] { return ctl->releaseCount.load(std::memory_order_acquire) == 2; })) + << "the failed write was never handled"; + EXPECT_EQ(fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "FAILURE reported while chunk 0's scatter could still write dst"; + + gate.open(); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "the scatter's ACK was not awaited"; + auto const result = fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); + EXPECT_TRUE(result.quiesced); + EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[0], bufs.descBytes, kArenaFill)) + << "chunk 0's scatter had not finished when the FAILURE resolved"; + EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[1], bufs.descBytes, 0)) << "chunk 1 was never sent"; + // Reuse dst, then shut the receiver down (joins its scatter workers, syncs the device): nothing may + // write it again. + ASSERT_EQ(cudaMemset(bufs.dst, kReusedFill, bufs.total), cudaSuccess); + receiver->tx->shutdown(); + EXPECT_TRUE(deviceBytesEqual(bufs.dst, 0, bufs.total, kReusedFill)) << "dst was written after it was reused"; + sender->tx->shutdown(); +} + +// Same as above, but the sent chunks are still QUEUED on the receiver: chunk 0 waits for an exec +// context in the only scatter worker and chunk 1 waits in the job queue, while chunk 2's write fails. +// Without an exec context no answer is possible, so the order is deterministic. +TEST(BounceTransportFailure, FailureWaitsForQueuedScatters) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + // The timeout never fires here: only the receiver's answers may resolve the failed request. + auto c = cfg(/*timeoutMs=*/30000); + c.scatterWorkerCount = 1; + c.maxInflightChunksPerRequest = 3; + capRegions(c, c.maxInflightChunksPerRequest); // three chunks, all granted at once + // Chunk i is descriptor i. Declared before the nodes, so it is freed after they shut down. + OwnedXferBufs bufs(/*nDescs=*/3, static_cast(c.maxChunkSizeBytes), /*seed=*/92); + auto ctl = std::make_shared(); + auto sender = makeFakeNode("queuedScatterSender", c, ctl); + auto receiver = bounce_test::makeNode("queuedScatterReceiver", c, 1024); + if (!sender || !receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender->tx->addPeer(receiver->name, receiver->ch->localEndpoint())); + ASSERT_TRUE(receiver->tx->addPeer(sender->name, sender->ch->localEndpoint())); + ASSERT_EQ(cudaMemset(receiver->arena->base(), kArenaFill, receiver->arena->bytes()), cudaSuccess); + + ExecHold senderGathers(*sender->exec, /*leaveFree=*/1); // post index == chunk index + ExecHold receiverScatters(*receiver->exec, /*leaveFree=*/0); + ASSERT_TRUE(senderGathers.complete()); + ASSERT_TRUE(receiverScatters.complete()); + ctl->scriptWrite(0, WriteScript::kComplete); + ctl->scriptWrite(1, WriteScript::kComplete); + // Polled in chunk order, writes 0 and 1 are Sent (DATA out) before write 2's failure is seen. + ctl->scriptWrite(2, WriteScript::kFail); + + auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, receiver->name); + // Writes 0 and 1 were released when their DATA went out; write 2 after failRequest. + ASSERT_TRUE(waitUntil([&] { return ctl->releaseCount.load(std::memory_order_acquire) == 3; })) + << "the failed write was never handled"; + EXPECT_EQ(fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "FAILURE reported while chunks 0 and 1 could still be scattered into dst"; + + receiverScatters.releaseAll(); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) + << "the scatters' ACKs were not awaited"; + auto const result = fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); + EXPECT_TRUE(result.quiesced); + for (std::uint32_t chunk = 0; chunk < 2; ++chunk) + { + EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[chunk], bufs.descBytes, kArenaFill)) + << "chunk " << chunk << " was not scattered before the FAILURE resolved"; + } + EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[2], bufs.descBytes, 0)) << "chunk 2 was never sent"; + // Reuse dst, then shut the receiver down (joins its scatter workers, syncs the device): nothing may + // write it again. + ASSERT_EQ(cudaMemset(bufs.dst, kReusedFill, bufs.total), cudaSuccess); + receiver->tx->shutdown(); + EXPECT_TRUE(deviceBytesEqual(bufs.dst, 0, bufs.total, kReusedFill)) << "dst was written after it was reused"; + sender->tx->shutdown(); +} + +// A NACK fails the request but answers only its own chunk: the request waits for the other Sent +// chunk's answer, and not for a chunk whose write is still in flight (its DATA was never sent, so the +// receiver will never scatter or answer it). +TEST(BounceTransportFailure, NackWaitsOnlyForOtherSentChunks) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); // waiting for chunk 2 would show up as a missed 5 s bound + c.maxInflightChunksPerRequest = 3; + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("nackWait", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ExecHold senderGathers(*peers->sender->exec, /*leaveFree=*/1); // post index == chunk index + ASSERT_TRUE(senderGathers.complete()); + peers->ctl->scriptWrite(0, WriteScript::kComplete); + peers->ctl->scriptWrite(1, WriteScript::kComplete); + peers->ctl->scriptWrite(2, WriteScript::kStall); + + ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/93); + ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); + peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/3); + ASSERT_TRUE(waitData(*peers->receiver, req.rid, 2)) << "chunks 0 and 1 never reached Sent"; + ASSERT_TRUE(waitUntil([&] { return peers->ctl->postCount.load(std::memory_order_acquire) == 3; })) + << "chunk 2's write was never posted"; + + peers->nack(req.rid, 1); + ASSERT_TRUE(syncSender(*peers)); + EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "resolved before chunk 0 was answered"; + peers->ack(req.rid, 0); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) + << "waited for chunk 2, whose DATA was never sent"; + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); + EXPECT_TRUE(result.quiesced); + EXPECT_EQ(peers->ctl->releaseCount.load(std::memory_order_acquire), 2u) + << "chunk 2's write must still be in flight"; + + // End the orphaned write so shutdown skips its 5 s write drain; its deferred cancel follows. + peers->ctl->scriptWrite(2, WriteScript::kComplete); + EXPECT_TRUE(waitCancel(*peers->receiver, req.rid)); +} + +// NACKs are answers too: a failed request whose remaining chunks are all NACKed (e.g. a receiver with a +// CUDA fault NACKs every queued scatter) resolves quiesced at once, not in-doubt at requestTimeoutMs. +TEST(BounceTransportFailure, NackAnswersEndTheWait) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); // ignoring the NACKs would show up as a missed 5 s bound + c.maxInflightChunksPerRequest = 3; + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("nackAnswers", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/106); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + peers->nack(req.rid, 2); // fails the request: chunks 0 and 1 are awaited + ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); + peers->nack(req.rid, 0); + peers->nack(req.rid, 1); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) + << "NACKs to a failed request were not taken as answers"; + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); + EXPECT_TRUE(result.quiesced); +} + +// Only an answer for an awaited chunk (same chunk index AND region) ends a failed request's wait. A +// duplicate, an answer for a chunk that is not awaited and a wrong-region answer are all ignored. +TEST(BounceTransportFailure, OnlyMatchingAnswersEndTheWait) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); + c.maxInflightChunksPerRequest = 3; + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("matchAnswers", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/102); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + peers->nack(req.rid, 2); // fails the request: chunks 0 and 1 are awaited + ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); + + peers->ack(req.rid, 1); // a real answer: only chunk 0 is awaited now + peers->ack(req.rid, 1); // duplicate + peers->ack(req.rid, 2); // chunk 2 is not awaited: its NACK already answered it + peers->receiver->sendTo(peers->senderName, b::encodeAck(req.rid, 0, fakeRegion(0) + 1)); // wrong region + ASSERT_TRUE(syncSender(*peers)); + EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "an answer that matches no awaited chunk ended the wait"; + + peers->ack(req.rid, 0); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); + EXPECT_TRUE(result.quiesced); +} + +// A failed request that still has unposted chunks (more chunks than maxInflightChunksPerRequest, the +// production norm) stays inert while it awaits answers: a GRANT that crossed its cancel must not make +// it gather, post or send DATA, and once it resolves it must hold no staging region. +TEST(BounceTransportFailure, FailedRequestWithUnpostedChunksStaysInert) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); // maxInflightChunksPerRequest = 2 + capRegions(c, c.maxInflightChunksPerRequest); // the sender's staging holds exactly two chunks + auto peers = makeScriptedPeers("unpostedChunks", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ExecHold senderGathers(*peers->sender->exec, /*leaveFree=*/1); // post index == chunk index + ASSERT_TRUE(senderGathers.complete()); + peers->ctl->scriptWrite(0, WriteScript::kComplete); + peers->ctl->scriptWrite(1, WriteScript::kStall); + + ScriptedRequest req(*peers, /*chunks=*/4, /*seed=*/107); + ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); + peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/2); + ASSERT_TRUE(waitData(*peers->receiver, req.rid, 1)); + ASSERT_TRUE(waitUntil([&] { return peers->ctl->postCount.load(std::memory_order_acquire) == 2; })); + peers->ctl->scriptWrite(1, WriteScript::kFail); + // Failed: chunk 0 is awaited, and chunks 2 and 3 were never posted. + ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); + + peers->grant(req.rid, /*firstChunk=*/2, /*endChunk=*/4); + ASSERT_TRUE(syncSender(*peers)); + EXPECT_EQ(peers->ctl->postCount.load(std::memory_order_acquire), 2u) << "the failed request posted a write"; + EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + peers->ack(req.rid, 0); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); + EXPECT_TRUE(result.quiesced); + + // No staging may stay held: a fresh 2-chunk request needs both regions to get both chunks Sent. + ScriptedRequest next(*peers, /*chunks=*/2, /*seed=*/108); + ASSERT_TRUE(grantUntilAllSent(*peers, next)) << "the failed request leaked a staging region"; + EXPECT_EQ(peers->ctl->postCount.load(std::memory_order_acquire), 4u) + << "the failed request posted a write (only the fresh request's two may follow its own two)"; +} + +// A receiver that stops answering is assumed dead requestTimeoutMs after its LAST answer: the request +// then resolves in-doubt. Each answer restarts that bound. Bounds are measured from the test's own +// send times, so they hold however the threads are scheduled. +TEST(BounceTransportFailure, SilentReceiverResolvesInDoubtAfterTimeoutSinceLastAnswer) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + constexpr auto kTimeout = std::chrono::milliseconds(2000); + // Room for thread scheduling on top of the bound; far below a second kTimeout. + constexpr auto kSchedulingSlack = std::chrono::seconds(3); + auto c = cfg(static_cast(kTimeout.count())); + c.maxInflightChunksPerRequest = 3; + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("silentReceiver", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/94); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + + // Fails the request; chunks 0 and 1 still owe an answer. + peers->nack(req.rid, 2); + // Let a quarter of the bound pass first: if the ACK below did not restart the deadline, the request + // would resolve kTimeout/4 early and fail EXPECT_GE(waited, kTimeout). + std::this_thread::sleep_for(kTimeout / 4); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "resolved before requestTimeoutMs passed since the failure"; + auto const lastAnswerSent = std::chrono::steady_clock::now(); + peers->ack(req.rid, 0); // the last answer: chunk 1 never gets one + + ASSERT_EQ(req.fut.wait_for(kTimeout + kSchedulingSlack), std::future_status::ready) + << "an unanswered chunk held the request past its bound"; + auto const waited = std::chrono::steady_clock::now() - lastAnswerSent; + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); + EXPECT_FALSE(result.quiesced) << "chunk 1 was never answered"; + EXPECT_GE(waited, kTimeout) << "the last answer did not restart the requestTimeoutMs bound"; + EXPECT_LT(waited, kTimeout + kSchedulingSlack); +} + +// The answer wait also restarts at the failure itself, however long the failing write took: the +// receiver gets a full requestTimeoutMs after the failure to answer a sent chunk. +TEST(BounceTransportFailure, AnswerWaitStartsAtTheFailure) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + constexpr auto kTimeout = std::chrono::milliseconds(1500); + auto c = cfg(static_cast(kTimeout.count())); + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("answerWaitStart", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ExecHold senderGathers(*peers->sender->exec, /*leaveFree=*/1); // post index == chunk index + ASSERT_TRUE(senderGathers.complete()); + peers->ctl->scriptWrite(0, WriteScript::kComplete); + peers->ctl->scriptWrite(1, WriteScript::kStall); + + ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/109); + ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); + peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/2); + ASSERT_TRUE(waitData(*peers->receiver, req.rid, 1)); + ASSERT_TRUE(waitUntil([&] { return peers->ctl->postCount.load(std::memory_order_acquire) == 2; })); + // Write 1 stalls for most of requestTimeoutMs (its post was the request's last progress), then fails. + std::this_thread::sleep_for(kTimeout * 6 / 10); + peers->ctl->scriptWrite(1, WriteScript::kFail); + ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); + // Counted from the stalled write's post, the bound would already be over here. + std::this_thread::sleep_for(kTimeout / 2); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "resolved in-doubt less than requestTimeoutMs after the failure"; + peers->ack(req.rid, 0); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); + EXPECT_TRUE(result.quiesced); +} + +// A kNoProgressTimeout means no progress for requestTimeoutMs, so a request with a sent chunk still +// unanswered resolves in-doubt at once instead of waiting another requestTimeoutMs. +TEST(BounceTransportFailure, NoProgressTimeoutWithUnansweredChunkResolvesInDoubtAtOnce) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + constexpr auto kTimeout = std::chrono::milliseconds(2000); + auto c = cfg(static_cast(kTimeout.count())); + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("noProgress", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/95); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + auto const lastAnswerSent = std::chrono::steady_clock::now(); + peers->ack(req.rid, 0); // the request's last progress; chunk 1 stays unanswered + + ASSERT_EQ(req.fut.wait_for(2 * kTimeout), std::future_status::ready) + << "the no-progress timeout waited another requestTimeoutMs for answers"; + auto const waited = std::chrono::steady_clock::now() - lastAnswerSent; + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kNoProgressTimeout); + EXPECT_FALSE(result.quiesced) << "chunk 1 was never answered"; + EXPECT_GE(waited, kTimeout); + EXPECT_LT(waited, 2 * kTimeout); +} + +// An abandoned flow (GRANT mispair) dies through checkTimeouts' kNoProgressTimeout. With a Sent chunk +// unanswered it resolves in-doubt within one requestTimeoutMs of its last progress, reporting the +// abandon cause; it must not wait another requestTimeoutMs for answers. +TEST(BounceTransportFailure, AbandonWithSentChunkResolvesInDoubtWithinOneTimeout) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + constexpr auto kTimeout = std::chrono::milliseconds(1500); + auto c = cfg(static_cast(kTimeout.count())); + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("abandonSent", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/110); + ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); + peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/1); + ASSERT_TRUE(waitData(*peers->receiver, req.rid, 1)); + auto const chunk0Sent = std::chrono::steady_clock::now(); // after the request's last progress + // Chunk 1's credit is half the chunk: a mispair, so the flow is abandoned. + std::vector const undersized{{/*addr=*/0, peers->chunkBytes / 2, /*devId=*/0, fakeRegion(1)}}; + peers->receiver->sendTo(peers->senderName, b::encodeGrant(req.rid, undersized)); + + ASSERT_EQ(req.fut.wait_for(kTimeout + kTimeout / 2), std::future_status::ready) + << "the abandoned flow waited another requestTimeoutMs for answers"; + EXPECT_LT(std::chrono::steady_clock::now() - chunk0Sent, kTimeout + kTimeout / 2); + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kProtocolError); + EXPECT_FALSE(result.quiesced) << "chunk 0 was never answered"; +} + +// A kNoProgressTimeout with every chunk Sent and none answered resolves in-doubt. ACKs arriving +// afterwards, enough to count every chunk, hit an erased request: the reactor drops them and keeps +// serving (the FIFO barrier below). +TEST(BounceTransportFailure, LateAcksAfterNoProgressTimeoutChangeNothing) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + constexpr auto kTimeout = std::chrono::milliseconds(1000); + auto c = cfg(static_cast(kTimeout.count())); + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("lateAcks", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/96); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + ASSERT_EQ(req.fut.wait_for(kTimeout + std::chrono::seconds(5)), std::future_status::ready); + EXPECT_EQ(req.fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(req.fut.get().reason, b::BounceFailReason::kNoProgressTimeout); + EXPECT_FALSE(req.fut.get().quiesced); + + peers->ack(req.rid, 0); + peers->ack(req.rid, 1); + EXPECT_TRUE(syncSender(*peers)) << "the reactor stopped after the late ACKs"; +} + +// forgetPeer while a failed request awaits answers must not fail it again: it keeps its first reason +// and keeps waiting, and a later answer resolves it with that ORIGINAL reason. +TEST(BounceTransportFailure, ForgetPeerWhileAwaitingAnswersKeepsWaiting) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("forgetAwaiting", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/98); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + // Never granted: forgetPeer fails it, which proves the forget was applied. + ScriptedRequest sentinel(*peers, /*chunks=*/1, /*seed=*/99); + + peers->nack(req.rid, 1); + // No write was in flight, so the cancel went out at once: the NACK failed the request. + ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); + peers->sender->tx->forgetPeer(peers->receiverName); + ASSERT_EQ(sentinel.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "forgetPeer not applied"; + EXPECT_EQ(sentinel.fut.get().reason, b::BounceFailReason::kPeerDropped); + EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "forgetPeer resolved a request that still awaits an answer"; + + peers->ack(req.rid, 0); // answers still reach the sender's ROUTER after forgetPeer + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); + EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack) << "forgetPeer failed the request a second time"; + EXPECT_TRUE(result.quiesced); +} + +// forgetPeer fails a request whose chunks were all Sent; it then waits for their answers. ACKs for +// every chunk must resolve it FAILURE (kPeerDropped, quiesced), never SUCCESS, even though the ACK +// count reaches numChunks. +TEST(BounceTransportFailure, AnswersAfterForgetPeerResolveFailureNotSuccess) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("forgetAllSent", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/100); + ASSERT_TRUE(grantUntilAllSent(*peers, req)); + // Never granted: forgetPeer fails it in the same pass, which proves the forget was applied. + ScriptedRequest sentinel(*peers, /*chunks=*/1, /*seed=*/101); + + peers->sender->tx->forgetPeer(peers->receiverName); + ASSERT_EQ(sentinel.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "forgetPeer not applied"; + EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "forgetPeer resolved a request whose sent chunks were unanswered"; + + peers->ack(req.rid, 0); + peers->ack(req.rid, 1); + ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto const result = req.fut.get(); + EXPECT_EQ(result.state, kvc::TransferState::kFAILURE) << "a failed request resolved SUCCESS"; + EXPECT_EQ(result.reason, b::BounceFailReason::kPeerDropped); + EXPECT_TRUE(result.quiesced); +} + +// Shutdown resolves every pending request. One already awaiting answers keeps its own reason, and any +// request with a sent chunk unanswered is in-doubt: the receiver may still scatter into its dst. +TEST(BounceTransportFailure, ShutdownResolvesUnansweredRequestsInDoubtWithTheirOwnReason) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); + c.maxInflightChunksPerRequest = 3; + capRegions(c, c.maxInflightChunksPerRequest); + auto peers = makeScriptedPeers("shutdownAwaiting", c); + if (!peers) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + ScriptedRequest failed(*peers, /*chunks=*/2, /*seed=*/104); + ASSERT_TRUE(grantUntilAllSent(*peers, failed)); + peers->nack(failed.rid, 1); // fails it: chunk 0 is awaited + ASSERT_TRUE(waitCancel(*peers->receiver, failed.rid)); + // Still live: its only chunk is Sent and unanswered. + ScriptedRequest live(*peers, /*chunks=*/1, /*seed=*/105); + ASSERT_TRUE(grantUntilAllSent(*peers, live)); + + peers->sender->tx->shutdown(); + ASSERT_EQ(failed.fut.wait_for(std::chrono::seconds(0)), std::future_status::ready); + EXPECT_EQ(failed.fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(failed.fut.get().reason, b::BounceFailReason::kPeerNack) << "shutdown replaced the original reason"; + EXPECT_FALSE(failed.fut.get().quiesced) << "chunk 0 was never answered"; + ASSERT_EQ(live.fut.wait_for(std::chrono::seconds(0)), std::future_status::ready); + EXPECT_EQ(live.fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(live.fut.get().reason, b::BounceFailReason::kShutdown); + EXPECT_FALSE(live.fut.get().quiesced) << "its sent chunk was never answered"; +} From 6d088796846999025740154e530824bcc12615f8 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:10:54 +0000 Subject: [PATCH 3/3] [TRTLLM-17105][fix] Prove a failed bounce transfer quiescent through isCompleted() A failed bounce request now records what it may still touch in caller memory: the sent chunks the receiver has not answered (it may still scatter them into the destination) and the gathers still running (they may still read the source). The record lives in a pure core, bounce/BounceResult.{h,cpp}, and only shrinks once the result is set. - wait() still returns only after the receiver has answered every sent chunk, at most request_timeout_ms after the failure or after the last answer, and at once after kNoProgressTimeout and at shutdown. - isCompleted() is true once the record is empty and never turns back to false. Answers arriving after an in-doubt result still count, and an orphaned gather finishing completes it. - Failed requests move from the active table to mFailed; settleFailedRequests, at the end of every IO tick, is the only place that resolves them and drops their entries. - The receiver NACKs a DATA it will not scatter unless a scatter still holds its region. A failed scatter launch syncs the stream before its NACK. kBounceVersion goes to 3, and BounceSender rejects a non-positive request_timeout_ms. - Tests: a host-only test runs FailedRequest through every event order and checks the contract after each event. Integration tests cover the reactor around it. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../tensorrt_llm/executor/transferAgent.h | 10 +- .../nixl_utils/bounce/BounceConfig.h | 6 +- .../nixl_utils/bounce/BounceMessage.h | 18 +- .../nixl_utils/bounce/BounceResult.cpp | 171 +++ .../nixl_utils/bounce/BounceResult.h | 159 +++ .../nixl_utils/bounce/BounceTransferStatus.h | 17 +- .../nixl_utils/bounce/BounceTransport.cpp | 437 ++++--- .../nixl_utils/bounce/BounceTransport.h | 174 ++- .../nixl_utils/bounce/CMakeLists.txt | 2 +- .../unit_tests/executor/bounce/CMakeLists.txt | 11 + .../executor/bounce/bounceAgentE2ETest.cpp | 9 +- .../bounce/bounceFailedRequestTest.cpp | 324 ++++++ .../bounce/bounceTransportFailureTest.cpp | 1000 ++++++++++------- 13 files changed, 1597 insertions(+), 741 deletions(-) create mode 100644 cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.cpp create mode 100644 cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h create mode 100644 cpp/tests/unit_tests/executor/bounce/bounceFailedRequestTest.cpp diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 73f672e26507..4e1aedc3902a 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -374,11 +374,11 @@ class TransferStatus { public: virtual ~TransferStatus() = default; - /// NIXL statuses return true only on success. Bounce statuses return true on success and on a - /// failure proven never to write the destination again (a source read by an orphaned local - /// gather may still be in flight); a failure without that proof stays false. Backends that - /// cannot distinguish failure from completion (mooncake) return true on any terminal state — - /// use wait() for the outcome. + /// True once the transfer is over and will not read its source or write its destination memory + /// again. After a failure it may turn true when the result is set or later, or never if the backend + /// cannot prove it. It does not turn back to false, except that a NIXL status reads false after release() or + /// once its agent is gone. Mooncake cannot tell failure from completion and returns true on any + /// terminal state. The outcome comes from wait(). [[nodiscard]] virtual bool isCompleted() const = 0; virtual TransferState wait(int64_t timeout_ms = -1) const = 0; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h index bd5acabd628e..e9688db82f32 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h @@ -166,9 +166,9 @@ struct BounceConfig int requestTimeoutMs{30000}; // request_timeout_ms; must be > 0 — the whole failure model // (abandoned-flow resolution, receiver lease, quarantine) hangs off // this timer, so applyParam rejects 0 (negatives already fail the - // strict unsigned parse) and keeps the current value. It also bounds - // how long a failed request waits for the receiver's answers (after - // the failure, and again after each answer) before resolving in-doubt. + // strict unsigned parse) and keeps the current value, and BounceSender + // rejects a non-positive one. It also bounds a failed request's answer + // wait (see FailedRequest). // Receiver-side lease on granted regions — DERIVED in deriveDependentTimeouts() as // 2 x requestTimeoutMs, not an independent knob (tests may still set the field directly). A // dead sender emits neither DATA nor a cancel, which is unobservable through the protocol diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h index 61ddcb4d5cdb..c62ad5bcdb23 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h @@ -97,10 +97,11 @@ static_assert(sizeof(BounceCreditEntry) == 24, "BounceCreditEntry must be 24 byt static_assert(sizeof(BounceScatterRun) == 36, "BounceScatterRun must be 36 bytes"); inline constexpr std::uint32_t kBounceMagic = 0x424E4332U; // 'B''N''C''2' -// v2: DATA scatter entries became strided RUNS (BounceScatterRun). The out-of-band capability -// handshake checks this version before bounce is selected; decodeHeader also rejects a mismatched -// control message defensively. -inline constexpr std::uint16_t kBounceVersion = 2U; +// v2: DATA scatter entries became strided RUNS (BounceScatterRun). v3: a chunk NACK, like an ACK, +// follows the chunk's last write into the destination, and a refused DATA is NACKed; the sender relies +// on both. The out-of-band capability handshake checks this version before bounce is selected; +// decodeHeader also rejects a mismatched control message defensively. +inline constexpr std::uint16_t kBounceVersion = 3U; // ---- encode (each returns a self-contained blob: header + payload) ---- /// WANT carries the per-chunk byte sizes the sender will write (the receiver allocates a region of @@ -122,10 +123,11 @@ inline constexpr std::uint16_t kBounceVersion = 2U; [[nodiscard]] std::string encodeData(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint32_t numChunks, std::uint64_t regionHandle, std::vector const& entries); [[nodiscard]] std::string encodeAck(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint64_t regionHandle); -/// NACK: the receiver KNOWS this chunk/request cannot complete (scatter failure, ungrantable chunk -/// size, plan overflow), so the sender fails it now instead of waiting out requestTimeoutMs. -/// Header-only like ACK. The sender fails the WHOLE request on any NACK, so chunkIdx/regionHandle are -/// informational only (0 when the receiver has no chunk context, e.g. a rejected WANT). +/// NACK: the receiver KNOWS this chunk/request cannot complete (scatter failure, a DATA it refused, +/// ungrantable chunk size, plan overflow), so the sender fails it now instead of waiting out +/// requestTimeoutMs. Header-only like ACK. A chunk NACK names (chunkIdx, regionHandle) like an ACK and +/// promises, like an ACK, that the chunk will not write the destination again. A request-level NACK (a +/// rejected WANT) carries 0/0; it precedes every grant, so it matches no sent chunk. [[nodiscard]] std::string encodeNack(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint64_t regionHandle); // ---- decode ---- diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.cpp new file mode 100644 index 000000000000..fc58e861ec82 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.cpp @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h" + +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +char const* toString(BounceFailReason reason) +{ + // Prefixed with the module name: these strings travel up through TransferStatus:: + // getLastStatusStr() into upper-layer error logs, where "RDMA write failed" alone would not + // say WHICH transport failed. + switch (reason) + { + case BounceFailReason::kNone: return "bounce: none"; + case BounceFailReason::kPlanRejected: return "bounce: plan rejected (request did not fit a transfer plan)"; + case BounceFailReason::kNoProgressTimeout: return "bounce: no GRANT/ACK progress within requestTimeoutMs"; + case BounceFailReason::kPeerDropped: return "bounce: peer dropped (forgetPeer/invalidateRemoteAgent)"; + case BounceFailReason::kGatherFailed: return "bounce: gather kernel failed (CUDA error)"; + case BounceFailReason::kWriteFailed: return "bounce: RDMA write failed"; + case BounceFailReason::kProtocolError: return "bounce: protocol error (GRANT mispair/plan overflow)"; + case BounceFailReason::kShutdown: return "bounce: transport shut down while pending"; + case BounceFailReason::kPeerNack: return "bounce: receiver reported it cannot complete the chunk/request (NACK)"; + } + return "bounce: unknown"; +} + +std::chrono::milliseconds answerWaitFor(BounceFailReason trigger, std::chrono::milliseconds requestTimeout) +{ + return trigger == BounceFailReason::kNoProgressTimeout ? std::chrono::milliseconds{0} : requestTimeout; +} + +OutstandingAccess::OutstandingAccess(std::vector unansweredChunks) + : mUnansweredChunks(std::move(unansweredChunks)) +{ +} + +OutstandingAccess::SourceRead::SourceRead(std::shared_ptr const& access) + : mAccess(access) +{ + std::lock_guard lk(access->mMu); + ++access->mSourceReads; +} + +OutstandingAccess::SourceRead::~SourceRead() +{ + if (mAccess != nullptr) // null once moved from + { + std::lock_guard lk(mAccess->mMu); + --mAccess->mSourceReads; + } +} + +bool OutstandingAccess::answer(std::uint32_t chunkIdx, std::uint64_t regionHandle) +{ + std::lock_guard lk(mMu); + auto const chunk = std::find_if(mUnansweredChunks.begin(), mUnansweredChunks.end(), + [&](SentChunk const& c) { return c.chunkIdx == chunkIdx && c.regionHandle == regionHandle; }); + if (chunk == mUnansweredChunks.end()) + { + return false; + } + mUnansweredChunks.erase(chunk); + return true; +} + +std::size_t OutstandingAccess::unansweredChunks() const +{ + std::lock_guard lk(mMu); + return mUnansweredChunks.size(); +} + +bool OutstandingAccess::quiesced() const +{ + std::lock_guard lk(mMu); + return mSourceReads == 0 && mUnansweredChunks.empty(); +} + +bool BounceResult::destinationQuiesced() const +{ + return outstanding == nullptr || outstanding->unansweredChunks() == 0; +} + +bool BounceResult::quiesced() const +{ + return outstanding == nullptr || outstanding->quiesced(); +} + +FailedRequest::FailedRequest(BounceFailReason reason, std::shared_ptr> promise, + std::shared_ptr outstanding, Clock::duration answerWait, Clock::time_point now) + : mReason(reason) + , mAnswerWait(answerWait) + , mRecord(outstanding) +{ + mAwaiting.emplace(AwaitingAnswers{std::move(promise), std::move(outstanding), now + answerWait}); +} + +bool FailedRequest::answer(std::uint32_t chunkIdx, std::uint64_t regionHandle, Clock::time_point now) +{ + auto const record = mRecord.lock(); + if (record == nullptr || !record->answer(chunkIdx, regionHandle)) + { + return false; + } + if (mAwaiting) + { + mAwaiting->deadline = now + mAnswerWait; // a live receiver gets a fresh wait for each remaining answer + } + return true; +} + +std::optional FailedRequest::resolveIfDue(Clock::time_point now) +{ + bool const due = mAwaiting && (mAwaiting->outstanding->unansweredChunks() == 0 || now >= mAwaiting->deadline); + return due ? resolve() : std::nullopt; +} + +std::optional FailedRequest::resolve() +{ + if (!mAwaiting) + { + return std::nullopt; + } + AwaitingAnswers awaiting = std::move(*mAwaiting); + mAwaiting.reset(); + // Count before the record moves into the result. + auto const unanswered = awaiting.outstanding->unansweredChunks(); + awaiting.promise->set_value({TransferState::kFAILURE, mReason, std::move(awaiting.outstanding)}); + return unanswered; +} + +BounceFailReason FailedRequest::reason() const +{ + return mReason; +} + +bool FailedRequest::resolved() const +{ + return !mAwaiting; +} + +bool FailedRequest::worthKeeping() const +{ + return !resolved() || unansweredChunks() > 0; +} + +std::size_t FailedRequest::unansweredChunks() const +{ + auto const record = mRecord.lock(); + return record == nullptr ? 0 : record->unansweredChunks(); +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h new file mode 100644 index 000000000000..d3fb7f38ee1c --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h @@ -0,0 +1,159 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/executor/transferAgent.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +/// Why a bounce request resolved kFAILURE — carried inside the request's result future and +/// surfaced to the upper layer through TransferStatus::getLastStatusStr(); the bare kFAILURE +/// enum alone tells an operator nothing. +enum class BounceFailReason : std::uint8_t +{ + kNone = 0, + kPlanRejected, // request did not fit a transfer plan (submit-time TLLM_CHECK) + kNoProgressTimeout, // no GRANT/ACK progress within requestTimeoutMs (peer unreachable/stalled) + kPeerDropped, // forgetPeer()/invalidateRemoteAgent while the request was in flight + kGatherFailed, // gather kernel launch / event-record / async poll failure (CUDA error) + kWriteFailed, // the RDMA write reported failure (getXferStatus == failed) + kProtocolError, // GRANT mispair or plan overflow — the flow was abandoned + kShutdown, // transport shut down while the request was still pending + kPeerNack, // receiver reported it cannot complete the chunk/request (NACK) +}; + +[[nodiscard]] char const* toString(BounceFailReason reason); + +/// How long a request failed by `trigger` holds its result back for the receiver's answers, restarted by +/// each answer. None after kNoProgressTimeout, which already saw no progress for requestTimeoutMs: the +/// result resolves at once, in-doubt if a sent chunk is unanswered (its DATA may be recent; a late answer +/// still ends it). +[[nodiscard]] std::chrono::milliseconds answerWaitFor( + BounceFailReason trigger, std::chrono::milliseconds requestTimeout); + +/// What a failed request may still do to caller memory after its result is set: gathers that may still +/// read its source, and sent chunks the receiver has not answered — it answers a chunk only after the +/// chunk's scatter has last written the destination. Once the result is set this only shrinks. Updated by +/// the sender's IO thread, read from any thread. +class OutstandingAccess +{ +public: + /// A chunk whose DATA is out, named the way the receiver's ACK/NACK names it. + struct SentChunk + { + std::uint32_t chunkIdx{}; + std::uint64_t regionHandle{}; + }; + + /// Recorded once, before the result is set: afterwards answers and SourceReads can only shrink it. + explicit OutstandingAccess(std::vector unansweredChunks = {}); + + /// One gather that may still read the source, until this is destroyed. It keeps the record alive. + class SourceRead + { + public: + explicit SourceRead(std::shared_ptr const& access); + SourceRead(SourceRead&& other) noexcept = default; + SourceRead(SourceRead const&) = delete; + SourceRead& operator=(SourceRead const&) = delete; + SourceRead& operator=(SourceRead&&) = delete; + ~SourceRead(); + + private: + std::shared_ptr mAccess; + }; + + /// The receiver answered (ACK or NACK) the chunk. False if it is not an unanswered chunk. + bool answer(std::uint32_t chunkIdx, std::uint64_t regionHandle); + + [[nodiscard]] std::size_t unansweredChunks() const; + [[nodiscard]] bool quiesced() const; + +private: + mutable std::mutex mMu; + std::uint32_t mSourceReads{0}; + std::vector mUnansweredChunks; +}; + +/// A bounce request's single result: the terminal state, on kFAILURE its cause, and what the request may +/// still access. The state and cause never change; the outstanding accesses only end. +struct BounceResult +{ + TransferState state{TransferState::kFAILURE}; + BounceFailReason reason{BounceFailReason::kNone}; + std::shared_ptr outstanding; // null: nothing outstanding + + [[nodiscard]] bool destinationQuiesced() const; + [[nodiscard]] bool quiesced() const; +}; + +/// A failed request's result, held back until every sent chunk is answered or the answer wait is over (a +/// caller that trusts wait() would otherwise reuse the destination while the receiver may still scatter into +/// it). Answers keep shrinking the record after that, so isCompleted() can turn true after wait() returns. +class FailedRequest +{ +public: + using Clock = std::chrono::steady_clock; + + /// `answerWait` is how long the receiver gets to answer, from now and again after each answer. + FailedRequest(BounceFailReason reason, std::shared_ptr> promise, + std::shared_ptr outstanding, Clock::duration answerWait, Clock::time_point now); + + /// The receiver answered (ACK or NACK) a sent chunk. False if it is not an unanswered chunk. + bool answer(std::uint32_t chunkIdx, std::uint64_t regionHandle, Clock::time_point now); + /// Resolve the result once every sent chunk is answered or the answer wait is over. If this call resolved + /// it: the number of sent chunks still unanswered then. + std::optional resolveIfDue(Clock::time_point now); + /// Resolve the result whatever is still unanswered (shutdown); returns like resolveIfDue. + std::optional resolve(); + + [[nodiscard]] BounceFailReason reason() const; + [[nodiscard]] bool resolved() const; + /// The result is held back, or an answer can still change it (a sent chunk is unanswered and the record + /// is alive). + [[nodiscard]] bool worthKeeping() const; + +private: + struct AwaitingAnswers + { + std::shared_ptr> promise; + std::shared_ptr outstanding; // keeps mRecord alive until the result is set + Clock::time_point deadline; + }; + + /// 0 once the record has expired. + [[nodiscard]] std::size_t unansweredChunks() const; + + BounceFailReason mReason; + Clock::duration mAnswerWait; + std::optional mAwaiting; // empty once resolved + // Once resolved, the result's holders and running orphaned gathers keep the record alive. + std::weak_ptr mRecord; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h index 3724e081ca20..614d7ce4404f 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h @@ -17,7 +17,7 @@ #pragma once -#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h" #include "tensorrt_llm/executor/transferAgent.h" #include @@ -39,18 +39,11 @@ class BounceTransferStatus final : public TransferStatus { } + /// The TransferStatus contract; see FailedRequest for when a failure gets there. [[nodiscard]] bool isCompleted() const override { - // Terminal AND the destination is quiesced (BounceResult::quiesced): kSUCCESS, or a kFAILURE the - // receiver has fully answered; an in-doubt kFAILURE stays false. Callers read the outcome from - // wait()/getLastStatusStr. get() cannot throw: every erase/shutdown path resolves the promise - // with set_value first, so broken_promise is unreachable. - if (!isReady()) - { - return false; - } - auto const& result = mFut.get(); - return result.state == TransferState::kSUCCESS || result.quiesced; + // Readiness first: the sender records every outstanding access before it sets the result. + return isReady() && mFut.get().quiesced(); } [[nodiscard]] TransferState wait(int64_t timeoutMs) const override @@ -84,7 +77,7 @@ class BounceTransferStatus final : public TransferStatus return {}; } std::string detail = toString(result.reason); - if (!result.quiesced) + if (!result.destinationQuiesced()) { detail += "; destination not proven quiescent"; } diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp index 3e0f68085c09..9fd19267a42a 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h" +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/tllmException.h" @@ -34,26 +35,6 @@ namespace tensorrt_llm::executor::kv_cache::bounce { -char const* toString(BounceFailReason reason) -{ - // Prefixed with the module name: these strings travel up through TransferStatus:: - // getLastStatusStr() into upper-layer error logs, where "RDMA write failed" alone would not - // say WHICH transport failed. - switch (reason) - { - case BounceFailReason::kNone: return "bounce: none"; - case BounceFailReason::kPlanRejected: return "bounce: plan rejected (request did not fit a transfer plan)"; - case BounceFailReason::kNoProgressTimeout: return "bounce: no GRANT/ACK progress within requestTimeoutMs"; - case BounceFailReason::kPeerDropped: return "bounce: peer dropped (forgetPeer/invalidateRemoteAgent)"; - case BounceFailReason::kGatherFailed: return "bounce: gather kernel failed (CUDA error)"; - case BounceFailReason::kWriteFailed: return "bounce: RDMA write failed"; - case BounceFailReason::kProtocolError: return "bounce: protocol error (GRANT mispair/plan overflow)"; - case BounceFailReason::kShutdown: return "bounce: transport shut down while pending"; - case BounceFailReason::kPeerNack: return "bounce: receiver reported it cannot complete the chunk/request (NACK)"; - } - return "bounce: unknown"; -} - namespace { constexpr char kSep = '\x1f'; // unit separator: agent names won't contain it @@ -354,17 +335,27 @@ void BounceReceiver::onData(std::string const& peer, BounceMsgHeader const& h, s // software leg of the sender's ackWait (the message copy cost scales with the entry count). BounceNvtxScope onDataScope(kNvtxOnData, "onData rid=%llu chunk=%u bytes=%zu", static_cast(h.requestId), h.chunkIdx, blob.size()); - std::vector entries; - if (!decodeScatter(blob, h, entries)) + // Claim the region for this DATA's scatter. A region a scatter already holds means a duplicate or stale + // DATA (DATA is not retransmitted): drop it without an answer, since that scatter may be this very + // chunk's and only its own answer may end the sender's wait. + auto const [claim, claimed] = mScattering.try_emplace(h.regionHandle, false); + if (!claimed) { + TLLM_LOG_WARNING( + "BounceTransport(%s): dropping duplicate/stale DATA for a region still scattering peer=%s rid=%llu " + "chunk=%u region=%llu", + mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, + static_cast(h.regionHandle)); return; } + // NACK a DATA this receiver will not scatter, so its sender stops waiting for the chunk: an undecodable one, + // or one whose region the flow no longer holds (cancelled/reclaimed; scattering could read another flow's data). + std::vector entries; auto const key = makeKey(peer, h.requestId); - // Drop a DATA whose region this flow no longer holds — it was cancelled/reclaimed (e.g. an empty - // WANT raced ahead of this DATA), and the region may have been re-granted to another flow. - // Scattering it would read a freed/re-owned region and corrupt that other flow's data. - if (!mCtx.scheduler.heldByFlow(key, h.regionHandle)) + if (!decodeScatter(blob, h, entries) || !mCtx.scheduler.heldByFlow(key, h.regionHandle)) { + mScattering.erase(claim); + mCtx.channel->sendTo(peer, encodeNack(h.requestId, h.chunkIdx, h.regionHandle)); return; } ScatterJob job; @@ -379,16 +370,6 @@ void BounceReceiver::onData(std::string const& peer, BounceMsgHeader const& h, s job.entries = std::move(entries); job.nvtxQueue = bounceRangeStart( kNvtxScatterQueue, "scatterQueue rid=%llu chunk=%u", static_cast(h.requestId), h.chunkIdx); - // DATA is not retransmitted by this protocol. An occupied region therefore indicates stale or - // malformed input; drop it instead of adding replay state or risking two readers of one grant. - if (!mScattering.emplace(job.offset, false).second) - { - TLLM_LOG_WARNING("BounceTransport(%s): dropping duplicate DATA peer=%s rid=%llu chunk=%u region=%llu", - mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, - static_cast(h.regionHandle)); - bounceRangeEnd(job.nvtxQueue); - return; - } { std::lock_guard lk(mJobMu); mJobs.emplace_back(std::move(job)); @@ -547,9 +528,8 @@ void BounceReceiver::checkTimeouts() mCtx.sendGrants(mCtx.scheduler.reapQuarantine()); if (mCtx.cfg.receiverFlowTimeoutMs <= 0) { - // Defensive only, mirroring BounceSender::checkTimeouts(): the config layer rejects - // request_timeout_ms <= 0 and the lease is derived from it, so this is reachable only - // through a directly-constructed BounceConfig. + // Defensive only: the lease is derived from requestTimeoutMs (> 0, checked by BounceSender), so a + // non-positive one comes only from a BounceConfig whose lease was set directly. return; } for (auto const& flow : mCtx.scheduler.staleFlows(std::chrono::milliseconds(mCtx.cfg.receiverFlowTimeoutMs))) @@ -761,9 +741,11 @@ void BounceReceiver::scatterWorkerLoop() } } cudaError_t syncErr = cudaSuccess; - if (launchErr == cudaSuccess) { - // Sync leg: the actual GPU wait (kernel queueing + run time on the exec stream). + // Sync leg: the actual GPU wait (kernel queueing + run time on the exec stream). Also after a + // failed launch: launchBatchedCopy reports cudaGetLastError(), which may be stale (e.g. from the + // warn-only cudaSetDevice above) with the scatter enqueued, and a NACK, like an ACK, must follow + // the chunk's last write into dst. BounceNvtxScope syncScope(kNvtxScatterSync, "scatterSync rid=%llu chunk=%u", static_cast(job.rid), job.chunkIdx); syncErr = cudaStreamSynchronize(ctx->stream); @@ -771,15 +753,6 @@ void BounceReceiver::scatterWorkerLoop() bool const ok = srcInBounds && launchErr == cudaSuccess && syncErr == cudaSuccess; if (!ok) { - if (launchErr != cudaSuccess) - { - // A failed launch can still have enqueued the scatter: launchBatchedCopy reports - // cudaGetLastError(), which may be a stale error from an earlier call on this thread - // (e.g. the warn-only cudaSetDevice above). The NACK below tells the sender this - // chunk will not write dst again, so drain the stream before sending it. (Also runs - // when nothing was launched — launchErr keeps its initial value — on an idle stream.) - (void) cudaStreamSynchronize(ctx->stream); - } // Clear any non-sticky error so the next launch does not inherit it; a sticky // (context-level) error cannot be cleared and will surface on the next CUDA call. (void) cudaGetLastError(); @@ -836,6 +809,8 @@ void BounceReceiver::scatterWorkerLoop() BounceSender::BounceSender(BounceContext& ctx) : mCtx(ctx) { + TLLM_CHECK_WITH_INFO( + mCtx.cfg.requestTimeoutMs > 0, "bounce requestTimeoutMs must be > 0, got %d", mCtx.cfg.requestTimeoutMs); } std::shared_future BounceSender::submit( @@ -860,14 +835,13 @@ std::shared_future BounceSender::submit( // silently absorbed as a slow success. Whether to retry is the caller's decision. TLLM_LOG_WARNING( "BounceTransport(%s): plan build for peer %s rejected: %s", mCtx.selfName.c_str(), peer.c_str(), e.what()); - // Nothing was sent, so nothing can write the destination. - promise->set_value({TransferState::kFAILURE, BounceFailReason::kPlanRejected, /*quiesced=*/true}); + promise->set_value({TransferState::kFAILURE, BounceFailReason::kPlanRejected}); return fut; } auto const numChunks = static_cast(plan.numChunks()); if (numChunks == 0) { - promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone, /*quiesced=*/true}); + promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); return fut; } @@ -941,10 +915,6 @@ void BounceSender::onGrant(std::string const& peer, BounceMsgHeader const& h, st mCtx.selfName.c_str(), peer.c_str(), req.peer.c_str(), static_cast(h.requestId)); return; } - if (req.failure) - { - return; // awaiting answers: nothing will be written; its cancel reclaims what this GRANT holds - } bounceRangeEnd(req.nvtxGrantWait); // first GRANT ends the credit-wait span (no-op on later GRANTs) bounceRangeEnd(req.nvtxCreditStarved); // a GRANT ends the current starvation period (pump may reopen one) for (auto const& credit : credits) @@ -1035,10 +1005,6 @@ void BounceSender::attachCredits(std::uint64_t rid, Request& req) void BounceSender::pumpRequest(std::uint64_t rid, Request& req) { - if (req.failure) - { - return; // awaiting answers: posted and pendingCredits stay empty - } if (req.abandonReason != BounceFailReason::kNone) { // Abandoned (GRANT mispair / plan overflow): stop attaching credits and launching gathers. @@ -1297,12 +1263,7 @@ bool BounceSender::drainGatherReady() } for (auto rid : toFail) { - auto it = mRequests.find(rid); - if (it != mRequests.end()) - { - failRequest(rid, it->second, BounceFailReason::kGatherFailed); - didWork = true; - } + didWork |= failRequest(rid, BounceFailReason::kGatherFailed, std::chrono::steady_clock::now()); } return didWork; } @@ -1363,12 +1324,7 @@ bool BounceSender::pollSenderHandles() } for (auto rid : toFail) { - auto it = mRequests.find(rid); - if (it != mRequests.end()) - { - failRequest(rid, it->second, BounceFailReason::kWriteFailed); - didWork = true; - } + didWork |= failRequest(rid, BounceFailReason::kWriteFailed, std::chrono::steady_clock::now()); } return didWork; } @@ -1379,10 +1335,7 @@ void BounceSender::onNack(std::string const& peer, BounceMsgHeader const& h) auto it = mRequests.find(h.requestId); if (it == mRequests.end()) { - // Expected late arrival: the request already resolved (e.g. an answer arriving after it - // resolved in-doubt on requestTimeoutMs, or a NACK racing kNoProgressTimeout). - TLLM_LOG_DEBUG("BounceTransport(%s): dropping NACK for unknown rid=%llu peer=%s chunk=%u", - mCtx.selfName.c_str(), static_cast(h.requestId), peer.c_str(), h.chunkIdx); + answerFailedRequest(peer, h); return; } Request& req = it->second; @@ -1393,19 +1346,11 @@ void BounceSender::onNack(std::string const& peer, BounceMsgHeader const& h) h.chunkIdx); return; } - if (req.failure) - { - onAnswerAfterFailure(req, h); - return; - } // The NACK answers its own chunk (the receiver never scattered it, or synced before NACKing), so - // recycle that chunk now: failRequest must not wait for an answer that already arrived. A + // recycle that chunk now: the failed request must not wait for an answer that already arrived. A // request-level NACK (rejected WANT) arrives before any chunk was sent and matches none. releaseAnsweredChunk(req, h); - // Same terminal path as a failed RDMA write (pollSenderHandles): failRequest recycles/defers the - // staging regions, sends the cancel that reclaims the receiver's regions and resolves the future - // once the receiver has answered every other sent chunk. - failRequest(h.requestId, req, BounceFailReason::kPeerNack); + failRequest(h.requestId, BounceFailReason::kPeerNack, std::chrono::steady_clock::now()); } void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) @@ -1419,6 +1364,7 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) auto it = mRequests.find(h.requestId); if (it == mRequests.end()) { + answerFailedRequest(peer, h); return; } Request& req = it->second; @@ -1429,13 +1375,6 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) h.chunkIdx); return; } - if (req.failure) - { - // A failed request resolves FAILURE even if every chunk ends up ACKed (e.g. forgetPeer after - // all were Sent): never take the SUCCESS path below. - onAnswerAfterFailure(req, h); - return; - } if (!releaseAnsweredChunk(req, h)) { TLLM_LOG_WARNING("BounceTransport(%s): dropping stale/invalid ACK peer=%s rid=%llu chunk=%u region=%llu", @@ -1449,7 +1388,9 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk was ACKed if (req.acked >= req.numChunks) { - resolve(req, {TransferState::kSUCCESS, BounceFailReason::kNone, /*quiesced=*/true}); + endPipelineSpans(req); + bounceRangeEnd(req.nvtxReq); + req.promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); mRequests.erase(it); } } @@ -1473,73 +1414,86 @@ bool BounceSender::releaseAnsweredChunk(Request& req, BounceMsgHeader const& h) return true; } -void BounceSender::onAnswerAfterFailure(Request& req, BounceMsgHeader const& h) +void BounceSender::answerFailedRequest(std::string const& peer, BounceMsgHeader const& h) { - auto const answered = std::find_if(req.awaiting.begin(), req.awaiting.end(), - [&h](AnswerKey const& k) { return k.chunkIdx == h.chunkIdx && k.regionHandle == h.regionHandle; }); - if (answered == req.awaiting.end()) + auto it = mFailed.find(h.requestId); + if (it == mFailed.end()) + { + // Expected late arrival: the request is gone (every sent chunk answered, nobody holds its + // result any more, or shutdown). + TLLM_LOG_DEBUG("BounceTransport(%s): dropping answer for unknown rid=%llu peer=%s chunk=%u", + mCtx.selfName.c_str(), static_cast(h.requestId), peer.c_str(), h.chunkIdx); + return; + } + FailedEntry& failed = it->second; + if (peer != failed.peer) + { + TLLM_LOG_WARNING("BounceTransport(%s): dropping wrong-peer answer peer=%s expected=%s rid=%llu chunk=%u", + mCtx.selfName.c_str(), peer.c_str(), failed.peer.c_str(), static_cast(h.requestId), + h.chunkIdx); + return; + } + if (!failed.request.worthKeeping()) + { + // Settled, only waiting for this tick's sweep to drop it: as good as gone. + TLLM_LOG_DEBUG("BounceTransport(%s): dropping answer for settled rid=%llu peer=%s chunk=%u", + mCtx.selfName.c_str(), static_cast(h.requestId), peer.c_str(), h.chunkIdx); + return; + } + if (!failed.request.answer(h.chunkIdx, h.regionHandle, std::chrono::steady_clock::now())) { TLLM_LOG_WARNING( "BounceTransport(%s): dropping stale/invalid answer for failed request peer=%s rid=%llu chunk=%u " "region=%llu", - mCtx.selfName.c_str(), req.peer.c_str(), static_cast(h.requestId), h.chunkIdx, + mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, static_cast(h.regionHandle)); - return; - } - // Its staging region was already released when the request failed; only the bookkeeping is left. - req.awaiting.erase(answered); - // A live receiver gets a fresh requestTimeoutMs for each remaining answer. - req.lastProgress = std::chrono::steady_clock::now(); - if (req.awaiting.empty()) - { - resolveFailure(h.requestId, req); - mRequests.erase(h.requestId); } } void BounceSender::checkTimeouts() { - if (mCtx.cfg.requestTimeoutMs <= 0) - { - // Defensive only: the config layer rejects request_timeout_ms <= 0 (abandoned-flow - // resolution and the receiver lease both hang off this timer), so this is reachable only - // through a directly-constructed BounceConfig. - return; - } + std::lock_guard lk(mReqMu); auto const now = std::chrono::steady_clock::now(); - auto const limit = std::chrono::milliseconds(mCtx.cfg.requestTimeoutMs); - std::vector stuck; + failStalledRequests(now); + settleFailedRequests(now); +} + +void BounceSender::failStalledRequests(std::chrono::steady_clock::time_point now) +{ + std::vector stalled; + for (auto const& [rid, req] : mRequests) { - std::lock_guard lk(mReqMu); - for (auto& [rid, req] : mRequests) + if (now - req.lastProgress > requestTimeout()) { - if (now - req.lastProgress > limit) - { - stuck.push_back(rid); - } + stalled.push_back(rid); } - for (auto rid : stuck) + } + for (auto rid : stalled) + { + // Peer never granted or stopped making progress (unreachable / not bounce-ready / congested): fail + // the request so wait() returns instead of hanging. + failRequest(rid, BounceFailReason::kNoProgressTimeout, now); + } +} + +void BounceSender::settleFailedRequests(std::chrono::steady_clock::time_point now) +{ + for (auto it = mFailed.begin(); it != mFailed.end();) + { + if (auto const unanswered = it->second.request.resolveIfDue(now)) { - auto it = mRequests.find(rid); - if (it == mRequests.end()) - { - continue; - } - Request& req = it->second; - if (req.failure) - { - // No answer for requestTimeoutMs while chunks are still awaited: in-doubt (see failRequest). - resolveFailure(rid, req); - mRequests.erase(it); - continue; - } - // Peer never granted or stopped making progress (unreachable / not bounce-ready / - // congested). Fail the request so wait() returns FAILURE instead of hanging. - failRequest(rid, req, BounceFailReason::kNoProgressTimeout); + onFailureResolved(it->first, it->second, *unanswered); } + it = it->second.request.worthKeeping() ? std::next(it) : mFailed.erase(it); } } +std::size_t BounceSender::failedRequestCount() +{ + std::lock_guard lk(mReqMu); + return mFailed.size(); +} + bool BounceSender::drainOrphanLocal() { if (mOrphanLocal.empty()) @@ -1635,7 +1589,7 @@ bool BounceSender::drainOrphanGather() mCtx.selfName.c_str(), static_cast(o.rid), static_cast(o.localOffset), static_cast(ageMs)); } - keep.push_back(o); // kernel may still be writing the region -> wait + keep.push_back(std::move(o)); // kernel may still be writing the region -> wait continue; } if (st != cudaSuccess) @@ -1657,23 +1611,34 @@ bool BounceSender::drainOrphanGather() return didWork; } -void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason reason) +bool BounceSender::failRequest(std::uint64_t rid, BounceFailReason trigger, std::chrono::steady_clock::time_point now) { - if (req.failure) + auto node = mRequests.extract(rid); + if (node.empty()) { - // Already awaiting answers (forget() reaches here again): keep the first reason and the answer - // deadline; staging and the cancel were handled at the first failure. - return; + return false; } - // An abandoned flow (GRANT mispair / plan overflow) reaches here through the timeout path; - // report the specific abandon cause, not the generic timeout. - BounceFailReason const reported = req.abandonReason != BounceFailReason::kNone ? req.abandonReason : reason; - req.failure = reported; + Request& req = node.mapped(); + BounceFailReason const cause = req.failureCause(trigger); // The one log line every sender-side failure passes through (timeout, peer drop, write/gather // failure) — keep it WARNING and keep the progress context. TLLM_LOG_WARNING("BounceTransport(%s): request FAILED rid=%llu peer=%s reason=\"%s\" chunks acked=%u posted=%u/%u", - mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(reported), req.acked, + mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(cause), req.acked, req.nextPost, req.numChunks); + auto const outstanding = unansweredSentChunks(req); + bool const writeInFlight = retireChunks(rid, req, outstanding, now); + cancelFlow(rid, req.peer, writeInFlight); + // The answer wait follows the trigger, not the reported cause: an abandoned flow reports + // kProtocolError but fails through the no-progress timeout. settleFailedRequests resolves it later + // in this same tick: every failRequest caller runs in tick() before checkTimeouts. + mFailed.emplace( + rid, toFailedEntry(std::move(req), cause, outstanding, answerWaitFor(trigger, requestTimeout()), now)); + return true; +} + +bool BounceSender::retireChunks(std::uint64_t rid, Request& req, std::shared_ptr const& outstanding, + std::chrono::steady_clock::time_point now) +{ // Release in-flight transfer handles and return each gather-staging region to the shared arena — // but only once nothing is still touching the region's memory, else recycling races a live DMA. // Nothing here blocks on the GPU or the NIC (this runs on the IO thread, commonly from onNack): @@ -1683,112 +1648,108 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason // gather would scribble a re-granted region — the write is not ordered against the new owner; // GatherFailed includes "launch OK, event-record failed", so its event cannot be trusted but // its stream can). Park it in mOrphanGather; drainOrphanGather() returns the ctx + region once - // cudaStreamQuery(ctx->stream) reports the stream idle. + // cudaStreamQuery(ctx->stream) reports the stream idle. Until then it may still read the + // source too. // - Gathered: gather done (event observed), ctx already returned -> recycle now. // - Sent: write landed (poll==kDone), xfer already released in pollSenderHandles, NIC done - // reading -> recycle now. Its DATA is out, though: the receiver may still scatter it into the - // destination, so remember it as awaiting the receiver's answer. - bool deferredWrite = false; + // reading -> recycle now. + bool writeInFlight = false; for (auto& p : req.posted) { - // Close this chunk's NVTX spans (whichever leg it died in); 0 handles are no-ops. - bounceRangeEnd(p.nvtxGather); - bounceRangeEnd(p.nvtxWrite); - bounceRangeEnd(p.nvtxAckWait); + endChunkSpans(p); if (p.state == PostState::Writing) { - mOrphanLocal.push_back( - OrphanLocal{std::move(p.xfer), p.localOffset, rid, std::chrono::steady_clock::now(), false}); - deferredWrite = true; + mOrphanLocal.push_back(OrphanLocal{ + .xfer = std::move(p.xfer), .offset = p.localOffset, .rid = rid, .since = now, .warned = false}); + writeInFlight = true; continue; // do NOT release xfer or recycle the region yet } if (p.ctx != nullptr) // Gathering or GatherFailed: the exec stream may still write the region { - mOrphanGather.push_back( - OrphanGather{p.ctx, p.localOffset, rid, std::chrono::steady_clock::now(), /*warned=*/false}); - p.ctx = nullptr; + mOrphanGather.push_back(OrphanGather{.ctx = p.ctx, + .localOffset = p.localOffset, + .rid = rid, + .since = now, + .warned = false, + .sourceRead = OutstandingAccess::SourceRead{outstanding}}); continue; // do NOT release ctx or recycle the region yet } - if (p.state == PostState::Sent) - { - req.awaiting.push_back(AnswerKey{p.chunkIdx, p.remoteHandle}); - } mCtx.sendGrants(mCtx.scheduler.releaseLocal(p.localOffset)); } - // Every chunk is now recycled, parked as an orphan or awaiting its answer, so the drain loops have - // nothing left to gather, post or poll for this request; its parked credits die with the cancel. - req.posted.clear(); - req.pendingCredits.clear(); + return writeInFlight; +} + +void BounceSender::cancelFlow(std::uint64_t rid, std::string const& peer, bool writeInFlight) +{ // Retract the credit request so the receiver stops holding/granting for it. If any RDMA write is // still in flight it is landing on the receiver's region; sending the cancel NOW would let the // receiver reclaim+re-grant that region under the write -> corruption. Defer the cancel until the // flow's writes drain (drainOrphanLocal sends it). With no in-flight write, send it immediately. // The cancel does not stop the receiver's queued scatters: they still run and answer. - if (deferredWrite) + if (writeInFlight) { - mPendingCancel[rid] = req.peer; + mPendingCancel[rid] = peer; } else { - mCtx.channel->sendTo(req.peer, encodeCancel(rid, mCtx.channel->localEndpoint())); - } - // Resolve at once when nothing is awaited; on kNoProgressTimeout (checkTimeouts, abandoned flows - // included), which already saw no progress for requestTimeoutMs, so an unanswered chunk is reported - // in-doubt rather than waited for again (its DATA may be recent); and when requestTimeoutMs <= 0 - // (direct test configs only), which leaves no bound to wait under. This tests the caller's `reason`, - // not `reported`: an abandoned flow reports kProtocolError but arrives here from checkTimeouts. - bool const waitForAnswers - = !req.awaiting.empty() && reason != BounceFailReason::kNoProgressTimeout && mCtx.cfg.requestTimeoutMs > 0; - if (waitForAnswers) - { - // Resolving now would let the caller reuse the destination while the receiver may still - // scatter into it. lastProgress restarts here and at each answer; checkTimeouts resolves - // in-doubt after requestTimeoutMs without one. The pipeline is over: only nvtxReq stays open. - bounceRangeEnd(req.nvtxGrantWait); - bounceRangeEnd(req.nvtxCreditStarved); - bounceRangeEnd(req.nvtxArenaStarved); - req.lastProgress = std::chrono::steady_clock::now(); - return; + mCtx.channel->sendTo(peer, encodeCancel(rid, mCtx.channel->localEndpoint())); } - resolveFailure(rid, req); - mRequests.erase(rid); } -void BounceSender::resolveFailure(std::uint64_t rid, Request& req) +BounceSender::FailedEntry BounceSender::toFailedEntry(Request&& req, BounceFailReason cause, + std::shared_ptr outstanding, std::chrono::milliseconds answerWait, + std::chrono::steady_clock::time_point now) { - // A sent chunk the receiver has not answered may still be scattered into the destination. - bool const quiesced = req.awaiting.empty(); - if (!quiesced) + endPipelineSpans(req); + return FailedEntry{.peer = req.peer, + .request = FailedRequest{cause, std::move(req.promise), std::move(outstanding), answerWait, now}, + .nvtxReq = req.nvtxReq}; +} + +std::shared_ptr BounceSender::unansweredSentChunks(Request const& req) +{ + std::vector sent; + for (auto const& p : req.posted) { - TLLM_LOG_WARNING( - "BounceTransport(%s): request rid=%llu peer=%s resolved FAILURE (\"%s\") with %zu sent chunk(s) " - "unanswered -> destination not proven quiescent (the receiver may still write it)", - mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(*req.failure), - req.awaiting.size()); + if (p.state == PostState::Sent) + { + sent.push_back({.chunkIdx = p.chunkIdx, .regionHandle = p.remoteHandle}); + } } - resolve(req, {TransferState::kFAILURE, *req.failure, quiesced}); + return std::make_shared(std::move(sent)); } -void BounceSender::resolve(Request& req, BounceResult const& result) +void BounceSender::endPipelineSpans(Request& req) { bounceRangeEnd(req.nvtxGrantWait); bounceRangeEnd(req.nvtxCreditStarved); bounceRangeEnd(req.nvtxArenaStarved); - bounceRangeEnd(req.nvtxReq); - try - { - req.promise->set_value(result); - } - catch (...) +} + +void BounceSender::endChunkSpans(Posted& p) +{ + // Whichever leg the chunk ended in; 0 handles are no-ops. + bounceRangeEnd(p.nvtxGather); + bounceRangeEnd(p.nvtxWrite); + bounceRangeEnd(p.nvtxAckWait); +} + +void BounceSender::onFailureResolved(std::uint64_t rid, FailedEntry& failed, std::size_t unanswered) +{ + bounceRangeEnd(failed.nvtxReq); + if (unanswered > 0) { - // set_value throws std::future_error ONLY if the promise is already satisfied — a benign - // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. + TLLM_LOG_WARNING( + "BounceTransport(%s): request rid=%llu peer=%s resolved FAILURE (\"%s\") with %zu sent chunk(s) " + "unanswered -> destination not proven quiescent (the receiver may still write it)", + mCtx.selfName.c_str(), static_cast(rid), failed.peer.c_str(), + toString(failed.request.reason()), unanswered); } } void BounceSender::forget(std::string const& peer) { - // Fail every request to the gone peer (failRequest leaves one already awaiting answers alone). + // Fail every active request to the gone peer. Answers to its sent chunks still reach our ROUTER. std::lock_guard lk(mReqMu); std::vector toFail; for (auto const& [rid, req] : mRequests) @@ -1800,19 +1761,15 @@ void BounceSender::forget(std::string const& peer) } for (auto rid : toFail) { - auto it = mRequests.find(rid); - if (it != mRequests.end()) - { - failRequest(rid, it->second, BounceFailReason::kPeerDropped); - } + failRequest(rid, BounceFailReason::kPeerDropped, std::chrono::steady_clock::now()); } } void BounceSender::failAll() { - // Fail any still-pending requests so no submit() future hangs, releasing their in-flight - // transfer handles first (same handle-leak fix as failRequest). Called after the device has been - // synced and the IO thread joined, so no lock contention — but keep mReqMu for consistency. + // Resolve every request, active or already failed, so no submit() future hangs. Called after the + // device has been synced and the IO thread joined, so no lock contention — but keep mReqMu for + // consistency. // cudaDeviceSynchronize covers only local kernels, NOT the NIC: TransferStatus::release() cannot // reliably cancel a posted NIXL write (UCX's release issues ucp_request_cancel, which acts only on // tag-matching receives, and returns success regardless; libfabric's release is a no-op), so an @@ -1868,13 +1825,13 @@ void BounceSender::failAll() } } std::vector> releaseRetry; + auto const now = std::chrono::steady_clock::now(); for (auto& [rid, req] : mRequests) { + auto const outstanding = unansweredSentChunks(req); for (auto& p : req.posted) { - bounceRangeEnd(p.nvtxGather); - bounceRangeEnd(p.nvtxWrite); - bounceRangeEnd(p.nvtxAckWait); + endChunkSpans(p); if (p.state == PostState::Writing && p.xfer != nullptr) { if (!p.xfer->release()) @@ -1887,20 +1844,22 @@ void BounceSender::failAll() mCtx.exec->release(p.ctx); // still-gathering chunk: return its borrowed exec context p.ctx = nullptr; } - if (p.state == PostState::Sent) - { - req.awaiting.push_back(AnswerKey{p.chunkIdx, p.remoteHandle}); - } } - if (!req.failure) + auto const cause = req.failureCause(BounceFailReason::kShutdown); + mFailed.emplace( + rid, toFailedEntry(std::move(req), cause, outstanding, /*answerWait=*/std::chrono::milliseconds{0}, now)); + } + mRequests.clear(); + // No answer can arrive any more, so every result resolves now; one with a sent chunk unanswered is + // in-doubt, since the caller's device sync drained our kernels, not the receiver's scatter. + for (auto& [rid, failed] : mFailed) + { + if (auto const unanswered = failed.request.resolve()) { - req.failure = BounceFailReason::kShutdown; // a request already awaiting answers keeps its reason + onFailureResolved(rid, failed, *unanswered); } - // Any sent chunk still unanswered makes the result in-doubt: the caller's device sync drained our - // kernels, not the receiver's scatter. - resolveFailure(rid, req); } - mRequests.clear(); + mFailed.clear(); // Release the handles of deferred writes (waited on above) without recycling their regions or // sending deferred control messages: no producer remains, and shutdown must not grant work // against an arena being torn down. diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h index b4fb96af7500..d2b8277434eb 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h @@ -20,6 +20,7 @@ #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ControlChannel.h" #include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h" @@ -36,7 +37,6 @@ #include #include #include -#include #include #include #include @@ -79,9 +79,12 @@ namespace tensorrt_llm::executor::kv_cache::bounce // availability have been observed. // [S] ACK(rid, chunk, h) -> releaseLocal() that chunk's region; acked++; // all acked -> resolve the request's future SUCCESS. -// (A request awaiting answers only matches them; see failRequest.) +// [S] NACK(rid, chunk, h) -> fail the request (kPeerNack); like an ACK, it answers its chunk. +// For a failed request, ACKs and NACKs end its unanswered sent chunks. // [R] WANT(rid, sizes[]) -> scheduler.onWant(peer:rid, sizes) -> send GRANT(s). -// [R] DATA(rid, chunk, h, plan) -> enqueue a scatter job over region `h`. +// [R] DATA(rid, chunk, h, plan) -> enqueue a scatter job over region `h`; a DATA it will not +// scatter (undecodable, or `h` no longer the flow's) -> NACK, +// unless region `h` is still scattering. // 2. poll every in-flight write (TransferStatus::wait(0)); on Done send DATA (data has // landed at the remote) and mark on-wire. THIS is why no notifMsg is needed. // 3. drain scatter completions posted by workers: release receiver regions and re-GRANT. @@ -93,54 +96,15 @@ namespace tensorrt_llm::executor::kv_cache::bounce // "peer\x1f rid" so multiple concurrent requests from one peer are independent flows. // // Lifetime: submit() returns a shared_future; it resolves SUCCESS on full ACK and -// FAILURE on transfer error, peer invalidation, shutdown, or request timeout. requestTimeoutMs must -// be > 0 (the config layer enforces this): abandoned-flow resolution and the receiver lease both -// hang off this timer, so disabling it would let a protocol-error flow hold wait() forever. -// A FAILURE first waits for the receiver's answer to every chunk whose DATA was sent, for at most -// requestTimeoutMs after the failure or after the receiver's last answer, so that -// BounceResult::quiesced can promise the destination is no longer written. +// FAILURE on transfer error, peer invalidation, shutdown, or request timeout (see +// BounceConfig::requestTimeoutMs). When a FAILURE resolves, and what it may still access, is +// FailedRequest's contract. // // The agent (data plane) & ControlChannel are injected (not owned) so the same reactor runs // unchanged in tests and production (both over the real NIXL/zmq stack; failure tests subclass // NixlTransferAgent and override postXferRequest to inject deterministic transfer faults). // ============================================================================ -/// Why a bounce request resolved kFAILURE — carried inside the request's result future and -/// surfaced to the upper layer through TransferStatus::getLastStatusStr(); the bare kFAILURE -/// enum alone tells an operator nothing. -enum class BounceFailReason : std::uint8_t -{ - kNone = 0, - kPlanRejected, // request did not fit a transfer plan (submit-time TLLM_CHECK) - kNoProgressTimeout, // no GRANT/ACK progress within requestTimeoutMs (peer unreachable/stalled) - kPeerDropped, // forgetPeer()/invalidateRemoteAgent while the request was in flight - kGatherFailed, // gather kernel launch / event-record / async poll failure (CUDA error) - kWriteFailed, // the RDMA write reported failure (getXferStatus == failed) - kProtocolError, // GRANT mispair or plan overflow — the flow was abandoned - kShutdown, // transport shut down while the request was still pending - kPeerNack, // receiver reported it cannot complete the chunk/request (NACK) -}; - -[[nodiscard]] char const* toString(BounceFailReason reason); - -/// A bounce request's single result: the terminal state plus, on kFAILURE, its cause. The future -/// is the ONE source of truth — no side channel to keep in sync. -struct BounceResult -{ - TransferState state{TransferState::kFAILURE}; - BounceFailReason reason{BounceFailReason::kNone}; - // The receiver will not write the destination again for this request: always on kSUCCESS, and on - // kFAILURE once the receiver has answered (ACK or NACK) every chunk whose DATA was sent, because - // only that chunk's scatter writes the destination and the answer follows it. A failed request - // waits for those answers. It resolves in-doubt (false) when no answer comes within - // requestTimeoutMs after the failure or after the receiver's last answer (assumed dead, as the - // receiver lease and quarantine assume), on kNoProgressTimeout with a sent chunk unanswered (no - // progress for requestTimeoutMs, though the last DATA may be recent), and at shutdown with a sent - // chunk unanswered (see BounceSender::failRequest and failAll). Covers the destination only: a - // gather orphaned by the failure may still be reading the source. - bool quiesced{false}; -}; - /// Shared dependencies used by both roles. Holds the injected channel/agent/arena/exec (borrowed, /// not owned), the one CreditScheduler that carves the shared arena for both roles, and /// sendGrants(). Owned by value by BounceTransport and referenced by the sender and receiver. @@ -289,7 +253,8 @@ class BounceSender explicit BounceSender(BounceContext& ctx); /// Submit a WRITE of (src -> dst) descriptors to `peer`. Returns a future that resolves - /// {kSUCCESS} once every chunk is scattered+ACKed, or {kFAILURE, reason} on error/shutdown. + /// {kSUCCESS} once every chunk is scattered+ACKed, or {kFAILURE, reason, outstanding} on + /// error/shutdown (see FailedRequest). /// Chunks are cut to cfg.maxChunkSizeBytes. NixlTransferAgent validates peer compatibility /// before calling this; direct users of BounceTransport must establish the same invariant /// themselves. @@ -300,7 +265,7 @@ class BounceSender void onAck(std::string const& peer, BounceMsgHeader const& h); /// The receiver reported it cannot complete a chunk/request: fail the request now (kPeerNack) /// instead of waiting out requestTimeoutMs. Unknown / wrong-peer NACKs are dropped like ACKs. - /// For a request that already failed, a NACK is just the answer it is waiting for. + /// For a failed request, a NACK answers one of its sent chunks, like an ACK. void onNack(std::string const& peer, BounceMsgHeader const& h); /// Issue the RDMA write for any chunk whose gather kernel has now completed (poll ctx->event, /// non-blocking). Decouples the IO thread from gather latency on a shared GPU. Returns true if it @@ -317,18 +282,19 @@ class BounceSender /// Retry parked credits for every request (called on the IO loop after ACKs free regions). void drainPendingPosts(); /// Fail requests that have made no progress within requestTimeoutMs (e.g. peer never granted), - /// and resolve in-doubt the failed requests whose receiver has sent no answer for that long. + /// resolve failed requests whose answer wait is over, and drop those no longer worth keeping. void checkTimeouts(); - /// A peer is gone: fail any in-flight request targeting it. Its wait() returns once its sent - /// chunks are answered (answers still arrive on our ROUTER) or the answer wait times out. + /// A peer is gone: fail every active request targeting it (see FailedRequest; answers to its sent + /// chunks still reach our ROUTER). void forget(std::string const& peer); - /// Shutdown: fail every still-pending request (a request already awaiting answers keeps its reason; - /// any request with an unanswered sent chunk resolves in-doubt), wait (bounded) for in-flight - /// outbound writes to reach a terminal state so the NIC is done reading the arena before it is - /// freed, then release each write's TransferStatus handle and return the parked gathers' exec - /// contexts. Recycles no regions (all are left to arena destruction) and sends no grants. Called - /// after the IO/workers are joined and local CUDA work has been synced. + /// Shutdown, called after the IO/workers are joined and local CUDA work has been synced: wait + /// (bounded) for in-flight outbound writes to reach a terminal state so the NIC is done reading the + /// arena before it is freed, then resolve every request (an active one fails with kShutdown unless + /// abandoned; see FailedRequest), release each write's TransferStatus handle and return the parked + /// gathers' exec contexts. Recycles no regions (all are left to arena destruction) and sends no grants. void failAll(); + /// Failed requests still kept for the receiver's answers (observability; takes the request lock). + [[nodiscard]] std::size_t failedRequestCount(); /// True while a local gather region is held or an orphaned write/gather is still in flight /// (drives the IO loop's 0ms busy-poll). Called on the IO thread. @@ -379,14 +345,6 @@ class BounceSender std::uint64_t nvtxAckWait{0}; // DATA sent -> ACK received }; - // A chunk whose DATA was sent, identified the way the receiver's ACK/NACK names it (the rid and - // peer are the owning request's). - struct AnswerKey - { - std::uint32_t chunkIdx{}; - std::uint64_t regionHandle{}; - }; - struct Request { std::string peer; @@ -410,14 +368,9 @@ class BounceSender // Set by the abandon sites (GRANT mispair / plan overflow): the request then dies through // the timeout path, but failRequest reports THIS more specific cause instead. BounceFailReason abandonReason{BounceFailReason::kNone}; - // Set once the request has failed. A failed request stays in mRequests only while awaiting - // answers, and then it only collects them: no more posts, gathers or credits. - std::optional failure; - std::vector awaiting; // sent chunks the receiver has not answered yet (once failed) - // Last time this request made forward progress (granted+posted a chunk, or got an ACK; for a - // failed request: the failure, or an answer). A request stuck with no progress for - // requestTimeoutMs (e.g. the peer never GRANTs because it is unreachable / not bounce-ready) - // is failed rather than hanging wait() forever. + // Last time this request made forward progress (granted+posted a chunk, or got an ACK). + // A request stuck with no progress for requestTimeoutMs (e.g. the peer never GRANTs because + // it is unreachable / not bounce-ready) is failed rather than hanging wait() forever. std::chrono::steady_clock::time_point lastProgress; // Cross-thread NVTX spans (0 == none; ended on failure paths too). std::uint64_t nvtxReq{0}; // submit -> promise resolved @@ -425,6 +378,20 @@ class BounceSender // Pipeline-starvation spans (can open/close several times per request, one range per period): std::uint64_t nvtxCreditStarved{0}; // every credit consumed, chunks left -> next GRANT (onGrant ends) std::uint64_t nvtxArenaStarved{0}; // credits parked: no local region / exec ctx -> parked credits drain + + [[nodiscard]] BounceFailReason failureCause(BounceFailReason trigger) const + { + return abandonReason != BounceFailReason::kNone ? abandonReason : trigger; + } + }; + + // A failed request whose chunks are retired (by failRequest or failAll): only its result and the + // receiver's answers to its sent chunks are left. + struct FailedEntry + { + std::string peer; + FailedRequest request; + std::uint64_t nvtxReq{0}; // submit -> result resolved }; // Regions of a FAILED request whose RDMA write was still in flight (state == Writing). @@ -456,6 +423,7 @@ class BounceSender std::uint64_t rid{}; std::chrono::steady_clock::time_point since{}; // when the gather was orphaned (age for the stall warning) bool warned{false}; // stall WARNING already emitted once + OutstandingAccess::SourceRead sourceRead; // ends when the orphan is dropped }; /// GRANT-mispair guard shared by attachCredits()/pumpRequest(): a credit smaller than its chunk @@ -477,31 +445,50 @@ class BounceSender /// (rest parked/retried). /// Caller MUST hold mReqMu. Runs attachCredits() first so credit pairing stays in chunk order. void pumpRequest(std::uint64_t rid, Request& req); - /// Fail `rid` with `reason` (unless a more specific one was tagged earlier), logging ONE warning - /// with the request's progress context — the single choke point that keeps every failure mode - /// observable. See the body for the region-recycling rules and for when it resolves at once. - /// Otherwise its sent chunks move to `awaiting` and it resolves once the receiver has answered them - /// all (quiesced) or has sent no answer for requestTimeoutMs (in-doubt; checkTimeouts). Idempotent: - /// a later failure (forget) keeps the first reason and answer deadline. - void failRequest(std::uint64_t rid, Request& req, BounceFailReason reason); + /// Fail the active request `rid`, logging ONE warning with its progress context — the single choke + /// point that keeps every failure mode observable: retire its chunks, cancel its flow and move it to + /// mFailed, where its result waits answerWaitFor(trigger) from `now` for answers. False if `rid` is not + /// active. + bool failRequest(std::uint64_t rid, BounceFailReason trigger, std::chrono::steady_clock::time_point now); + /// Orphan the chunks whose write or gather may still run (a running gather keeps a SourceRead on + /// `outstanding`) and recycle the rest. True if a write is still in flight. + bool retireChunks(std::uint64_t rid, Request& req, std::shared_ptr const& outstanding, + std::chrono::steady_clock::time_point now); + /// Retract the flow's credits on the receiver, once no RDMA write of it is in flight. + void cancelFlow(std::uint64_t rid, std::string const& peer, bool writeInFlight); + /// The failed form of `req`, holding its result back for `answerWait`. Ends `req`'s pipeline spans. + static FailedEntry toFailedEntry(Request&& req, BounceFailReason cause, + std::shared_ptr outstanding, std::chrono::milliseconds answerWait, + std::chrono::steady_clock::time_point now); /// Recycle the Sent chunk an ACK/NACK answers (same chunk index and region handle): end its /// ack-wait span, return its staging region and drop it from `posted`. False if none matches. bool releaseAnsweredChunk(Request& req, BounceMsgHeader const& h); - /// An ACK or NACK for a request awaiting answers: it only tells us that chunk's scatter will not - /// write the destination again. Once every sent chunk is answered, resolves the request kFAILURE - /// (quiesced) and erases it. - void onAnswerAfterFailure(Request& req, BounceMsgHeader const& h); - /// Resolve a failed request with its `failure` reason: quiesced iff `awaiting` is empty, else - /// in-doubt with a WARNING. The caller erases the request. - void resolveFailure(std::uint64_t rid, Request& req); - /// Deliver the request's one result: end its request-level NVTX spans and set the promise. - static void resolve(Request& req, BounceResult const& result); + /// An ACK or NACK for a request that is no longer active: the answer to a failed request's sent chunk. + void answerFailedRequest(std::string const& peer, BounceMsgHeader const& h); + void failStalledRequests(std::chrono::steady_clock::time_point now); + /// The one place that resolves failed requests whose answer wait is over and drops those no longer + /// worth keeping. + void settleFailedRequests(std::chrono::steady_clock::time_point now); + /// Log and trace a failed request's result that was just resolved with `unanswered` sent chunks. + void onFailureResolved(std::uint64_t rid, FailedEntry& failed, std::size_t unanswered); + /// A record of `req`'s Sent chunks: their DATA is out, so the receiver may scatter them into the + /// destination until it answers. + static std::shared_ptr unansweredSentChunks(Request const& req); + static void endPipelineSpans(Request& req); + static void endChunkSpans(Posted& p); + + [[nodiscard]] std::chrono::milliseconds requestTimeout() const + { + return std::chrono::milliseconds(mCtx.cfg.requestTimeoutMs); + } BounceContext& mCtx; - // request table (IO thread + submit() from app threads) + // Request tables (IO thread + submit() from app threads): active requests, and failed ones whose + // result still waits for answers or is resolved in-doubt with someone holding it. std::mutex mReqMu; std::unordered_map mRequests; + std::unordered_map mFailed; std::atomic mNextRid{1}; std::vector mOrphanLocal; @@ -562,13 +549,20 @@ class BounceTransport void forgetPeer(std::string const& peer); /// Submit a WRITE of (src -> dst) descriptors to `peer`. Returns a future that resolves - /// {kSUCCESS} once every chunk is scattered+ACKed, or {kFAILURE, reason} on error/shutdown. + /// {kSUCCESS} once every chunk is scattered+ACKed, or {kFAILURE, reason, outstanding} on + /// error/shutdown (see FailedRequest). [[nodiscard]] std::shared_future submit( TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer) { return mSender.submit(srcDescs, dstDescs, peer); } + /// Failed requests still kept for the receiver's answers (observability; takes the request lock). + [[nodiscard]] std::size_t failedRequestCount() + { + return mSender.failedRequestCount(); + } + /// This side's EFFECTIVE per-chunk cap: cfg.maxChunkSizeBytes AFTER the constructor's clamp to /// the usable arena capacity. Advertised in the capability handshake; bounce engages toward a /// peer only when both sides advertise the SAME value (see NixlTransferAgent::loadRemoteAgent), diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CMakeLists.txt index 3c48d9d5f470..d7bbf9e8ca39 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CMakeLists.txt @@ -16,7 +16,7 @@ add_library( tllm_executor_bounce_core STATIC EXCLUDE_FROM_ALL BuddyAllocator.cpp CreditScheduler.cpp BounceTransferPlan.cpp - BounceMessage.cpp) + BounceMessage.cpp BounceResult.cpp) add_library(tllm::executor_bounce_core ALIAS tllm_executor_bounce_core) target_link_libraries(tllm_executor_bounce_core PUBLIC tllm::common_logging Threads::Threads) diff --git a/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt b/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt index 8b6a6e835527..39618f36f577 100644 --- a/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt +++ b/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt @@ -50,6 +50,16 @@ add_tllm_gtest( tllm::executor_bounce_core LABELS "unit;executor;bounce;host") +add_tllm_gtest( + bounceFailedRequestTest + OWNER + executor + SOURCES + bounceFailedRequestTest.cpp + LINK_LIBRARIES + tllm::executor_bounce_core + LABELS + "unit;executor;bounce;host") # Header-only: knob resolution (agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* # env > default), byte-size suffixes, garbage fallback, derived-timeout # invariants. @@ -141,6 +151,7 @@ if(NIXL_ROOT) ${BOUNCE_DIR}/BuddyAllocator.cpp ${BOUNCE_DIR}/CreditScheduler.cpp ${BOUNCE_DIR}/BounceMessage.cpp + ${BOUNCE_DIR}/BounceResult.cpp ${BOUNCE_DIR}/BounceArena.cpp ${BOUNCE_DIR}/ExecPool.cpp ${BOUNCE_DIR}/GatherScatterKernel.cu diff --git a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp index 1e7cb13e2e29..8ab5d71fdb89 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp @@ -232,7 +232,14 @@ TEST(BounceAgentE2E, FailureWithNothingSentIsCompleted) ASSERT_NE(status, nullptr); EXPECT_EQ(a->getBounceSubmitCount(), 1u) << "the request must take the bounce path"; EXPECT_EQ(waitTerminal(status, 30), kvc::TransferState::kFAILURE); - EXPECT_TRUE(status->isCompleted()) << "nothing was sent, so the destination is quiescent"; + // An eager gather the failure orphaned may still be reading the source; it ends shortly after. + auto const completeBy = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!status->isCompleted() && std::chrono::steady_clock::now() < completeBy) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(status->isCompleted()) + << "nothing was sent, so the request stops accessing memory once its gathers end"; auto const detail = status->getLastStatusStr(); EXPECT_NE(detail.find("no GRANT/ACK progress"), std::string::npos) << detail; EXPECT_EQ(detail.find("not proven quiescent"), std::string::npos) << detail; diff --git a/cpp/tests/unit_tests/executor/bounce/bounceFailedRequestTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceFailedRequestTest.cpp new file mode 100644 index 000000000000..cfaf2f408159 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceFailedRequestTest.cpp @@ -0,0 +1,324 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; +namespace kvc = tensorrt_llm::executor::kv_cache; + +namespace +{ + +using Clock = b::FailedRequest::Clock; +constexpr auto kAnswerWait = std::chrono::seconds(10); +constexpr std::uint64_t kFirstRegion = 100; // chunk i was granted region kFirstRegion + i +constexpr auto kReason = b::BounceFailReason::kWriteFailed; + +b::BounceTransferStatus resolvedStatus(b::BounceResult const& result) +{ + std::promise promise; + promise.set_value(result); + return b::BounceTransferStatus{promise.get_future().share()}; +} + +// Everything that can happen to a failed request: the receiver's answers, time passing, its orphaned +// gather ending, the result's last holder letting go, and shutdown. +enum class Event +{ + kAnswerChunk0, + kAnswerChunk1, + kCrossedAnswer, // names chunk 1 with chunk 0's region + kTimePasses, // to just before the answer deadline + kAnswerWaitEnds, // to the answer deadline + kGatherEnds, + kHolderLetsGo, + kShutdown, // last +}; + +// Indexed by Event. +constexpr std::array kEventNames{ + "answer0", "answer1", "crossedAnswer", "timePasses", "answerWaitEnds", "gatherEnds", "holderLetsGo", "shutdown"}; +static_assert(kEventNames.size() == static_cast(Event::kShutdown) + 1, "one name per Event"); + +// One failed request driven the way BounceSender drives it — failRequest and answerFailedRequest feed it, +// the per-tick settleFailedRequests (sweep) resolves and drops it, failAll resolves it — next to a model of +// what must hold: which sent chunks are still unanswered, whether its gather still runs, whether anyone +// holds the result, and when the answer wait ends. +class FailedRequestModel +{ +public: + FailedRequestModel(std::uint32_t sentChunks, bool gatherRunning, Clock::duration answerWait) + : mAnswerWait(answerWait) + { + std::vector sent; + for (std::uint32_t chunk = 0; chunk < sentChunks; ++chunk) + { + sent.push_back({.chunkIdx = chunk, .regionHandle = kFirstRegion + chunk}); + mUnanswered.insert(chunk); + } + auto outstanding = std::make_shared(std::move(sent)); + if (gatherRunning) + { + mGather.emplace(outstanding); + } + auto promise = std::make_shared>(); + mHolder.emplace(promise->get_future().share()); + mDeadline = mNow + mAnswerWait; + mFailedRequest.emplace(kReason, std::move(promise), std::move(outstanding), mAnswerWait, mNow); + sweep(); + } + + void apply(Event event) + { + switch (event) + { + case Event::kAnswerChunk0: answer(0, kFirstRegion); break; + case Event::kAnswerChunk1: answer(1, kFirstRegion + 1); break; + case Event::kCrossedAnswer: answer(1, kFirstRegion); break; + case Event::kTimePasses: mNow = std::max(mNow, mDeadline - std::chrono::milliseconds(1)); break; + case Event::kAnswerWaitEnds: mNow = std::max(mNow, mDeadline); break; + case Event::kGatherEnds: mGather.reset(); break; + case Event::kHolderLetsGo: mHolder.reset(); break; + case Event::kShutdown: shutdown(); break; + } + sweep(); // every tick ends with settleFailedRequests + } + + void check() + { + if (mFailedRequest) + { + EXPECT_EQ(mFailedRequest->resolved(), mResolved); + bool const worthKeeping = !mResolved || (!mUnanswered.empty() && outstandingAlive()); + EXPECT_EQ(mFailedRequest->worthKeeping(), worthKeeping); + } + if (!mHolder) + { + return; + } + EXPECT_EQ(mHolder->wait(0), mResolved ? kvc::TransferState::kFAILURE : kvc::TransferState::kIN_PROGRESS); + // Completed exactly when resolved and nothing can still read the source or write the destination. + bool const completed = mHolder->isCompleted(); + EXPECT_EQ(completed, mResolved && mUnanswered.empty() && !mGather.has_value()); + EXPECT_TRUE(completed || !mWasCompleted) << "isCompleted() went back to false"; + mWasCompleted = completed; + if (mResolved) + { + std::string expected = b::toString(kReason); + if (!mUnanswered.empty()) + { + expected += "; destination not proven quiescent"; + } + EXPECT_EQ(mHolder->getLastStatusStr(), expected); + } + } + +private: + // The resolved record lives on through the result's holder and the running gather's source read. + [[nodiscard]] bool outstandingAlive() const + { + return !mResolved || mHolder.has_value() || mGather.has_value(); + } + + // answerFailedRequest: a matching answer ends its chunk and restarts the answer wait. + void answer(std::uint32_t chunk, std::uint64_t region) + { + if (!mFailedRequest) + { + return; + } + bool const matches = outstandingAlive() && region == kFirstRegion + chunk && mUnanswered.count(chunk) > 0; + EXPECT_EQ(mFailedRequest->answer(chunk, region, mNow), matches); + if (matches) + { + mUnanswered.erase(chunk); + if (!mResolved) + { + mDeadline = mNow + mAnswerWait; + } + } + } + + // settleFailedRequests: resolve when due, drop once not worth keeping. + void sweep() + { + if (!mFailedRequest) + { + return; + } + bool const due = !mResolved && (mUnanswered.empty() || mNow >= mDeadline); + auto const unanswered = mFailedRequest->resolveIfDue(mNow); + EXPECT_EQ(unanswered.has_value(), due) << "resolved at the wrong moment"; + if (unanswered) + { + EXPECT_EQ(*unanswered, mUnanswered.size()); + } + mResolved = mResolved || due; + if (!mFailedRequest->worthKeeping()) + { + mFailedRequest.reset(); + } + } + + // failAll: whatever is unanswered stays so, since no answer can arrive any more. + void shutdown() + { + if (!mFailedRequest) + { + return; + } + auto const unanswered = mFailedRequest->resolve(); + EXPECT_EQ(unanswered.has_value(), !mResolved); + if (unanswered) + { + EXPECT_EQ(*unanswered, mUnanswered.size()); + } + mResolved = true; + mFailedRequest.reset(); + } + + Clock::duration mAnswerWait; + Clock::time_point mNow{Clock::time_point{} + std::chrono::hours(1)}; + Clock::time_point mDeadline; + std::set mUnanswered; + std::optional mGather; + std::optional mHolder; + std::optional mFailedRequest; + bool mResolved{false}; + bool mWasCompleted{false}; +}; + +std::string describe( + std::uint32_t sentChunks, bool gatherRunning, Clock::duration answerWait, std::vector const& order) +{ + std::ostringstream out; + out << "sentChunks=" << sentChunks << " gatherRunning=" << gatherRunning + << " answerWaitMs=" << std::chrono::duration_cast(answerWait).count() << " order:"; + for (auto const event : order) + { + out << ' ' << kEventNames.at(static_cast(event)); + } + return out.str(); +} + +} // namespace + +// Every order of every event that can reach a failed request with up to two sent chunks and one running +// gather, with the usual answer wait and with none (kNoProgressTimeout); the model's checks run after each +// event. Time passes twice, so an answer can fall between the two. Stops at the first order that breaks one. +TEST(BounceFailedRequest, EveryEventOrderKeepsTheContract) +{ + std::array const answerWaits{kAnswerWait, Clock::duration::zero()}; + for (std::uint32_t sentChunks = 0; sentChunks <= 2; ++sentChunks) + { + for (auto const answerWait : answerWaits) + { + for (bool const gatherRunning : {false, true}) + { + std::vector order{Event::kCrossedAnswer, Event::kTimePasses, Event::kTimePasses, + Event::kAnswerWaitEnds, Event::kHolderLetsGo, Event::kShutdown}; + if (sentChunks >= 1) + { + order.push_back(Event::kAnswerChunk0); + } + if (sentChunks >= 2) + { + order.push_back(Event::kAnswerChunk1); + } + if (gatherRunning) + { + order.push_back(Event::kGatherEnds); + } + std::sort(order.begin(), order.end()); + do + { + FailedRequestModel model(sentChunks, gatherRunning, answerWait); + model.check(); + for (auto const event : order) + { + model.apply(event); + model.check(); + } + if (::testing::Test::HasFailure()) + { + FAIL() << describe(sentChunks, gatherRunning, answerWait, order); + } + } while (std::next_permutation(order.begin(), order.end())); + } + } + } +} + +// The status contract over hand-built results that left nothing outstanding (the model test covers in-flight +// and outstanding ones): completed, with the failure cause as the status string. +TEST(BounceTransferStatus, IsCompletedOnlyForAResolvedQuiescedResult) +{ + auto const success = resolvedStatus({.state = kvc::TransferState::kSUCCESS}); + EXPECT_TRUE(success.isCompleted()); + EXPECT_EQ(success.getLastStatusStr(), ""); + + auto const quiescedFailure + = resolvedStatus({.state = kvc::TransferState::kFAILURE, .reason = b::BounceFailReason::kPlanRejected}); + EXPECT_TRUE(quiescedFailure.isCompleted()); + EXPECT_EQ(quiescedFailure.getLastStatusStr(), b::toString(b::BounceFailReason::kPlanRejected)); +} + +// A gather's read of the source ends when its SourceRead is destroyed, and only once: a moved-from one +// ends nothing. +TEST(BounceOutstandingAccess, SourceReadEndsWhenDestroyed) +{ + auto access = std::make_shared(); + std::optional read{std::in_place, access}; + EXPECT_FALSE(access->quiesced()); + + std::optional moved{std::move(*read)}; + read.reset(); + EXPECT_FALSE(access->quiesced()) << "the moved-from SourceRead ended the read"; + moved.reset(); + EXPECT_TRUE(access->quiesced()); +} + +// Every trigger waits the full request timeout for the receiver's answers, except kNoProgressTimeout, which +// resolves at once. +TEST(BounceFailedRequest, AnswerWaitFollowsTheTrigger) +{ + constexpr std::chrono::milliseconds kRequestTimeout{30000}; + for (auto const trigger : + {b::BounceFailReason::kNone, b::BounceFailReason::kPlanRejected, b::BounceFailReason::kNoProgressTimeout, + b::BounceFailReason::kPeerDropped, b::BounceFailReason::kGatherFailed, b::BounceFailReason::kWriteFailed, + b::BounceFailReason::kProtocolError, b::BounceFailReason::kShutdown, b::BounceFailReason::kPeerNack}) + { + auto const expected + = trigger == b::BounceFailReason::kNoProgressTimeout ? std::chrono::milliseconds{0} : kRequestTimeout; + EXPECT_EQ(b::answerWaitFor(trigger, kRequestTimeout), expected) << b::toString(trigger); + } +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp index fe60c05ab86d..735099ceee07 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp @@ -23,6 +23,10 @@ // individual writes (complete / stall / fail): those use a FakeXferAgent (a NixlTransferAgent subclass // overriding postXferRequest), because real NIXL cannot be made to do that. // +// The last section checks the FAILURE contract around the failed-request core (bounceFailedRequestTest): +// which chunks a failure records, which messages answer them, and that dst is not written after the +// status completes. +// // The happy-path / concurrency / bidirectional / multi-agent coverage lives in bounceAgentE2ETest, // which drives the SAME pipeline through the production entry point (NixlTransferAgent:: // submitTransferRequests) with one-directional AgentDesc bootstrap, so it is not duplicated here. @@ -38,6 +42,7 @@ #include #include #include +#include #include #include #include @@ -243,29 +248,91 @@ bool waitGrant(b::ZmqControlChannel& ch, std::uint64_t rid, std::chrono::millise }); } -// Wait up to `budget` for a NACK for `rid`. -bool waitNack(b::ZmqControlChannel& ch, std::uint64_t rid, std::chrono::milliseconds budget) +// Wait up to 5 s for a NACK naming (rid, chunkIdx, regionHandle). +bool waitChunkNack(b::ZmqControlChannel& ch, std::uint64_t rid, std::uint32_t chunkIdx, std::uint64_t regionHandle) { - return pumpChannel(ch, budget, + return pumpChannel(ch, std::chrono::seconds(5), [&](b::BounceMsgHeader const& h, std::string const&) - { return static_cast(h.msgType) == b::BounceMsgType::kNACK && h.requestId == rid; }); + { + return static_cast(h.msgType) == b::BounceMsgType::kNACK && h.requestId == rid + && h.chunkIdx == chunkIdx && h.regionHandle == regionHandle; + }); } -// Count ACKs for `rid` over the FULL `budget` window (a negative-assertion helper: always runs to -// the deadline so extra/duplicate ACKs have time to show up). -int countAcks(b::ZmqControlChannel& ch, std::chrono::milliseconds budget, std::uint64_t rid) +// The ACKs and NACKs seen for one rid. +struct AnswerCounts { - int n = 0; - pumpChannel(ch, budget, - [&](b::BounceMsgHeader const& h, std::string const&) + std::uint64_t rid{}; + int acks{0}; + int nacks{0}; + + void add(b::BounceMsgHeader const& h) + { + if (h.requestId != rid) { - if (static_cast(h.msgType) == b::BounceMsgType::kACK && h.requestId == rid) + return; + } + auto const type = static_cast(h.msgType); + acks += type == b::BounceMsgType::kACK ? 1 : 0; + nacks += type == b::BounceMsgType::kNACK ? 1 : 0; + } +}; + +// FIFO barrier against `node`'s reactor. It NACKs a WANT for a zero-byte chunk while dispatching it and reads +// one peer's messages in order, so once that NACK arrives it has dispatched everything `ch` sent before; and +// it sends `ch` everything on one DEALER, so `onEarlier` sees each message it sent `ch` before the NACK. +template +bool reactorBarrier(b::ZmqControlChannel& ch, std::string const& node, Fn&& onEarlier) +{ + static std::atomic nextRid{std::uint64_t{1} << 40}; // never one of a test's own rids + auto const rid = nextRid.fetch_add(1); + ch.sendTo(node, b::encodeWant(rid, {0}, ch.localEndpoint())); + return pumpChannel(ch, std::chrono::seconds(5), + [&](b::BounceMsgHeader const& h, std::string const& blob) + { + if (static_cast(h.msgType) == b::BounceMsgType::kNACK && h.requestId == rid) { - ++n; + return true; } + onEarlier(h, blob); + return false; + }); +} + +bool reactorBarrier(b::ZmqControlChannel& ch, std::string const& node) +{ + return reactorBarrier(ch, node, [](b::BounceMsgHeader const&, std::string const&) {}); +} + +// A DATA whose header claims one run more than its payload holds: it passes the header check but not +// decodeScatter. +std::string undecodableData(std::uint64_t rid, std::uint32_t chunkIdx, std::uint32_t numChunks, std::uint64_t region, + b::BounceScatterRun const& run) +{ + auto blob = b::encodeData(rid, chunkIdx, numChunks, region, {run}); + b::BounceMsgHeader header{}; + EXPECT_TRUE(b::decodeHeader(blob, header)); + header.count += 1; + std::memcpy(blob.data(), &header, sizeof(header)); + return blob; +} + +// Count answers until the first ACK for `counts.rid` arrives (up to 5 s: it follows a scatter, whose first +// launch also loads the kernel module), then for one more second, so an extra answer has time to show up. +void countAnswersPastFirstAck(b::ZmqControlChannel& ch, AnswerCounts& counts) +{ + pumpChannel(ch, std::chrono::seconds(5), + [&](b::BounceMsgHeader const& h, std::string const&) + { + counts.add(h); + return counts.acks > 0; + }); + pumpChannel(ch, std::chrono::seconds(1), + [&](b::BounceMsgHeader const& h, std::string const&) + { + counts.add(h); return false; }); - return n; } // ---- Helpers for the FAILURE-contract tests (last section) ---- @@ -287,6 +354,12 @@ bool waitUntil(Pred&& done, std::chrono::milliseconds budget = std::chrono::seco return true; } +// Wait until the fake agent has had `n` writes posted. +bool waitPosts(XferControls const& ctl, std::uint64_t n) +{ + return waitUntil([&] { return ctl.postCount.load(std::memory_order_acquire) == n; }); +} + // Takes exec contexts away from a node's reactor and gives them back on destruction, so a failed // ASSERT cannot strand them. Holding every receiver context stops scatters before they launch. // Holding all but one sender context serializes the gathers, so writes post in chunk order and a @@ -318,17 +391,17 @@ class ExecHold ExecHold& operator=(ExecHold const&) = delete; // Whether it holds every context but the `leaveFree` it was asked to leave (i.e. the pool was idle). - [[nodiscard]] bool complete() const + [[nodiscard]] bool acquiredAll() const { return mHeld.size() + mLeaveFree == mPool.size(); } - [[nodiscard]] b::ExecCtx* front() const + [[nodiscard]] b::ExecCtx* first() const { return mHeld.front(); } - void releaseFront() + void releaseFirst() { mPool.release(mHeld.front()); mHeld.erase(mHeld.begin()); @@ -465,6 +538,71 @@ constexpr unsigned char kArenaFill = 0xA5; // What the caller writes into dst once the request has resolved. constexpr unsigned char kReusedFill = 0x3C; +// Reuse dst, then shut the receiver down (joins its scatter workers, syncs the device): nothing may write +// dst again. +void expectDstUntouchedAfterReuse(bounce_test::Node& receiver, OwnedXferBufs const& bufs) +{ + ASSERT_EQ(cudaMemset(bufs.dst, kReusedFill, bufs.total), cudaSuccess); + receiver.tx->shutdown(); + EXPECT_TRUE(deviceBytesEqual(bufs.dst, 0, bufs.total, kReusedFill)) << "dst was written after it was reused"; +} + +// A FakeXferAgent sender and a real receiver whose scatters write dst, for `chunks` one-chunk descriptors +// (chunk i is descriptor i). +struct GatedScatterPair +{ + GatedScatterPair(std::string const& name, b::BounceConfig const& c, std::uint32_t chunks, std::uint32_t seed) + : bufs(chunks, static_cast(c.maxChunkSizeBytes), seed) + , sender(makeFakeNode(name + "Sender", c, ctl)) + , receiver(bounce_test::makeNode(name + "Receiver", c, 1024)) + { + } + + // False when the NIXL agent/backend is unavailable (the caller GTEST_SKIPs). + [[nodiscard]] bool available() const + { + return sender != nullptr && receiver != nullptr; + } + + // Wire the peers and prefill the receiver arena with kArenaFill. Leave the sender one exec context, so + // its gathers and writes post in chunk order, and hold every receiver context, so no scatter starts + // until the test lets one. + [[nodiscard]] bool arm() + { + if (!sender->tx->addPeer(receiver->name, receiver->ch->localEndpoint()) + || !receiver->tx->addPeer(sender->name, sender->ch->localEndpoint()) + || cudaMemset(receiver->arena->base(), kArenaFill, receiver->arena->bytes()) != cudaSuccess) + { + return false; + } + senderGathers.emplace(*sender->exec, /*leaveFree=*/1); + receiverScatters.emplace(*receiver->exec, /*leaveFree=*/0); + return senderGathers->acquiredAll() && receiverScatters->acquiredAll(); + } + + // Give the receiver's scatter worker only a context behind a closed StreamGate: the first scatter waits + // there until gate->open(). + [[nodiscard]] bool gateFirstScatter() + { + gate.emplace(receiverScatters->first()->stream); + if (!gate->installed()) + { + return false; + } + receiverScatters->releaseFirst(); + return true; + } + + OwnedXferBufs bufs; // first, so it is freed after the nodes shut down + std::shared_ptr ctl{std::make_shared()}; + std::unique_ptr sender; + std::unique_ptr receiver; + // Last, so they let go of the exec contexts (and open the gate) before the nodes go away. + std::optional senderGathers; + std::optional receiverScatters; + std::optional gate; +}; + // Wait for the sender's (non-cancel) WANT and store its rid in `rid`. bool waitWant(b::ZmqControlChannel& ch, std::uint64_t& rid, std::chrono::milliseconds budget = std::chrono::seconds(5)) { @@ -499,17 +637,20 @@ bool waitData(b::ZmqControlChannel& ch, std::uint64_t rid, std::size_t chunks, }); } -// Wait for the sender's cancel (an empty WANT) for `rid`, i.e. proof that it failed the request. +// Whether a message is the sender's cancel (an empty WANT) for `rid`. +bool isCancelFor(b::BounceMsgHeader const& h, std::string const& blob, std::uint64_t rid) +{ + std::vector chunkBytes; + std::string endpoint; + return static_cast(h.msgType) == b::BounceMsgType::kWANT && h.requestId == rid + && b::decodeWant(blob, h, chunkBytes, endpoint) && b::isCancelWant(chunkBytes); +} + +// Wait for the sender's cancel for `rid`, i.e. proof that it failed the request. bool waitCancel(b::ZmqControlChannel& ch, std::uint64_t rid, std::chrono::milliseconds budget = std::chrono::seconds(5)) { - return pumpChannel(ch, budget, - [&](b::BounceMsgHeader const& h, std::string const& blob) - { - std::vector chunkBytes; - std::string endpoint; - return static_cast(h.msgType) == b::BounceMsgType::kWANT && h.requestId == rid - && b::decodeWant(blob, h, chunkBytes, endpoint) && b::isCancelWant(chunkBytes); - }); + return pumpChannel( + ch, budget, [&](b::BounceMsgHeader const& h, std::string const& blob) { return isCancelFor(h, blob, rid); }); } // The region handle a scripted receiver grants to chunk `chunkIdx`. @@ -522,14 +663,18 @@ std::uint64_t fakeRegion(std::uint32_t chunkIdx) // plays the receiver by hand. Writes are fake, so the credits carry fake addresses. struct ScriptedPeers { - // Declared first, so destroyed last: freed only after the sender has shut down. + ~ScriptedPeers() + { + sender.reset(); // shut down before the channel and the request buffers go away + } + std::deque requestBufs; std::shared_ptr ctl{std::make_shared()}; std::string senderName; std::string receiverName; std::uint32_t chunkBytes{}; std::unique_ptr receiver; - std::unique_ptr sender; // declared last: shut down before the channel goes away + std::unique_ptr sender; // GRANT chunks [firstChunk, endChunk) of `rid` in one message; chunk i gets fakeRegion(i). void grant(std::uint64_t rid, std::uint32_t firstChunk, std::uint32_t endChunk) const @@ -537,7 +682,7 @@ struct ScriptedPeers std::vector credits; for (std::uint32_t i = firstChunk; i < endChunk; ++i) { - credits.push_back({/*addr=*/0, chunkBytes, /*devId=*/0, fakeRegion(i)}); + credits.push_back({.addr = 0, .len = chunkBytes, .devId = 0, .regionHandle = fakeRegion(i)}); } receiver->sendTo(senderName, b::encodeGrant(rid, credits)); } @@ -589,7 +734,7 @@ struct ScriptedRequest }; // Learn `req`'s rid from its WANT, grant all its chunks at once and wait until every one is Sent. -bool grantUntilAllSent(ScriptedPeers& peers, ScriptedRequest& req) +bool grantAllAndWaitSent(ScriptedPeers& peers, ScriptedRequest& req) { if (!waitWant(*peers.receiver, req.rid)) { @@ -599,27 +744,6 @@ bool grantUntilAllSent(ScriptedPeers& peers, ScriptedRequest& req) return waitData(*peers.receiver, req.rid, req.chunks); } -// FIFO barrier: once this returns true, the sender's reactor has dispatched every message the scripted -// receiver sent before it. A never-granted sentinel request is NACKed and must resolve; the reactor reads -// one peer's messages in order. -bool syncSender(ScriptedPeers& peers) -{ - ScriptedRequest sentinel(peers, /*chunks=*/1, /*seed=*/0); - if (!waitWant(*peers.receiver, sentinel.rid)) - { - return false; - } - peers.nack(sentinel.rid, 0); - return sentinel.fut.wait_for(std::chrono::seconds(5)) == std::future_status::ready; -} - -// A status over a future already resolved with `result`. -b::BounceTransferStatus resolvedStatus(b::BounceResult const& result) -{ - std::promise promise; - promise.set_value(result); - return b::BounceTransferStatus{promise.get_future().share()}; -} } // namespace // A peer that never GRANTs must fail the request on requestTimeoutMs, not hang. The WANT is @@ -643,7 +767,6 @@ TEST(BounceTransportFailure, NoGrantTimesOutNotHang) ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "request hung with no grant"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kNoProgressTimeout); - EXPECT_TRUE(fut.get().quiesced) << "nothing was sent, so nothing can write dst"; t->tx->shutdown(); bounce_test::freeXferBufs(bufs); @@ -654,7 +777,9 @@ TEST(BounceTransportFailure, NoGrantTimesOutNotHang) TEST(BounceTransportFailure, WriteFailureFailsRequest) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/5000); capRegions(c, c.maxInflightChunksPerRequest); auto ctl = std::make_shared(); @@ -662,7 +787,9 @@ TEST(BounceTransportFailure, WriteFailureFailsRequest) auto A = makeFakeNode("feA", c, ctl); auto B = bounce_test::makeNode("feB", c, 1024); if (!A || !B) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } A->tx->addPeer("feB", B->ch->localEndpoint()); B->tx->addPeer("feA", A->ch->localEndpoint()); @@ -671,7 +798,7 @@ TEST(BounceTransportFailure, WriteFailureFailsRequest) ASSERT_EQ(fut.wait_for(std::chrono::seconds(10)), std::future_status::ready) << "write failure hung"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kWriteFailed); - EXPECT_TRUE(fut.get().quiesced) << "no write completed, so no DATA was sent"; + EXPECT_TRUE(fut.get().destinationQuiesced()) << "no write completed, so no DATA was sent"; EXPECT_GE(ctl->postCount.load(std::memory_order_acquire), 1u); A->tx->shutdown(); @@ -689,7 +816,9 @@ TEST(BounceTransportFailure, WriteFailureFailsRequest) TEST(BounceTransportFailure, ReceiverNackFailsSenderBeforeTimeout) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } constexpr int kTimeoutMs = 20000; // generous: the NACK bound below is timeout/2, the hang guard is 10 s auto senderCfg = cfg(kTimeoutMs); senderCfg.maxChunkSizeBytes = 8192; @@ -697,7 +826,9 @@ TEST(BounceTransportFailure, ReceiverNackFailsSenderBeforeTimeout) auto A = bounce_test::makeNode("nackA", senderCfg, 1024); auto B = bounce_test::makeNode("nackB", receiverCfg, 1024); if (!A || !B) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } bounce_test::wirePair(*A, *B); auto bufs = bounce_test::makeXferBufs(/*nDescs=*/1, /*descBytes=*/8192, /*seed=*/12); @@ -709,7 +840,6 @@ TEST(BounceTransportFailure, ReceiverNackFailsSenderBeforeTimeout) auto const elapsed = std::chrono::steady_clock::now() - start; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kPeerNack); - EXPECT_TRUE(fut.get().quiesced) << "the WANT was rejected, so no DATA was sent"; EXPECT_LT(elapsed, std::chrono::milliseconds(kTimeoutMs / 2)) << "NACK did not short-circuit the timeout"; // The eager gather of the failed chunk holds an ExecCtx + staging region until failRequest (or, @@ -717,8 +847,8 @@ TEST(BounceTransportFailure, ReceiverNackFailsSenderBeforeTimeout) // ctx assertion is a leak guard, not parked-branch coverage: it cannot tell the parked // (Gathering) branch from the immediate (Gathered) one — an 8 KiB eager gather finishes before // the ZMQ round trip, and if the NACK wins the race nothing was posted at all — so the - // parked-gather branch is not deterministically exercised here (it would need a gather-stall - // hook). + // parked-gather branch is not exercised here; StatusCompletesOnlyOnceOrphanedGatherFinishes holds a + // gather behind a StreamGate to cover it. auto const releaseDeadline = std::chrono::steady_clock::now() + std::chrono::seconds(1); while (A->exec->freeCount() != execBaseline && std::chrono::steady_clock::now() < releaseDeadline) { @@ -733,26 +863,27 @@ TEST(BounceTransportFailure, ReceiverNackFailsSenderBeforeTimeout) // DATA has no retransmission path in this protocol. Receiving it twice for one granted region must // therefore produce one scatter and one ACK, rather than adding replay state for an impossible -// normal-flow event. +// normal-flow event. The duplicate gets no NACK either: it would answer a chunk whose first copy's +// scatter may still be writing dst. TEST(BounceTransportFailure, DuplicateDataProducesOneScatterAndAck) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/5000); capRegions(c, c.maxInflightChunksPerRequest); b::ZmqControlChannel sender("dupDataSender"); auto receiver = bounce_test::makeNode("dupDataReceiver", c, 1024); // receiver posts no writes if (!receiver) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ASSERT_TRUE(sender.addPeer("dupDataReceiver", receiver->ch->localEndpoint())); // Hold every context so both DATA messages reach the reactor before the first scatter can finish. - std::vector heldExecContexts; - while (auto* ctx = receiver->exec->tryAcquire()) - { - heldExecContexts.push_back(ctx); - } - ASSERT_EQ(heldExecContexts.size(), receiver->exec->size()); + ExecHold receiverScatters(*receiver->exec, /*leaveFree=*/0); + ASSERT_TRUE(receiverScatters.acquiredAll()); constexpr std::uint64_t rid = 17; sender.sendTo("dupDataReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); @@ -767,20 +898,115 @@ TEST(BounceTransportFailure, DuplicateDataProducesOneScatterAndAck) auto const data = b::encodeData(rid, /*chunkIdx=*/0, /*numChunks=*/1, credits.front().regionHandle, {run}); sender.sendTo("dupDataReceiver", data); sender.sendTo("dupDataReceiver", data); - // FIFO sentinel instead of a sleep: the reactor is a single FIFO consumer, so a GRANT for a - // fresh rid sent after both DATA messages proves it already consumed them (the second region is - // free under capRegions(2), so the sentinel WANT is grantable while rid=17 holds its region). - constexpr std::uint64_t sentinelRid = 18; - sender.sendTo("dupDataReceiver", b::encodeWant(sentinelRid, {256}, sender.localEndpoint())); - std::vector sentinelCredits; - ASSERT_TRUE(waitGrant(sender, sentinelRid, std::chrono::seconds(5), sentinelCredits)) - << "sentinel WANT never granted; DATA messages may not have reached the reactor"; - for (auto* ctx : heldExecContexts) + AnswerCounts answers{.rid = rid}; + ASSERT_TRUE(reactorBarrier( + sender, "dupDataReceiver", [&](b::BounceMsgHeader const& h, std::string const&) { answers.add(h); })) + << "the DATA messages never reached the reactor"; + receiverScatters.releaseAll(); + + countAnswersPastFirstAck(sender, answers); + EXPECT_EQ(answers.acks, 1); + EXPECT_EQ(answers.nacks, 0); + + receiver->tx->shutdown(); + EXPECT_EQ(cudaFree(dst), cudaSuccess); +} + +// A DATA the receiver will not scatter is NACKed, naming that DATA's chunk and region, so its sender stops +// waiting for the chunk instead of reporting it in-doubt: first one that does not decode, then one whose +// flow was cancelled, so its region is no longer the flow's. +TEST(BounceTransportFailure, RefusedDataIsNacked) +{ + if (!bounce_test::hasCuda()) { - receiver->exec->release(ctx); + GTEST_SKIP() << "no CUDA device"; + } + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel sender("refusedDataSender"); + auto receiver = bounce_test::makeNode("refusedDataReceiver", c, 1024); + if (!receiver) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + ASSERT_TRUE(sender.addPeer("refusedDataReceiver", receiver->ch->localEndpoint())); + + constexpr std::uint64_t rid = 19; + sender.sendTo("refusedDataReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, rid, std::chrono::seconds(5), credits)); + ASSERT_EQ(credits.size(), 1); + auto const region = credits.front().regionHandle; + + void* dst = nullptr; + ASSERT_EQ(cudaMalloc(&dst, 256), cudaSuccess); + b::BounceScatterRun const run{.bounceOffset = 0, + .dstAddr = reinterpret_cast(dst), + .dstStride = 0, + .bounceStride = 0, + .pieceSize = 256, + .count = 1}; + sender.sendTo("refusedDataReceiver", undecodableData(rid, /*chunkIdx=*/0, /*numChunks=*/2, region, run)); + EXPECT_TRUE(waitChunkNack(sender, rid, /*chunkIdx=*/0, region)) << "the undecodable DATA was not NACKed"; + + sender.sendTo("refusedDataReceiver", b::encodeCancel(rid, sender.localEndpoint())); + sender.sendTo("refusedDataReceiver", b::encodeData(rid, /*chunkIdx=*/1, /*numChunks=*/2, region, {run})); + EXPECT_TRUE(waitChunkNack(sender, rid, /*chunkIdx=*/1, region)) << "the DATA after the cancel was not NACKed"; + + receiver->tx->shutdown(); + EXPECT_EQ(cudaFree(dst), cudaSuccess); +} + +// A DATA refused while a scatter still holds its region gets no NACK: that scatter may be writing this very +// chunk. Here the flow is cancelled while the first copy's scatter waits in the queue, so a second copy, and +// an undecodable one, find the region no longer the flow's but still scattering. The first copy still runs +// and answers. +TEST(BounceTransportFailure, RefusedDataGetsNoNackWhileItsRegionScatters) +{ + if (!bounce_test::hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; } + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel sender("scatteringRefusedSender"); + auto receiver = bounce_test::makeNode("scatteringRefusedReceiver", c, 1024); + if (!receiver) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + ASSERT_TRUE(sender.addPeer("scatteringRefusedReceiver", receiver->ch->localEndpoint())); + ExecHold receiverScatters(*receiver->exec, /*leaveFree=*/0); + ASSERT_TRUE(receiverScatters.acquiredAll()); - EXPECT_EQ(countAcks(sender, std::chrono::seconds(2), rid), 1); + constexpr std::uint64_t rid = 23; + sender.sendTo("scatteringRefusedReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, rid, std::chrono::seconds(5), credits)); + ASSERT_EQ(credits.size(), 1); + + void* dst = nullptr; + ASSERT_EQ(cudaMalloc(&dst, 256), cudaSuccess); + b::BounceScatterRun const run{.bounceOffset = 0, + .dstAddr = reinterpret_cast(dst), + .dstStride = 0, + .bounceStride = 0, + .pieceSize = 256, + .count = 1}; + auto const data = b::encodeData(rid, /*chunkIdx=*/0, /*numChunks=*/1, credits.front().regionHandle, {run}); + sender.sendTo("scatteringRefusedReceiver", data); + sender.sendTo("scatteringRefusedReceiver", b::encodeCancel(rid, sender.localEndpoint())); + sender.sendTo("scatteringRefusedReceiver", data); + sender.sendTo("scatteringRefusedReceiver", + undecodableData(rid, /*chunkIdx=*/0, /*numChunks=*/1, credits.front().regionHandle, run)); + AnswerCounts answers{.rid = rid}; + ASSERT_TRUE(reactorBarrier( + sender, "scatteringRefusedReceiver", [&](b::BounceMsgHeader const& h, std::string const&) { answers.add(h); })); + receiverScatters.releaseAll(); + + countAnswersPastFirstAck(sender, answers); + EXPECT_EQ(answers.acks, 1) << "the first copy's scatter did not answer"; + EXPECT_EQ(answers.nacks, 0) << "the refused copy was NACKed while its region still scattered"; receiver->tx->shutdown(); EXPECT_EQ(cudaFree(dst), cudaSuccess); @@ -980,8 +1206,8 @@ TEST(BounceTransportFailure, GrantLeaseExpiryReclaimsSilentSendersRegion) ASSERT_EQ(credit2.size(), 1u); EXPECT_EQ(credit2.front().regionHandle, credit1.front().regionHandle); // the one region, recycled - // Late DATA for the EXPIRED grant must be dropped (flow reclaimed -> heldByFlow false): exactly - // one ACK total, and it belongs to rid=2's legitimate DATA. + // Late DATA for the EXPIRED grant must not be scattered (flow reclaimed -> heldByFlow false; the + // receiver NACKs it): exactly one ACK total, and it belongs to rid=2's legitimate DATA. void* dst = nullptr; ASSERT_EQ(cudaMalloc(&dst, 256), cudaSuccess); b::BounceScatterRun run{0, reinterpret_cast(dst), 0, 0, 256, 1}; @@ -1097,7 +1323,6 @@ TEST(BounceTransportFailure, CreditMispairAbandonIsTerminalAndIgnoresLaterGrants ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "abandoned flow never resolved"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kProtocolError); - EXPECT_TRUE(fut.get().quiesced) << "no write was posted, so no DATA was sent"; // The later valid GRANT must not have resurrected the abandoned flow into posting a write. EXPECT_EQ(ctl->postCount.load(std::memory_order_acquire), 0u) << "abandoned flow posted an RDMA write"; @@ -1109,20 +1334,23 @@ TEST(BounceTransportFailure, CreditMispairAbandonIsTerminalAndIgnoresLaterGrants TEST(BounceTransportFailure, ShutdownFailsInflight) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/0); // timeout disabled -> only shutdown can resolve it + } + auto c = cfg(/*timeoutMs=*/600000); // never fires here: only shutdown can resolve it auto t = bounce_test::makeNode("sdSolo", c, 1024); if (!t) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); - auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "nobody"); // stuck (no grant, no timeout) + auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "nobody"); // stuck: never granted EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); // still pending t->tx->shutdown(); ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "shutdown left request hanging"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); - EXPECT_TRUE(fut.get().quiesced) << "never granted, so no DATA was sent"; bounce_test::freeXferBufs(bufs); } @@ -1133,8 +1361,10 @@ TEST(BounceTransportFailure, ShutdownFailsInflight) TEST(BounceTransportFailure, ShutdownReleaseFailureDoesNotHangOrLoseHandle) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/0); + } + auto c = cfg(/*timeoutMs=*/600000); // never fires here: only shutdown can resolve it capRegions(c, c.maxInflightChunksPerRequest); b::ZmqControlChannel peer("shutdownReleasePeer"); auto ctl = std::make_shared(); @@ -1142,7 +1372,9 @@ TEST(BounceTransportFailure, ShutdownReleaseFailureDoesNotHangOrLoseHandle) ctl->releaseSucceeds.store(false, std::memory_order_release); auto sender = makeFakeNode("shutdownReleaseSender", c, ctl); if (!sender) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ASSERT_TRUE(sender->tx->addPeer("shutdownReleasePeer", peer.localEndpoint())); ASSERT_TRUE(peer.addPeer("shutdownReleaseSender", sender->ch->localEndpoint())); @@ -1166,22 +1398,26 @@ TEST(BounceTransportFailure, ShutdownReleaseFailureDoesNotHangOrLoseHandle) ASSERT_EQ(fut.wait_for(std::chrono::seconds(1)), std::future_status::ready); EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); - EXPECT_TRUE(fut.get().quiesced) << "the write never completed, so no DATA was sent"; + EXPECT_TRUE(fut.get().destinationQuiesced()) << "the write never completed, so no DATA was sent"; // failAll attempts the failed release exactly twice (initial + bounded retry), never spins. EXPECT_EQ(ctl->releaseCount.load(std::memory_order_acquire), 2u); bounce_test::freeXferBufs(bufs); } -// forgetPeer() (the invalidateRemoteAgent path) must fail any in-flight request to the gone peer, -// even with the request timeout disabled, modeling a peer that drops out mid-transfer. +// forgetPeer() (the invalidateRemoteAgent path) must fail any in-flight request to the gone peer +// well before the request timeout, modeling a peer that drops out mid-transfer. TEST(BounceTransportFailure, ForgetPeerFailsInflightRequest) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/0); // timeout disabled -> only forgetPeer can resolve it + } + auto c = cfg(/*timeoutMs=*/600000); // never fires here: only forgetPeer can resolve it auto t = bounce_test::makeNode("fpSolo", c, 1024); if (!t) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "gonePeer"); // stuck: never granted @@ -1263,7 +1499,9 @@ TEST(BounceTransportFailure, ForgetPeerInFlightRecovers) TEST(BounceTransportFailure, ForgetPeerFreesOrphanedRegionsToOtherPeersOnly) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/30000); c.scatterWorkerCount = 1; // exactly one job can be dequeued; the other must stay in mJobs c.quarantineMs = 30000; // the old-order symptom: the mis-granted region is parked this long @@ -1272,16 +1510,11 @@ TEST(BounceTransportFailure, ForgetPeerFreesOrphanedRegionsToOtherPeersOnly) b::ZmqControlChannel other("fgOther"); auto receiver = bounce_test::makeNode("fgReceiver", c, 1024); // receiver posts no writes if (!receiver) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ASSERT_TRUE(gone.addPeer("fgReceiver", receiver->ch->localEndpoint())); ASSERT_TRUE(other.addPeer("fgReceiver", receiver->ch->localEndpoint())); - // A zero-size WANT is NACKed inline by dispatch(); on the same DEALER it proves every earlier - // message was processed, and the tick that dispatched it runs drainForgets() before recv()ing again. - auto syncReactor = [&](std::uint64_t rid) - { - gone.sendTo("fgReceiver", b::encodeWant(rid, {0}, gone.localEndpoint())); - return waitNack(gone, rid, std::chrono::seconds(5)); - }; std::vector heldExecContexts; while (auto* ctx = receiver->exec->tryAcquire()) @@ -1310,10 +1543,12 @@ TEST(BounceTransportFailure, ForgetPeerFreesOrphanedRegionsToOtherPeersOnly) } gone.sendTo("fgReceiver", b::encodeWant(/*rid=*/1, {}, gone.localEndpoint())); // cancel -> orphans gone.sendTo("fgReceiver", b::encodeWant(/*rid=*/2, {chunk, chunk}, gone.localEndpoint())); // pending: arena full - ASSERT_TRUE(syncReactor(/*rid=*/3)) << "cancel / pending WANT not processed"; + ASSERT_TRUE(reactorBarrier(gone, "fgReceiver")) << "cancel / pending WANT not processed"; receiver->tx->forgetPeer("fgGone"); - ASSERT_TRUE(syncReactor(/*rid=*/4)) << "forget not applied"; // also re-adds fgGone's DEALER: a mis-grant is visible + // The tick that dispatches this barrier runs drainForgets() before it recv()s again; the barrier also + // re-adds fgGone's DEALER, so a mis-grant would be visible. + ASSERT_TRUE(reactorBarrier(gone, "fgReceiver")) << "forget not applied"; other.sendTo("fgReceiver", b::encodeWant(/*rid=*/10, {chunk}, other.localEndpoint())); std::vector otherCredits; @@ -1447,63 +1682,58 @@ TEST(BounceTransportFailure, ConcurrentMultiThreadedSubmit) // ============================================================================ // A FAILURE resolves only after the receiver has answered every chunk whose DATA was sent // ---------------------------------------------------------------------------- -// Only the receiver's scatter of a received DATA writes dst, and the receiver answers (ACK or NACK) -// after that scatter's last write. So a failed request waits for those answers before it resolves -// kFAILURE with quiesced == true; a receiver that sends no answer for requestTimeoutMs is assumed dead -// and the request resolves in-doubt (quiesced == false). These tests order events with exec-context -// holds, a stream gate, scripted writes and FIFO barriers, not by sleeping. +// The answer-wait rules themselves are checked over every event order in bounceFailedRequestTest. These +// tests check the reactor around them: which chunks a failure records, which messages answer them, and +// that nothing writes dst once the status completes. They order events with exec-context holds, a stream +// gate, scripted writes and FIFO barriers, and sleep only where elapsed time is what they measure. // ============================================================================ -// The status contract the agent hands out (pure: hand-built futures, no NIXL). isCompleted() is true -// exactly when the request is over AND nothing will write the destination again. -TEST(BounceTransferStatus, CompletedOnlyOnceTheDestinationIsQuiesced) -{ - std::promise pending; - b::BounceTransferStatus const inFlight{pending.get_future().share()}; - EXPECT_FALSE(inFlight.isCompleted()); - EXPECT_EQ(inFlight.wait(0), kvc::TransferState::kIN_PROGRESS); - EXPECT_EQ(inFlight.getLastStatusStr(), ""); - - auto const success = resolvedStatus({kvc::TransferState::kSUCCESS, b::BounceFailReason::kNone, /*quiesced=*/true}); - EXPECT_TRUE(success.isCompleted()); - EXPECT_EQ(success.wait(0), kvc::TransferState::kSUCCESS); - EXPECT_EQ(success.getLastStatusStr(), ""); - - auto const quiescedFailure - = resolvedStatus({kvc::TransferState::kFAILURE, b::BounceFailReason::kWriteFailed, /*quiesced=*/true}); - EXPECT_TRUE(quiescedFailure.isCompleted()) << "a destination proven quiescent may be reused"; - EXPECT_EQ(quiescedFailure.wait(0), kvc::TransferState::kFAILURE); - EXPECT_EQ(quiescedFailure.getLastStatusStr(), b::toString(b::BounceFailReason::kWriteFailed)); - - auto const inDoubt - = resolvedStatus({kvc::TransferState::kFAILURE, b::BounceFailReason::kNoProgressTimeout, /*quiesced=*/false}); - EXPECT_FALSE(inDoubt.isCompleted()) << "the receiver may still write the destination"; - EXPECT_EQ(inDoubt.wait(0), kvc::TransferState::kFAILURE); - EXPECT_EQ(inDoubt.getLastStatusStr(), - std::string(b::toString(b::BounceFailReason::kNoProgressTimeout)) + "; destination not proven quiescent"); - - pending.set_value({kvc::TransferState::kSUCCESS, b::BounceFailReason::kNone, /*quiesced=*/true}); - EXPECT_TRUE(inFlight.isCompleted()); -} - -// A request rejected at plan time never sent anything, so its FAILURE is quiesced. -TEST(BounceTransportFailure, PlanRejectionIsQuiesced) +// A request fails before any DATA was sent while its eager gather is still running. Nothing can write +// dst, so the result is quiesced at once, but the gather still reads the source, so the status completes +// only once that gather finishes. The gather is held by a stream gate on the only free exec context. +TEST(BounceTransportFailure, StatusCompletesOnlyOnceOrphanedGatherFinishes) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/5000); - // One descriptor larger than maxChunkSizeBytes: BounceTransferPlan::build rejects it. - OwnedXferBufs bufs(/*nDescs=*/1, /*descBytes=*/2 * static_cast(c.maxChunkSizeBytes), /*seed=*/81); - auto t = bounce_test::makeNode("prSolo", c, 1024); + } + auto c = cfg(/*timeoutMs=*/30000); // the timeout never fires here: forgetPeer fails the request + // One chunk. Declared before the node, so it is freed after the node shuts down. + OwnedXferBufs bufs(/*nDescs=*/1, static_cast(c.maxChunkSizeBytes), /*seed=*/94); + auto t = bounce_test::makeNode("orphanGatherSolo", c, 1024); if (!t) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + // A bound ROUTER that no transport ever recv()s on: the WANT is delivered but never granted. + b::ZmqControlChannel ghost("ghostPeer"); + t->tx->addPeer("ghostPeer", ghost.localEndpoint()); - auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "nobody"); - ASSERT_EQ(fut.wait_for(std::chrono::seconds(0)), std::future_status::ready) - << "plan rejection resolves in submit()"; + ExecHold senderGathers(*t->exec, /*leaveFree=*/0); + ASSERT_TRUE(senderGathers.acquiredAll()); + StreamGate gate(senderGathers.first()->stream); + ASSERT_TRUE(gate.installed()); + senderGathers.releaseFirst(); // the eager gather can get only the gated context + + auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "ghostPeer"); + b::BounceTransferStatus const status{fut}; + ASSERT_TRUE(waitUntil([&] { return t->exec->freeCount() == 0; })) << "the eager gather never started"; + t->tx->forgetPeer("ghostPeer"); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) + << "forgetPeer did not fail the request"; EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); - EXPECT_EQ(fut.get().reason, b::BounceFailReason::kPlanRejected); - EXPECT_TRUE(fut.get().quiesced); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kPeerDropped); + EXPECT_TRUE(fut.get().destinationQuiesced()) << "nothing was sent, so nothing can write dst"; + // Watched for a while, not once: an orphan drain that ended the read on an early poll shows up here. + auto const gatherHeldUntil = std::chrono::steady_clock::now() + std::chrono::milliseconds(100); + while (std::chrono::steady_clock::now() < gatherHeldUntil) + { + ASSERT_FALSE(status.isCompleted()) << "completed while the orphaned gather could still read the source"; + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + + gate.open(); + EXPECT_TRUE(waitUntil([&] { return status.isCompleted(); })) << "the finished gather was never observed"; t->tx->shutdown(); } @@ -1514,59 +1744,88 @@ TEST(BounceTransportFailure, PlanRejectionIsQuiesced) TEST(BounceTransportFailure, FailureWaitsForRunningScatter) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } // The timeout never fires here: only the receiver's answers may resolve the failed request. auto c = cfg(/*timeoutMs=*/30000); c.scatterWorkerCount = 1; capRegions(c, c.maxInflightChunksPerRequest); // two chunks, both granted at once - // Chunk i is descriptor i. Declared before the nodes, so it is freed after they shut down. - OwnedXferBufs bufs(/*nDescs=*/2, static_cast(c.maxChunkSizeBytes), /*seed=*/91); - auto ctl = std::make_shared(); - auto sender = makeFakeNode("runScatterSender", c, ctl); - auto receiver = bounce_test::makeNode("runScatterReceiver", c, 1024); - if (!sender || !receiver) + GatedScatterPair pair("runScatter", c, /*chunks=*/2, /*seed=*/91); + if (!pair.available()) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; - ASSERT_TRUE(sender->tx->addPeer(receiver->name, receiver->ch->localEndpoint())); - ASSERT_TRUE(receiver->tx->addPeer(sender->name, sender->ch->localEndpoint())); - ASSERT_EQ(cudaMemset(receiver->arena->base(), kArenaFill, receiver->arena->bytes()), cudaSuccess); - - ExecHold senderGathers(*sender->exec, /*leaveFree=*/1); // post index == chunk index - ExecHold receiverScatters(*receiver->exec, /*leaveFree=*/0); - ASSERT_TRUE(senderGathers.complete()); - ASSERT_TRUE(receiverScatters.complete()); - StreamGate gate(receiverScatters.front()->stream); - ASSERT_TRUE(gate.installed()); - receiverScatters.releaseFront(); // the worker can get only the gated context - ctl->scriptWrite(0, WriteScript::kComplete); - ctl->scriptWrite(1, WriteScript::kStall); + } + ASSERT_TRUE(pair.arm()); + ASSERT_TRUE(pair.gateFirstScatter()); + pair.ctl->scriptWrite(0, WriteScript::kComplete); + pair.ctl->scriptWrite(1, WriteScript::kStall); - auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, receiver->name); + auto fut = pair.sender->tx->submit(pair.bufs.srcDescs, pair.bufs.dstDescs, pair.receiver->name); // The worker took the gated context: chunk 0 is Sent and its scatter waits behind the gate. - ASSERT_TRUE(waitUntil([&] { return receiver->exec->freeCount() == 0; })) << "chunk 0 never reached the scatter"; - ASSERT_TRUE(waitUntil([&] { return ctl->postCount.load(std::memory_order_acquire) == 2; })); - ctl->scriptWrite(1, WriteScript::kFail); + ASSERT_TRUE(waitUntil([&] { return pair.receiver->exec->freeCount() == 0; })) + << "chunk 0 never reached the scatter"; + ASSERT_TRUE(waitPosts(*pair.ctl, 2)); + pair.ctl->scriptWrite(1, WriteScript::kFail); // Write 0's handle was released when its DATA went out; write 1's is released by the orphan // drain that runs only after failRequest. - ASSERT_TRUE(waitUntil([&] { return ctl->releaseCount.load(std::memory_order_acquire) == 2; })) + ASSERT_TRUE(waitUntil([&] { return pair.ctl->releaseCount.load(std::memory_order_acquire) == 2; })) << "the failed write was never handled"; - EXPECT_EQ(fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(100)), std::future_status::timeout) << "FAILURE reported while chunk 0's scatter could still write dst"; - gate.open(); + pair.gate->open(); ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "the scatter's ACK was not awaited"; auto const result = fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); - EXPECT_TRUE(result.quiesced); - EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[0], bufs.descBytes, kArenaFill)) + EXPECT_TRUE(result.destinationQuiesced()); + EXPECT_TRUE(deviceBytesEqual(pair.bufs.dst, pair.bufs.off[0], pair.bufs.descBytes, kArenaFill)) << "chunk 0's scatter had not finished when the FAILURE resolved"; - EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[1], bufs.descBytes, 0)) << "chunk 1 was never sent"; - // Reuse dst, then shut the receiver down (joins its scatter workers, syncs the device): nothing may - // write it again. - ASSERT_EQ(cudaMemset(bufs.dst, kReusedFill, bufs.total), cudaSuccess); - receiver->tx->shutdown(); - EXPECT_TRUE(deviceBytesEqual(bufs.dst, 0, bufs.total, kReusedFill)) << "dst was written after it was reused"; - sender->tx->shutdown(); + EXPECT_TRUE(deviceBytesEqual(pair.bufs.dst, pair.bufs.off[1], pair.bufs.descBytes, 0)) << "chunk 1 was never sent"; + expectDstUntouchedAfterReuse(*pair.receiver, pair.bufs); + pair.sender->tx->shutdown(); +} + +// End to end on a real receiver: a failure resolved in-doubt while chunk 0's scatter is held completes once +// that scatter runs and answers, and nothing writes dst after the status says so. +TEST(BounceTransportFailure, LateAnswerCompletesAnInDoubtFailureAfterItsLastWrite) +{ + if (!bounce_test::hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + constexpr auto kTimeout = std::chrono::milliseconds(2000); + auto c = cfg(static_cast(kTimeout.count())); + c.scatterWorkerCount = 1; + capRegions(c, c.maxInflightChunksPerRequest); + GatedScatterPair pair("lateAnswer", c, /*chunks=*/2, /*seed=*/113); + if (!pair.available()) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + ASSERT_TRUE(pair.arm()); + ASSERT_TRUE(pair.gateFirstScatter()); + pair.ctl->scriptWrite(0, WriteScript::kComplete); + pair.ctl->scriptWrite(1, WriteScript::kStall); + + auto fut = pair.sender->tx->submit(pair.bufs.srcDescs, pair.bufs.dstDescs, pair.receiver->name); + b::BounceTransferStatus const status{fut}; + ASSERT_TRUE(waitUntil([&] { return pair.receiver->exec->freeCount() == 0; })) + << "chunk 0 never reached the scatter"; + ASSERT_TRUE(waitPosts(*pair.ctl, 2)); + pair.ctl->scriptWrite(1, WriteScript::kFail); + ASSERT_EQ(fut.wait_for(kTimeout + std::chrono::seconds(5)), std::future_status::ready) + << "the answer wait did not end"; + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kWriteFailed); + EXPECT_FALSE(status.isCompleted()) << "completed while chunk 0's scatter could still write dst"; + + pair.gate->open(); + ASSERT_TRUE(waitUntil([&] { return status.isCompleted(); })) << "chunk 0's late ACK did not complete it"; + EXPECT_TRUE(deviceBytesEqual(pair.bufs.dst, pair.bufs.off[0], pair.bufs.descBytes, kArenaFill)) + << "completed before chunk 0's scatter finished"; + expectDstUntouchedAfterReuse(*pair.receiver, pair.bufs); + pair.sender->tx->shutdown(); } // Same as above, but the sent chunks are still QUEUED on the receiver: chunk 0 waits for an exec @@ -1575,75 +1834,69 @@ TEST(BounceTransportFailure, FailureWaitsForRunningScatter) TEST(BounceTransportFailure, FailureWaitsForQueuedScatters) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } // The timeout never fires here: only the receiver's answers may resolve the failed request. auto c = cfg(/*timeoutMs=*/30000); c.scatterWorkerCount = 1; c.maxInflightChunksPerRequest = 3; capRegions(c, c.maxInflightChunksPerRequest); // three chunks, all granted at once - // Chunk i is descriptor i. Declared before the nodes, so it is freed after they shut down. - OwnedXferBufs bufs(/*nDescs=*/3, static_cast(c.maxChunkSizeBytes), /*seed=*/92); - auto ctl = std::make_shared(); - auto sender = makeFakeNode("queuedScatterSender", c, ctl); - auto receiver = bounce_test::makeNode("queuedScatterReceiver", c, 1024); - if (!sender || !receiver) + GatedScatterPair pair("queuedScatter", c, /*chunks=*/3, /*seed=*/92); + if (!pair.available()) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; - ASSERT_TRUE(sender->tx->addPeer(receiver->name, receiver->ch->localEndpoint())); - ASSERT_TRUE(receiver->tx->addPeer(sender->name, sender->ch->localEndpoint())); - ASSERT_EQ(cudaMemset(receiver->arena->base(), kArenaFill, receiver->arena->bytes()), cudaSuccess); - - ExecHold senderGathers(*sender->exec, /*leaveFree=*/1); // post index == chunk index - ExecHold receiverScatters(*receiver->exec, /*leaveFree=*/0); - ASSERT_TRUE(senderGathers.complete()); - ASSERT_TRUE(receiverScatters.complete()); - ctl->scriptWrite(0, WriteScript::kComplete); - ctl->scriptWrite(1, WriteScript::kComplete); + } + ASSERT_TRUE(pair.arm()); + pair.ctl->scriptWrite(0, WriteScript::kComplete); + pair.ctl->scriptWrite(1, WriteScript::kComplete); // Polled in chunk order, writes 0 and 1 are Sent (DATA out) before write 2's failure is seen. - ctl->scriptWrite(2, WriteScript::kFail); + pair.ctl->scriptWrite(2, WriteScript::kFail); - auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, receiver->name); + auto fut = pair.sender->tx->submit(pair.bufs.srcDescs, pair.bufs.dstDescs, pair.receiver->name); // Writes 0 and 1 were released when their DATA went out; write 2 after failRequest. - ASSERT_TRUE(waitUntil([&] { return ctl->releaseCount.load(std::memory_order_acquire) == 3; })) + ASSERT_TRUE(waitUntil([&] { return pair.ctl->releaseCount.load(std::memory_order_acquire) == 3; })) << "the failed write was never handled"; - EXPECT_EQ(fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(100)), std::future_status::timeout) << "FAILURE reported while chunks 0 and 1 could still be scattered into dst"; - receiverScatters.releaseAll(); + pair.receiverScatters->releaseAll(); ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "the scatters' ACKs were not awaited"; auto const result = fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); - EXPECT_TRUE(result.quiesced); + EXPECT_TRUE(result.destinationQuiesced()); for (std::uint32_t chunk = 0; chunk < 2; ++chunk) { - EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[chunk], bufs.descBytes, kArenaFill)) + EXPECT_TRUE(deviceBytesEqual(pair.bufs.dst, pair.bufs.off[chunk], pair.bufs.descBytes, kArenaFill)) << "chunk " << chunk << " was not scattered before the FAILURE resolved"; } - EXPECT_TRUE(deviceBytesEqual(bufs.dst, bufs.off[2], bufs.descBytes, 0)) << "chunk 2 was never sent"; - // Reuse dst, then shut the receiver down (joins its scatter workers, syncs the device): nothing may - // write it again. - ASSERT_EQ(cudaMemset(bufs.dst, kReusedFill, bufs.total), cudaSuccess); - receiver->tx->shutdown(); - EXPECT_TRUE(deviceBytesEqual(bufs.dst, 0, bufs.total, kReusedFill)) << "dst was written after it was reused"; - sender->tx->shutdown(); + EXPECT_TRUE(deviceBytesEqual(pair.bufs.dst, pair.bufs.off[2], pair.bufs.descBytes, 0)) << "chunk 2 was never sent"; + expectDstUntouchedAfterReuse(*pair.receiver, pair.bufs); + pair.sender->tx->shutdown(); } // A NACK fails the request but answers only its own chunk: the request waits for the other Sent // chunk's answer, and not for a chunk whose write is still in flight (its DATA was never sent, so the -// receiver will never scatter or answer it). -TEST(BounceTransportFailure, NackWaitsOnlyForOtherSentChunks) +// receiver will never scatter or answer it). The cancel waits for that write: sent earlier, it would let +// the receiver re-grant a region the write still lands in. +TEST(BounceTransportFailure, NackWaitsOnlyForOtherSentChunksAndDefersCancelUntilWriteEnds) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/30000); // waiting for chunk 2 would show up as a missed 5 s bound c.maxInflightChunksPerRequest = 3; capRegions(c, c.maxInflightChunksPerRequest); auto peers = makeScriptedPeers("nackWait", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ExecHold senderGathers(*peers->sender->exec, /*leaveFree=*/1); // post index == chunk index - ASSERT_TRUE(senderGathers.complete()); + ASSERT_TRUE(senderGathers.acquiredAll()); peers->ctl->scriptWrite(0, WriteScript::kComplete); peers->ctl->scriptWrite(1, WriteScript::kComplete); peers->ctl->scriptWrite(2, WriteScript::kStall); @@ -1652,11 +1905,13 @@ TEST(BounceTransportFailure, NackWaitsOnlyForOtherSentChunks) ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/3); ASSERT_TRUE(waitData(*peers->receiver, req.rid, 2)) << "chunks 0 and 1 never reached Sent"; - ASSERT_TRUE(waitUntil([&] { return peers->ctl->postCount.load(std::memory_order_acquire) == 3; })) - << "chunk 2's write was never posted"; + ASSERT_TRUE(waitPosts(*peers->ctl, 3)) << "chunk 2's write was never posted"; + bool cancelled = false; + auto const noteCancel = [&](b::BounceMsgHeader const& h, std::string const& blob) + { cancelled = cancelled || isCancelFor(h, blob, req.rid); }; peers->nack(req.rid, 1); - ASSERT_TRUE(syncSender(*peers)); + ASSERT_TRUE(reactorBarrier(*peers->receiver, peers->senderName, noteCancel)); EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) << "resolved before chunk 0 was answered"; peers->ack(req.rid, 0); @@ -1665,65 +1920,51 @@ TEST(BounceTransportFailure, NackWaitsOnlyForOtherSentChunks) auto const result = req.fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); - EXPECT_TRUE(result.quiesced); + EXPECT_TRUE(result.destinationQuiesced()); EXPECT_EQ(peers->ctl->releaseCount.load(std::memory_order_acquire), 2u) << "chunk 2's write must still be in flight"; + ASSERT_TRUE(reactorBarrier(*peers->receiver, peers->senderName, noteCancel)); + EXPECT_FALSE(cancelled) << "the cancel went out while chunk 2's write could still land in its region"; - // End the orphaned write so shutdown skips its 5 s write drain; its deferred cancel follows. peers->ctl->scriptWrite(2, WriteScript::kComplete); - EXPECT_TRUE(waitCancel(*peers->receiver, req.rid)); -} - -// NACKs are answers too: a failed request whose remaining chunks are all NACKed (e.g. a receiver with a -// CUDA fault NACKs every queued scatter) resolves quiesced at once, not in-doubt at requestTimeoutMs. -TEST(BounceTransportFailure, NackAnswersEndTheWait) -{ - if (!bounce_test::hasCuda()) - GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/30000); // ignoring the NACKs would show up as a missed 5 s bound - c.maxInflightChunksPerRequest = 3; - capRegions(c, c.maxInflightChunksPerRequest); - auto peers = makeScriptedPeers("nackAnswers", c); - if (!peers) - GTEST_SKIP() << "NIXL agent/backend unavailable"; - - ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/106); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); - peers->nack(req.rid, 2); // fails the request: chunks 0 and 1 are awaited - ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); - peers->nack(req.rid, 0); - peers->nack(req.rid, 1); - ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) - << "NACKs to a failed request were not taken as answers"; - auto const result = req.fut.get(); - EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); - EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); - EXPECT_TRUE(result.quiesced); + EXPECT_TRUE(waitCancel(*peers->receiver, req.rid)) << "the deferred cancel never followed the write"; } -// Only an answer for an awaited chunk (same chunk index AND region) ends a failed request's wait. A -// duplicate, an answer for a chunk that is not awaited and a wrong-region answer are all ignored. +// Only an answer for an awaited chunk (same chunk index AND region, from the request's peer) ends a failed +// request's wait, and a NACK answers like an ACK. A duplicate, an answer for a chunk that is not awaited, a +// wrong-region answer and an answer from another peer are all ignored. TEST(BounceTransportFailure, OnlyMatchingAnswersEndTheWait) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/30000); c.maxInflightChunksPerRequest = 3; capRegions(c, c.maxInflightChunksPerRequest); auto peers = makeScriptedPeers("matchAnswers", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/102); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); + ASSERT_TRUE(grantAllAndWaitSent(*peers, req)); peers->nack(req.rid, 2); // fails the request: chunks 0 and 1 are awaited ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); - peers->ack(req.rid, 1); // a real answer: only chunk 0 is awaited now - peers->ack(req.rid, 1); // duplicate - peers->ack(req.rid, 2); // chunk 2 is not awaited: its NACK already answered it - peers->receiver->sendTo(peers->senderName, b::encodeAck(req.rid, 0, fakeRegion(0) + 1)); // wrong region - ASSERT_TRUE(syncSender(*peers)); + peers->nack(req.rid, 1); // a real answer: only chunk 0 is awaited now + peers->nack(req.rid, 1); // duplicate + peers->ack(req.rid, 2); // chunk 2 is not awaited: its NACK already answered it + auto const ungrantedRegion = fakeRegion(req.chunks); + peers->receiver->sendTo(peers->senderName, b::encodeAck(req.rid, 0, ungrantedRegion)); + peers->receiver->sendTo(peers->senderName, b::encodeAck(req.rid, /*chunkIdx=*/1, fakeRegion(0))); // crossed + b::ZmqControlChannel intruder("matchAnswersIntruder"); + ASSERT_TRUE(intruder.addPeer(peers->senderName, peers->sender->ch->localEndpoint())); + intruder.sendTo(peers->senderName, b::encodeAck(req.rid, 0, fakeRegion(0))); + // Each channel needs its own barrier: the reactor orders messages per peer only. + ASSERT_TRUE(reactorBarrier(intruder, peers->senderName)); + ASSERT_TRUE(reactorBarrier(*peers->receiver, peers->senderName)); EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) << "an answer that matches no awaited chunk ended the wait"; @@ -1732,7 +1973,8 @@ TEST(BounceTransportFailure, OnlyMatchingAnswersEndTheWait) auto const result = req.fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); - EXPECT_TRUE(result.quiesced); + EXPECT_TRUE(result.destinationQuiesced()); + EXPECT_EQ(peers->sender->tx->failedRequestCount(), 0u) << "kept a failed request whose chunks are all answered"; } // A failed request that still has unposted chunks (more chunks than maxInflightChunksPerRequest, the @@ -1741,14 +1983,18 @@ TEST(BounceTransportFailure, OnlyMatchingAnswersEndTheWait) TEST(BounceTransportFailure, FailedRequestWithUnpostedChunksStaysInert) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/30000); // maxInflightChunksPerRequest = 2 capRegions(c, c.maxInflightChunksPerRequest); // the sender's staging holds exactly two chunks auto peers = makeScriptedPeers("unpostedChunks", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ExecHold senderGathers(*peers->sender->exec, /*leaveFree=*/1); // post index == chunk index - ASSERT_TRUE(senderGathers.complete()); + ASSERT_TRUE(senderGathers.acquiredAll()); peers->ctl->scriptWrite(0, WriteScript::kComplete); peers->ctl->scriptWrite(1, WriteScript::kStall); @@ -1756,48 +2002,80 @@ TEST(BounceTransportFailure, FailedRequestWithUnpostedChunksStaysInert) ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/2); ASSERT_TRUE(waitData(*peers->receiver, req.rid, 1)); - ASSERT_TRUE(waitUntil([&] { return peers->ctl->postCount.load(std::memory_order_acquire) == 2; })); + ASSERT_TRUE(waitPosts(*peers->ctl, 2)); peers->ctl->scriptWrite(1, WriteScript::kFail); // Failed: chunk 0 is awaited, and chunks 2 and 3 were never posted. ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); peers->grant(req.rid, /*firstChunk=*/2, /*endChunk=*/4); - ASSERT_TRUE(syncSender(*peers)); - EXPECT_EQ(peers->ctl->postCount.load(std::memory_order_acquire), 2u) << "the failed request posted a write"; + ASSERT_TRUE(reactorBarrier(*peers->receiver, peers->senderName)); EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout); peers->ack(req.rid, 0); ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); auto const result = req.fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); - EXPECT_TRUE(result.quiesced); + EXPECT_TRUE(result.destinationQuiesced()); // No staging may stay held: a fresh 2-chunk request needs both regions to get both chunks Sent. ScriptedRequest next(*peers, /*chunks=*/2, /*seed=*/108); - ASSERT_TRUE(grantUntilAllSent(*peers, next)) << "the failed request leaked a staging region"; + ASSERT_TRUE(grantAllAndWaitSent(*peers, next)) << "the failed request leaked a staging region"; EXPECT_EQ(peers->ctl->postCount.load(std::memory_order_acquire), 4u) << "the failed request posted a write (only the fresh request's two may follow its own two)"; } +// A request that fails while its eagerly gathered chunk still waits for its credit returns that chunk's +// staging region: a fresh request that needs all of the sender's staging still gets every chunk Sent. +TEST(BounceTransportFailure, FailureRecyclesAGatheredChunkStillWaitingForItsCredit) +{ + if (!bounce_test::hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + auto c = cfg(/*timeoutMs=*/30000); + capRegions(c, c.maxInflightChunksPerRequest); // the sender's staging holds exactly two chunks + auto peers = makeScriptedPeers("gatheredRecycle", c); + if (!peers) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + + // submit() launches the eager gather; its exec context comes back once the reactor saw it finish. + ScriptedRequest gathered(*peers, /*chunks=*/1, /*seed=*/115); + ASSERT_TRUE(waitWant(*peers->receiver, gathered.rid)); + ASSERT_TRUE(waitUntil([&] { return peers->sender->exec->freeCount() == peers->sender->exec->size(); })) + << "the eager gather never finished"; + peers->nack(gathered.rid, 0); + ASSERT_EQ(gathered.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + + ScriptedRequest next(*peers, /*chunks=*/2, /*seed=*/116); + ASSERT_TRUE(grantAllAndWaitSent(*peers, next)) << "the failed request kept its gathered chunk's staging"; +} + // A receiver that stops answering is assumed dead requestTimeoutMs after its LAST answer: the request -// then resolves in-doubt. Each answer restarts that bound. Bounds are measured from the test's own -// send times, so they hold however the threads are scheduled. +// then resolves in-doubt. Each answer restarts that bound, and an answer after it still ends its chunk's +// write. Bounds are measured from the test's own send times, so they hold however the threads are scheduled. TEST(BounceTransportFailure, SilentReceiverResolvesInDoubtAfterTimeoutSinceLastAnswer) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } constexpr auto kTimeout = std::chrono::milliseconds(2000); - // Room for thread scheduling on top of the bound; far below a second kTimeout. - constexpr auto kSchedulingSlack = std::chrono::seconds(3); + // Room for thread scheduling on top of the bound; less than a second kTimeout, so a wait that ran + // twice still fails EXPECT_LT below. + constexpr auto kSchedulingSlack = kTimeout * 9 / 10; auto c = cfg(static_cast(kTimeout.count())); c.maxInflightChunksPerRequest = 3; capRegions(c, c.maxInflightChunksPerRequest); auto peers = makeScriptedPeers("silentReceiver", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ScriptedRequest req(*peers, /*chunks=*/3, /*seed=*/94); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); + ASSERT_TRUE(grantAllAndWaitSent(*peers, req)); // Fails the request; chunks 0 and 1 still owe an answer. peers->nack(req.rid, 2); @@ -1807,7 +2085,7 @@ TEST(BounceTransportFailure, SilentReceiverResolvesInDoubtAfterTimeoutSinceLastA ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) << "resolved before requestTimeoutMs passed since the failure"; auto const lastAnswerSent = std::chrono::steady_clock::now(); - peers->ack(req.rid, 0); // the last answer: chunk 1 never gets one + peers->ack(req.rid, 0); // the last answer within the wait ASSERT_EQ(req.fut.wait_for(kTimeout + kSchedulingSlack), std::future_status::ready) << "an unanswered chunk held the request past its bound"; @@ -1815,9 +2093,13 @@ TEST(BounceTransportFailure, SilentReceiverResolvesInDoubtAfterTimeoutSinceLastA auto const result = req.fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack); - EXPECT_FALSE(result.quiesced) << "chunk 1 was never answered"; + EXPECT_FALSE(result.destinationQuiesced()) << "chunk 1 was not answered within the wait"; EXPECT_GE(waited, kTimeout) << "the last answer did not restart the requestTimeoutMs bound"; EXPECT_LT(waited, kTimeout + kSchedulingSlack); + + b::BounceTransferStatus const status{req.fut}; + peers->ack(req.rid, 1); + EXPECT_TRUE(waitUntil([&] { return status.isCompleted(); })) << "the late answer was dropped"; } // The answer wait also restarts at the failure itself, however long the failing write took: the @@ -1825,15 +2107,19 @@ TEST(BounceTransportFailure, SilentReceiverResolvesInDoubtAfterTimeoutSinceLastA TEST(BounceTransportFailure, AnswerWaitStartsAtTheFailure) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; - constexpr auto kTimeout = std::chrono::milliseconds(1500); + } + constexpr auto kTimeout = std::chrono::milliseconds(5000); auto c = cfg(static_cast(kTimeout.count())); capRegions(c, c.maxInflightChunksPerRequest); auto peers = makeScriptedPeers("answerWaitStart", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ExecHold senderGathers(*peers->sender->exec, /*leaveFree=*/1); // post index == chunk index - ASSERT_TRUE(senderGathers.complete()); + ASSERT_TRUE(senderGathers.acquiredAll()); peers->ctl->scriptWrite(0, WriteScript::kComplete); peers->ctl->scriptWrite(1, WriteScript::kStall); @@ -1841,7 +2127,7 @@ TEST(BounceTransportFailure, AnswerWaitStartsAtTheFailure) ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); peers->grant(req.rid, /*firstChunk=*/0, /*endChunk=*/2); ASSERT_TRUE(waitData(*peers->receiver, req.rid, 1)); - ASSERT_TRUE(waitUntil([&] { return peers->ctl->postCount.load(std::memory_order_acquire) == 2; })); + ASSERT_TRUE(waitPosts(*peers->ctl, 2)); // Write 1 stalls for most of requestTimeoutMs (its post was the request's last progress), then fails. std::this_thread::sleep_for(kTimeout * 6 / 10); peers->ctl->scriptWrite(1, WriteScript::kFail); @@ -1855,51 +2141,27 @@ TEST(BounceTransportFailure, AnswerWaitStartsAtTheFailure) auto const result = req.fut.get(); EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); EXPECT_EQ(result.reason, b::BounceFailReason::kWriteFailed); - EXPECT_TRUE(result.quiesced); -} - -// A kNoProgressTimeout means no progress for requestTimeoutMs, so a request with a sent chunk still -// unanswered resolves in-doubt at once instead of waiting another requestTimeoutMs. -TEST(BounceTransportFailure, NoProgressTimeoutWithUnansweredChunkResolvesInDoubtAtOnce) -{ - if (!bounce_test::hasCuda()) - GTEST_SKIP() << "no CUDA device"; - constexpr auto kTimeout = std::chrono::milliseconds(2000); - auto c = cfg(static_cast(kTimeout.count())); - capRegions(c, c.maxInflightChunksPerRequest); - auto peers = makeScriptedPeers("noProgress", c); - if (!peers) - GTEST_SKIP() << "NIXL agent/backend unavailable"; - - ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/95); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); - auto const lastAnswerSent = std::chrono::steady_clock::now(); - peers->ack(req.rid, 0); // the request's last progress; chunk 1 stays unanswered - - ASSERT_EQ(req.fut.wait_for(2 * kTimeout), std::future_status::ready) - << "the no-progress timeout waited another requestTimeoutMs for answers"; - auto const waited = std::chrono::steady_clock::now() - lastAnswerSent; - auto const result = req.fut.get(); - EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); - EXPECT_EQ(result.reason, b::BounceFailReason::kNoProgressTimeout); - EXPECT_FALSE(result.quiesced) << "chunk 1 was never answered"; - EXPECT_GE(waited, kTimeout); - EXPECT_LT(waited, 2 * kTimeout); + EXPECT_TRUE(result.destinationQuiesced()); } // An abandoned flow (GRANT mispair) dies through checkTimeouts' kNoProgressTimeout. With a Sent chunk // unanswered it resolves in-doubt within one requestTimeoutMs of its last progress, reporting the -// abandon cause; it must not wait another requestTimeoutMs for answers. +// abandon cause; it must not wait another requestTimeoutMs for answers. The sender keeps it for late +// answers only while someone holds its result. TEST(BounceTransportFailure, AbandonWithSentChunkResolvesInDoubtWithinOneTimeout) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; - constexpr auto kTimeout = std::chrono::milliseconds(1500); + } + constexpr auto kTimeout = std::chrono::milliseconds(3000); auto c = cfg(static_cast(kTimeout.count())); capRegions(c, c.maxInflightChunksPerRequest); auto peers = makeScriptedPeers("abandonSent", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/110); ASSERT_TRUE(waitWant(*peers->receiver, req.rid)); @@ -1907,138 +2169,112 @@ TEST(BounceTransportFailure, AbandonWithSentChunkResolvesInDoubtWithinOneTimeout ASSERT_TRUE(waitData(*peers->receiver, req.rid, 1)); auto const chunk0Sent = std::chrono::steady_clock::now(); // after the request's last progress // Chunk 1's credit is half the chunk: a mispair, so the flow is abandoned. - std::vector const undersized{{/*addr=*/0, peers->chunkBytes / 2, /*devId=*/0, fakeRegion(1)}}; + std::vector const undersized{ + {.addr = 0, .len = peers->chunkBytes / 2, .devId = 0, .regionHandle = fakeRegion(1)}}; peers->receiver->sendTo(peers->senderName, b::encodeGrant(req.rid, undersized)); ASSERT_EQ(req.fut.wait_for(kTimeout + kTimeout / 2), std::future_status::ready) << "the abandoned flow waited another requestTimeoutMs for answers"; EXPECT_LT(std::chrono::steady_clock::now() - chunk0Sent, kTimeout + kTimeout / 2); - auto const result = req.fut.get(); - EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); - EXPECT_EQ(result.reason, b::BounceFailReason::kProtocolError); - EXPECT_FALSE(result.quiesced) << "chunk 0 was never answered"; -} - -// A kNoProgressTimeout with every chunk Sent and none answered resolves in-doubt. ACKs arriving -// afterwards, enough to count every chunk, hit an erased request: the reactor drops them and keeps -// serving (the FIFO barrier below). -TEST(BounceTransportFailure, LateAcksAfterNoProgressTimeoutChangeNothing) -{ - if (!bounce_test::hasCuda()) - GTEST_SKIP() << "no CUDA device"; - constexpr auto kTimeout = std::chrono::milliseconds(1000); - auto c = cfg(static_cast(kTimeout.count())); - capRegions(c, c.maxInflightChunksPerRequest); - auto peers = makeScriptedPeers("lateAcks", c); - if (!peers) - GTEST_SKIP() << "NIXL agent/backend unavailable"; - - ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/96); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); - ASSERT_EQ(req.fut.wait_for(kTimeout + std::chrono::seconds(5)), std::future_status::ready); EXPECT_EQ(req.fut.get().state, kvc::TransferState::kFAILURE); - EXPECT_EQ(req.fut.get().reason, b::BounceFailReason::kNoProgressTimeout); - EXPECT_FALSE(req.fut.get().quiesced); + EXPECT_EQ(req.fut.get().reason, b::BounceFailReason::kProtocolError); + EXPECT_FALSE(req.fut.get().destinationQuiesced()) << "chunk 0 was never answered"; - peers->ack(req.rid, 0); - peers->ack(req.rid, 1); - EXPECT_TRUE(syncSender(*peers)) << "the reactor stopped after the late ACKs"; + EXPECT_EQ(peers->sender->tx->failedRequestCount(), 1u) << "dropped an in-doubt request whose result is held"; + req.fut = {}; + EXPECT_TRUE(waitUntil([&] { return peers->sender->tx->failedRequestCount() == 0; })) + << "kept an in-doubt request whose result nobody holds"; } -// forgetPeer while a failed request awaits answers must not fail it again: it keeps its first reason -// and keeps waiting, and a later answer resolves it with that ORIGINAL reason. -TEST(BounceTransportFailure, ForgetPeerWhileAwaitingAnswersKeepsWaiting) +// forgetPeer fails only active requests. One that already failed keeps its reason and its answer wait; a +// fully sent active one fails with kPeerDropped and then waits for its chunks' answers, which still reach +// the sender's ROUTER. Both resolve FAILURE once answered, never SUCCESS, though every chunk ends up ACKed. +TEST(BounceTransportFailure, ForgetPeerFailsOnlyActiveRequestsAndLaterAnswersStillCount) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/30000); capRegions(c, c.maxInflightChunksPerRequest); - auto peers = makeScriptedPeers("forgetAwaiting", c); + auto peers = makeScriptedPeers("forgetPeer", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } - ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/98); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); + ScriptedRequest failed(*peers, /*chunks=*/2, /*seed=*/98); + ASSERT_TRUE(grantAllAndWaitSent(*peers, failed)); + peers->nack(failed.rid, 1); // fails it: chunk 0 is awaited + ASSERT_TRUE(waitCancel(*peers->receiver, failed.rid)); + ScriptedRequest active(*peers, /*chunks=*/2, /*seed=*/100); + ASSERT_TRUE(grantAllAndWaitSent(*peers, active)); // Never granted: forgetPeer fails it, which proves the forget was applied. ScriptedRequest sentinel(*peers, /*chunks=*/1, /*seed=*/99); - peers->nack(req.rid, 1); - // No write was in flight, so the cancel went out at once: the NACK failed the request. - ASSERT_TRUE(waitCancel(*peers->receiver, req.rid)); peers->sender->tx->forgetPeer(peers->receiverName); ASSERT_EQ(sentinel.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "forgetPeer not applied"; EXPECT_EQ(sentinel.fut.get().reason, b::BounceFailReason::kPeerDropped); - EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) - << "forgetPeer resolved a request that still awaits an answer"; - - peers->ack(req.rid, 0); // answers still reach the sender's ROUTER after forgetPeer - ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); - auto const result = req.fut.get(); - EXPECT_EQ(result.state, kvc::TransferState::kFAILURE); - EXPECT_EQ(result.reason, b::BounceFailReason::kPeerNack) << "forgetPeer failed the request a second time"; - EXPECT_TRUE(result.quiesced); -} - -// forgetPeer fails a request whose chunks were all Sent; it then waits for their answers. ACKs for -// every chunk must resolve it FAILURE (kPeerDropped, quiesced), never SUCCESS, even though the ACK -// count reaches numChunks. -TEST(BounceTransportFailure, AnswersAfterForgetPeerResolveFailureNotSuccess) -{ - if (!bounce_test::hasCuda()) - GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/30000); - capRegions(c, c.maxInflightChunksPerRequest); - auto peers = makeScriptedPeers("forgetAllSent", c); - if (!peers) - GTEST_SKIP() << "NIXL agent/backend unavailable"; - - ScriptedRequest req(*peers, /*chunks=*/2, /*seed=*/100); - ASSERT_TRUE(grantUntilAllSent(*peers, req)); - // Never granted: forgetPeer fails it in the same pass, which proves the forget was applied. - ScriptedRequest sentinel(*peers, /*chunks=*/1, /*seed=*/101); - - peers->sender->tx->forgetPeer(peers->receiverName); - ASSERT_EQ(sentinel.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "forgetPeer not applied"; - EXPECT_EQ(req.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + EXPECT_EQ(failed.fut.wait_for(std::chrono::seconds(0)), std::future_status::timeout) + << "forgetPeer resolved a failed request that still awaits an answer"; + EXPECT_EQ(active.fut.wait_for(std::chrono::milliseconds(100)), std::future_status::timeout) << "forgetPeer resolved a request whose sent chunks were unanswered"; - peers->ack(req.rid, 0); - peers->ack(req.rid, 1); - ASSERT_EQ(req.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); - auto const result = req.fut.get(); - EXPECT_EQ(result.state, kvc::TransferState::kFAILURE) << "a failed request resolved SUCCESS"; - EXPECT_EQ(result.reason, b::BounceFailReason::kPeerDropped); - EXPECT_TRUE(result.quiesced); + peers->ack(failed.rid, 0); + peers->ack(active.rid, 0); + peers->ack(active.rid, 1); + ASSERT_EQ(failed.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + ASSERT_EQ(active.fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_EQ(failed.fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(failed.fut.get().reason, b::BounceFailReason::kPeerNack) << "forgetPeer failed the request a second time"; + EXPECT_TRUE(failed.fut.get().destinationQuiesced()); + EXPECT_EQ(active.fut.get().state, kvc::TransferState::kFAILURE) << "a failed request resolved SUCCESS"; + EXPECT_EQ(active.fut.get().reason, b::BounceFailReason::kPeerDropped); + EXPECT_TRUE(active.fut.get().destinationQuiesced()); } -// Shutdown resolves every pending request. One already awaiting answers keeps its own reason, and any -// request with a sent chunk unanswered is in-doubt: the receiver may still scatter into its dst. +// Shutdown resolves every pending request. One already awaiting answers keeps its own reason, an abandoned +// one keeps its abandon cause, and any request with a sent chunk unanswered is in-doubt: the receiver may +// still scatter into its dst. TEST(BounceTransportFailure, ShutdownResolvesUnansweredRequestsInDoubtWithTheirOwnReason) { if (!bounce_test::hasCuda()) + { GTEST_SKIP() << "no CUDA device"; + } auto c = cfg(/*timeoutMs=*/30000); c.maxInflightChunksPerRequest = 3; capRegions(c, c.maxInflightChunksPerRequest); auto peers = makeScriptedPeers("shutdownAwaiting", c); if (!peers) + { GTEST_SKIP() << "NIXL agent/backend unavailable"; + } ScriptedRequest failed(*peers, /*chunks=*/2, /*seed=*/104); - ASSERT_TRUE(grantUntilAllSent(*peers, failed)); + ASSERT_TRUE(grantAllAndWaitSent(*peers, failed)); peers->nack(failed.rid, 1); // fails it: chunk 0 is awaited ASSERT_TRUE(waitCancel(*peers->receiver, failed.rid)); // Still live: its only chunk is Sent and unanswered. ScriptedRequest live(*peers, /*chunks=*/1, /*seed=*/105); - ASSERT_TRUE(grantUntilAllSent(*peers, live)); + ASSERT_TRUE(grantAllAndWaitSent(*peers, live)); + // Abandoned, not yet failed: its credit is half its chunk. + ScriptedRequest abandoned(*peers, /*chunks=*/1, /*seed=*/114); + ASSERT_TRUE(waitWant(*peers->receiver, abandoned.rid)); + std::vector const undersized{ + {.addr = 0, .len = peers->chunkBytes / 2, .devId = 0, .regionHandle = fakeRegion(0)}}; + peers->receiver->sendTo(peers->senderName, b::encodeGrant(abandoned.rid, undersized)); + ASSERT_TRUE(reactorBarrier(*peers->receiver, peers->senderName)); peers->sender->tx->shutdown(); ASSERT_EQ(failed.fut.wait_for(std::chrono::seconds(0)), std::future_status::ready); EXPECT_EQ(failed.fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(failed.fut.get().reason, b::BounceFailReason::kPeerNack) << "shutdown replaced the original reason"; - EXPECT_FALSE(failed.fut.get().quiesced) << "chunk 0 was never answered"; + EXPECT_FALSE(failed.fut.get().destinationQuiesced()) << "chunk 0 was never answered"; ASSERT_EQ(live.fut.wait_for(std::chrono::seconds(0)), std::future_status::ready); EXPECT_EQ(live.fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(live.fut.get().reason, b::BounceFailReason::kShutdown); - EXPECT_FALSE(live.fut.get().quiesced) << "its sent chunk was never answered"; + EXPECT_FALSE(live.fut.get().destinationQuiesced()) << "its sent chunk was never answered"; + ASSERT_EQ(abandoned.fut.wait_for(std::chrono::seconds(0)), std::future_status::ready); + EXPECT_EQ(abandoned.fut.get().reason, b::BounceFailReason::kProtocolError) << "shutdown replaced the abandon cause"; }