diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 247b7c0c5c06..4e1aedc3902a 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. + /// 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 f5acdefb0abe..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,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, 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 new file mode 100644 index 000000000000..614d7ce4404f --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferStatus.h @@ -0,0 +1,96 @@ +/* + * 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/BounceResult.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)) + { + } + + /// The TransferStatus contract; see FailedRequest for when a failure gets there. + [[nodiscard]] bool isCompleted() const override + { + // 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 + { + 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.destinationQuiesced()) + { + 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..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))) @@ -684,8 +664,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 +677,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) @@ -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); @@ -789,8 +771,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( @@ -827,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( @@ -1279,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; } @@ -1345,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; } @@ -1361,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 failed / was cancelled (e.g. a scatter NACK - // racing the timeout or a second chunk's NACK after the first one failed it). - 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; @@ -1375,9 +1346,11 @@ 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. - failRequest(h.requestId, req, BounceFailReason::kPeerNack); + // The NACK answers its own chunk (the receiver never scattered it, or synced before NACKing), so + // 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); + failRequest(h.requestId, BounceFailReason::kPeerNack, std::chrono::steady_clock::now()); } void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) @@ -1391,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; @@ -1401,22 +1375,7 @@ 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) - { - // 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; - } - } - 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,57 +1388,112 @@ 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); + endPipelineSpans(req); 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. - } + req.promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); mRequests.erase(it); } } -void BounceSender::checkTimeouts() +bool BounceSender::releaseAnsweredChunk(Request& req, BounceMsgHeader const& h) { - if (mCtx.cfg.requestTimeoutMs <= 0) + // 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()) { - // 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 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::answerFailedRequest(std::string const& peer, BounceMsgHeader const& h) +{ + 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(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, + static_cast(h.regionHandle)); + } +} + +void BounceSender::checkTimeouts() +{ + 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()) - { - // 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); - } + 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()) @@ -1575,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) @@ -1597,19 +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) { - // 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) + auto node = mRequests.extract(rid); + if (node.empty()) { - reason = req.abandonReason; + return false; } + 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(reason), 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): @@ -1619,64 +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. - bool deferredWrite = false; + 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 } mCtx.sendGrants(mCtx.scheduler.releaseLocal(p.localOffset)); } + 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. - if (deferredWrite) + // The cancel does not stop the receiver's queued scatters: they still run and answer. + if (writeInFlight) { - mPendingCancel[rid] = req.peer; + mPendingCancel[rid] = peer; } else { - mCtx.channel->sendTo(req.peer, encodeCancel(rid, mCtx.channel->localEndpoint())); + mCtx.channel->sendTo(peer, encodeCancel(rid, mCtx.channel->localEndpoint())); } +} + +BounceSender::FailedEntry BounceSender::toFailedEntry(Request&& req, BounceFailReason cause, + std::shared_ptr outstanding, std::chrono::milliseconds answerWait, + std::chrono::steady_clock::time_point now) +{ + 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) + { + if (p.state == PostState::Sent) + { + sent.push_back({.chunkIdx = p.chunkIdx, .regionHandle = p.remoteHandle}); + } + } + return std::make_shared(std::move(sent)); +} + +void BounceSender::endPipelineSpans(Request& req) +{ bounceRangeEnd(req.nvtxGrantWait); bounceRangeEnd(req.nvtxCreditStarved); bounceRangeEnd(req.nvtxArenaStarved); - bounceRangeEnd(req.nvtxReq); - try - { - req.promise->set_value({TransferState::kFAILURE, reason}); - } - 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); } - 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 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) @@ -1688,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 @@ -1756,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()) @@ -1776,21 +1845,21 @@ void BounceSender::failAll() p.ctx = nullptr; } } - bounceRangeEnd(req.nvtxGrantWait); - bounceRangeEnd(req.nvtxCreditStarved); - bounceRangeEnd(req.nvtxArenaStarved); - bounceRangeEnd(req.nvtxReq); - try - { - req.promise->set_value({TransferState::kFAILURE, BounceFailReason::kShutdown}); - } - catch (...) + 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()) { - // 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. + onFailureResolved(rid, failed, *unanswered); } } - 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 42f12ee0820c..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" @@ -78,8 +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. +// [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. @@ -91,41 +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. +// 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}; -}; - /// 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. @@ -274,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. @@ -285,6 +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 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 @@ -300,16 +281,20 @@ 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), + /// 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 so its wait() returns. + /// 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, 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. @@ -393,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). @@ -424,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 @@ -445,16 +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); - /// 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. - 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 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; @@ -515,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/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..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 @@ -150,11 +161,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 7fa215f170e6..8ab5d71fdb89 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) { @@ -193,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(); @@ -200,6 +202,52 @@ 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); + // 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; + + 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) @@ -212,7 +260,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/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/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&) diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp index 6b90b9dd7da2..735099ceee07 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp @@ -19,9 +19,13 @@ // (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 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:: @@ -29,12 +33,25 @@ #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 #include namespace b = tensorrt_llm::executor::kv_cache::bounce; @@ -42,6 +59,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 +76,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 +106,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 +119,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 +153,7 @@ class FakeXferStatus final : public kvc::TransferStatus private: std::shared_ptr mCtl; + std::uint64_t mPostIndex{}; bool mReleased{false}; }; @@ -113,8 +172,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: @@ -189,30 +248,502 @@ 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; + }); +} + +// The ACKs and NACKs seen for one rid. +struct AnswerCounts +{ + std::uint64_t rid{}; + int acks{0}; + int nacks{0}; + + void add(b::BounceMsgHeader const& h) + { + if (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) + { + 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 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) +// 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) { - int n = 0; - pumpChannel(ch, budget, + 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&) { - if (static_cast(h.msgType) == b::BounceMsgType::kACK && h.requestId == rid) + counts.add(h); + return false; + }); +} + +// ---- 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; +} + +// 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 +// 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) { - ++n; + 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 acquiredAll() const + { + return mHeld.size() + mLeaveFree == mPool.size(); + } + + [[nodiscard]] b::ExecCtx* first() const + { + return mHeld.front(); + } + + void releaseFirst() + { + 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; + +// 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)) +{ + 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; }); - return n; } + +// 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; + }); +} + +// 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) { return isCancelFor(h, blob, rid); }); +} + +// 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 +{ + ~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; + + // 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, .len = chunkBytes, .devId = 0, .regionHandle = 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 grantAllAndWaitSent(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); +} + } // namespace // A peer that never GRANTs must fail the request on requestTimeoutMs, not hang. The WANT is @@ -246,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(); @@ -254,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()); @@ -263,6 +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().destinationQuiesced()) << "no write completed, so no DATA was sent"; EXPECT_GE(ctl->postCount.load(std::memory_order_acquire), 1u); A->tx->shutdown(); @@ -280,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; @@ -288,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); @@ -307,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) { @@ -323,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())); @@ -357,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); @@ -570,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}; @@ -698,14 +1334,18 @@ 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"; @@ -721,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(); @@ -730,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())); @@ -754,21 +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().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 @@ -850,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 @@ -859,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()) @@ -897,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; @@ -1030,3 +1678,603 @@ TEST(BounceTransportFailure, ConcurrentMultiThreadedSubmit) bounce_test::freeXferBufs(bb); } } + +// ============================================================================ +// A FAILURE resolves only after the receiver has answered every chunk whose DATA was sent +// ---------------------------------------------------------------------------- +// 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. +// ============================================================================ + +// 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=*/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()); + + 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::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(); +} + +// 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 + GatedScatterPair pair("runScatter", c, /*chunks=*/2, /*seed=*/91); + 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); + // The worker took the gated context: chunk 0 is Sent and its scatter waits behind the gate. + 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 pair.ctl->releaseCount.load(std::memory_order_acquire) == 2; })) + << "the failed write was never handled"; + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(100)), std::future_status::timeout) + << "FAILURE reported while chunk 0's scatter could still write dst"; + + 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.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(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 +// 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 + GatedScatterPair pair("queuedScatter", c, /*chunks=*/3, /*seed=*/92); + if (!pair.available()) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + 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. + pair.ctl->scriptWrite(2, WriteScript::kFail); + + 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 pair.ctl->releaseCount.load(std::memory_order_acquire) == 3; })) + << "the failed write was never handled"; + 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"; + + 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.destinationQuiesced()); + for (std::uint32_t chunk = 0; chunk < 2; ++chunk) + { + 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(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). 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.acquiredAll()); + 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(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(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); + 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.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"; + + peers->ctl->scriptWrite(2, WriteScript::kComplete); + 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, 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(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->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"; + + 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.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 +// 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.acquiredAll()); + 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(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(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.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(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, 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; 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(grantAllAndWaitSent(*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 within the wait + + 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.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 +// 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(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.acquiredAll()); + 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(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); + 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.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. 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(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)); + 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, .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); + EXPECT_EQ(req.fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(req.fut.get().reason, b::BounceFailReason::kProtocolError); + EXPECT_FALSE(req.fut.get().destinationQuiesced()) << "chunk 0 was never answered"; + + 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 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("forgetPeer", c); + if (!peers) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + + 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->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(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(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, 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(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(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().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().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"; +}