From c53c82f95a083f9972cb197325e3f7ee093d032e Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:20:46 +0800 Subject: [PATCH] fix: require physical quiescence before retiring C++ bounce KV Separate logical transfer failure from physical completion with a wire-v3 drain fence, retained scatter/gather/RDMA ownership, and bounded fail-stop containment. Connect native proof to Python session retirement, capability negotiation, and the final executor resource-release boundary. Preserve default Python bounce selection and qualify the opt-in native lifecycle profile. Add CUDA fault-injection, native status-wrapper, lifecycle and retirement regressions. This provides the ownership prerequisite for NVIDIA/TensorRT-LLM#20049. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../tensorrt_llm/executor/transferAgent.h | 13 + .../nixl_utils/agentBindings.cpp | 2 + .../nixl_utils/bounce/BounceConfig.h | 22 +- .../nixl_utils/bounce/BounceMessage.cpp | 15 + .../nixl_utils/bounce/BounceMessage.h | 41 +- .../nixl_utils/bounce/BounceTransport.cpp | 851 ++++++++++-------- .../nixl_utils/bounce/BounceTransport.h | 107 ++- .../nixl_utils/bounce/CreditScheduler.cpp | 69 ++ .../nixl_utils/bounce/CreditScheduler.h | 55 +- .../nixl_utils/transferAgent.cpp | 31 +- .../nixl_utils/transferAgent.h | 2 + .../executor/bounce/bounceAgentE2ETest.cpp | 88 ++ .../bounce/bounceMessageCodecTest.cpp | 27 + .../executor/bounce/bounceTestNixlNode.h | 2 +- .../bounce/bounceTransportFailureTest.cpp | 478 ++++++++-- .../executor/bounce/creditSchedulerTest.cpp | 32 + docs/source/features/disagg-serving.md | 6 + .../_torch/disaggregation/base/agent.py | 34 + .../disaggregation/kv_cache_transceiver.py | 4 + .../_torch/disaggregation/native/peer.py | 3 +- .../_torch/disaggregation/native/rank_info.py | 17 + .../disaggregation/native/retirement.py | 25 +- .../_torch/disaggregation/native/transfer.py | 146 ++- .../_torch/disaggregation/nixl/_agent_cpp.py | 8 + .../orchestration/coordinator.py | 2 +- .../_torch/disaggregation/transceiver.py | 88 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 3 + .../test_disagg_coordinator_transfers.py | 13 +- .../test_cache_transceiver_single_process.py | 63 +- .../test_kv_transfer_fail_stop.py | 13 +- .../disaggregated/test_retirement_deadline.py | 10 +- .../test_transceiver_bounded_polling.py | 2 + .../test_transfer_late_settlement.py | 34 +- .../test_transfer_ownership_regressions.py | 683 +++++++++++++- .../disaggregated/test_unsubmitted_session.py | 15 +- 35 files changed, 2348 insertions(+), 656 deletions(-) diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 247b7c0c5c06..de17538cd9e0 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -393,6 +393,19 @@ class TransferStatus { return {}; } + + /// True only after no future source read or destination write is possible. + /// Logical failure, cancellation and handle release are not this proof. + [[nodiscard]] virtual bool isQuiesced() const + { + return false; + } + + /// Request cancellation without waiting; acceptance is not a quiescence proof. + [[nodiscard]] virtual bool requestCancel() + { + return false; + } }; struct BaseAgentConfig diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index 6b7b8eb2457a..da1cb68d64f9 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -188,6 +188,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) // the GIL here is safe and necessary for correct behavior when the concrete // subclass type is not directly registered (e.g., agents created via factory). nb::class_(m, "TransferStatus") + .def("is_quiesced", &kvc::TransferStatus::isQuiesced, nb::call_guard()) + .def("request_cancel", &kvc::TransferStatus::requestCancel, nb::call_guard()) .def("is_completed", &kvc::TransferStatus::isCompleted, nb::call_guard()) .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard()) // Failure detail for the last terminal state (empty if unavailable). Named to match the 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..5070ff365b12 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 @@ -130,7 +130,7 @@ namespace tensorrt_llm::executor::kv_cache::bounce return parsed * mult; } -/// POD config for the bounce v2 pipeline. There is no `enabled` field: the on/off switch is +/// POD config for the bounce pipeline. There is no `enabled` field: the on/off switch is /// CacheTransceiverConfig's agent_bounce_buffer_enable + kv_cache_bounce_size_mb (which the Python /// frontend folds into the agent's arena size in MiB: 0 = off, >0 = arena size), and at runtime /// "bounce on" simply means the owning agent built its bounce state (mBounce != nullptr). @@ -163,24 +163,22 @@ struct BounceConfig // means no minimum. std::size_t minDescriptorCount{1024}; // min_descriptor_count std::size_t maxAverageDescriptorSizeBytes{16ULL << 10}; // max_average_descriptor_size - 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. + int requestTimeoutMs{30000}; // request_timeout_ms; must be > 0 for logical failure detection + // and receiver lease fencing. Expiry initiates physical drain; it + // cannot itself establish quiescence or permit memory reuse. // 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 // alone — so a flow whose grants see no progress (no GRANT sent, no DATA received) for this - // long is reclaimed and its regions quarantined (below) before reuse. The lease must EXCEED + // long is fenced and its regions retained until explicit physical drain. The lease must EXCEED // the peers' requestTimeoutMs (a live sender abandons + cancels first, so only // dead/unreachable peers ever hit this) — the 2x derivation relies on both ends running the // SAME request_timeout_ms, which the capability handshake enforces (strict equality; a // mismatched peer falls back to standard NIXL). int receiverFlowTimeoutMs{60000}; - // How long a receiver-reclaimed, possibly-still-being-written region stays out of the arena - // before reuse — DERIVED in deriveDependentTimeouts() as requestTimeoutMs (equal to the peer's - // by the handshake). A one-sided RDMA write cannot be aborted, so time is the only barrier - // against re-granting a region a gone peer's NIC may still be writing. + // Legacy scheduler quarantine interval, retained for direct scheduler callers/tests and sweep + // cadence. The v3 transport uses permanent per-flow quarantine until DRAIN proves that all + // remote writes stopped; this interval never permits release of those regions. int quarantineMs{30000}; bool disableFabricMemory{false}; // disable_fabric_memory // enable_eager_gather: launch a chunk's gather at submit() time, before the receiver's GRANT @@ -195,8 +193,8 @@ struct BounceConfig bool useZeroCopyArguments{true}; /// Re-derive the lease/quarantine values from the one user-visible timeout (see the field - /// comments): the lease must exceed the peers' request timeout, and time is the only write - /// barrier for a reclaimed region. The config layer guarantees requestTimeoutMs > 0 (applyParam + /// comments): the lease exceeds the peers' request timeout but is only a fencing deadline, + /// never a physical completion barrier. The config layer guarantees requestTimeoutMs > 0 (applyParam /// rejects 0); the > 0 guard below only matters for directly-constructed configs. Must be /// re-run whenever requestTimeoutMs changes: fromEnv() always calls it, fromParams() only when /// the dict actually provides request_timeout_ms — so a caller-tweaked receiverFlowTimeoutMs / diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp index 3abd1463c1a3..780ef8e0bf8e 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp @@ -118,6 +118,21 @@ std::string encodeCancel(std::uint64_t requestId, std::string const& endpoint) return encodeWant(requestId, {}, endpoint); } +std::string encodeDrain(std::uint64_t requestId, std::string const& endpoint) +{ + auto blob = encodeWant(requestId, {}, endpoint); + BounceMsgHeader header{}; + std::memcpy(&header, blob.data(), sizeof(header)); + header.msgType = static_cast(BounceMsgType::kDRAIN); + std::memcpy(blob.data(), &header, sizeof(header)); + return blob; +} + +std::string encodeDrainAck(std::uint64_t requestId) +{ + return encodeHeaderOnly(makeHeader(BounceMsgType::kDRAIN_ACK, requestId, 0, 0, 0, 0, 0)); +} + std::string encodeGrant(std::uint64_t requestId, std::vector const& credits) { auto const bytes = static_cast(credits.size() * sizeof(BounceCreditEntry)); 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..478ef4be9641 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 @@ -24,9 +24,9 @@ namespace tensorrt_llm::executor::kv_cache::bounce { -/// Bounce v2 control-plane wire format. Carried over the ControlChannel (zmq by default). +/// Bounce v3 control-plane wire format. Carried over the ControlChannel (zmq by default). /// Pure encode/decode — no IO — so it is fully unit-testable. NOTE: this is NOT the NIXL -/// notifMsg (v2 does not use notifMsg at all); it is our own message set. +/// notifMsg (bounce does not use notifMsg); it is our own message set. /// /// ENDIANNESS: fields are memcpy'd raw (host byte order), so this assumes both peers are /// little-endian — true for all NVIDIA GPU hosts (x86_64 / aarch64 LE). It is NOT a portable @@ -35,11 +35,13 @@ namespace tensorrt_llm::executor::kv_cache::bounce /// treats any unknown `msgType` as a no-op (default branch), so a bogus type can't misroute. enum class BounceMsgType : std::uint16_t { - kWANT = 1, // sender -> receiver: per-chunk byte sizes[] for requestId (empty list cancels) - kGRANT = 2, // receiver -> sender: credits[] for requestId - kDATA = 3, // sender -> receiver: region written; scatter entries[] (after getXferStatus SUCCESS) - kACK = 4, // receiver -> sender: chunk scattered; region freed - kNACK = 5, // receiver -> sender: chunk/request cannot be completed; sender fails the request + kWANT = 1, // sender -> receiver: per-chunk byte sizes[] for requestId (empty list cancels) + kGRANT = 2, // receiver -> sender: credits[] for requestId + kDATA = 3, // sender -> receiver: region written; scatter entries[] (after getXferStatus SUCCESS) + kACK = 4, // receiver -> sender: chunk scattered; region freed + kDRAIN = 6, // sender -> receiver: no more RDMA; fence and drain every scatter + kDRAIN_ACK = 7, // receiver -> sender: flow fenced and all scatter physically complete + kNACK = 5, // receiver -> sender: chunk/request cannot be completed; sender fails the request }; #pragma pack(push, 1) @@ -51,7 +53,7 @@ struct BounceMsgHeader std::uint32_t magic; // kMagic std::uint16_t version; // kBounceVersion std::uint16_t msgType; // BounceMsgType - std::uint64_t requestId; // WANT/GRANT/DATA/ACK + std::uint64_t requestId; // all control messages std::uint64_t regionHandle; // DATA/ACK: the arena region offset of this chunk; 0 elsewhere (64-bit -> arena may exceed 4 GiB) std::uint32_t chunkIdx; // DATA/ACK @@ -97,10 +99,10 @@ 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 +// v3: explicit DRAIN/DRAIN_ACK proves physical quiescence after logical failure. 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; +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 @@ -112,12 +114,14 @@ inline constexpr std::uint16_t kBounceVersion = 2U; /// decode via decodeWant. [[nodiscard]] std::string encodeWant( std::uint64_t requestId, std::vector const& chunkBytes, std::string const& endpoint); -/// Cancel/retract a request: a WANT with an EMPTY chunk list (the receiver frees everything it -/// allocated/held for `requestId`). This is the ONLY meaning of an empty WANT — submit() never sends -/// a zero-chunk WANT (a 0-chunk transfer resolves SUCCESS without any WANT) — so an empty chunk list -/// is unambiguously a cancel. Thin named wrapper over encodeWant to make the intent explicit at call -/// sites; the wire form is identical (still a WANT), so it reuses the receiver's onWant/reclaim path. +/// Cancel/retract a request with an EMPTY WANT. This fences future WANT/DATA for the flow but +/// supplies no physical completion evidence: unknown writes remain quarantined and queued scatters +/// remain owned until completion. A zero-chunk transfer succeeds without sending WANT. Use DRAIN +/// only after every write into the peer arena is proved stopped; its acknowledgment proves that +/// the receiver has fenced the flow and physically completed all admitted scatters. [[nodiscard]] std::string encodeCancel(std::uint64_t requestId, std::string const& endpoint); +[[nodiscard]] std::string encodeDrain(std::uint64_t requestId, std::string const& endpoint); +[[nodiscard]] std::string encodeDrainAck(std::uint64_t requestId); [[nodiscard]] std::string encodeGrant(std::uint64_t requestId, std::vector const& credits); [[nodiscard]] std::string encodeData(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint32_t numChunks, std::uint64_t regionHandle, std::vector const& entries); @@ -184,10 +188,9 @@ struct BounceHandshake // This agent's effective per-chunk cap after the arena-capacity clamp. Peers compare this field // exactly, which also keeps their scatter-plan capacities consistent. std::uint64_t maxChunkSizeBytes{0}; - // This agent's request_timeout_ms. Peers compare it exactly: the receiver's region lease (2x its - // value) must exceed the sender's timeout. The quarantine (1x) starts at lease expiry OR at a - // receiver-side peer teardown (forget()/invalidateRemoteAgent); on the lease path it needs no - // extra margin of its own. + // This agent's request_timeout_ms. Peers compare it exactly so the receiver lease (2x this + // value) gives live senders time to initiate drain. Lease expiry fences a flow; it never proves + // that a remote write stopped and cannot release quarantined storage. std::int32_t requestTimeoutMs{0}; // Control-channel address: a zmq endpoint (kZMQ). std::string endpoint; 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..ed5398ce1f6f 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 @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -49,6 +50,7 @@ char const* toString(BounceFailReason reason) 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::kCancelled: return "bounce: cancellation requested"; case BounceFailReason::kPeerNack: return "bounce: receiver reported it cannot complete the chunk/request (NACK)"; } return "bounce: unknown"; @@ -197,25 +199,51 @@ cudaError_t launchPrepared(ExecCtx* ctx, std::size_t n, bool zeroCopy) // BounceContext // ============================================================================ -void BounceContext::sendGrants(std::vector const& grants) +bool BounceContext::addPeer(std::string const& peer, std::string const& endpoint) { - if (grants.empty()) + std::lock_guard lock(peerMutex); + auto const it = peerEndpoints.find(peer); + if (it != peerEndpoints.end() && it->second != endpoint) { - return; + TLLM_LOG_ERROR("Bounce peer %s reused its lifetime identity with a different endpoint", peer.c_str()); + return false; } - // grants are keyed by flow id ("peerrid"); split back to (agent name, rid) to address them. - std::unordered_map> byFlow; - for (auto const& g : grants) + if (!channel->addPeer(peer, endpoint)) + { + return false; + } + peerEndpoints.emplace(peer, endpoint); + return true; +} + +void BounceContext::sendGrants(std::vector const& grants) +{ + try { - // Carry OUR (receiver) device id so the sender writes the remote desc to the right GPU - // even if the two agents don't share a device index. `regionHandle` (our arena offset) is - // echoed back in DATA so we can locate + free the region. - byFlow[g.flow].push_back(BounceCreditEntry{g.addr, g.len, static_cast(deviceId), g.offset}); + if (grants.empty()) + { + return; + } + // grants are keyed by flow id ("peerrid"); split back to (agent name, rid) to address them. + std::unordered_map> byFlow; + for (auto const& g : grants) + { + // Carry OUR (receiver) device id so the sender writes the remote desc to the right GPU + // even if the two agents don't share a device index. `regionHandle` (our arena offset) is + // echoed back in DATA so we can locate + free the region. + byFlow[g.flow].push_back(BounceCreditEntry{g.addr, g.len, static_cast(deviceId), g.offset}); + } + for (auto const& [flow, creds] : byFlow) + { + auto [peer, rid] = splitKey(flow); + channel->sendTo(peer, encodeGrant(rid, creds)); + } } - for (auto const& [flow, creds] : byFlow) + catch (std::exception const& error) { - auto [peer, rid] = splitKey(flow); - channel->sendTo(peer, encodeGrant(rid, creds)); + // Allocation ownership already changed. A lost control message can time out; + // replaying its completed allocation release could corrupt a new owner. + TLLM_LOG_WARNING("Bounce grant publication failed: %s", error.what()); } } @@ -282,7 +310,7 @@ void BounceReceiver::onWant(std::string const& peer, BounceMsgHeader const& h, s bool reversePathReady = false; try { - reversePathReady = mCtx.channel->addPeer(peer, endpoint); + reversePathReady = mCtx.addPeer(peer, endpoint); } catch (tensorrt_llm::common::TllmException const& e) { @@ -295,16 +323,15 @@ void BounceReceiver::onWant(std::string const& peer, BounceMsgHeader const& h, s auto const key = makeKey(peer, h.requestId); if (isCancelWant(chunkBytes)) { - // Explicit cancel/abort (the sender failed or retracted): precisely free this flow's - // granted-but-unwritten regions now — otherwise they stay held until peer loss and a - // long-running receiver leaks up to one request's in-flight allocation cap per failed rid. - // Any region whose scatter is still running is deferred (flagged orphaned in mScattering, - // freed on completion). - // Immediate free (no quarantine) is safe HERE because a cancel is sender-initiated: the - // sender defers it until its last in-flight RDMA write reached a terminal state - // (mPendingCancel / drainOrphanLocal), so nothing can still be writing these regions. + // Empty WANT is a legacy cancellation signal, not a proof that DMA stopped. + closeFlow(key); reclaimAndFlagDeferred( - [&](auto const& busy, auto& deferred) { return mCtx.scheduler.reclaimFlow(key, busy, deferred); }); + [&](auto const& busy, auto& deferred) { return mCtx.scheduler.quarantineFlow(key, busy, deferred); }); + return; + } + if (mClosedFlows.count(key) != 0 || mCtx.shutdownRequested.load(std::memory_order_acquire)) + { + mCtx.channel->sendTo(peer, encodeNack(h.requestId, 0, 0)); return; } if (!reversePathReady) @@ -363,7 +390,7 @@ void BounceReceiver::onData(std::string const& peer, BounceMsgHeader const& h, s // 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 (mClosedFlows.count(key) != 0 || !mCtx.scheduler.heldByFlow(key, h.regionHandle)) { return; } @@ -389,6 +416,7 @@ void BounceReceiver::onData(std::string const& peer, BounceMsgHeader const& h, s bounceRangeEnd(job.nvtxQueue); return; } + ++mOutstanding[key]; { std::lock_guard lk(mJobMu); mJobs.emplace_back(std::move(job)); @@ -396,6 +424,68 @@ void BounceReceiver::onData(std::string const& peer, BounceMsgHeader const& h, s mJobCv.notify_one(); } +void BounceReceiver::closeFlow(std::string const& key) +{ + constexpr std::size_t kMaxClosedFlows = 1048576; + if (mClosedFlows.count(key) == 0 && mClosedFlows.size() >= kMaxClosedFlows) + { + TLLM_LOG_ERROR("Bounce receiver exhausted its retained flow fences; terminating before unsafe reuse"); + std::abort(); + } + mClosedFlows.insert(key); +} + +void BounceReceiver::onDrain(std::string const& peer, BounceMsgHeader const& h, std::string const& blob) +{ + std::vector chunks; + std::string endpoint; + if (!decodeWant(blob, h, chunks, endpoint) || !chunks.empty()) + { + return; + } + if (!mCtx.addPeer(peer, endpoint)) + { + return; // a changed peer incarnation cannot attest an old flow's drain + } + auto const key = makeKey(peer, h.requestId); + closeFlow(key); + mDrainRequested.insert(key); + // DRAIN attests that no sender DMA can still touch these grants. Already queued + // scatters remain counted until their streams have physically finished. + reclaimAndFlagDeferred( + [&](auto const& busy, auto& deferred) { return mCtx.scheduler.reclaimFlow(key, busy, deferred); }); + mCtx.sendGrants(mCtx.scheduler.releaseQuarantinedFlow(key)); + maybeAcknowledgeDrain(key); +} + +void BounceReceiver::maybeAcknowledgeDrain(std::string const& key) +{ + if (mDrainRequested.count(key) == 0 || mOutstanding.count(key) != 0) + { + return; + } + auto const [peer, rid] = splitKey(key); + mCtx.channel->sendTo(peer, encodeDrainAck(rid)); +} + +void BounceReceiver::beginShutdown() +{ + mCtx.scheduler.closeGrantAdmission(); + for (auto const& key : mCtx.scheduler.flowKeys()) + { + closeFlow(key); + auto const [peer, rid] = splitKey(key); + reclaimAndFlagDeferred( + [&](auto const& busy, auto& deferred) { return mCtx.scheduler.quarantineFlow(key, busy, deferred); }); + mCtx.channel->sendTo(peer, encodeNack(rid, 0, 0)); + } +} + +bool BounceReceiver::hasUnprovenAccess() const +{ + return !mScattering.empty() || mCtx.scheduler.hasUnprovenRemoteWrites(); +} + std::unordered_set BounceReceiver::scatteringRegions() const { std::unordered_set busy; @@ -434,6 +524,11 @@ bool BounceReceiver::drainScatterDone() BounceNvtxScope drainScope(kNvtxDoneDrain, "doneDrain n=%zu", done.size()); for (auto& d : done) { + auto outstanding = mOutstanding.find(d.key); + if (outstanding != mOutstanding.end() && --outstanding->second == 0) + { + mOutstanding.erase(outstanding); + } // The ACK was already sent by the scatter worker itself (latency: it is on the sender's // ackWait critical path); only the region bookkeeping happens here, on the IO thread. // Worker finished reading this region. Was its flow reclaimed (peer gone / cancel) mid-scatter? @@ -453,63 +548,23 @@ bool BounceReceiver::drainScatterDone() { mCtx.sendGrants(mCtx.scheduler.onScatterDone(d.key, d.offset)); } + maybeAcknowledgeDrain(d.key); } return true; } void BounceReceiver::forget(std::string const& peer) { - // Reclaim every flow this peer was granted ("peer\x1f rid"). - // (1) Drop this peer's not-yet-started scatter jobs — no point scattering for a gone peer; their - // incoming regions are quarantined (or freed) by the reclaim below. A queued job whose - // flow was ALREADY reclaimed (cancel / lease expiry flagged it orphaned in mScattering) sits - // in the scheduler's orphan set; collect those offsets here and free them in (4), AFTER the - // peer's flows are dropped — freeOrphanRegion runs schedule(), which would otherwise grant the - // freed region straight back to a flow of the very peer being forgotten. - std::vector orphanedOffsets; + auto const prefix = peer + kSep; + for (auto const& key : mCtx.scheduler.flowKeys()) { - std::lock_guard lk(mJobMu); - std::deque keep; - for (auto& j : mJobs) + if (key.compare(0, prefix.size(), prefix) != 0) { - if (j.peer == peer) - { - bounceRangeEnd(j.nvtxQueue); // job dropped, close its queue-wait span - auto it = mScattering.find(j.offset); - if (it != mScattering.end()) - { - if (it->second) - { - orphanedOffsets.push_back(j.offset); - } - mScattering.erase(it); - } - } - else - { - keep.push_back(std::move(j)); - } + continue; } - mJobs.swap(keep); - } - // (2) Regions of this peer still scattering are reads already RUNNING in a worker — they must not - // be re-granted until the worker finishes. reclaimByPrefix defers those; we flag them orphaned - // in mScattering so drainScatterDone frees them via freeOrphanRegion on completion. - // (3) Every OTHER held region may STILL be RDMA-written by the peer: forget() is - // receiver-initiated (invalidateRemoteAgent), so — unlike an explicit cancel — no - // sender-side drain guarantees those writes ended, and a one-sided write cannot be aborted. - // quarantineFor > 0 keeps them out of the arena until checkTimeouts()'s reapQuarantine. - reclaimAndFlagDeferred( - [&](auto const& busy, auto& deferred) - { - return mCtx.scheduler.reclaimByPrefix( - peer + kSep, busy, deferred, std::chrono::milliseconds(std::max(0, mCtx.cfg.quarantineMs))); - }); - // (4) Now that none of this peer's flows is in the ring, free the orphaned regions collected in - // (1); the re-grants can only reach other peers. - for (auto const off : orphanedOffsets) - { - mCtx.sendGrants(mCtx.scheduler.freeOrphanRegion(off)); // mirrors drainScatterDone + closeFlow(key); + reclaimAndFlagDeferred( + [&](auto const& busy, auto& deferred) { return mCtx.scheduler.quarantineFlow(key, busy, deferred); }); } } @@ -521,7 +576,7 @@ void BounceReceiver::checkTimeouts() return; // throttled: no need to scan on every ~1ms tick } // Drain time-box is otherwise only re-evaluated on scheduler events; a quiet system needs this - // regardless of whether the lease/quarantine sweeps below are enabled. + // regardless of whether the lease sweep below is enabled. mCtx.sendGrants(mCtx.scheduler.pollDrain()); // Sweep granularity follows the smallest ENABLED timeout (a tenth of it, clamped to // [50ms, 1s]) instead of a fixed constant: the 60s/30s defaults yield a 1s sweep, while a @@ -534,16 +589,16 @@ void BounceReceiver::checkTimeouts() } if (smallestMs <= 0) { - // Lease disabled AND quarantine disabled: nothing is ever quarantined (reclaims free - // immediately) and no flow can go stale — nothing to sweep for. + // No time-based scheduler sweeps are configured. Explicit DRAIN still governs release + // of regions pinned by cancellation or peer loss. mNextSweep = now + std::chrono::seconds(1); return; } mNextSweep = now + std::clamp( std::chrono::milliseconds(smallestMs / 10), std::chrono::milliseconds(50), std::chrono::milliseconds(1000)); - // Quarantine over for any expired region: no write posted before its flow's lease expired can - // plausibly still be in flight, so it may re-enter circulation (and may grant a waiter). + // Service the scheduler's legacy timed quarantine API. This transport only populates + // permanent flow quarantine, which this call cannot release. mCtx.sendGrants(mCtx.scheduler.reapQuarantine()); if (mCtx.cfg.receiverFlowTimeoutMs <= 0) { @@ -556,20 +611,16 @@ void BounceReceiver::checkTimeouts() { // The flow's sender went silent after taking grants: it is dead or unreachable (a LIVE // sender either makes progress or abandons via requestTimeoutMs + cancel well before this - // lease expires). Reclaim the whole flow — regions a worker still reads defer as orphans, - // all others quarantine (the silent peer's NIC may still be writing them). + // lease expires). Fence the flow; active scatter regions remain owned until completion + // and all unknown-write regions stay quarantined until explicit DRAIN. auto const [peer, rid] = splitKey(flow); TLLM_LOG_WARNING( "BounceTransport(%s): flow lease expired (no progress within %d ms) peer=%s rid=%llu -> reclaiming " - "(regions quarantined for %d ms before reuse)", - mCtx.selfName.c_str(), mCtx.cfg.receiverFlowTimeoutMs, peer.c_str(), static_cast(rid), - mCtx.cfg.quarantineMs); + "(unknown-write regions retained until DRAIN)", + mCtx.selfName.c_str(), mCtx.cfg.receiverFlowTimeoutMs, peer.c_str(), static_cast(rid)); + closeFlow(flow); reclaimAndFlagDeferred( - [&](auto const& busy, auto& deferred) - { - return mCtx.scheduler.reclaimFlow( - flow, busy, deferred, std::chrono::milliseconds(std::max(0, mCtx.cfg.quarantineMs))); - }); + [&](auto const& busy, auto& deferred) { return mCtx.scheduler.quarantineFlow(flow, busy, deferred); }); } } @@ -583,10 +634,12 @@ void BounceReceiver::recoverScatterJob( static_cast(nackSent), static_cast(donePushed)); if (ctx != nullptr) { - // A scatter kernel may still be reading the region: drain the stream before the ScatterDone - // below lets the scheduler re-grant it. Best effort — the error (if any) is already logged. - (void) cudaStreamSynchronize(ctx->stream); - (void) cudaGetLastError(); + if (cudaStreamSynchronize(ctx->stream) != cudaSuccess) + { + // Retain the context, region and outstanding claim. Shutdown's independent + // guard and the owning Python deadline contain an unprovable GPU failure. + return; + } mCtx.exec->release(ctx); } // An ACK that was attempted (delivered or not) must never be followed by a NACK: a duplicated or @@ -663,7 +716,9 @@ void BounceReceiver::scatterWorkerLoop() } if (ctx == nullptr) { - break; // shutting down + std::lock_guard lk(mDoneMu); + mDone.push_back(ScatterDone{job.key, job.offset}); + continue; // never launched: cancellation is physical completion } auto const n = static_cast(job.entries.size()); // Validate every scatter SOURCE stays inside THIS flow's granted region before launching. @@ -761,8 +816,8 @@ void BounceReceiver::scatterWorkerLoop() } } cudaError_t syncErr = cudaSuccess; - if (launchErr == cudaSuccess) { + // Even a failed launch may follow queued asynchronous argument copies. // Sync leg: the actual GPU wait (kernel queueing + run time on the exec stream). BounceNvtxScope syncScope(kNvtxScatterSync, "scatterSync rid=%llu chunk=%u", static_cast(job.rid), job.chunkIdx); @@ -782,7 +837,15 @@ void BounceReceiver::scatterWorkerLoop() mCtx.selfName.c_str(), static_cast(srcInBounds), static_cast(launchErr), static_cast(syncErr), static_cast(job.rid), job.chunkIdx); } - mCtx.exec->release(ctx); // kernel done (or failed) -> return the context + if (syncErr != cudaSuccess) + { + // No physical proof: keep the exec context and the flow claim rooted. + nackSent = true; + mCtx.channel->sendTo(job.peer, encodeNack(job.rid, job.chunkIdx, job.offset)); + ctx = nullptr; + continue; + } + mCtx.exec->release(ctx); // stream is physically idle ctx = nullptr; // released: the exception handler must not release it again // ACK straight from the worker (ControlChannel::sendTo is thread-safe): the data IS at its // final dst here, and skipping the done-queue -> IO-thread hop shaves its drain latency off @@ -829,9 +892,14 @@ BounceSender::BounceSender(BounceContext& ctx) { } -std::shared_future BounceSender::submit( - TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer) +std::shared_future BounceSender::submit(TransferDescs const& srcDescs, TransferDescs const& dstDescs, + std::string const& peer, std::shared_ptr state) +try { + if (!state) + { + state = std::make_shared(); + } auto promise = std::make_shared>(); auto fut = promise->get_future().share(); BounceTransferPlan plan; @@ -851,12 +919,14 @@ std::shared_future BounceSender::submit( // silently absorbed as a slow success. Whether to retry is the caller's decision. TLLM_LOG_WARNING( "BounceTransport(%s): plan build for peer %s rejected: %s", mCtx.selfName.c_str(), peer.c_str(), e.what()); + state->quiesced.store(true, std::memory_order_release); promise->set_value({TransferState::kFAILURE, BounceFailReason::kPlanRejected}); return fut; } auto const numChunks = static_cast(plan.numChunks()); if (numChunks == 0) { + state->quiesced.store(true, std::memory_order_release); promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); return fut; } @@ -870,14 +940,26 @@ std::shared_future BounceSender::submit( } std::uint64_t const rid = mNextRid.fetch_add(1, std::memory_order_relaxed); + if (rid == 0) + { + TLLM_LOG_ERROR("Bounce request identifier exhausted; terminating before identity reuse"); + std::abort(); + } std::uint64_t const planBytes = plan.totalBytes(); { std::lock_guard lk(mReqMu); + if (mCtx.shutdownRequested.load(std::memory_order_acquire)) + { + state->quiesced.store(true, std::memory_order_release); + promise->set_value({TransferState::kFAILURE, BounceFailReason::kShutdown}); + return fut; + } Request req; req.peer = peer; req.numChunks = numChunks; req.plan = std::move(plan); req.promise = promise; + req.physical = std::move(state); req.lastProgress = std::chrono::steady_clock::now(); req.nvtxReq = bounceRangeStart(kNvtxRequest, "req rid=%llu chunks=%u bytes=%llu", static_cast(rid), numChunks, static_cast(planBytes)); @@ -885,10 +967,17 @@ std::shared_future BounceSender::submit( req.nvtxGrantWait = bounceRangeStart(kNvtxGrantWait, "grantWait rid=%llu", static_cast(rid)); mRequests.emplace(rid, std::move(req)); + // Publication is ordered with failure/cancellation under mReqMu. A failed + // request must never publish a new WANT after its DRAIN was acknowledged. + try + { + mCtx.channel->sendTo(peer, encodeWant(rid, chunkBytes, mCtx.channel->localEndpoint())); + } + catch (std::exception const&) + { + failRequest(rid, mRequests.at(rid), BounceFailReason::kProtocolError); + } } - // Ask the receiver to grant a region for each chunk of this request flow. The WANT carries our - // own control endpoint so the receiver can addPeer us and send GRANT/ACK back (self-bootstrap). - mCtx.channel->sendTo(peer, encodeWant(rid, chunkBytes, mCtx.channel->localEndpoint())); if (mCtx.cfg.enableEagerGather) { // Overlap the WANT->GRANT control round-trip with the gather: launch this request's first @@ -909,6 +998,12 @@ std::shared_future BounceSender::submit( return fut; } +catch (std::bad_alloc const&) +{ + TLLM_LOG_ERROR("Bounce sender allocation failure after admission; terminating before memory reuse"); + std::abort(); +} + void BounceSender::onGrant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob) { std::vector credits; @@ -1021,6 +1116,11 @@ void BounceSender::attachCredits(std::uint64_t rid, Request& req) void BounceSender::pumpRequest(std::uint64_t rid, Request& req) { + if (req.physical->cancelRequested.load(std::memory_order_acquire) + || mCtx.shutdownRequested.load(std::memory_order_acquire)) + { + return; + } if (req.abandonReason != BounceFailReason::kNone) { // Abandoned (GRANT mispair / plan overflow): stop attaching credits and launching gathers. @@ -1083,6 +1183,31 @@ void BounceSender::pumpRequest(std::uint64_t rid, Request& req) req.pendingCredits.pop_front(); req.nextCredit += 1; } + Posted posted; + posted.chunkIdx = chunkIdx; + posted.localOffset = *localOff; + posted.ctx = ctx; + posted.hasCredit = haveCredit; + if (haveCredit) + { + posted.remoteHandle = credit.regionHandle; + posted.remoteAddr = credit.addr; + posted.remoteDevId = credit.devId; + } + posted.writeBytes = static_cast(chunk.packedBytes); + posted.state = PostState::GatherFailed; + try + { + req.posted.push_back(std::move(posted)); + } + catch (std::exception const&) + { + mCtx.exec->release(ctx); + mCtx.sendGrants(mCtx.scheduler.releaseLocal(*localOff)); + throw; + } + auto& p = req.posted.back(); + ++req.nextPost; auto const nDesc = static_cast(chunk.srcPtrs.size()); // Covers plan-array prep + gather launch + event record (the synchronous launch cost; // the gather's GPU time is the async `gather` span ended in drainGatherReady). @@ -1109,8 +1234,6 @@ void BounceSender::pumpRequest(std::uint64_t rid, Request& req) TLLM_LOG_WARNING( "BounceTransport(%s): rid=%llu chunk=%u plan entries %zu > exec capacity %zu; abandoning flow", mCtx.selfName.c_str(), static_cast(rid), chunkIdx, nTotal, maxEntries); - mCtx.exec->release(ctx); - mCtx.sendGrants(mCtx.scheduler.releaseLocal(*localOff)); req.abandonReason = BounceFailReason::kProtocolError; // Abandon: pumpRequest ignores the request from now on; once its in-flight chunks // drain, lastProgress freezes and checkTimeouts fails it within one requestTimeoutMs. @@ -1148,29 +1271,12 @@ void BounceSender::pumpRequest(std::uint64_t rid, Request& req) mCtx.selfName.c_str(), static_cast(gatherErr), static_cast(recordErr), static_cast(rid), chunkIdx); } - Posted p; - p.chunkIdx = chunkIdx; - p.localOffset = *localOff; - p.ctx = ctx; - p.hasCredit = haveCredit; - if (haveCredit) - { - p.remoteHandle = credit.regionHandle; - p.remoteAddr = credit.addr; - p.remoteDevId = credit.devId; - } - p.writeBytes = static_cast(chunk.packedBytes); - // Gather in flight; the write is issued later by drainGatherReady once the event signals. If - // the gather launch/record failed, go straight to GatherFailed so drainGatherReady fails the - // request without ever trusting the (un)recorded event. p.state = gatherFailed ? PostState::GatherFailed : PostState::Gathering; if (!gatherFailed) { p.nvtxGather = bounceRangeStart(kNvtxGather, "gather rid=%llu chunk=%u bytes=%llu", static_cast(rid), chunkIdx, static_cast(chunk.packedBytes)); } - req.posted.push_back(std::move(p)); - req.nextPost += 1; req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk's gather launched } // Reconcile the pipeline-starvation NVTX spans after every pump pass (perf visibility only): @@ -1217,6 +1323,10 @@ bool BounceSender::drainGatherReady() std::vector toFail; for (auto& [rid, req] : mRequests) { + if (req.physical->cancelRequested.load(std::memory_order_acquire)) + { + continue; + } for (auto& p : req.posted) { if (p.state == PostState::Writing || p.state == PostState::Sent) @@ -1270,7 +1380,17 @@ bool BounceSender::drainGatherReady() {MemoryDesc{reinterpret_cast(mCtx.arena->at(p.localOffset)), p.writeBytes, static_cast(mCtx.deviceId)}}}; TransferDescs const dst{MemoryType::kVRAM, {MemoryDesc{p.remoteAddr, p.writeBytes, p.remoteDevId}}}; - p.xfer = mCtx.agent.postXferRequest(TransferOp::kWRITE, src, dst, req.peer, std::nullopt); + p.state = PostState::Writing; + try + { + p.xfer = mCtx.agent.postXferRequest(TransferOp::kWRITE, src, dst, req.peer, std::nullopt); + } + catch (std::exception const&) + { + req.abandonReason = BounceFailReason::kWriteFailed; + toFail.push_back(rid); + break; // an escaped submission with no handle cannot be proved stopped + } } p.state = PostState::Writing; didWork = true; @@ -1304,7 +1424,16 @@ bool BounceSender::pollSenderHandles() } // wait(0) is one non-blocking status query: IN_PROGRESS / SUCCESS / FAILURE. // A failed post left p.xfer null -> report FAILURE. - TransferState const st = p.xfer != nullptr ? p.xfer->wait(0) : TransferState::kFAILURE; + TransferState st = TransferState::kFAILURE; + try + { + st = p.xfer != nullptr ? p.xfer->wait(0) : TransferState::kFAILURE; + } + catch (std::exception const&) + { + toFail.push_back(rid); + break; + } if (st == TransferState::kSUCCESS) { // End the write span BEFORE building the DATA message so nixlWrite measures only the @@ -1320,12 +1449,30 @@ bool BounceSender::pollSenderHandles() BounceNvtxScope dataScope(kNvtxDataSend, "dataSend rid=%llu chunk=%u n=%u bytes=%zu", static_cast(rid), p.chunkIdx, nRuns, static_cast(nRuns) * sizeof(BounceScatterRun)); - mCtx.channel->sendTo( - req.peer, encodeData(rid, p.chunkIdx, req.numChunks, p.remoteHandle, chunk.scatterRuns)); + try + { + mCtx.channel->sendTo( + req.peer, encodeData(rid, p.chunkIdx, req.numChunks, p.remoteHandle, chunk.scatterRuns)); + } + catch (std::exception const&) + { + toFail.push_back(rid); + break; + } + } + bool released = false; + try + { + released = p.xfer->release(); } - if (!p.xfer->release()) + catch (std::exception const&) { - // The write is terminal, so progress is safe; the status object keeps the + toFail.push_back(rid); + break; + } + if (!released) + { + // The write succeeded, so progress is safe; the status object keeps the // handle and its destructor retries the backend release. TLLM_LOG_WARNING("BounceTransport(%s): terminal write handle release deferred rid=%llu chunk=%u", mCtx.selfName.c_str(), static_cast(rid), p.chunkIdx); @@ -1410,8 +1557,9 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) 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)); + auto const offset = pit->localOffset; req.posted.erase(pit); + mCtx.sendGrants(mCtx.scheduler.releaseLocal(offset)); found = true; break; } @@ -1435,6 +1583,7 @@ void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) bounceRangeEnd(req.nvtxReq); try { + req.physical->quiesced.store(true, std::memory_order_release); req.promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); } catch (...) @@ -1482,123 +1631,64 @@ void BounceSender::checkTimeouts() bool BounceSender::drainOrphanLocal() { - if (mOrphanLocal.empty()) + bool progress = false; + for (auto it = mOrphanLocal.begin(); it != mOrphanLocal.end();) { - return false; - } - bool didWork = false; - std::vector keep; - keep.reserve(mOrphanLocal.size()); - for (auto& o : mOrphanLocal) - { - if (o.xfer != nullptr && o.xfer->wait(0) == TransferState::kIN_PROGRESS) - { - // Warn once when an orphaned write outlives requestTimeoutMs: nothing gives up on it - // (the NIC may still read the region), but an operator should know the region is held. - auto const ageMs - = std::chrono::duration_cast(std::chrono::steady_clock::now() - o.since) - .count(); - if (!o.warned && ageMs >= mCtx.cfg.requestTimeoutMs) - { - o.warned = true; - TLLM_LOG_WARNING( - "BounceTransport(%s): orphaned RDMA write rid=%llu region offset=%llu has not reached " - "a terminal state after %lld ms; staging region held until it does", - mCtx.selfName.c_str(), static_cast(o.rid), - static_cast(o.offset), static_cast(ageMs)); - } - keep.push_back(std::move(o)); // write still in flight -> the NIC may still read the region; wait + if (it->xfer == nullptr || !it->xfer->isQuiesced()) + { + ++it; continue; } - // Terminal (Done or Failed): the NIC is finished with the region (source AND the receiver's - // destination) -> recycle the local source now. - if (o.xfer != nullptr && !o.xfer->release()) + auto orphan = std::move(*it); + it = mOrphanLocal.erase(it); + // Remove the owner before any release/publication can throw. Physical proof + // permits allocation reuse even when releasing a backend handle fails. + try { - // The write is terminal, so recycling is safe; only the backend handle remains retained - // (the status object's destructor retries the release). - TLLM_LOG_WARNING("BounceTransport(%s): terminal orphan handle release deferred rid=%llu", - mCtx.selfName.c_str(), static_cast(o.rid)); + if (!orphan.xfer->release()) + { + TLLM_LOG_WARNING("Bounce terminal RDMA handle release deferred rid=%llu", + static_cast(orphan.rid)); + } } - mCtx.sendGrants(mCtx.scheduler.releaseLocal(o.offset)); - didWork = true; - } - mOrphanLocal.swap(keep); - // Send any deferred cancel whose flow now has NO in-flight write left: the receiver may safely - // reclaim its regions (the writes have landed/failed, no more DMA targets them). - for (auto it = mPendingCancel.begin(); it != mPendingCancel.end();) - { - std::uint64_t const rid = it->first; - bool const stillInFlight = std::any_of( - mOrphanLocal.begin(), mOrphanLocal.end(), [rid](OrphanLocal const& o) { return o.rid == rid; }); - if (stillInFlight) + catch (std::exception const& error) { - ++it; - continue; + TLLM_LOG_WARNING("Bounce terminal RDMA handle release failed: %s", error.what()); } - mCtx.channel->sendTo(it->second, encodeCancel(rid, mCtx.channel->localEndpoint())); - it = mPendingCancel.erase(it); - didWork = true; + mCtx.sendGrants(mCtx.scheduler.releaseLocal(orphan.offset)); + progress = true; } - return didWork; + return progress; } bool BounceSender::drainOrphanGather() { - if (mOrphanGather.empty()) + bool progress = false; + for (auto it = mOrphanGather.begin(); it != mOrphanGather.end();) { - return false; - } - bool didWork = false; - std::vector keep; - keep.reserve(mOrphanGather.size()); - for (auto& o : mOrphanGather) - { - // Non-blocking poll of the exec STREAM, not the event: a GatherFailed chunk may have a - // successfully launched kernel whose cudaEventRecord failed, so only the stream tells us when - // the region is no longer being written. A held ExecCtx stream carries only this chunk's work. - cudaError_t const st = cudaStreamQuery(o.ctx->stream); - if (st == cudaErrorNotReady) - { - // Warn once when an orphaned gather outlives requestTimeoutMs: nothing gives up on it - // (the kernel may still write the region), but an operator should know the region is held. - // A never-completing gather keeps busy() true (0 ms poll with the 50 us backoff) until - // process exit. - auto const ageMs - = std::chrono::duration_cast(std::chrono::steady_clock::now() - o.since) - .count(); - if (!o.warned && ageMs >= mCtx.cfg.requestTimeoutMs) - { - o.warned = true; - TLLM_LOG_WARNING( - "BounceTransport(%s): orphaned gather rid=%llu region offset=%llu has not completed after " - "%lld ms; staging region held until it does", - 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 - continue; - } - if (st != cudaSuccess) + // An event may never have been recorded. Only an idle stream proves that + // both source reads and staging/plan-buffer writes have stopped. + if (cudaStreamQuery(it->ctx->stream) != cudaSuccess) { - // A non-NotReady error here is most likely a sticky async kernel fault, which - // cudaGetLastError cannot clear (the clear only helps a non-sticky one): the WARNING is the - // signal. The region is recycled only because the request already failed — nothing - // consumes its content. For a sticky launch failure this WARNING duplicates the one - // pumpRequest already emitted (acceptable: this second one names the region). - (void) cudaGetLastError(); - TLLM_LOG_WARNING("BounceTransport(%s): orphaned gather rid=%llu ended with CUDA error %s", - mCtx.selfName.c_str(), static_cast(o.rid), cudaGetErrorString(st)); + ++it; + continue; } - mCtx.exec->release(o.ctx); - mCtx.sendGrants(mCtx.scheduler.releaseLocal(o.localOffset)); - didWork = true; + auto const orphan = *it; + it = mOrphanGather.erase(it); + mCtx.exec->release(orphan.ctx); + mCtx.sendGrants(mCtx.scheduler.releaseLocal(orphan.localOffset)); + progress = true; } - mOrphanGather.swap(keep); - return didWork; + return progress; } void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason reason) { + req.physical->cancelRequested.store(true, std::memory_order_release); + if (req.abandonReason == BounceFailReason::kNone) + { + req.abandonReason = reason; + } // An abandoned flow (GRANT mispair / plan overflow) reaches here through the timeout path; // report the specific abandon cause, not the generic timeout. if (req.abandonReason != BounceFailReason::kNone) @@ -1614,7 +1704,7 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason // 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): // - Writing: the RDMA write may still be reading the region as its source. Defer to mOrphanLocal; - // drainOrphanLocal() releases the xfer + region once poll() is terminal. + // drainOrphanLocal() releases the xfer + region only after isQuiesced() proves it stopped. // - Gathering / GatherFailed: our gather kernel may still be WRITING the region (an abandoned // 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 @@ -1623,9 +1713,14 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason // - 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; - for (auto& p : req.posted) - { + // Allocate every owner record before changing resource ownership. + mOrphanLocal.reserve(mOrphanLocal.size() + req.posted.size()); + mOrphanGather.reserve(mOrphanGather.size() + req.posted.size()); + mDrains.try_emplace(rid, Drain{req.peer, req.physical}); + while (!req.posted.empty()) + { + auto p = std::move(req.posted.back()); + req.posted.pop_back(); // Close this chunk's NVTX spans (whichever leg it died in); 0 handles are no-ops. bounceRangeEnd(p.nvtxGather); bounceRangeEnd(p.nvtxWrite); @@ -1634,7 +1729,6 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason { mOrphanLocal.push_back( OrphanLocal{std::move(p.xfer), p.localOffset, rid, std::chrono::steady_clock::now(), false}); - deferredWrite = 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 @@ -1646,18 +1740,6 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason } mCtx.sendGrants(mCtx.scheduler.releaseLocal(p.localOffset)); } - // 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) - { - mPendingCancel[rid] = req.peer; - } - else - { - mCtx.channel->sendTo(req.peer, encodeCancel(rid, mCtx.channel->localEndpoint())); - } bounceRangeEnd(req.nvtxGrantWait); bounceRangeEnd(req.nvtxCreditStarved); bounceRangeEnd(req.nvtxArenaStarved); @@ -1674,156 +1756,118 @@ void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason mRequests.erase(rid); } -void BounceSender::forget(std::string const& peer) +bool BounceSender::hasUnprovenAccess() +{ + std::lock_guard lock(mReqMu); + return !mRequests.empty() || !mDrains.empty() || !mOrphanLocal.empty() || !mOrphanGather.empty(); +} + +void BounceSender::cancelAll() { - // Fail any in-flight request targeting the gone peer so its wait() returns. std::lock_guard lk(mReqMu); - std::vector toFail; - for (auto const& [rid, req] : mRequests) + for (auto& [rid, request] : mRequests) { - if (req.peer == peer) - { - toFail.push_back(rid); - } + request.physical->cancelRequested.store(true, std::memory_order_release); } - for (auto rid : toFail) +} + +void BounceSender::onDrainAck(std::string const& peer, BounceMsgHeader const& h) +{ + std::lock_guard lk(mReqMu); + auto it = mDrains.find(h.requestId); + if (it != mDrains.end() && it->second.peer == peer) { - auto it = mRequests.find(rid); - if (it != mRequests.end()) - { - failRequest(rid, it->second, BounceFailReason::kPeerDropped); - } + it->second.acknowledged = true; } } -void BounceSender::failAll() +bool BounceSender::progressDrains() { - // 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. - // 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 - // outbound write may still be READING the arena when ~NixlBounceState frees it. The only - // protection is to wait for those writes to reach a terminal state — bounded by a fixed - // kShutdownWriteDrainMs so a dead peer cannot hang teardown — and warn once if some are - // still in flight past the deadline. The once-retry of failed releases below is kept only as a - // cheap guard for other backends. This bounded wait runs while NixlTransferAgent::shutdown holds - // mLock, which is acceptable at teardown. It covers OUTBOUND writes only: inbound writes into - // regions we granted can neither be waited on nor revoked — deregistration invalidates the rkey and - // the process is exiting (the same exposure as KV-pool deregistration). std::lock_guard lk(mReqMu); + std::vector cancelled; + for (auto const& [rid, request] : mRequests) { - std::vector inFlight; - for (auto& [rid, req] : mRequests) + if (request.physical->cancelRequested.load(std::memory_order_acquire)) { - for (auto& p : req.posted) - { - if (p.state == PostState::Writing && p.xfer != nullptr) - { - inFlight.push_back(p.xfer.get()); - } - } + cancelled.push_back(rid); } - for (auto& o : mOrphanLocal) + } + for (auto const rid : cancelled) + { + failRequest(rid, mRequests.at(rid), BounceFailReason::kCancelled); + } + bool progress = !cancelled.empty(); + auto const now = std::chrono::steady_clock::now(); + constexpr auto kDrainRetry = std::chrono::milliseconds(100); + for (auto it = mDrains.begin(); it != mDrains.end();) + { + auto const rid = it->first; + auto& drain = it->second; + if (mRequests.count(rid) != 0) { - if (o.xfer != nullptr) - { - inFlight.push_back(o.xfer.get()); - } + ++it; + continue; // failure cleanup has not yet transferred every physical owner } - // Fixed, independent of requestTimeoutMs: a dead peer's write never completes, so waiting - // longer buys nothing; on a live link a default 32 MiB chunk lands well within 5 s; and a tiny - // request_timeout_ms must not shorten a legitimate in-flight write. - constexpr int kShutdownWriteDrainMs = 5000; - auto const deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(kShutdownWriteDrainMs); - while (!inFlight.empty()) - { - inFlight.erase(std::remove_if(inFlight.begin(), inFlight.end(), - [](TransferStatus* x) { return x->wait(0) != TransferState::kIN_PROGRESS; }), - inFlight.end()); - if (inFlight.empty() || std::chrono::steady_clock::now() >= deadline) - { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(1)); + bool const writing = std::any_of( + mOrphanLocal.begin(), mOrphanLocal.end(), [rid](auto const& orphan) { return orphan.rid == rid; }); + bool const gathering = std::any_of( + mOrphanGather.begin(), mOrphanGather.end(), [rid](auto const& orphan) { return orphan.rid == rid; }); + if (!writing && !gathering && drain.acknowledged) + { + drain.physical->quiesced.store(true, std::memory_order_release); + it = mDrains.erase(it); + progress = true; + continue; } - if (!inFlight.empty()) + if (!writing && !drain.acknowledged && now >= drain.retryAt) { - TLLM_LOG_WARNING( - "BounceTransport(%s): %zu outbound writes still in flight at shutdown; arena teardown may race the NIC", - mCtx.selfName.c_str(), inFlight.size()); + drain.retryAt = now + kDrainRetry; + mCtx.channel->sendTo(drain.peer, encodeDrain(rid, mCtx.channel->localEndpoint())); + progress = true; } + ++it; } - std::vector> releaseRetry; - for (auto& [rid, req] : mRequests) + return progress; +} + +void BounceSender::forget(std::string const& peer) +{ + // Fail any in-flight request targeting the gone peer so its wait() returns. + std::lock_guard lk(mReqMu); + std::vector toFail; + for (auto const& [rid, req] : mRequests) { - for (auto& p : req.posted) - { - bounceRangeEnd(p.nvtxGather); - bounceRangeEnd(p.nvtxWrite); - bounceRangeEnd(p.nvtxAckWait); - if (p.state == PostState::Writing && p.xfer != nullptr) - { - if (!p.xfer->release()) - { - releaseRetry.push_back(std::move(p.xfer)); - } - } - if (p.ctx != nullptr) - { - mCtx.exec->release(p.ctx); // still-gathering chunk: return its borrowed exec context - 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 (...) + if (req.peer == peer) { - // 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. + toFail.push_back(rid); } } - mRequests.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. - for (auto& o : mOrphanLocal) + for (auto rid : toFail) { - if (o.xfer != nullptr && !o.xfer->release()) + auto it = mRequests.find(rid); + if (it != mRequests.end()) { - releaseRetry.push_back(std::move(o.xfer)); + failRequest(rid, it->second, BounceFailReason::kPeerDropped); } } - mOrphanLocal.clear(); - // Parked gathers: the caller's cudaDeviceSynchronize already drained every stream, so return the - // exec contexts unconditionally. Their regions are left to arena destruction like every other - // chunk's: the ring is torn down and grants are discarded, so recycling would be pointless. - for (auto const& o : mOrphanGather) +} + +void BounceSender::failAll() +{ + std::lock_guard lk(mReqMu); + if (!mOrphanLocal.empty() || !mOrphanGather.empty() || mCtx.scheduler.localHeldCount() != 0) { - mCtx.exec->release(o.ctx); + TLLM_LOG_ERROR("Bounce sender teardown reached unproved source access; terminating before memory reuse"); + std::abort(); } - mOrphanGather.clear(); - // Retry failed releases once after all producers/futures have been stopped. A persistent - // backend failure gets a final bounded attempt from the status object's destructor. - for (auto& status : releaseRetry) + for (auto& [rid, request] : mRequests) { - if (!status->release()) - { - TLLM_LOG_WARNING( - "BounceTransport(%s): NIXL handle still retained after shutdown retry", mCtx.selfName.c_str()); - } + request.promise->set_value({TransferState::kFAILURE, BounceFailReason::kShutdown}); } - // Deferred cancels aren't sent at shutdown: in-flight RDMA writes aren't drained by the device - // sync (they're NIXL, not CUDA-stream), so a cancel could still race a write. The receiver - // reclaims those regions via its own teardown / peer-loss path. - mPendingCancel.clear(); + mRequests.clear(); + // Remote drain may remain unproved after local memory is safe. Retained statuses + // keep quiesced=false; transport destruction does not manufacture remote proof. + mDrains.clear(); } // ============================================================================ @@ -1883,31 +1927,61 @@ BounceTransport::~BounceTransport() void BounceTransport::shutdown() { + std::lock_guard shutdownLock(mShutdownMu); bool expected = false; - if (!mCtx.stop.compare_exchange_strong(expected, true)) + if (!mCtx.shutdownRequested.compare_exchange_strong(expected, true)) { - return; // already shut down + return; } - mReceiver.wake(); // wake scatter workers so they observe stop + + struct GuardState + { + std::mutex mutex; + std::condition_variable cv; + bool done{false}; + }; + + auto guardState = std::make_shared(); + // This guard precedes every join and CUDA wait. A stuck kernel/NIC must terminate + // the process before any arena or exec allocation can be destroyed underneath it. + std::thread guard( + [guardState] + { + std::unique_lock lock(guardState->mutex); + if (!guardState->cv.wait_for(lock, std::chrono::seconds(5), [&] { return guardState->done; })) + { + TLLM_LOG_ERROR("Bounce shutdown could not prove physical quiescence; terminating before memory reuse"); + std::abort(); + } + }); + while (!mCtx.shutdownQuiesced.load(std::memory_order_acquire)) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + mCtx.stop.store(true, std::memory_order_release); + mReceiver.wake(); if (mIoThread.joinable()) { mIoThread.join(); } mReceiver.joinWorkers(); - // Threads are joined, but in-flight gather/scatter kernels may still be queued on ExecPool - // streams referencing the arena. Drain the device BEFORE we release contexts and let the caller - // tear down ExecPool/BounceArena — otherwise ~ExecPool's cudaStreamDestroy + ~BounceArena's - // cudaFree could race a kernel still reading the arena. (Teardown-only; a full-device sync is - // acceptable here.) Warn-only (we're in teardown / a dtor path -> must not throw); the user - // should still see a GPU fault. shutdown() may run on a thread whose current device isn't ours. - TLLM_CUDA_CHECK_WARN(cudaSetDevice(mCtx.deviceId)); - TLLM_CUDA_CHECK_WARN(cudaDeviceSynchronize()); + if (cudaSetDevice(mCtx.deviceId) != cudaSuccess || cudaDeviceSynchronize() != cudaSuccess) + { + TLLM_LOG_ERROR("Bounce shutdown could not synchronize CUDA; terminating before memory reuse"); + std::abort(); + } mSender.failAll(); + { + std::lock_guard lock(guardState->mutex); + guardState->done = true; + } + guardState->cv.notify_one(); + guard.join(); } bool BounceTransport::addPeer(std::string const& peer, std::string const& endpoint) { - return mCtx.channel->addPeer(peer, endpoint); + return mCtx.addPeer(peer, endpoint); } std::string BounceTransport::localHandshakeBlob() const @@ -1929,10 +2003,9 @@ std::string BounceTransport::localHandshakeBlob() const bool BounceTransport::registerPeerHandshake(std::string const& peer, std::string const& blob) { - // Treat registration as replacement, not an additive update. A peer may be reloaded under the - // same name with bounce disabled or with changed settings; clear the previously validated route - // first so any missing/malformed/incompatible replacement immediately falls back to NIXL. - mCtx.channel->removePeer(peer); + // Revalidate admission on every metadata load. Missing, malformed or incompatible metadata + // clears the bounce capability, while the immutable old control route remains available for + // outstanding drains. A new transport incarnation must have a fresh peer name. { std::lock_guard lk(mPeerMu); mHandshakedPeers.erase(peer); @@ -1960,11 +2033,9 @@ bool BounceTransport::registerPeerHandshake(std::string const& peer, std::string // STRICT equality. maxChunkSizeBytes is compared on the effective (post-clamp) values both // sides advertise; each side already clamped its own value to its usable arena capacity, so // equality also guarantees our chunks always fit the peer's arena and its scatter scratch - // (sized for its own maxChunkSizeBytes). requestTimeoutMs must match because the receiver's - // lease (2x ITS value) must exceed the SENDER's timeout; the quarantine (1x) starts at lease expiry - // or at a receiver-side peer teardown (forgetPeer via reclaimByPrefix) and on the lease path needs - // no extra margin (see BounceConfig::deriveDependentTimeouts). Local-only knobs - // (worker/stream counts, the + // (sized for its own maxChunkSizeBytes). Matching requestTimeoutMs gives live senders time to + // initiate drain before the receiver lease (2x this value) fences the flow. Lease expiry never + // supplies physical completion evidence. Local-only knobs (worker/stream counts, the // zero-copy-argument path, granularity, arena size) intentionally do NOT have to match. if (handshake.wireVersion != kBounceVersion || handshake.controlKind != localControlKind || handshake.maxChunkSizeBytes != mCtx.cfg.maxChunkSizeBytes @@ -1996,7 +2067,7 @@ bool BounceTransport::registerPeerHandshake(std::string const& peer, std::string bool peerAdded = false; try { - peerAdded = mCtx.channel->addPeer(peer, handshake.endpoint); + peerAdded = mCtx.addPeer(peer, handshake.endpoint); } catch (std::exception const& e) { @@ -2086,6 +2157,11 @@ void BounceTransport::ioLoop() { tick(peer, blob); } + catch (std::bad_alloc const&) + { + TLLM_LOG_ERROR("Bounce reactor allocation failure; terminating before memory reuse"); + std::abort(); + } catch (std::exception const& e) { TLLM_LOG_WARNING("BounceTransport(%s): IO tick failed: %s", mCtx.selfName.c_str(), e.what()); @@ -2103,6 +2179,11 @@ void BounceTransport::tick(std::string& peer, std::string& blob) // waiting on a control message), sleep up to 1ms so the IO thread doesn't busy-spin a core. // A request merely waiting for a GRANT (nothing posted yet) is NOT "busy", so a stalled / // unreachable peer won't spin a core until requestTimeoutMs. Both checks are IO-thread-only. + if (mCtx.shutdownRequested.load(std::memory_order_acquire)) + { + mReceiver.beginShutdown(); + mSender.cancelAll(); + } bool const busy = mSender.busy() || mReceiver.busy(); int const timeoutMs = busy ? 0 : 1; bool work = false; @@ -2111,6 +2192,7 @@ void BounceTransport::tick(std::string& peer, std::string& blob) dispatch(peer, blob); work = true; } + work |= mSender.progressDrains(); work |= mSender.drainGatherReady(); // post writes for chunks whose gather kernel just finished work |= mSender.pollSenderHandles(); work |= mReceiver.drainScatterDone(); @@ -2119,7 +2201,12 @@ void BounceTransport::tick(std::string& peer, std::string& blob) drainForgets(); mSender.drainPendingPosts(); // retry credits parked when the arena was full (onAck freed regions) mSender.checkTimeouts(); - mReceiver.checkTimeouts(); // expire grant leases of silent senders + free post-quarantine regions + mReceiver.checkTimeouts(); // fence expired flows; unknown writes remain quarantined + if (mCtx.shutdownRequested.load(std::memory_order_acquire) && !mSender.hasUnprovenAccess() + && !mReceiver.hasUnprovenAccess()) + { + mCtx.shutdownQuiesced.store(true, std::memory_order_release); + } // Idle backoff: when there IS in-flight work (busy → 0ms poll) but nothing actually advanced // this pass — the classic case being a gather stalled behind unrelated model kernels (gather // event stays NotReady) — keep latency low for the first few spins, then sleep briefly so we @@ -2151,6 +2238,8 @@ void BounceTransport::dispatch(std::string const& peer, std::string const& blob) } switch (static_cast(h.msgType)) { + case BounceMsgType::kDRAIN: mReceiver.onDrain(peer, h, blob); break; + case BounceMsgType::kDRAIN_ACK: mSender.onDrainAck(peer, h); break; case BounceMsgType::kWANT: mReceiver.onWant(peer, h, blob); break; // [R] case BounceMsgType::kGRANT: mSender.onGrant(peer, h, blob); break; // [S] case BounceMsgType::kDATA: mReceiver.onData(peer, h, blob); break; // [R] 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..53922764f929 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 @@ -54,7 +54,7 @@ namespace tensorrt_llm::executor::kv_cache::bounce { // ============================================================================ -// Bounce v2 reactor +// Bounce v3 reactor // ---------------------------------------------------------------------------- // One BounceTransport per agent. Protocol progress is driven by ONE IO thread plus M scatter // workers. submit() may also prepare eager gathers on application threads, so the sender request @@ -91,7 +91,9 @@ 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 +// FAILURE on transfer error, peer invalidation, shutdown, or request timeout. FAILURE does not +// release caller memory: BounceRequestState independently records physical quiescence after local +// gathers/writes and remote scatter drain. 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. // @@ -113,19 +115,27 @@ enum class BounceFailReason : std::uint8_t 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 + kCancelled, // caller requested cancellation; physical drain remains independent 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. +/// A bounce request's logical result: terminal state plus, on kFAILURE, its cause. The future +/// never changes a failure into success when the separate physical state later becomes quiesced. struct BounceResult { TransferState state{TransferState::kFAILURE}; BounceFailReason reason{BounceFailReason::kNone}; }; +/// Independent, monotonic physical evidence retained by statuses after logical failure. +struct BounceRequestState +{ + std::atomic quiesced{false}; + std::atomic cancelRequested{false}; +}; + /// Shared dependencies used by both roles. Holds the injected channel/agent/arena/exec (borrowed, /// not owned), the one CreditScheduler that carves the shared arena for both roles, and /// sendGrants(). Owned by value by BounceTransport and referenced by the sender and receiver. @@ -147,7 +157,7 @@ class BounceContext cfg.maxInflightChunksPerRequest) { // Receiver drain mode must give up well before the request timeout so a silent peer cannot - // stall every other flow for the full lease + quarantine (see CreditScheduler::mDrainTimeout). + // stall every other flow while its grants remain pinned (see CreditScheduler::mDrainTimeout). scheduler.setDrainTimeout(std::chrono::milliseconds(std::max(1, cfg.requestTimeoutMs / 2))); } @@ -158,6 +168,9 @@ class BounceContext /// GRANT control message to each peer. Used by BOTH roles (the receiver grants incoming regions; /// the sender's releaseLocal/reclaim can free arena bytes that re-grant a waiting remote flow). void sendGrants(std::vector const& grants); + bool addPeer(std::string const& peer, std::string const& endpoint); + std::mutex peerMutex; + std::unordered_map peerEndpoints; std::string selfName; BounceConfig cfg; @@ -168,6 +181,8 @@ class BounceContext ExecPool* exec{}; // gather/scatter exec contexts (streams/scratch), borrowed per kernel CreditScheduler scheduler; // shared region allocator; internally synchronized std::atomic stop{false}; + std::atomic shutdownRequested{false}; + std::atomic shutdownQuiesced{false}; }; /// Receiver role ([R]): WANT -> grant regions, DATA -> scatter into the caller's KV, then ACK. Owns @@ -187,18 +202,21 @@ class BounceReceiver void onWant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); void onData(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); + void onDrain(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); + [[nodiscard]] bool hasUnprovenAccess() const; + void beginShutdown(); + /// Drain scatter bookkeeping completions posted by workers, release regions, and re-grant newly /// available space. Successful workers have already sent ACK. Returns true if any were drained. bool drainScatterDone(); /// A peer is gone: reclaim every receiver-side flow of that peer (deferring regions a worker is - /// still reading, quarantining regions the peer may still be RDMA-writing) and drop its - /// not-yet-started scatter jobs. + /// still reading, quarantining regions the peer may still be RDMA-writing). Queued and dequeued + /// scatter jobs retain ownership until physically complete. void forget(std::string const& peer); - /// Time-driven reclamation (IO thread, called from tick()): reclaim flows the scheduler reports - /// idle beyond cfg.receiverFlowTimeoutMs — a dead/unreachable sender emits neither DATA nor a - /// cancel, so without this its granted regions leak forever — and return quarantined regions to - /// the arena once their cfg.quarantineMs deadline passes (scheduler.reapQuarantine). + /// Time-driven fencing (IO thread, called from tick()): close flows idle beyond + /// cfg.receiverFlowTimeoutMs and retain their unknown-write regions until explicit DRAIN. + /// Elapsed time alone never releases regions admitted by this transport. void checkTimeouts(); /// True while any scatter is enqueued-or-running (drives the IO loop's 0ms busy-poll). @@ -236,8 +254,12 @@ class BounceReceiver std::function(std::unordered_set const&, std::vector&)> const& reclaim); void scatterWorkerLoop(); - /// Exception recovery for one scatter job (worker thread): release a still-held exec context, - /// NACK unless an ACK or NACK was already attempted, and push the ScatterDone unless already pushed. + void closeFlow(std::string const& key); + void maybeAcknowledgeDrain(std::string const& key); + + /// Exception recovery for one scatter job (worker thread): synchronize any held exec context + /// before releasing it or posting ScatterDone. CUDA errors retain ownership. Send NACK unless + /// an ACK or NACK was already attempted. void recoverScatterJob( ScatterJob const& job, ExecCtx* ctx, bool acked, bool nackSent, bool donePushed, char const* what); @@ -260,6 +282,10 @@ class BounceReceiver // drainScatterDone frees the region via freeOrphanRegion (flow already gone) or onScatterDone // (normal completion). std::unordered_map mScattering; + // Counts include jobs waiting in the queue and jobs dequeued before acquiring an ExecCtx. + std::unordered_map mOutstanding; + std::unordered_set mClosedFlows; + std::unordered_set mDrainRequested; std::chrono::steady_clock::time_point mNextSweep{}; // checkTimeouts() throttle (IO-thread-only) }; @@ -278,11 +304,16 @@ class BounceSender /// Chunks are cut to cfg.maxChunkSizeBytes. NixlTransferAgent validates peer compatibility /// before calling this; direct users of BounceTransport must establish the same invariant /// themselves. - [[nodiscard]] std::shared_future submit( - TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer); + [[nodiscard]] std::shared_future submit(TransferDescs const& srcDescs, TransferDescs const& dstDescs, + std::string const& peer, std::shared_ptr state = {}); void onGrant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); void onAck(std::string const& peer, BounceMsgHeader const& h); + void onDrainAck(std::string const& peer, BounceMsgHeader const& h); + bool progressDrains(); + void cancelAll(); + [[nodiscard]] bool hasUnprovenAccess(); + /// 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. void onNack(std::string const& peer, BounceMsgHeader const& h); @@ -292,7 +323,8 @@ class BounceSender bool drainGatherReady(); bool pollSenderHandles(); /// Free local regions whose failed request still had an RDMA write in flight, once that write - /// reaches a terminal state (so the NIC is done reading the source). Returns true if any freed. + /// supplies explicit physical completion evidence. FAILURE alone is insufficient. Returns true + /// if any regions were freed. bool drainOrphanLocal(); /// Free the staging regions of failed requests whose gather may still be running, once their exec /// stream is idle (non-blocking cudaStreamQuery), and return the exec context. Returns true if @@ -304,11 +336,9 @@ class BounceSender void checkTimeouts(); /// A peer is gone: fail any in-flight request targeting it so its wait() returns. 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. + /// Final shutdown assertion: every request, orphan and remote drain must already have settled. + /// Called after the guarded reactor drain, worker joins and successful CUDA synchronization; + /// any remaining owner requires fail-stop before arena destruction. void failAll(); /// True while a local gather region is held or an orphaned write/gather is still in flight @@ -380,6 +410,8 @@ class BounceSender // here and retries each loop iteration as ACKs free regions. FIFO: paired with chunks in order. std::deque pendingCredits; std::shared_ptr> promise; + std::shared_ptr physical; + // Set by the abandon sites (GRANT mispair / plan overflow): the request then dies through // the timeout path, but failRequest reports THIS more specific cause instead. BounceFailReason abandonReason{BounceFailReason::kNone}; @@ -397,7 +429,7 @@ class BounceSender // Regions of a FAILED request whose RDMA write was still in flight (state == Writing). // The NIC may still be reading the source region, so recycling is deferred until the write reaches - // a terminal state — drainOrphanLocal() polls and only then releases the xfer handle + returns the + // a physically quiesced state — drainOrphanLocal() checks evidence before releasing the handle + // region. IO-thread-only (no lock). (xfer handle, arena offset). struct OrphanLocal { @@ -411,7 +443,7 @@ class BounceSender // Staging regions of a FAILED request whose gather may still be running (state == Gathering, or // GatherFailed — a failed cudaEventRecord after a successful launch still leaves a kernel on the // stream). A held ExecCtx stream carries only this chunk's work, so recycling is deferred until - // cudaStreamQuery(ctx->stream) is no longer NotReady — drainOrphanGather() polls it (non-blocking) + // cudaStreamQuery(ctx->stream) returns cudaSuccess — drainOrphanGather() polls it (non-blocking) // and only then returns the exec context + region. Parking (instead of a cudaStreamSynchronize in // failRequest) keeps the IO thread off the GPU on the NACK path. A never-completing gather keeps // busy() true (0 ms poll with the 50 us backoff) until process exit. It also pins one of the @@ -459,12 +491,16 @@ class BounceSender std::vector mOrphanLocal; std::vector mOrphanGather; - // A failed request's cancel (empty WANT) must be DEFERRED while any of its RDMA writes are still - // in flight: those writes are landing on the RECEIVER's regions, and an early cancel would let the - // receiver reclaim + re-grant those regions under the in-flight write -> cross-node corruption. - // drainOrphanLocal() sends the cancel once the flow's last in-flight write reaches terminal. - // rid -> peer. IO-thread-only. - std::unordered_map mPendingCancel; + + struct Drain + { + std::string peer; + std::shared_ptr physical; + std::chrono::steady_clock::time_point retryAt{}; + bool acknowledged{false}; + }; + + std::unordered_map mDrains; }; /// The thin reactor: owns the shared context, the sender + receiver, and the IO thread. Routes each @@ -472,6 +508,9 @@ class BounceSender class BounceTransport { public: + /// @param selfName Unique for this transport lifetime, including process restarts. Production + /// uses an instance UUID plus rank. Custom callers must never reuse a name for a new incarnation; + /// the transport rejects replacing an existing name with a different endpoint. /// @param arena ONE shared data buffer serving BOTH roles: receiver (remote senders' RDMA-write /// targets, granted as variable regions by the scheduler) and sender (local gather staging, via /// acquireLocal). The caller must already have NIXL-registered it with `agent`. @@ -501,7 +540,8 @@ class BounceTransport /// equality of wireVersion, controlKind and the effective maxChunkSizeBytes — both sides /// advertise post-clamp values, so equality also guarantees every chunk fits the peer's arena /// and scatter scratch. Registration replaces any previous result for the same peer; an empty, - /// incompatible or unparsable replacement clears the old route/capability and returns false, so + /// incompatible or unparsable replacement clears admission capability and returns false. The old + /// control route remains available for outstanding drains, so /// the caller keeps that peer on the standard per-desc NIXL path. Thread-safe. bool registerPeerHandshake(std::string const& peer, std::string const& blob); @@ -516,10 +556,10 @@ class BounceTransport /// 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. - [[nodiscard]] std::shared_future submit( - TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer) + [[nodiscard]] std::shared_future submit(TransferDescs const& srcDescs, TransferDescs const& dstDescs, + std::string const& peer, std::shared_ptr state = {}) { - return mSender.submit(srcDescs, dstDescs, peer); + return mSender.submit(srcDescs, dstDescs, peer, std::move(state)); } /// This side's EFFECTIVE per-chunk cap: cfg.maxChunkSizeBytes AFTER the constructor's clamp to @@ -549,6 +589,7 @@ class BounceTransport BounceSender mSender; std::thread mIoThread; + std::mutex mShutdownMu; // IO-loop idle backoff: consecutive "busy poll but nothing happened" iterations. When in-flight // work exists (so the poll timeout is 0ms) but a gather is stalled behind unrelated GPU kernels, // the loop would otherwise spin a core at 100% on cudaEventQuery. After a threshold of no-progress diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp index ad54147a0461..0dd5c847ecb4 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp @@ -163,6 +163,10 @@ void CreditScheduler::issueGrant( std::vector CreditScheduler::schedule() { std::vector grants; + if (mGrantAdmissionClosed) + { + return grants; + } // Re-sweep as long as the previous sweep granted something; stop when a whole sweep makes no // progress or the ring is empty. while (!mRing.empty()) @@ -350,6 +354,71 @@ void CreditScheduler::dropFlow(std::string const& flow, std::unordered_set CreditScheduler::quarantineFlow( + std::string const& flow, std::unordered_set const& busy, std::vector& deferredOut) +{ + std::lock_guard lk(mMu); + auto it = mFlows.find(flow); + if (it == mFlows.end()) + { + return {}; + } + for (auto const offset : it->second.held) + { + if (busy.count(offset) != 0) + { + deferredOut.push_back(offset); + mOrphans.insert(offset); + } + else + { + mUnproven[flow].insert(offset); + } + } + mFlows.erase(it); + dropFromRing(flow); + return schedule(); +} + +std::vector CreditScheduler::releaseQuarantinedFlow(std::string const& flow) +{ + std::lock_guard lk(mMu); + auto it = mUnproven.find(flow); + if (it != mUnproven.end()) + { + for (auto const offset : it->second) + { + mArena.free(offset); + } + mUnproven.erase(it); + } + return schedule(); +} + +std::vector CreditScheduler::flowKeys() const +{ + std::lock_guard lk(mMu); + std::vector keys; + for (auto const& [key, state] : mFlows) + { + keys.push_back(key); + } + return keys; +} + +void CreditScheduler::closeGrantAdmission() +{ + std::lock_guard lk(mMu); + mGrantAdmissionClosed = true; +} + +bool CreditScheduler::hasUnprovenRemoteWrites() const +{ + std::lock_guard lk(mMu); + return !mUnproven.empty() + || std::any_of(mFlows.begin(), mFlows.end(), [](auto const& entry) { return !entry.second.held.empty(); }); +} + std::vector CreditScheduler::reclaimByPrefix(std::string const& prefix, std::unordered_set const& busy, std::vector& deferredOut, std::chrono::milliseconds quarantineFor) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h index 163e96ebf25f..6777f5b731af 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h @@ -93,40 +93,43 @@ class CreditScheduler /// A region finished scattering on the receiver -> free it and re-schedule. Idempotent. [[nodiscard]] std::vector onScatterDone(std::string const& flow, std::uint64_t offset); - /// Reclaim every flow whose id starts with `prefix` (drop all flows of a gone peer). DEFERS - /// freeing any held region whose offset is in `busy` (a scatter is still reading it): those are - /// appended to `deferredOut` instead of freed, and the caller MUST later call freeOrphanRegion() - /// for each once its scatter completes. `quarantineFor` > 0 marks a RECEIVER-initiated reclaim - /// (peer loss / lease expiry): unlike a sender cancel there is no sender-side drain guaranteeing - /// the RDMA writes into the granted regions have ended, and a one-sided write cannot be aborted - /// — so every non-busy held region is QUARANTINED for that duration (kept out of the arena; - /// reapQuarantine() frees it once the deadline passes) instead of freed immediately. + /// Reclaim every flow whose id starts with `prefix`. Regions in `busy` remain owned as + /// orphans until the caller proves scatter completion and calls freeOrphanRegion(). This + /// low-level API requires independent evidence that remote writes have stopped. A positive + /// `quarantineFor` delays reuse but cannot establish that evidence; the v3 transport uses + /// quarantineFlow() for every flow whose remote writes remain unproven. [[nodiscard]] std::vector reclaimByPrefix(std::string const& prefix, std::unordered_set const& busy, std::vector& deferredOut, std::chrono::milliseconds quarantineFor = std::chrono::milliseconds{0}); + /// Fence a flow without assuming remote RDMA stopped. Its non-scattering allocations + /// remain held until releaseQuarantinedFlow receives explicit sender drain evidence. + [[nodiscard]] std::vector quarantineFlow(std::string const& flow, + std::unordered_set const& busy, std::vector& deferredOut); + [[nodiscard]] std::vector releaseQuarantinedFlow(std::string const& flow); + [[nodiscard]] std::vector flowKeys() const; + [[nodiscard]] bool hasUnprovenRemoteWrites() const; + void closeGrantAdmission(); + /// Free a region deferred by reclaimByPrefix (its in-flight scatter has finished) + re-schedule. [[nodiscard]] std::vector freeOrphanRegion(std::uint64_t offset); - /// Cancel ONE flow (explicit abort / empty WANT): free its held regions and drop it. Any held - /// region in `busy` (a scatter is still reading it) is deferred to `deferredOut` instead of freed - /// (caller later calls freeOrphanRegion). Frees the granted-but-unwritten regions a failed sender - /// would otherwise leak until peer loss. The immediate-free default is safe ONLY when the SENDER - /// initiated the reclaim (its cancel is sent after draining in-flight writes); receiver-initiated - /// reclaims pass `quarantineFor` > 0 (same semantics as reclaimByPrefix). + /// Reclaim one flow with the same ownership requirements as reclaimByPrefix(). The v3 + /// transport calls this only after DRAIN proves remote writes stopped, and retains every busy + /// scatter in `deferredOut` until completion. An empty WANT alone cannot justify this call. [[nodiscard]] std::vector reclaimFlow(std::string const& flow, std::unordered_set const& busy, std::vector& deferredOut, std::chrono::milliseconds quarantineFor = std::chrono::milliseconds{0}); /// Flows HOLDING at least one region with no progress (no WANT refresh, no grant issued, no /// scatter completed) for longer than `idleLimit`. A dead sender emits neither DATA nor a cancel - /// — unobservable through the protocol alone — so the receiver reclaims these via - /// reclaimFlow(quarantineFor > 0). Pending-only flows tie up no memory and are never reported: + /// — unobservable through the protocol alone — so the receiver fences these via + /// quarantineFlow(). Pending-only flows tie up no memory and are never reported: /// they may legitimately queue behind a full arena for a long time. [[nodiscard]] std::vector staleFlows(std::chrono::milliseconds idleLimit) const; - /// Free every quarantined region whose deadline has passed (no write posted before its flow's - /// lease expired can plausibly still be in flight) + re-schedule. Returns the resulting grants. + /// Release legacy timed-quarantine regions and re-schedule. This cannot release permanent + /// flow quarantine and cannot establish remote-write completion evidence. [[nodiscard]] std::vector reapQuarantine(); /// Time-driven re-evaluation of the drain time-box: schedule() only runs on scheduler events, so @@ -240,6 +243,7 @@ class CreditScheduler // except acquireLocal() from submit() app threads (eager gather staging). mutable std::mutex mMu; + bool mGrantAdmissionClosed{false}; BuddyAllocator mArena; // the single shared region allocator (byte offsets) std::uint64_t mBaseAddr{}; // device addr of offset 0 std::uint32_t mMaxInflightChunksPerRequest{}; // per-request in-flight allocation cap @@ -254,26 +258,27 @@ class CreditScheduler std::size_t mEagerBudgetBytes{0}; std::unordered_set mOrphans; // regions deferred by flow/peer reclamation (busy scatter), awaiting freeOrphanRegion - // Regions reclaimed by a RECEIVER-initiated teardown while possibly still being RDMA-written by - // the peer (granted, DATA never arrived), mapped to their reuse deadline. They stay allocated in - // the arena (so schedule() can never re-grant them) until reapQuarantine() passes the deadline. + // Legacy timed quarantine for direct scheduler callers with independent completion evidence. + // The v3 transport uses mUnproven instead: elapsed time cannot retire unknown remote writes. std::unordered_map mQuarantined; + std::unordered_map> mUnproven; + std::function mClock; // injectable for tests // Receiver-only anti-starvation barrier. Every successful remote grant advances mGrantSequence. // A flow whose head allocation keeps failing while enough other grants pass it becomes mDrainFlow; // schedule() then pauses NEW remote grants until that one head fits. Local acquireLocal() remains // untouched, avoiding a cross-direction circular wait in the shared-arena case. The drain is // time-boxed by mDrainTimeout (see setDrainTimeout): a silent peer holding regions must not stall - // every other peer until its lease + quarantine expire, so an overlong drain is abandoned and the + // every other peer while its regions remain pinned, so an overlong drain is abandoned and the // normal round-robin sweep resumes. The expiry is checked inside schedule() (on scheduler // events) and by pollDrain() (from the periodic timeout sweep, whose granularity is 50 ms..1 s, so // expiry may be observed up to one sweep late) — it therefore also fires when no event arrives. // Abandoning restarts the head's bypass count, so it may re-latch once the threshold is crossed // again: under sustained traffic a permanently blocked head alternates between letting // ~threshold grants through and blocking for up to mDrainTimeout. The time-box bounds each stall; - // it does not remove the episode — aggregate throughput over a lease + quarantine window stays - // low (~threshold grants per mDrainTimeout) until the silent peer's regions are lease-reclaimed - // and reaped, or the blocked head's own sender cancels at requestTimeoutMs (a pending-only head + // it does not remove the episode — throughput remains low (~threshold grants per mDrainTimeout) + // until the silent peer's regions acquire physical drain evidence, or the blocked head's own + // sender cancels at requestTimeoutMs (a pending-only head // holds no regions and is never lease-reclaimed itself). static constexpr std::uint64_t kMinimumBypassGrants{8}; static constexpr std::uint64_t kBypassRounds{2}; 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..db1ba1bddc89 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -94,8 +94,10 @@ namespace class BounceTransferStatus final : public TransferStatus { public: - explicit BounceTransferStatus(std::shared_future fut) + explicit BounceTransferStatus( + std::shared_future fut, std::shared_ptr physical) : mFut(std::move(fut)) + , mPhysical(std::move(physical)) { } @@ -126,6 +128,17 @@ class BounceTransferStatus final : public TransferStatus return TransferState::kIN_PROGRESS; } + [[nodiscard]] bool isQuiesced() const override + { + return mPhysical->quiesced.load(std::memory_order_acquire); + } + + [[nodiscard]] bool requestCancel() override + { + mPhysical->cancelRequested.store(true, std::memory_order_release); + return true; + } + /// Failure cause recorded by the bounce transport (empty while in flight / on success). [[nodiscard]] std::string getLastStatusStr() const override { @@ -139,6 +152,7 @@ class BounceTransferStatus final : public TransferStatus private: std::shared_future mFut; + std::shared_ptr mPhysical; }; } // namespace @@ -743,6 +757,11 @@ std::string NixlTransferStatus::getLastStatusStr() const return queryStatus() == NIXL_SUCCESS; } +bool NixlTransferStatus::isQuiesced() const +{ + return mQuiesced.load(std::memory_order_acquire) || queryStatus() == NIXL_SUCCESS; +} + nixl_status_t NixlTransferStatus::queryStatus() const { auto const query = [this]() @@ -760,6 +779,10 @@ nixl_status_t NixlTransferStatus::queryStatus() const return NIXL_ERR_INVALID_PARAM; } auto const status = agent->getXferStatus(mHandle); + if (status == NIXL_SUCCESS) + { + mQuiesced.store(true, std::memory_order_release); + } mLastStatus.store(static_cast(status), std::memory_order_relaxed); return status; }; @@ -1048,8 +1071,10 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) // Programmatic check: getBounceSubmitCount() / bounce_submit_count on the Python binding. 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)); + auto physical = std::make_shared(); + auto fut = mBounce->transport->submit( + request.getSrcDescs(), request.getDstDescs(), request.getRemoteName(), physical); + return std::make_unique(std::move(fut), std::move(physical)); } #endif diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index 90794052d6f2..656de5397d78 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -80,6 +80,7 @@ class NixlTransferStatus final : public TransferStatus NixlTransferStatus& operator=(NixlTransferStatus&&) = delete; [[nodiscard]] bool isCompleted() const override; + [[nodiscard]] bool isQuiesced() const override; [[nodiscard]] TransferState wait(int64_t timeout_ms = -1) const override; @@ -95,6 +96,7 @@ class NixlTransferStatus final : public TransferStatus std::weak_ptr mWeakAgent; nixlXferReqH* mHandle{}; mutable std::atomic mLastStatus{0}; + mutable std::atomic mQuiesced{false}; bool const mSynchronizeHandleAccess; mutable std::mutex mHandleMutex; }; diff --git a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp index 7fa215f170e6..96aea91874a3 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp @@ -136,6 +136,52 @@ std::unique_ptr tryMakeAgent(kvc::BaseAgentConfig cfg, s return nullptr; } } + +class FailedPhysicalStatus final : public kvc::TransferStatus +{ +public: + explicit FailedPhysicalStatus(std::shared_ptr> stopped) + : mStopped(std::move(stopped)) + { + } + + bool isCompleted() const override + { + return false; + } + + kvc::TransferState wait(int64_t) const override + { + return kvc::TransferState::kFAILURE; + } + + bool isQuiesced() const override + { + return mStopped->load(std::memory_order_acquire); + } + +private: + std::shared_ptr> mStopped; +}; + +class FailedPhysicalAgent final : public kvc::NixlTransferAgent +{ +public: + FailedPhysicalAgent(kvc::BaseAgentConfig const& config, std::shared_ptr> stopped) + : kvc::NixlTransferAgent(config) + , mStopped(std::move(stopped)) + { + } + + std::unique_ptr postXferRequest(kvc::TransferOp, kvc::TransferDescs const&, + kvc::TransferDescs const&, std::string const&, std::optional const&) override + { + return std::make_unique(mStopped); + } + +private: + std::shared_ptr> mStopped; +}; } // namespace // BaseAgentConfig::agentBufferSizeMb (derived from CacheTransceiverConfig's @@ -194,6 +240,9 @@ TEST(BounceAgentE2E, SubmitTransferRequestsUsesBounce) EXPECT_EQ(waitTerminal(status, 30), kvc::TransferState::kSUCCESS) << "bounce transfer via submitTransferRequests did not succeed"; EXPECT_TRUE(verifyXferBufs(bufs)); + EXPECT_TRUE(status->isQuiesced()); + (void) status->release(); + EXPECT_TRUE(status->isQuiesced()); a->shutdown(); b->shutdown(); @@ -248,7 +297,9 @@ TEST(BounceAgentE2E, EffectiveChunkCapFallsBackToStandardNixl) auto const bucket = static_cast(kvc::bounce::BounceRejectReason::kDescriptorShape); EXPECT_EQ(a->getBounceRejectCounts()[bucket], 1u); EXPECT_TRUE(verifyXferBufs(bufs)); + EXPECT_TRUE(status->isQuiesced()); EXPECT_TRUE(status->release()); + EXPECT_TRUE(status->isQuiesced()); status.reset(); a->deregisterMemory(req.getSrcDescs()); @@ -423,3 +474,40 @@ TEST(BounceAgentE2E, MultiAgentManySendersToOneReceiver) freeXferBufs(x); } } + +TEST(BounceAgentE2E, FailedBounceStatusSeparatesFailureFromPhysicalCompletion) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + std::string skipMessage; + auto receiver = tryMakeAgent(makeBounceConfig("proofAgentReceiver"), skipMessage); + if (!receiver) + { + GTEST_SKIP() << "NIXL unavailable: " << skipMessage; + } + auto stopped = std::make_shared>(false); + auto sender = std::make_unique(makeBounceConfig("proofAgentSender"), stopped); + sender->loadRemoteAgent("proofAgentReceiver", receiver->getLocalAgentDesc()); + auto buffers = makeXferBufs(8, 256, 19); + auto status = sender->submitTransferRequests(makeReq(buffers, "proofAgentReceiver")); + ASSERT_NE(status, nullptr); + ASSERT_EQ(waitTerminal(status, 10), kvc::TransferState::kFAILURE); + EXPECT_FALSE(status->isCompleted()); + EXPECT_FALSE(status->isQuiesced()); + EXPECT_TRUE(status->requestCancel()); + EXPECT_FALSE(status->isQuiesced()); + stopped->store(true, std::memory_order_release); + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!status->isQuiesced() && std::chrono::steady_clock::now() < deadline) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(status->isQuiesced()); + EXPECT_FALSE(status->isCompleted()); + EXPECT_EQ(status->wait(0), kvc::TransferState::kFAILURE); + sender->shutdown(); + receiver->shutdown(); + freeXferBufs(buffers); +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp index 8e775f65a7df..ce446a331ef8 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp @@ -279,3 +279,30 @@ TEST(BounceMessageCodec, HandshakeTruncatedRejected) EXPECT_FALSE(b::decodeHandshake(blob.substr(0, blob.size() - 4), out)); // endpoint cut short EXPECT_FALSE(b::decodeHandshake(blob.substr(0, 10), out)); // fixed part cut short } + +TEST(BounceMessageCodec, DrainRoundTripAndAcknowledgment) +{ + auto const blob = b::encodeDrain(42, "tcp://127.0.0.1:12345"); + b::BounceMsgHeader header{}; + ASSERT_TRUE(b::decodeHeader(blob, header)); + EXPECT_EQ(static_cast(header.msgType), b::BounceMsgType::kDRAIN); + EXPECT_EQ(header.requestId, 42u); + std::vector chunks; + std::string endpoint; + ASSERT_TRUE(b::decodeWant(blob, header, chunks, endpoint)); + EXPECT_TRUE(chunks.empty()); + EXPECT_EQ(endpoint, "tcp://127.0.0.1:12345"); + ASSERT_TRUE(b::decodeHeader(b::encodeDrainAck(42), header)); + EXPECT_EQ(static_cast(header.msgType), b::BounceMsgType::kDRAIN_ACK); + EXPECT_EQ(header.requestId, 42u); + EXPECT_EQ(header.payloadBytes, 0u); +} + +TEST(BounceMessageCodec, PreviousWireVersionCannotAttestDrain) +{ + auto blob = b::encodeDrainAck(42); + auto const oldVersion = static_cast(b::kBounceVersion - 1); + std::memcpy(blob.data() + offsetof(b::BounceMsgHeader, version), &oldVersion, sizeof(oldVersion)); + b::BounceMsgHeader header{}; + EXPECT_FALSE(b::decodeHeader(blob, header)); +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h b/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h index 2f4dc8c74a53..bf961d6c1aaa 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h +++ b/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h @@ -185,7 +185,7 @@ inline std::unique_ptr makeNode(std::string const& name, b::BounceConfig c return nullptr; } n->arena = std::make_unique(cfg.arenaSizeBytes, 0, /*allowFabric=*/false); - n->exec = std::make_unique(cfg.maxInflightChunksPerRequest + 4, maxDescs, 0, cfg.useZeroCopyArguments); + n->exec = std::make_unique(cfg.copyStreamCount, maxDescs, 0, cfg.useZeroCopyArguments); if (!n->agent->registerRegionImpl(n->arena->base(), n->arena->bytes(), /*deviceId=*/0)) { return nullptr; diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp index 6b90b9dd7da2..7aee630d2311 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp @@ -33,7 +33,11 @@ #include #include +#include +#include +#include #include +#include #include #include @@ -46,9 +50,11 @@ namespace // statuses it hands out. struct XferControls { - std::atomic failWrites{false}; // every posted write polls as FAILURE - std::atomic allowTerminal{true}; // false -> posted writes stay IN_PROGRESS - std::atomic releaseSucceeds{true}; // false -> release() fails (handle retained) + std::atomic throwAfterPost{false}; + std::atomic failWrites{false}; // every posted write polls as FAILURE + std::atomic physicallyStopped{true}; // independent physical proof for injected failures + std::atomic allowTerminal{true}; // false -> posted writes stay IN_PROGRESS + std::atomic releaseSucceeds{true}; // false -> release() fails (handle retained) std::atomic postCount{0}; std::atomic releaseCount{0}; }; @@ -68,6 +74,11 @@ class FakeXferStatus final : public kvc::TransferStatus return wait(0) != kvc::TransferState::kIN_PROGRESS; } + [[nodiscard]] bool isQuiesced() const override + { + return mCtl->physicallyStopped.load(std::memory_order_acquire) && wait(0) != kvc::TransferState::kIN_PROGRESS; + } + [[nodiscard]] kvc::TransferState wait(int64_t) const override { if (mCtl->failWrites.load(std::memory_order_acquire)) @@ -114,6 +125,10 @@ class FakeXferAgent final : public kvc::NixlTransferAgent kvc::TransferDescs const&, std::string const&, std::optional const&) override { mCtl->postCount.fetch_add(1, std::memory_order_relaxed); + if (mCtl->throwAfterPost.load(std::memory_order_acquire)) + { + throw std::runtime_error("injected ambiguous backend admission"); + } return std::make_unique(mCtl); } @@ -173,6 +188,46 @@ bool pumpChannel(b::ZmqControlChannel& ch, std::chrono::milliseconds budget, Fn& return false; } +bool drainFlow(b::ZmqControlChannel& sender, std::string const& receiver, std::uint64_t rid) +{ + sender.sendTo(receiver, b::encodeDrain(rid, sender.localEndpoint())); + return pumpChannel(sender, std::chrono::seconds(5), + [rid](b::BounceMsgHeader const& header, std::string const&) { + return header.requestId == rid + && static_cast(header.msgType) == b::BounceMsgType::kDRAIN_ACK; + }); +} + +bool acknowledgeDrain(b::ZmqControlChannel& receiver, std::string const& sender) +{ + return pumpChannel(receiver, std::chrono::seconds(5), + [&](b::BounceMsgHeader const& header, std::string const& blob) + { + if (static_cast(header.msgType) != b::BounceMsgType::kDRAIN) + { + return false; + } + std::vector chunks; + std::string endpoint; + if (!b::decodeWant(blob, header, chunks, endpoint) || !receiver.addPeer(sender, endpoint)) + { + return false; + } + receiver.sendTo(sender, b::encodeDrainAck(header.requestId)); + return true; + }); +} + +bool waitPhysical(std::shared_ptr const& state) +{ + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!state->quiesced.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return state->quiesced.load(std::memory_order_acquire); +} + // Wait up to `budget` for a GRANT for `rid`, decoding its credits into `out`. bool waitGrant(b::ZmqControlChannel& ch, std::uint64_t rid, std::chrono::milliseconds budget, std::vector& out) @@ -213,6 +268,28 @@ int countAcks(b::ZmqControlChannel& ch, std::chrono::milliseconds budget, std::u }); return n; } + +struct StreamGate +{ + std::atomic entered{false}; + std::atomic open{false}; + + ~StreamGate() + { + open.store(true, std::memory_order_release); + } +}; + +void CUDART_CB waitAtGate(void* data) +{ + auto& gate = *static_cast(data); + gate.entered.store(true, std::memory_order_release); + while (!gate.open.load(std::memory_order_acquire)) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } +} + } // namespace // A peer that never GRANTs must fail the request on requestTimeoutMs, not hang. The WANT is @@ -237,6 +314,7 @@ TEST(BounceTransportFailure, NoGrantTimesOutNotHang) EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); EXPECT_EQ(fut.get().reason, b::BounceFailReason::kNoProgressTimeout); + ASSERT_TRUE(acknowledgeDrain(ghost, "ngSolo")); t->tx->shutdown(); bounce_test::freeXferBufs(bufs); } @@ -371,7 +449,7 @@ TEST(BounceTransportFailure, DuplicateDataProducesOneScatterAndAck) } EXPECT_EQ(countAcks(sender, std::chrono::seconds(2), rid), 1); - + ASSERT_TRUE(drainFlow(sender, "dupDataReceiver", sentinelRid)); receiver->tx->shutdown(); EXPECT_EQ(cudaFree(dst), cudaSuccess); } @@ -405,6 +483,7 @@ TEST(BounceTransportFailure, DuplicateWantIsDroppedWithoutRegrant) std::vector extra; EXPECT_FALSE(waitGrant(sender, rid, std::chrono::milliseconds(500), extra)) << "duplicate WANT was re-granted"; + ASSERT_TRUE(drainFlow(sender, "dupWantReceiver", rid)); receiver->tx->shutdown(); } @@ -441,14 +520,14 @@ TEST(BounceTransportFailure, MalformedWantChunkSizeIsRejectedNotWedged) std::vector credits; ASSERT_TRUE(waitGrant(sender, rid, std::chrono::seconds(5), credits)) << "legit WANT starved, rid=" << rid; // Cancel to release the granted region and drop the flow from the ring. - sender.sendTo("badWantReceiver", b::encodeWant(rid, {}, sender.localEndpoint())); + sender.sendTo("badWantReceiver", b::encodeDrain(rid, sender.localEndpoint())); } // The receiver must still grant fresh legitimate demand — pre-fix the drain latch wedges this. sender.sendTo("badWantReceiver", b::encodeWant(/*rid=*/200, {256}, sender.localEndpoint())); std::vector credits; EXPECT_TRUE(waitGrant(sender, 200, std::chrono::seconds(5), credits)) << "receiver wedged after zero-size WANT"; - sender.sendTo("badWantReceiver", b::encodeWant(200, {}, sender.localEndpoint())); + sender.sendTo("badWantReceiver", b::encodeDrain(200, sender.localEndpoint())); // Oversized mode: pre-fix this grant monopolizes the whole arena, so the follow-up legitimate // WANT starves. @@ -465,6 +544,7 @@ TEST(BounceTransportFailure, MalformedWantChunkSizeIsRejectedNotWedged) EXPECT_FALSE(waitGrant(sender, 1, std::chrono::milliseconds(200), badCredits)) << "zero-size WANT was granted"; EXPECT_FALSE(waitGrant(sender, 2, std::chrono::milliseconds(200), badCredits)) << "oversized WANT was granted"; + ASSERT_TRUE(drainFlow(sender, "badWantReceiver", 201)); receiver->tx->shutdown(); } @@ -530,16 +610,15 @@ TEST(BounceTransportFailure, OversizedScatterRunListIsRejectedNotCounted) std::vector credits2; EXPECT_TRUE(waitGrant(sender, 2, std::chrono::seconds(5), credits2)) << "region pinned by oversized run list"; + ASSERT_TRUE(drainFlow(sender, "bigRunReceiver", 2)); receiver->tx->shutdown(); EXPECT_EQ(cudaFree(dst), cudaSuccess); } -// A sender that takes a GRANT and then dies emits neither DATA nor a cancel, so nothing -// event-driven can ever reclaim its region — the receiver's grant lease must. After -// receiverFlowTimeoutMs of silence the flow is reclaimed and its region quarantined for -// quarantineMs, after which it must serve the next waiting flow; a late DATA for the expired -// grant must be dropped (no scatter, no ACK). -TEST(BounceTransportFailure, GrantLeaseExpiryReclaimsSilentSendersRegion) +// Lease expiry fences a silent sender but cannot prove its NIC stopped. Its region must remain +// unavailable beyond the old timed quarantine, until an explicit DRAIN permits reuse. Late DATA +// for the fenced flow must be dropped even after another flow acquires the same region. +TEST(BounceTransportFailure, GrantLeaseExpiryRetainsRegionUntilExplicitDrain) { if (!bounce_test::hasCuda()) GTEST_SKIP() << "no CUDA device"; @@ -565,8 +644,10 @@ TEST(BounceTransportFailure, GrantLeaseExpiryReclaimsSilentSendersRegion) std::vector credit2; EXPECT_FALSE(waitGrant(sender, 2, std::chrono::milliseconds(200), credit2)) << "region re-granted before lease expiry"; - // ...but once the silent flow's lease expires and the quarantine passes, it must be. - ASSERT_TRUE(waitGrant(sender, 2, std::chrono::seconds(10), credit2)) << "silent sender's region never reclaimed"; + // Time alone is not a DMA-completion proof, including after the old quarantine. + EXPECT_FALSE(waitGrant(sender, 2, std::chrono::milliseconds(1000), credit2)); + sender.sendTo("glReceiver", b::encodeDrain(1, sender.localEndpoint())); + ASSERT_TRUE(waitGrant(sender, 2, std::chrono::seconds(5), credit2)); ASSERT_EQ(credit2.size(), 1u); EXPECT_EQ(credit2.front().regionHandle, credit1.front().regionHandle); // the one region, recycled @@ -690,77 +771,61 @@ TEST(BounceTransportFailure, CreditMispairAbandonIsTerminalAndIgnoresLaterGrants // The later valid GRANT must not have resurrected the abandoned flow into posting a write. EXPECT_EQ(ctl->postCount.load(std::memory_order_acquire), 0u) << "abandoned flow posted an RDMA write"; + ASSERT_TRUE(acknowledgeDrain(peer, "mispairSender")); sender->tx->shutdown(); bounce_test::freeXferBufs(bufs); } // shutdown() with an in-flight (stuck) request must resolve its future FAILURE, never leave wait() hanging. -TEST(BounceTransportFailure, ShutdownFailsInflight) +TEST(BounceTransportFailure, ShutdownWithoutRemoteProofFailsStop) { if (!bounce_test::hasCuda()) GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/0); // timeout disabled -> 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) - EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); // still pending - t->tx->shutdown(); - ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "shutdown left request hanging"; - EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); - EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); - - bounce_test::freeXferBufs(bufs); + GTEST_FLAG_SET(death_test_style, "threadsafe"); + EXPECT_EXIT( + { + auto node = bounce_test::makeNode("shutdownUnproved", cfg(30000), 1024); + auto buffers = bounce_test::makeXferBufs(8, 256, 1); + auto future = node->tx->submit(buffers.srcDescs, buffers.dstDescs, "missingPeer"); + // MPI installs a SIGABRT handler that converts abort into exit(1). + // Restore the signal only after fixture admission has succeeded. + std::signal(SIGABRT, SIG_DFL); + std::fprintf(stderr, "awaiting physical quiescence guard\n"); + node->tx->shutdown(); + }, + ::testing::KilledBySignal(SIGABRT), "awaiting physical quiescence guard"); } -// Shutdown uses the status object's bounded release instead of polling forever. If the backend -// cannot abort, the future still fails promptly; the release is attempted once in failAll and once -// in its retry pass (the production NixlTransferStatus destructor makes a final attempt). -TEST(BounceTransportFailure, ShutdownReleaseFailureDoesNotHangOrLoseHandle) +TEST(BounceTransportFailure, AmbiguousSubmissionCannotReleaseSourceAtShutdown) { if (!bounce_test::hasCuda()) GTEST_SKIP() << "no CUDA device"; - auto c = cfg(/*timeoutMs=*/0); - capRegions(c, c.maxInflightChunksPerRequest); - b::ZmqControlChannel peer("shutdownReleasePeer"); - auto ctl = std::make_shared(); - ctl->allowTerminal.store(false, std::memory_order_release); - 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())); - - auto bufs = bounce_test::makeXferBufs(/*nDescs=*/1, /*descBytes=*/256, /*seed=*/5); - auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, "shutdownReleasePeer"); - b::BounceCreditEntry credit{/*addr=*/0, /*len=*/256, /*devId=*/0, /*regionHandle=*/91}; - peer.sendTo("shutdownReleaseSender", b::encodeGrant(/*requestId=*/1, {credit})); - - auto const postDeadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); - while (std::chrono::steady_clock::now() < postDeadline && ctl->postCount.load(std::memory_order_acquire) == 0) - { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - } - ASSERT_EQ(ctl->postCount.load(std::memory_order_acquire), 1u); - - auto const start = std::chrono::steady_clock::now(); - sender->tx->shutdown(); - // A write that never completes is drained for the fixed kShutdownWriteDrainMs (5 s) — intended - // teardown behaviour — so the bound only asserts "bounded, no hang", not "instant". - EXPECT_LT(std::chrono::steady_clock::now() - start, std::chrono::seconds(10)); - 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); - // 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); + GTEST_FLAG_SET(death_test_style, "threadsafe"); + EXPECT_EXIT( + { + auto controls = std::make_shared(); + controls->throwAfterPost.store(true); + auto sender = makeFakeNode("throwSender", cfg(30000), controls); + b::ZmqControlChannel peer("throwPeer"); + sender->tx->addPeer("throwPeer", peer.localEndpoint()); + peer.addPeer("throwSender", sender->ch->localEndpoint()); + auto buffers = bounce_test::makeXferBufs(1, 256, 1); + auto state = std::make_shared(); + auto future = sender->tx->submit(buffers.srcDescs, buffers.dstDescs, "throwPeer", state); + peer.sendTo("throwSender", b::encodeGrant(1, {{0, 256, 0, 1}})); + if (future.wait_for(std::chrono::seconds(5)) != std::future_status::ready || state->quiesced.load()) + { + std::_Exit(2); + } + // MPI installs a SIGABRT handler that converts abort into exit(1). + // Restore the signal only after fixture admission has succeeded. + std::signal(SIGABRT, SIG_DFL); + std::fprintf(stderr, "awaiting physical quiescence guard\n"); + sender->tx->shutdown(); + }, + ::testing::KilledBySignal(SIGABRT), "awaiting physical quiescence guard"); } -// 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. TEST(BounceTransportFailure, ForgetPeerFailsInflightRequest) { if (!bounce_test::hasCuda()) @@ -770,6 +835,8 @@ TEST(BounceTransportFailure, ForgetPeerFailsInflightRequest) if (!t) GTEST_SKIP() << "NIXL agent/backend unavailable"; + b::ZmqControlChannel gone("gonePeer"); + ASSERT_TRUE(t->tx->addPeer("gonePeer", gone.localEndpoint())); auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "gonePeer"); // stuck: never granted EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); // still pending @@ -780,6 +847,8 @@ TEST(BounceTransportFailure, ForgetPeerFailsInflightRequest) // forgetPeer for an unrelated/unknown peer must be a harmless no-op (no crash). t->tx->forgetPeer("someoneElse"); + ASSERT_TRUE(t->tx->addPeer("gonePeer", gone.localEndpoint())); + ASSERT_TRUE(acknowledgeDrain(gone, "fpSolo")); t->tx->shutdown(); bounce_test::freeXferBufs(bufs); } @@ -802,7 +871,8 @@ TEST(BounceTransportFailure, ForgetPeerInFlightRecovers) bounce_test::wirePair(*A, *B); auto bufs = bounce_test::makeXferBufs(/*nDescs=*/24, /*descBytes=*/600, /*seed=*/1); - auto fut = A->tx->submit(bufs.srcDescs, bufs.dstDescs, "fpB"); + auto physical = std::make_shared(); + auto fut = A->tx->submit(bufs.srcDescs, bufs.dstDescs, "fpB", physical); A->tx->forgetPeer("fpB"); // drop the peer (queued; applied on A's IO thread) ASSERT_EQ(fut.wait_for(std::chrono::seconds(10)), std::future_status::ready) << "request hung after forgetPeer"; auto const res = fut.get(); @@ -812,15 +882,16 @@ TEST(BounceTransportFailure, ForgetPeerInFlightRecovers) { EXPECT_EQ(res.reason, b::BounceFailReason::kPeerDropped); } - bounce_test::freeXferBufs(bufs); - // Recovery + no-leak: forgetPeer is a one-shot event; the scheduler/request reclaim is drained by - // the time fut resolved, so these new flows aren't reclaimed. forgetPeer ALSO drops the + // Logical failure may precede physical drain. Re-establish the original control route and + // wait for physical evidence before freeing buffers or testing recovery. forgetPeer drops the // control-channel DEALER to fpB synchronously (on this thread), so re-establish it with addPeer // before recovering — deterministic because forgetPeer's removePeer happens-before this addPeer // (no async removePeer can race/erase the freshly re-added dealer). NIXL metadata persists, so // only the dealer is re-added (no loadRemoteAgent / full re-wire). A->tx->addPeer("fpB", B->ch->localEndpoint()); + ASSERT_TRUE(waitPhysical(physical)); + bounce_test::freeXferBufs(bufs); for (int k = 0; k < 5; ++k) { auto rb @@ -837,23 +908,16 @@ TEST(BounceTransportFailure, ForgetPeerInFlightRecovers) B->tx->shutdown(); } -// forget() must drop the gone peer's flows BEFORE freeing the orphaned regions of its queued (not yet -// started) scatter jobs: freeOrphanRegion() runs schedule(), so with free-first the region is -// granted straight back to a pending flow of the very peer being forgotten, and the reclaim that -// follows then quarantines it (quarantineMs) instead of returning it to the arena. Setup: the -// receiver's ExecPool is drained from the test so the single scatter worker spins in tryAcquire on -// job A while job B stays queued (mScattering set for both); a cancel flags both regions orphaned; a -// second WANT from the same peer is left pending (arena full); then forgetPeer. Correct order: the -// pending flow is gone before region B is freed, so a second peer's WANT gets it at once. Old order: -// region B goes to the forgotten peer's pending flow and is quarantined for 30 s, so the second peer -// starves (and the forgotten peer sees a GRANT after forgetPeer). +// Peer loss fences every pending flow before any queued or dequeued scatter can release a +// region. Held execution contexts keep both scatters pending until explicitly released. Their +// regions may then serve another peer, never a pending flow belonging to the forgotten peer. 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 + c.quarantineMs = 30000; // elapsed time must not release unknown remote writes capRegions(c, c.maxInflightChunksPerRequest); // arena holds exactly TWO regions b::ZmqControlChannel gone("fgGone"); b::ZmqControlChannel other("fgOther"); @@ -904,8 +968,8 @@ TEST(BounceTransportFailure, ForgetPeerFreesOrphanedRegionsToOtherPeersOnly) other.sendTo("fgReceiver", b::encodeWant(/*rid=*/10, {chunk}, other.localEndpoint())); std::vector otherCredits; - EXPECT_TRUE(waitGrant(other, 10, std::chrono::seconds(2), otherCredits)) - << "orphaned region not returned to the arena by forgetPeer (granted back to the forgotten peer?)"; + EXPECT_FALSE(waitGrant(other, 10, std::chrono::milliseconds(200), otherCredits)) + << "queued/dequeued scatter still owns both regions"; std::vector stray; EXPECT_FALSE(waitGrant(gone, 2, std::chrono::milliseconds(200), stray)) << "forgotten peer was granted a region"; @@ -913,6 +977,10 @@ TEST(BounceTransportFailure, ForgetPeerFreesOrphanedRegionsToOtherPeersOnly) { receiver->exec->release(ctx); } + ASSERT_TRUE(waitGrant(other, 10, std::chrono::seconds(5), otherCredits)); + ASSERT_TRUE(drainFlow(gone, "fgReceiver", 1)); + ASSERT_TRUE(drainFlow(gone, "fgReceiver", 2)); + ASSERT_TRUE(drainFlow(other, "fgReceiver", 10)); receiver->tx->shutdown(); EXPECT_EQ(cudaFree(dst), cudaSuccess); } @@ -1030,3 +1098,239 @@ TEST(BounceTransportFailure, ConcurrentMultiThreadedSubmit) bounce_test::freeXferBufs(bb); } } + +TEST(BounceTransportFailure, DrainWaitsForDequeuedScatterAndRejectsLateData) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto config = cfg(30000); + config.copyStreamCount = 1; + config.scatterWorkerCount = 1; + auto receiver = bounce_test::makeNode("drainReceiver", config, 1024); + if (!receiver) + GTEST_SKIP() << "NIXL unavailable"; + b::ZmqControlChannel sender("drainSender"); + ASSERT_TRUE(sender.addPeer(receiver->name, receiver->ch->localEndpoint())); + auto buffers = bounce_test::makeXferBufs(1, 256, 1); + auto* context = receiver->exec->tryAcquire(); + ASSERT_NE(context, nullptr); + sender.sendTo(receiver->name, b::encodeWant(1, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, 1, std::chrono::seconds(5), credits)); + ASSERT_EQ(cudaMemset(reinterpret_cast(credits[0].addr), 0x5A, 256), cudaSuccess); + auto const data = b::encodeData( + 1, 0, 1, credits[0].regionHandle, {{0, reinterpret_cast(buffers.dst), 0, 0, 256, 1}}); + sender.sendTo(receiver->name, data); + sender.sendTo(receiver->name, b::encodeDrain(1, sender.localEndpoint())); + EXPECT_FALSE(pumpChannel(sender, std::chrono::milliseconds(200), + [](auto const& h, auto const&) + { return static_cast(h.msgType) == b::BounceMsgType::kDRAIN_ACK; })); + receiver->exec->release(context); + ASSERT_TRUE(drainFlow(sender, receiver->name, 1)); + unsigned char value = 0; + ASSERT_EQ(cudaMemcpy(&value, buffers.dst, 1, cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(value, 0x5A); + ASSERT_EQ(cudaMemset(buffers.dst, 0x7B, 256), cudaSuccess); + sender.sendTo(receiver->name, b::encodeWant(1, {256}, sender.localEndpoint())); + sender.sendTo(receiver->name, data); + ASSERT_TRUE(drainFlow(sender, receiver->name, 1)); + ASSERT_EQ(cudaMemcpy(&value, buffers.dst, 1, cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(value, 0x7B); + receiver->tx->shutdown(); + bounce_test::freeXferBufs(buffers); +} + +TEST(BounceTransportFailure, DrainWaitsForSubmittedCudaScatterWork) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto config = cfg(30000); + config.copyStreamCount = 1; + config.scatterWorkerCount = 1; + auto receiver = bounce_test::makeNode("cudaDrainReceiver", config, 1024); + if (!receiver) + GTEST_SKIP() << "NIXL unavailable"; + b::ZmqControlChannel sender("cudaDrainSender"); + ASSERT_TRUE(sender.addPeer(receiver->name, receiver->ch->localEndpoint())); + auto buffers = bounce_test::makeXferBufs(1, 256, 1); + sender.sendTo(receiver->name, b::encodeWant(1, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, 1, std::chrono::seconds(5), credits)); + ASSERT_EQ(cudaMemset(reinterpret_cast(credits[0].addr), 0x5A, 256), cudaSuccess); + StreamGate gate; + auto* context = receiver->exec->tryAcquire(); + ASSERT_NE(context, nullptr); + ASSERT_EQ(cudaLaunchHostFunc(context->stream, waitAtGate, &gate), cudaSuccess); + receiver->exec->release(context); + sender.sendTo(receiver->name, + b::encodeData( + 1, 0, 1, credits[0].regionHandle, {{0, reinterpret_cast(buffers.dst), 0, 0, 256, 1}})); + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (receiver->exec->freeCount() != 0 && std::chrono::steady_clock::now() < deadline) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(receiver->exec->freeCount(), 0u); + sender.sendTo(receiver->name, b::encodeDrain(1, sender.localEndpoint())); + EXPECT_FALSE(pumpChannel(sender, std::chrono::milliseconds(200), + [](auto const& h, auto const&) + { return static_cast(h.msgType) == b::BounceMsgType::kDRAIN_ACK; })); + gate.open.store(true, std::memory_order_release); + ASSERT_TRUE(drainFlow(sender, receiver->name, 1)); + unsigned char value = 0; + ASSERT_EQ(cudaMemcpy(&value, buffers.dst, 1, cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(value, 0x5A); + ASSERT_EQ(cudaMemset(buffers.dst, 0x7B, 256), cudaSuccess); + sender.sendTo(receiver->name, + b::encodeData( + 1, 0, 1, credits[0].regionHandle, {{0, reinterpret_cast(buffers.dst), 0, 0, 256, 1}})); + ASSERT_TRUE(drainFlow(sender, receiver->name, 1)); + ASSERT_EQ(cudaMemcpy(&value, buffers.dst, 1, cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(value, 0x7B); + receiver->tx->shutdown(); + bounce_test::freeXferBufs(buffers); +} + +TEST(BounceTransportFailure, RemoteDrainAckCannotRetireRunningSourceGather) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto config = cfg(30000); + config.copyStreamCount = 1; + auto sender = bounce_test::makeNode("gatherDrainSender", config, 1024); + if (!sender) + GTEST_SKIP() << "NIXL unavailable"; + b::ZmqControlChannel receiver("gatherDrainPeer"); + ASSERT_TRUE(sender->tx->addPeer("gatherDrainPeer", receiver.localEndpoint())); + auto buffers = bounce_test::makeXferBufs(1, 256, 1); + StreamGate gate; + auto* context = sender->exec->tryAcquire(); + ASSERT_NE(context, nullptr); + ASSERT_EQ(cudaLaunchHostFunc(context->stream, waitAtGate, &gate), cudaSuccess); + sender->exec->release(context); + auto state = std::make_shared(); + auto future = sender->tx->submit(buffers.srcDescs, buffers.dstDescs, "gatherDrainPeer", state); + state->cancelRequested.store(true, std::memory_order_release); + ASSERT_EQ(future.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_EQ(future.get().state, kvc::TransferState::kFAILURE); + ASSERT_TRUE(acknowledgeDrain(receiver, sender->name)); + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + EXPECT_FALSE(state->quiesced.load()); + gate.open.store(true, std::memory_order_release); + ASSERT_TRUE(waitPhysical(state)); + EXPECT_EQ(future.get().state, kvc::TransferState::kFAILURE); + sender->tx->shutdown(); + bounce_test::freeXferBufs(buffers); +} + +TEST(BounceTransportFailure, FailedRdmaNeedsProofAndDrainAcknowledgmentRetries) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto controls = std::make_shared(); + controls->failWrites.store(true); + controls->physicallyStopped.store(false); + auto sender = makeFakeNode("proofSender", cfg(30000), controls); + if (!sender) + GTEST_SKIP() << "NIXL unavailable"; + b::ZmqControlChannel receiver("proofReceiver"); + ASSERT_TRUE(sender->tx->addPeer("proofReceiver", receiver.localEndpoint())); + ASSERT_TRUE(receiver.addPeer(sender->name, sender->ch->localEndpoint())); + auto buffers = bounce_test::makeXferBufs(1, 256, 1); + auto state = std::make_shared(); + auto future = sender->tx->submit(buffers.srcDescs, buffers.dstDescs, "proofReceiver", state); + receiver.sendTo(sender->name, b::encodeGrant(1, {{0, 256, 0, 1}})); + ASSERT_EQ(future.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_FALSE(state->quiesced.load()); + EXPECT_FALSE(pumpChannel(receiver, std::chrono::milliseconds(200), + [](auto const& h, auto const&) + { return static_cast(h.msgType) == b::BounceMsgType::kDRAIN; })); + EXPECT_EQ(controls->releaseCount.load(), 0u); + controls->physicallyStopped.store(true); + // Drop the first DRAIN, then drop the first ACK: each subsequent DRAIN is retryable. + ASSERT_TRUE(pumpChannel(receiver, std::chrono::seconds(5), + [](auto const& h, auto const&) + { return static_cast(h.msgType) == b::BounceMsgType::kDRAIN; })); + ASSERT_TRUE(pumpChannel(receiver, std::chrono::seconds(5), + [](auto const& h, auto const&) + { return static_cast(h.msgType) == b::BounceMsgType::kDRAIN; })); + EXPECT_FALSE(state->quiesced.load()); + ASSERT_TRUE(acknowledgeDrain(receiver, sender->name)); + ASSERT_TRUE(waitPhysical(state)); + sender->tx->shutdown(); + bounce_test::freeXferBufs(buffers); +} + +TEST(BounceTransportFailure, ChangedEndpointDrainCannotReleaseOldIncarnation) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto config = cfg(30000); + capRegions(config, 1); + auto receiver = bounce_test::makeNode("identityReceiver", config, 1024); + if (!receiver) + GTEST_SKIP() << "NIXL unavailable"; + b::ZmqControlChannel sender("identitySender"); + b::ZmqControlChannel replacement("otherIncarnation"); + ASSERT_TRUE(sender.addPeer(receiver->name, receiver->ch->localEndpoint())); + sender.sendTo(receiver->name, b::encodeWant(1, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, 1, std::chrono::seconds(5), credits)); + sender.sendTo(receiver->name, b::encodeDrain(1, replacement.localEndpoint())); + sender.sendTo(receiver->name, b::encodeWant(2, {256}, sender.localEndpoint())); + EXPECT_FALSE(waitGrant(sender, 2, std::chrono::milliseconds(200), credits)); + sender.sendTo(receiver->name, b::encodeDrain(1, sender.localEndpoint())); + ASSERT_TRUE(waitGrant(sender, 2, std::chrono::seconds(5), credits)); + ASSERT_TRUE(drainFlow(sender, receiver->name, 2)); + receiver->tx->shutdown(); +} + +TEST(BounceTransportFailure, ShutdownGuardStartsBeforeWaitingForBlockedScatter) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + GTEST_FLAG_SET(death_test_style, "threadsafe"); + EXPECT_EXIT( + { + auto config = cfg(30000); + config.copyStreamCount = 1; + config.scatterWorkerCount = 1; + auto receiver = bounce_test::makeNode("blockedReceiver", config, 1024); + b::ZmqControlChannel sender("blockedSender"); + sender.addPeer(receiver->name, receiver->ch->localEndpoint()); + auto buffers = bounce_test::makeXferBufs(1, 256, 1); + sender.sendTo(receiver->name, b::encodeWant(1, {256}, sender.localEndpoint())); + std::vector credits; + if (!waitGrant(sender, 1, std::chrono::seconds(5), credits)) + { + std::_Exit(2); + } + StreamGate gate; + auto* context = receiver->exec->tryAcquire(); + if (context == nullptr || cudaLaunchHostFunc(context->stream, waitAtGate, &gate) != cudaSuccess) + { + std::_Exit(2); + } + receiver->exec->release(context); + sender.sendTo(receiver->name, + b::encodeData(1, 0, 1, credits[0].regionHandle, + {{0, reinterpret_cast(buffers.dst), 0, 0, 256, 1}})); + sender.sendTo(receiver->name, b::encodeDrain(1, sender.localEndpoint())); + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while ((!gate.entered.load() || receiver->exec->freeCount() != 0) + && std::chrono::steady_clock::now() < deadline) + { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + if (!gate.entered.load() || receiver->exec->freeCount() != 0) + { + std::_Exit(2); + } + // MPI installs a SIGABRT handler that converts abort into exit(1). + // Restore the signal only after fixture admission has succeeded. + std::signal(SIGABRT, SIG_DFL); + std::fprintf(stderr, "awaiting physical quiescence guard\n"); + receiver->tx->shutdown(); + }, + ::testing::KilledBySignal(SIGABRT), "awaiting physical quiescence guard"); +} diff --git a/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp b/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp index 19c29bc47fca..0ae6924f80ef 100644 --- a/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp @@ -866,3 +866,35 @@ TEST(CreditScheduler, FourThreadsPerSenderCannotStarveLargeChunk) { expectLargeChunkProgressWithBoundedSenderConcurrency(/*numSenders=*/4, /*threadsPerSender=*/4); } + +TEST(CreditScheduler, UnprovedFlowNeverExpiresAndExplicitDrainReleasesOnlyItsRegions) +{ + auto s = makeSched(/*nRegions=*/2, /*maxInflightChunksPerRequest=*/2); + auto grants = s.onWant("lost", want(2)); + ASSERT_EQ(grants.size(), 2u); + std::vector deferred; + EXPECT_TRUE(s.quarantineFlow("lost", {grants[0].offset}, deferred).empty()); + ASSERT_EQ(deferred.size(), 1u); + EXPECT_TRUE(s.hasUnprovenRemoteWrites()); + EXPECT_TRUE(s.reapQuarantine().empty()); + EXPECT_EQ(s.freeBytes(), 0u); + EXPECT_TRUE(s.onWant("next", want(2)).empty()); + auto fromDrain = s.releaseQuarantinedFlow("lost"); + ASSERT_EQ(fromDrain.size(), 1u); + EXPECT_EQ(fromDrain[0].offset, grants[1].offset); + auto fromScatter = s.freeOrphanRegion(deferred[0]); + ASSERT_EQ(fromScatter.size(), 1u); + EXPECT_EQ(fromScatter[0].offset, grants[0].offset); +} + +TEST(CreditScheduler, ShutdownClosesGrantAdmissionBeforeCompletionReleasesRegions) +{ + auto s = makeSched(/*nRegions=*/1, /*maxInflightChunksPerRequest=*/1); + auto first = s.onWant("first", want(1)); + ASSERT_EQ(first.size(), 1u); + EXPECT_TRUE(s.onWant("next", want(1)).empty()); + s.closeGrantAdmission(); + EXPECT_TRUE(s.onScatterDone("first", first[0].offset).empty()); + EXPECT_TRUE(s.onWant("late", want(1)).empty()); + EXPECT_EQ(freeRegions(s), 1u); +} diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index 2e06df609aed..9ec2b4976cd4 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -161,6 +161,12 @@ cache_transceiver_config: `agent_bounce_buffer_enable` (default `false`) switches to the C++ transfer-agent bounce implementation, which uses one buffer of `kv_cache_bounce_size_mb` MiB shared by send and receive (prefer a power of two; the usable capacity rounds down to one). By default it engages only for KV writes with many small descriptors (at least 1024, averaging at most 16 KiB), i.e. head-mismatch layouts such as context TP != generation TP or DP; head-matched layouts (MLA, symmetric TP) stay on the standard NIXL path unless `agent_bounce_params` relaxes the gate (lower `min_descriptor_count`, where `0` means no minimum; raise `max_average_descriptor_size`, e.g. `4MB`; `max_average_descriptor_size: 0` routes every outbound write to standard NIXL, acting as an outbound gate off switch — it disables outbound routing only: the arena is still allocated, the handshake still advertised and inbound transfers still served; set `agent_bounce_buffer_enable: false` to turn the feature off). Set it consistently on the context and generation workers; in particular `request_timeout_ms` and the effective chunk cap must match on both sides, otherwise the pair falls back to standard NIXL with a WARNING when the peer is loaded (transfers are then routed silently). The compared chunk cap is the effective one, min(`max_chunk_size`, the arena's usable capacity): `kv_cache_bounce_size_mb` matters only when its usable capacity (rounded down to a power of two) is below `max_chunk_size`; both sides must then clamp to the same chunk cap, otherwise the pair falls back to standard NIXL. +The C++ bounce path retains source and destination KV after failure until physical completion is proved. Cancellation closes admission, drains local gathers and RDMA, and exchanges a `DRAIN`/`DRAIN_ACK` fence with the receiver after its queued and running scatters stop. A logical failure or timeout never authorizes memory reuse. This also applies to operations routed to standard NIXL: success proves completion, but an unproven backend failure retains its memory. If proof does not arrive within the retirement grace period, the worker uses the MPI/launcher hard-stop path; the entire MPI job may terminate. Use separate launcher jobs when independent availability is required. + +For this lifecycle protocol, enabling C++ bounce requires both endpoints to use matching ownership-capable builds, the C++ NIXL binding, a finite positive `kv_transfer_timeout_ms`, `TRTLLM_DISAGG_NO_RETRY=1`, PP1/CP1, and asynchronous generation-first requests with a non-negative `disagg_request_id`. Layerwise and pipelined transfer are not supported by this initial profile. Multiple TP writers and attention DP remain supported. The MPI communicator must match the executor's mapping and support `MPI_THREAD_MULTIPLE`. Unsupported profiles fail before transfer admission; an ownership mismatch fails before request addresses are published. These checks do not change selection of the default Python bounce implementation. + +At the transfer-agent API, `TransferStatus.is_quiesced()` is independent of `is_completed()`: a failed bounce operation can eventually be quiescent while remaining failed. `request_cancel()` only requests asynchronous drain; its return value is not completion evidence. Standalone C++ callers must retain their memory until quiescence, use a unique agent name for every transport incarnation, and avoid replacing a live peer's endpoint under the same name. The production Python transceiver generates UUID-based instance names. Standalone transport shutdown drains outstanding work and terminates the process if it cannot establish quiescence within its bounded shutdown window; it never frees a live bounce arena after merely logging a warning. + For example, you could launch two context servers and one generation server as follows: ```bash diff --git a/tensorrt_llm/_torch/disaggregation/base/agent.py b/tensorrt_llm/_torch/disaggregation/base/agent.py index 6d0e22e279a8..7c2c4f3b5e24 100644 --- a/tensorrt_llm/_torch/disaggregation/base/agent.py +++ b/tensorrt_llm/_torch/disaggregation/base/agent.py @@ -1,3 +1,18 @@ +# 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. + import os from abc import ABC, abstractmethod from dataclasses import dataclass @@ -83,6 +98,14 @@ class RegMemoryDescs: class TransferStatus(ABC): + def is_quiesced(self) -> bool: + """Whether all memory access has stopped; failure alone is not proof.""" + return False + + def request_cancel(self) -> bool: + """Request asynchronous cancellation without implying physical completion.""" + return False + @abstractmethod def is_completed(self) -> bool: ... @@ -181,3 +204,14 @@ def _try_load_cpp_binding(): def use_pure_python_transfer_agent() -> bool: return _use_pure_python_transfer_agent + + +def supports_native_quiescence() -> bool: + """Require the native proof API, not the conservative Python defaults.""" + return ( + not use_pure_python_transfer_agent() + and _cpp_binding is not None + and all( + hasattr(_cpp_binding.TransferStatus, name) for name in ("is_quiesced", "request_cancel") + ) + ) diff --git a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py index 48d048f4bcc4..bebcf099cee8 100644 --- a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py @@ -302,6 +302,10 @@ def has_retired_send_session(self, req: LlmRequest) -> bool: """Whether the send session closed before its final slice.""" return False + def require_request_quiescence(self, req: LlmRequest) -> None: + """Guard final executor resource release when the backend tracks owners.""" + return + @property def consumes_transfer_buffer(self) -> bool: """Return whether this runtime consumes the C++ CacheTransBuffer budget.""" diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 72bd85972e0c..e83eeb3d3dab 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -24,7 +24,7 @@ from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxTransferLayout from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy -from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo, validate_ownership_peer from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, MapperKind, PoolView from tensorrt_llm._torch.disaggregation.resource.utils import ( @@ -72,6 +72,7 @@ def __init__(self, self_rank_info: RankInfo, self_extractor: KVRegionExtractorV1 ] = {} # peer_key -> {(self_lg, self_pi) -> (peer_lg, peer_pi)} def register(self, peer_name: str, peer_rank: int, peer_ri: RankInfo): + validate_ownership_peer(self._ri, peer_ri) assert self._self_ext_cache is not None if not self._check_peer_compatible(peer_ri): raise ValueError( diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 8fbbdc6f9583..687fb6349314 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -25,6 +25,19 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes +NATIVE_OWNERSHIP_PROTOCOL_VERSION = 1 + + +def validate_ownership_peer(local: "RankInfo", peer: "RankInfo") -> None: + """Reject incompatible retirement semantics before publishing any request.""" + local_version = getattr(local, "ownership_protocol_version", 0) + peer_version = getattr(peer, "ownership_protocol_version", 0) + if local_version != peer_version or local_version not in (0, NATIVE_OWNERSHIP_PROTOCOL_VERSION): + raise ValueError( + "native KV ownership protocol mismatch: both endpoints must enable " + f"the same physical retirement protocol (local={local_version}, peer={peer_version})" + ) + @dataclass class RankInfo: @@ -48,6 +61,7 @@ class RankInfo: attention: Optional[AttentionInfo] = None aux_meta: Optional[AuxBufferMeta] = None page_table: Optional[KVCachePageTable] = None + ownership_protocol_version: int = 0 @property def tp_size_per_dp_group(self) -> int: @@ -57,6 +71,9 @@ def tp_size_per_dp_group(self) -> int: def to_bytes(self) -> bytes: data = asdict(self) + # Preserve the legacy schema when neither endpoint opts into ownership. + if self.ownership_protocol_version == 0: + data.pop("ownership_protocol_version") data["attention"] = self.attention.to_dict() if self.attention is not None else None data["aux_meta"] = self.aux_meta.to_dict() if self.aux_meta is not None else None data["page_table"] = self.page_table.to_dict() if self.page_table is not None else None diff --git a/tensorrt_llm/_torch/disaggregation/native/retirement.py b/tensorrt_llm/_torch/disaggregation/native/retirement.py index 455376d42905..f9a0709c2580 100644 --- a/tensorrt_llm/_torch/disaggregation/native/retirement.py +++ b/tensorrt_llm/_torch/disaggregation/native/retirement.py @@ -15,7 +15,7 @@ @dataclass(frozen=True) class QuiescenceFatalEvent: - """Immutable evidence that an exposed owner missed its retirement deadline.""" + """Immutable evidence of unproven retirement or a missed drain deadline.""" request_id: int direction: Literal["send", "receive"] @@ -302,6 +302,29 @@ def request_shutdown(self) -> None: for owner in self._owners: owner.request_drain("shutdown") + def reject_unproven_retirement( + self, request_id: int, direction: Literal["send", "receive"] + ) -> QuiescenceFatalEvent: + """Latch fatal containment before an executor bypasses normal drain. + + Unlike an ordinary cancellation, final resource teardown cannot wait + for later polling. Its caller must hard-stop without releasing roots. + """ + with self.lock: + self._admission_closed = True + if self.fatal is None: + now = self.clock() + self.fatal = QuiescenceFatalEvent( + request_id, + direction, + "executor attempted resource release before physical quiescence", + now, + now, + now, + ) + self.wake.set() + return self.fatal + @property def admission_closed(self) -> bool: """Return the irreversible shutdown gate, without entering any session lock.""" diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 13612283b77d..f36375b3a856 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -69,7 +69,11 @@ ) from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager -from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo +from tensorrt_llm._torch.disaggregation.native.rank_info import ( + NATIVE_OWNERSHIP_PROTOCOL_VERSION, + RankInfo, + validate_ownership_peer, +) from tensorrt_llm._torch.disaggregation.native.retirement import ( QuiescenceFatalEvent, RetirementDeadline, @@ -406,7 +410,7 @@ def finish_publication(self) -> None: self._settle_if_drained_locked() def abort_publication(self, published_writers: set[int]) -> None: - """Close a failed fan-out around the writers whose sends succeeded.""" + """Close a failed fan-out around every writer whose send was attempted.""" with self._lock: published = frozenset(published_writers) if self._writer_cohort is not None and not published.issubset(self._writer_cohort): @@ -689,6 +693,7 @@ def __init__(self, params: DisaggregatedParams): self._perf_timer = PerfTimer() if perf_log_manager.enabled else None self._physical_lock = threading.Lock() self._physical_operations: dict[int, _PhysicalOperation] = {} + self._cancel_requested = False def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: """Bind session outcome and deadline before exposing the task to workers. @@ -752,6 +757,9 @@ def begin_backend_submission( operation = self._require_physical_operation_locked( peer_rank, (_PhysicalOperationState.ADMITTED,) ) + if self._cancel_requested: + operation.state = _PhysicalOperationState.NOT_SUBMITTED + raise _TransferNotSubmittedError("source operation was cancelled before submission") if self._retirement is not None and not self._retirement.expose(operation): operation.state = _PhysicalOperationState.NOT_SUBMITTED raise _TransferNotSubmittedError("source retirement admission is closed") @@ -765,6 +773,35 @@ def record_backend_submission(self, peer_rank: int, status: object) -> None: ) operation.status = status operation.state = _PhysicalOperationState.SUBMITTED + cancel_requested = self._cancel_requested + # Submission may return after cancellation took its status snapshot. + if cancel_requested: + self._request_status_cancel(status) + + @staticmethod + def _request_status_cancel(status: object) -> None: + """Invoke the backend outside every ownership/session lock.""" + cancel = getattr(status, "request_cancel", None) + if cancel is None: + return + try: + cancel() + except Exception as error: + # Cancellation is best effort. Retained ownership and the watchdog + # still prevent reuse if the backend cannot provide completion. + logger.warning(f"Unable to request physical transfer drain: {error}") + + def request_physical_cancel(self) -> None: + """Fence new operations and asynchronously drain retained status handles.""" + with self._physical_lock: + self._cancel_requested = True + statuses = [ + operation.status + for operation in self._physical_operations.values() + if operation.status is not None + ] + for status in statuses: + self._request_status_cancel(status) def mark_physical_operation_in_doubt(self, peer_rank: int) -> None: with self._physical_lock: @@ -811,7 +848,7 @@ def retire_backend_done_physical_operation(self, peer_rank: int) -> bool: return True def poll_in_doubt_physical_operation(self, peer_rank: int) -> bool: - """Retire once, only after a fresh DONE query on the retained status. + """Retire once, only after fresh physical proof from the retained status. Polling does not change the task's logical outcome. Keep strong local roots across the query, and reject a result if its operation changed. @@ -824,7 +861,8 @@ def poll_in_doubt_physical_operation(self, peer_rank: int) -> bool: if status is None: return False try: - completed = status.is_completed() + proof = getattr(status, "is_quiesced", None) + completed = proof() if proof is not None else False except Exception: # A backend query failure is not evidence that its accessors stopped. return False @@ -1175,6 +1213,8 @@ def _submit_transfer( # SUBMITTING already retains the request. Neither backend admission # nor completion may hold the lock needed by deadline containment. status = self._agent.submit_transfer_requests(request) + if status is None: + raise RuntimeError("backend submission returned no physical operation handle") task.record_backend_submission(peer_rank, status) except Exception as error: with self._ownership_poison_lock: @@ -1182,21 +1222,24 @@ def _submit_transfer( task.mark_physical_operation_in_doubt(peer_rank) return False, str(error) try: - if not status.wait(): - # A non-success query is a logical transfer failure, but the - # current binding does not expose a backend-defined proof that - # the request handle has released every memory accessor. Treat - # it as physically unresolved until that contract is available. + # A native wait must not hide the Python request timeout. Once this + # slice expires, request C++ drain while the watchdog retains KV. + timeout_ms = ( + max(1, int(task._retirement.timeout_s * 1000)) + if task._retirement is not None + else None + ) + if not status.wait(timeout_ms=timeout_ms): + # Failure is immutable, but the retained handle can later prove + # physical completion. Quarantine this operation, not unrelated + # sessions whose memory is independently owned. detail = getattr(status, "last_status_str", lambda: "")() - error = RuntimeError(f"NIXL transfer outcome is ambiguous: {detail}") - with self._ownership_poison_lock: - self._ownership_poisoned = error - task.mark_physical_operation_in_doubt(peer_rank) + task.mark_physical_operation_in_doubt(peer_rank) + task._request_status_cancel(status) return False, detail except Exception as error: - with self._ownership_poison_lock: - self._ownership_poisoned = error - task.mark_physical_operation_in_doubt(peer_rank) + task.mark_physical_operation_in_doubt(peer_rank) + task._request_status_cancel(status) return False, str(error) if not task.retire_backend_done_physical_operation(peer_rank): return False, "quiescence deadline expired before backend completion" @@ -1963,6 +2006,7 @@ def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): torch.cuda.set_device(self._device_id) CUASSERT(cudart.cudaSetDevice(self._device_id)) ri: RankInfo = RankInfo.from_bytes(message[1]) + validate_ownership_peer(self._registrar.self_rank_info, ri) agent_name = ri.instance_name + str(ri.instance_rank) logger.debug(f"Loading remote transfer agent descriptor for peer '{agent_name}'") # NIXL serializes descriptor load against submit internally. This lock @@ -2058,6 +2102,10 @@ def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): # _sessions_lock prevents a race between session lookup and req_info save. # session.lock atomically saves peer info and snapshots tasks against send(). info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) + # Validate before saving a request or dispatching even an existing task. + # An unregistered/rejected peer must never enter the pending-info cache. + peer_info = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) + validate_ownership_peer(self._registrar.self_rank_info, peer_info) with self._sessions_lock: session = self._get_session(info.unique_rid) fenced = ( @@ -2613,6 +2661,10 @@ def cancel_local(self, by_peer: bool = False, *, report_unsubmitted_aux: bool = req_infos = list((self._sender._get_req_info(self.disagg_request_id) or {}).values()) if report_unsubmitted_aux: aux_failures = self._claim_unsubmitted_aux_failures_locked(req_infos) + tasks = self.kv_tasks + ([self.aux_task] if self.aux_task is not None else []) + if self._enforce_physical_ownership: + for task in tasks: + task.request_physical_cancel() self._report_unsubmitted_aux_failures(aux_failures) return True @@ -2753,6 +2805,10 @@ def set_exception(self, reason: str = "") -> None: self.aux_task.fail(self._exception) req_infos = list((self._sender._get_req_info(self.disagg_request_id) or {}).values()) aux_failures = self._claim_unsubmitted_aux_failures_locked(req_infos) + tasks = self.kv_tasks + ([self.aux_task] if self.aux_task is not None else []) + if self._enforce_physical_ownership: + for task in tasks: + task.request_physical_cancel() self._report_unsubmitted_aux_failures(aux_failures) @property @@ -3315,11 +3371,9 @@ def publish_requests() -> None: if task._perf_timer is not None: task._perf_timer.record_task_start(rank) published_writers.add(rank) - try: - self._request_sender_data(peer_infos.sender_endpoints[rank], payload) - except Exception: - published_writers.discard(rank) - raise + # A send may escape before raising. The attempted writer keeps + # its destination claim until settlement or session fencing. + self._request_sender_data(peer_infos.sender_endpoints[rank], payload) # Gen-first ADP publishes the destination to every eligible DP group, # but the qualified immutable-request, no-retry/no-reroute profile @@ -3385,6 +3439,7 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: # (handled in dispatch_task) so only requests targeting this peer # fail, and cached so later requests fail fast. try: + validate_ownership_peer(self._registrar.self_rank_info, sender_info) MambaPolicy.validate_peer_compatible( self._registrar.self_rank_info, sender_info, @@ -3415,7 +3470,9 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: return sender_info else: - return self._sender_ep_instance_map[info_endpoint] + sender_info = self._sender_ep_instance_map[info_endpoint] + validate_ownership_peer(self._registrar.self_rank_info, sender_info) + return sender_info def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> set[str]: """Notify senders and return the endpoints whose cancel message was sent.""" @@ -3664,11 +3721,15 @@ def _poison_receiver_ownership(self, error: Exception) -> None: def _record_ownership_evidence_error(self, error: Exception) -> None: """Record a fatal ownership-evidence error and close receiver admission.""" + self._record_unproven_transfer(error) + self._poison_receiver_ownership(error) + + def _record_unproven_transfer(self, error: Exception) -> None: + """Fail this session while retaining owners until their drain completes.""" self._logical_outcomes.fail(error) self._exception = error if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR - self._poison_receiver_ownership(error) @property def status(self) -> SessionStatus: @@ -3757,10 +3818,8 @@ def try_begin_transfer( ): raise RuntimeError("publication did not queue the complete writer cohort") except Exception: - # ZMQ delivers multipart messages atomically. A successful - # send transfers responsibility to that writer; a failed - # send does not. Retain only the successfully queued prefix - # and wait for each of those writers to settle. + # Retain every attempted publication, including a send + # which may have escaped before raising. self._receiver._bounce.abort_publication( (self.disagg_request_id, task.slice_id), published_writers, @@ -3925,7 +3984,7 @@ def process_kv_agent_result( f"slice={receiver_slice_id} peer_rank={peer_rank}" ) task.fail(error) - self._record_ownership_evidence_error(error) + self._record_unproven_transfer(error) return if status == AgentResult.FAILED_QUIESCED: if not self._enforce_physical_ownership: @@ -4084,7 +4143,7 @@ def process_aux_agent_result(self, peer_rank: int, status: AgentResult): f"{self.request_id} peer_rank={peer_rank}" ) self._aux_status = TaskStatus.ERROR - self._record_ownership_evidence_error(error) + self._record_unproven_transfer(error) return if status == AgentResult.FAILED_QUIESCED: if self._aux_physical_owner is None: @@ -4514,6 +4573,10 @@ class TransferWorker: def __init__(self, config: TransferWorkerConfig): self._config = config self._retirement_watchdog = None + if config.agent_buffer_size_mb > 0 and ( + not config.enforce_physical_ownership or config.quiescence_fatal_callback is None + ): + raise ValueError("C++ bounce requires physical ownership and qualified fail-stop") if config.quiescence_fatal_callback is not None: if not config.enforce_physical_ownership: raise ValueError("deadline retirement requires physical ownership") @@ -4533,6 +4596,9 @@ def __init__(self, config: TransferWorkerConfig): config.device_id, self._aux_buffer.meta if self._aux_buffer is not None else None, ) + self._rank_info.ownership_protocol_version = ( + NATIVE_OWNERSHIP_PROTOCOL_VERSION if config.enforce_physical_ownership else 0 + ) self._setup_peer_infrastructure(kvm) self._setup_transfer_engine() if self._retirement_watchdog is not None: @@ -4554,6 +4620,23 @@ def request_shutdown(self) -> None: if watchdog is not None: watchdog.request_shutdown() + # Arm containment before any native cancel. Do not wait for a + # receiver's publication lock to send peer notifications here: + # shutdown must reach its ownership preflight even if publication + # is blocked. Ordinary cancellation still notifies all writers. + # Both directions keep their roots until the ordinary retirement + # checks observe proof; cancellation itself never releases memory. + for endpoint in (getattr(self, "_sender", None), getattr(self, "_receiver", None)): + if endpoint is None: + continue + with endpoint._sessions_lock: + sessions = list(endpoint._sessions.values()) + for session in sessions: + if isinstance(session, weakref.ReferenceType): + session = session() + if session is not None: + session.cancel_local() + def populate_instance_and_rank_info(self, endpoints: list[str], layer_num_per_pp: list[int]): assert self._rank_info is not None self._rank_info.sender_endpoints = endpoints @@ -4599,6 +4682,11 @@ def _setup_peer_infrastructure(self, kvm: KVCacheManager): self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor) def _setup_transfer_engine(self): + if self._config.enforce_physical_ownership: + from tensorrt_llm._torch.disaggregation.base.agent import supports_native_quiescence + + if not supports_native_quiescence(): + raise ValueError("physical ownership requires C++ quiescence-capable bindings") torch.cuda.set_device(self._config.device_id) CUASSERT(cudart.cudaSetDevice(self._config.device_id)) mapping = self._config.kv_cache_manager.mapping diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py index 68a8dce2ce59..29bad46f6c21 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py @@ -43,6 +43,14 @@ def is_completed(self) -> bool: """Check if transfer is completed (releases GIL).""" return self._cpp_status.is_completed() + def is_quiesced(self) -> bool: + """Poll physical completion independently of the logical outcome.""" + return self._cpp_status.is_quiesced() + + def request_cancel(self) -> bool: + """Ask the backend to drain; callers must still poll is_quiesced().""" + return self._cpp_status.request_cancel() + @nvtx_range("BindingsNixlTransferStatus.wait") def wait(self, timeout_ms=None) -> bool: """Wait for transfer to complete (releases GIL).""" diff --git a/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py b/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py index acd4d55a5930..949d2c0c75d1 100644 --- a/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py +++ b/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py @@ -410,7 +410,7 @@ def send_completed_context(self, requests: List[LlmRequest]) -> None: """Start async KV sends for finished context-only requests.""" # Do not send more chunks after an in-flight cancellation. cancel_pending_ids = set(self._registry.canceled_request_ids()) - bridge_enabled = getattr(self._transceiver, "_fp4_mla_bridge_enabled", False) is True + bridge_enabled = getattr(self._transceiver, "_enforce_physical_ownership", False) is True for req in requests: if not req.is_context_only_request or req.is_finished_due_to_cancellation: continue diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index b2b4eb46b5ba..337182762e4f 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -12,6 +12,7 @@ # 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. +import math import os import threading import time @@ -105,8 +106,8 @@ def _fail_unproven_kv_transfer( ) -> None: """Terminate the executor world without running ordinary resource cleanup. - Called by the independent retirement watchdog after fatal expiry is sticky - and transfer admission is closed. The existing MPI/launcher crash path + Called after the watchdog or final executor release guard latches fatal + containment and closes transfer admission. The MPI/launcher crash path makes the endpoint unhealthy; request failure is not memory-release proof. The communicator identifies the executor, not an isolated failure domain: Open MPI aborts the entire job, including other executors in split worlds. @@ -230,6 +231,31 @@ def _validate_fp4_mla_bridge_profile( return True +def _validate_agent_bounce_ownership_profile( + mapping: Mapping, config: CacheTransceiverConfig +) -> bool: + """Qualify the lifecycle assumptions required by native bounce draining.""" + if not config.agent_bounce_buffer_enable or config.kv_cache_bounce_size_mb <= 0: + return False + requirements = { + "C++ NIXL binding": not use_pure_python_transfer_agent(), + "finite positive kv_transfer_timeout_ms": ( + config.kv_transfer_timeout_ms is not None + and math.isfinite(config.kv_transfer_timeout_ms) + and config.kv_transfer_timeout_ms > 0 + ), + "TRTLLM_DISAGG_NO_RETRY=1": os.getenv("TRTLLM_DISAGG_NO_RETRY", "0") == "1", + "PP1/CP1": mapping.pp_size == 1 and mapping.cp_size == 1, + "non-pipelined transfer": not config.enable_pipelined_transfer, + "async transfer": os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") != "1", + "non-layerwise transfer": os.getenv("TRTLLM_DISAGG_LAYERWISE") != "1", + } + missing = [name for name, supported in requirements.items() if not supported] + if missing: + raise ValueError("C++ bounce physical ownership requires " + ", ".join(missing)) + return True + + class KvCacheTransceiverV2(KvCacheTransceiver): @property def consumes_transfer_buffer(self) -> bool: @@ -261,10 +287,16 @@ def __init__( else None ) self._check_compatible() - enforce_physical_ownership = _validate_fp4_mla_bridge_profile( + self._fp4_mla_bridge_enabled = _validate_fp4_mla_bridge_profile( mapping, kv_cache_manager, cache_transceiver_config ) - self._fp4_mla_bridge_enabled = enforce_physical_ownership + self._agent_bounce_ownership_enabled = _validate_agent_bounce_ownership_profile( + mapping, cache_transceiver_config + ) + enforce_physical_ownership = ( + self._fp4_mla_bridge_enabled or self._agent_bounce_ownership_enabled + ) + self._enforce_physical_ownership = enforce_physical_ownership self._retirement_mpi_comm = ( _retirement_executor_comm(mapping) if enforce_physical_ownership else None ) @@ -321,10 +353,10 @@ def __init__( ) if enforce_physical_ownership: logger.info( - "FP4 MLA KV ownership bridge ENABLED: " + "KV physical ownership ENABLED: " f"rank={rank}/{mapping.world_size}, backend=NIXL, runtime=PYTHON, " - "request_schedule_required=GENERATION_FIRST, attention_dp=True, pp=1, cp=1, " - "retry=False, async=True, layerwise=False, bounce_mb=0, " + "request_schedule_required=GENERATION_FIRST, pp=1, cp=1, " + "retry=False, async=True, layerwise=False, " "unproven_retirement=fail_stop, termination_scope=mpi_job_or_launcher, " f"kv_transfer_timeout_ms={self.kv_transfer_timeout_ms}" ) @@ -625,7 +657,7 @@ def _need_aux_transfer(req: LlmRequest) -> bool: return params is not None and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST def _validate_bridge_req(self, req: LlmRequest, synchronous: bool = False) -> bool: - if not getattr(self, "_fp4_mla_bridge_enabled", False): + if not getattr(self, "_enforce_physical_ownership", False): return True params = req.py_disaggregated_params rid = None if params is None else params.disagg_request_id @@ -637,7 +669,7 @@ def _validate_bridge_req(self, req: LlmRequest, synchronous: bool = False) -> bo or rid < 0 ): logger.error( - "FP4 MLA lifecycle bridge requires async generation-first requests " + "KV physical ownership requires async generation-first requests " "with a non-negative integer disagg_request_id" ) req.state = LlmRequestState.DISAGG_TRANS_ERROR @@ -989,6 +1021,34 @@ def has_retired_send_session(self, req: LlmRequest) -> bool: """Whether req's send session was torn down before its last slice.""" return req.py_kv_send_session_retired and get_unique_rid(req) not in self._send_sessions + def require_request_quiescence(self, req: LlmRequest) -> None: + """Protect KV release even when an unrelated executor error bypasses polling.""" + if not getattr(self, "_enforce_physical_ownership", False): + return + rid = get_unique_rid(req) + for direction, sessions in ( + ("send", self._send_sessions), + ("receive", self._recv_sessions), + ): + session = sessions.get(rid) + if session is None: + continue + # Close future admission before checking physical claims. This path + # is an executor teardown, not the normal asynchronous cancellation. + session.cancel_local() + if self._ownership_blocks_retirement(session) or session.close() is False: + watchdog = self._transfer_worker._retirement_watchdog + assert watchdog is not None + event = watchdog.reject_unproven_retirement(req.py_request_id, direction) + self._fail_unproven_transfer(event) + # A test/faulty callback must never authorize the caller's free. + raise RuntimeError("executor resource release refused: KV is not quiescent") + if direction == "send": + self._retire_send_session(rid, req, session_already_closed=True) + else: + self._recv_sessions.pop(rid, None) + self._recv_reqs.pop(rid, None) + def _build_prefill_extent( self, req: LlmRequest, @@ -1057,7 +1117,7 @@ def respond_and_send_async(self, req: LlmRequest) -> None: session = self._get_or_create_send_session(req) if session is None: return - bridge_enabled = getattr(self, "_fp4_mla_bridge_enabled", False) + bridge_enabled = getattr(self, "_enforce_physical_ownership", False) if bridge_enabled: # Root the request before backend admission so its source pages # outlive every ownership-enabled NIXL operation. @@ -1091,6 +1151,7 @@ def respond_and_send_async(self, req: LlmRequest) -> None: except Exception as error: if bridge_enabled: cast(Any, session).set_exception(f"transfer admission failed: {error}") + return raise @nvtx_range("KvCacheTransceiverV2.request_and_receive_sync") @@ -1200,12 +1261,15 @@ def request_and_receive_async(self, req: LlmRequest) -> None: # The handle that comes back is the contract's answer about this piece. What retires # the request is the sweep over the session tables, as it was before. fetches.fetch(extent, expected_write_bytes=chunk_bytes) - except Exception: + except Exception as error: # No session means no publication and nothing the sweep could ever pair the request # with, so the registration made here is undone here and the request goes terminal. if self._legacy_session(fetches) is None: del self._recv_reqs[rid] req.state = LlmRequestState.DISAGG_TRANS_ERROR + elif getattr(self, "_enforce_physical_ownership", False): + self._legacy_session(fetches).fail_admission(error) + return raise finally: # The session exists even when publication failed, and the legacy sweep owns it. @@ -1295,7 +1359,7 @@ def check_context_transfer_status( # DP ranks (entries that will never have a TxSession created for them). self._transfer_worker.sweep_stale_req_infos() - if getattr(self, "_fp4_mla_bridge_enabled", False): + if getattr(self, "_enforce_physical_ownership", False): # CtxTransferStatus has no cancellation channel. In the qualified # no-retry bridge, a globally quiesced cancellation is a terminal # send failure that the executor must consume to release its diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8f8d158cf0a2..ae265ec8b534 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -8266,6 +8266,9 @@ def _terminate_request(self, request: LlmRequest) -> None: def _free_request_resources(self, request: LlmRequest) -> None: """Release execution resources without removing response routing.""" + transceiver = getattr(self, "kv_cache_transceiver", None) + if transceiver is not None: + transceiver.require_request_quiescence(request) self.resource_manager.free_resources(request) self._prefetched_request_ids.discard(request.py_request_id) self.disagg.forget_request(request.py_request_id) diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py index 2e9fb0275eb3..4ba07328d52f 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py @@ -75,18 +75,19 @@ def test_send_releases_the_index_slot_on_the_target_and_draft_kv_managers() -> N @pytest.mark.parametrize( - ("bridge_enabled", "has_inflight", "releases_claim"), + ("ownership_enabled", "has_inflight", "releases_claim"), [(True, False, True), (True, True, False), (False, False, False)], ) -def test_bridge_rejection_releases_the_claim_only_without_a_physical_owner( - bridge_enabled: bool, has_inflight: bool, releases_claim: bool +def test_owned_transfer_rejection_releases_the_claim_only_without_a_physical_owner( + ownership_enabled: bool, has_inflight: bool, releases_claim: bool ) -> None: - """The FP4 MLA bridge may reject a send before any transfer session exists. + """An ownership-enabled transceiver may reject a send before any transfer session exists. Only then is there no physical accessor for the reap to poll, so the claim is released on the spot; a rejection with a live session, or a failure - without the bridge, stays claimed until the reap and the error path.""" + without ownership enforcement, stays claimed until the reap and the error path.""" h = _Harness(kv_transfer_timeout_ms=1000) - h.transceiver._fp4_mla_bridge_enabled = bridge_enabled + h.transceiver._fp4_mla_bridge_enabled = False + h.transceiver._enforce_physical_ownership = ownership_enabled h.transceiver.has_inflight_transfer = lambda _req: has_inflight h.transceiver.respond_and_send_async = lambda req: setattr( req, "state", LlmRequestState.DISAGG_TRANS_ERROR diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 0c6e83bdae42..fc9ea3d493b4 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -82,6 +82,7 @@ from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.bindings.internal.batch_manager import ReqIdsSet from tensorrt_llm.bindings.internal.testing import simulate_prefill_completion_only_use_for_testing +from tensorrt_llm.disaggregated_params import DisaggScheduleStyle from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, CacheTransceiverConfig, KvCacheConfig AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType @@ -1299,6 +1300,8 @@ def run_transfer_test( if expect_cpp_bounce else None, ) + if expect_cpp_bounce: + config.kv_transfer_timeout_ms = 30_000 ctx_tcs = create_instance_transceivers( ctx_tp, ctx_pp, ctx_enable_dp, ctx_managers, config, is_mla ) @@ -1341,7 +1344,14 @@ def run_transfer_test( is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, ) - ctx_request.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + schedule_style = ( + DisaggScheduleStyle.GENERATION_FIRST + if expect_cpp_bounce + else DisaggScheduleStyle.CONTEXT_FIRST + ) + ctx_request.py_disaggregated_params = DisaggregatedParams( + disagg_request_id=unique_rid, schedule_style=schedule_style + ) gen_request = LlmRequest( request_id=gen_rid, @@ -1358,6 +1368,7 @@ def run_transfer_test( ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, + schedule_style=schedule_style, ) if gen_scratch_tokens > 0: # Mark the request as MTP so its draft length matches the @@ -1421,6 +1432,19 @@ def run_transfer_test( gen_kv_caches[rank].append(kv) gen_tcs[rank].request_and_receive_async(req) + # Generation-first seals the AUX peer snapshot only after the scheduler + # observes every REQUEST_DATA. Model that admission boundary here too. + if expect_cpp_bounce: + ready_deadline = time.monotonic() + 10 + for rank in range(ctx_world): + for _, req in ctx_handle_map[rank]: + rid = get_unique_rid(req) + while not ctx_tcs[rank]._transfer_worker.has_all_peer_req_infos_for_send(rid): + assert time.monotonic() < ready_deadline, ( + "generation-first peer publication did not complete" + ) + time.sleep(0.01) + # 6. ctx send after for rank in range(ctx_world): for req_idx, req in ctx_handle_map[rank]: @@ -1990,10 +2014,10 @@ def test_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_pp() -> None: pytest.param(1, 1, 1, 1, False, True, False, id="v2_mha"), pytest.param(2, 1, 2, 1, True, True, False, id="v2_mla_tp2"), pytest.param(1, 1, 1, 1, True, False, True, id="v1_dsa_indexer"), - pytest.param(2, 1, 1, 2, False, True, False, id="v2_ctx_tp2_gen_pp2"), ], ) def test_python_nixl_cache_transceiver_uses_cpp_bounce( + monkeypatch: pytest.MonkeyPatch, ctx_tp: int, ctx_pp: int, gen_tp: int, @@ -2014,6 +2038,21 @@ def test_python_nixl_cache_transceiver_uses_cpp_bounce( from tensorrt_llm.tensorrt_llm_transfer_agent_binding import TransferStatus assert hasattr(TransferStatus, "get_last_status_str") + assert hasattr(TransferStatus, "is_quiesced") + assert hasattr(TransferStatus, "request_cancel") + monkeypatch.setenv("TRTLLM_DISAGG_NO_RETRY", "1") + # ThreadSafeDistributed models several logical ranks inside one process. + # Keep the real ownership/watchdog path while isolating its MPI kill route. + fatal_events = [] + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.transceiver._retirement_executor_comm", + lambda _mapping: None, + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_fail_unproven_transfer", + lambda _self, event: fatal_events.append(event), + ) run_transfer_test( ctx_tp=ctx_tp, @@ -2028,6 +2067,26 @@ def test_python_nixl_cache_transceiver_uses_cpp_bounce( enable_indexer_k_cache=enable_indexer_k_cache, expect_cpp_bounce=True, ) + assert not fatal_events + + +@pytest.mark.cpu_only +def test_cpp_bounce_rejects_pipeline_parallelism_before_agent_creation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.disaggregation.transceiver import ( + _validate_agent_bounce_ownership_profile, + ) + + monkeypatch.setenv("TRTLLM_DISAGG_NO_RETRY", "1") + config = CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + kv_cache_bounce_size_mb=64, + agent_bounce_buffer_enable=True, + ) + with pytest.raises(ValueError, match="PP1/CP1"): + _validate_agent_bounce_ownership_profile(Mapping(world_size=2, pp_size=2), config) # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py index 6ad0e67290ba..59a145a81295 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py +++ b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py @@ -329,7 +329,7 @@ def test_peer_loss_fails_closed_only_for_unsettled_exposure(exposure: str) -> No def test_failed_sender_write_retains_roots_until_fatal_deadline() -> None: - """A persistent backend error closes sender admission, not physical ownership.""" + """A persistent backend error retains ownership until bounded containment.""" now = [10.0] contain = Mock() watchdog = RetirementWatchdog(contain, clock=lambda: now[0]) @@ -338,7 +338,7 @@ def test_failed_sender_write_retains_roots_until_fatal_deadline() -> None: task.bind_logical_outcomes(transfer_mod._LogicalOutcomes(owner)) assert task.begin_physical_operation(7) request = object() - status = Mock(wait=Mock(return_value=False), is_completed=Mock(return_value=False)) + status = Mock(wait=Mock(return_value=False), is_quiesced=Mock(return_value=False)) status.last_status_str.return_value = "ERROR" sender = object.__new__(transfer_mod.Sender) sender._shutdown = True # No worker threads or transport were created by this fixture. @@ -355,13 +355,12 @@ def test_failed_sender_write_retains_roots_until_fatal_deadline() -> None: assert operation.status is status assert not task.resources_drained - # Sender poisoning rejects a different request even before grace expires. + # A retained proof-capable handle does not poison independent requests. later = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=13)) assert later.begin_physical_operation(8) - with pytest.raises(transfer_mod._TransferNotSubmittedError, match="before backend submission"): - sender._submit_transfer(later, 8, object()) - sender._agent.submit_transfer_requests.assert_called_once_with(request) - assert later._physical_operations[8].state is transfer_mod._PhysicalOperationState.NOT_SUBMITTED + sender._agent.submit_transfer_requests.return_value = Mock(wait=Mock(return_value=True)) + assert sender._submit_transfer(later, 8, object()) == (True, None) + assert sender._ownership_poisoned is None assert later.resources_drained now[0] = 10.5 diff --git a/tests/unittest/disaggregated/test_retirement_deadline.py b/tests/unittest/disaggregated/test_retirement_deadline.py index ce530f83c5cc..a5f7ab7e7d11 100644 --- a/tests/unittest/disaggregated/test_retirement_deadline.py +++ b/tests/unittest/disaggregated/test_retirement_deadline.py @@ -195,9 +195,9 @@ def test_blocked_backend_cannot_delay_fatal_deadline(blocked_call: str) -> None: task.bind_logical_outcomes(transfer_mod._LogicalOutcomes(retirement)) assert task.begin_physical_operation(7) request = object() - status = Mock(wait=Mock(return_value=True), is_completed=Mock(return_value=True)) + status = Mock(wait=Mock(return_value=True), is_quiesced=Mock(return_value=True)) - def block() -> object: + def block(timeout_ms: int | None = None) -> object: """Block only the external backend boundary, never production state logic.""" entered.set() assert release.wait(5) @@ -214,7 +214,7 @@ def block() -> object: task.begin_backend_submission(7, request) task.record_backend_submission(7, status) task.mark_physical_operation_in_doubt(7) - status.is_completed.side_effect = block + status.is_quiesced.side_effect = block work = partial(task.poll_in_doubt_physical_operation, 7) else: if blocked_call == "submit": @@ -351,6 +351,10 @@ def _receiver(watchdog: RetirementWatchdog) -> transfer_mod.Receiver: receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None receiver._bounce = Mock() + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) + ) + receiver._get_or_connect_dealer = Mock() return receiver diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index f9f86fed340b..2330de0c4f30 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -834,6 +834,7 @@ def _construct_worker_config(monkeypatch, cache_config) -> TransferWorkerConfig: # Everything the constructor can reach, not only what it reaches with these values: the # world-size and helix reads sit behind an env check, a monkeypatch and `cp_size == 1`. mapping = SimpleNamespace( + rank=0, cp_size=1, world_size=1, pp_size=1, @@ -929,6 +930,7 @@ def test_transceiver_routes_bounce_capacity_to_one_implementation( expected_python_bounce, expected_buffer_size_mb: int, ) -> None: + monkeypatch.setenv("TRTLLM_DISAGG_NO_RETRY", "1") worker_config = _construct_worker_config( monkeypatch, _make_cache_config( diff --git a/tests/unittest/disaggregated/test_transfer_late_settlement.py b/tests/unittest/disaggregated/test_transfer_late_settlement.py index 4dc9c93b1f40..37bfc436749d 100644 --- a/tests/unittest/disaggregated/test_transfer_late_settlement.py +++ b/tests/unittest/disaggregated/test_transfer_late_settlement.py @@ -77,9 +77,9 @@ def test_late_settlement_requires_positive_retained_status(query: object) -> Non status = None if query is None else Mock() if status is not None: if isinstance(query, Exception): - status.is_completed.side_effect = query + status.is_quiesced.side_effect = query else: - status.is_completed.return_value = query + status.is_quiesced.return_value = query task = _in_doubt_task(status) operation = task._physical_operations[7] request = operation.request @@ -101,7 +101,7 @@ def replace() -> bool: task._physical_operations[7].status = replacement return True - status.is_completed.side_effect = replace + status.is_quiesced.side_effect = replace assert not task.poll_in_doubt_physical_operation(7) assert not task.resources_drained assert task._physical_operations[7].status is replacement @@ -116,7 +116,7 @@ def done() -> bool: barrier.wait() return True - status.is_completed.side_effect = done + status.is_quiesced.side_effect = done task = _in_doubt_task(status) original_error = RuntimeError("original failure") task.fail(original_error) @@ -142,7 +142,7 @@ def poll(index: int) -> None: @pytest.mark.cpu_only def test_late_done_does_not_retire_active_sibling() -> None: - task = _in_doubt_task(Mock(is_completed=Mock(return_value=True))) + task = _in_doubt_task(Mock(is_quiesced=Mock(return_value=True))) assert task.begin_physical_operation(8) assert task.poll_in_doubt_physical_operation(7) assert not task.resources_drained @@ -240,7 +240,7 @@ def test_sender_late_done_reports_physical_failure_only( sender._sessions[401] = session if not auxiliary: task.expected_transfers = 1 - status = Mock(wait=Mock(return_value=False), is_completed=Mock(return_value=False)) + status = Mock(wait=Mock(return_value=False), is_quiesced=Mock(return_value=False)) status.last_status_str.return_value = "ERROR" sender._agent = SimpleNamespace(name="nixl", submit_transfer_requests=lambda request: status) request = SimpleNamespace(op="WRITE", remote_name="gen") @@ -274,7 +274,7 @@ def test_sender_late_done_reports_physical_failure_only( assert not task.resources_drained sender._poll_in_doubt_transfers(0) assert dealer.send.call_count == 1 - status.is_completed.return_value = True + status.is_quiesced.return_value = True sender._poll_in_doubt_transfers(0) sender._poll_in_doubt_transfers(0) @@ -283,7 +283,7 @@ def test_sender_late_done_reports_physical_failure_only( assert task._exception is original_error if handle is not None: assert isinstance(handle.poll(), Failed) - assert sender._ownership_poisoned is not None + assert sender._ownership_poisoned is None assert not sender._pending_settlements[0] assert dealer.send.call_count == 2 messages = [call.args[0] for call in dealer.send.call_args_list] @@ -303,7 +303,9 @@ def test_sender_late_done_reports_physical_failure_only( @pytest.mark.cpu_only @pytest.mark.parametrize("auxiliary", [False, True]) -def test_receiver_late_settlement_retains_logical_failure_and_quarantine(auxiliary: bool) -> None: +def test_receiver_late_settlement_retains_logical_failure_without_poisoning_worker( + auxiliary: bool, +) -> None: receiver = object.__new__(transfer_mod.Receiver) receiver._enforce_physical_ownership = True receiver._sessions_lock, receiver._sessions = threading.Lock(), {} @@ -340,7 +342,7 @@ def report(status: transfer_mod.AgentResult) -> None: assert session.resources_drained() assert session.status is SessionStatus.ERROR assert session.exception is error - assert receiver._ownership_poisoned is not None + assert receiver._ownership_poisoned is None receiver._bounce.record_failure.assert_called_once_with((401, 0), 7) assert session.close() @@ -349,7 +351,7 @@ def report(status: transfer_mod.AgentResult) -> None: def test_settlement_retries_in_order_without_releasing_twice() -> None: sender = _sender() task = transfer_mod.KVSendTask(_chunk(), _params(), slice_id=0) - status = Mock(is_completed=Mock(return_value=True)) + status = Mock(is_quiesced=Mock(return_value=True)) assert task.begin_physical_operation(7) task.begin_backend_submission(7, object()) task.record_backend_submission(7, status) @@ -372,7 +374,7 @@ def test_settlement_retries_in_order_without_releasing_twice() -> None: dealer.send.side_effect = [RuntimeError("initial send"), None, RuntimeError("late send"), None] sender._get_result_dealer = Mock(return_value=dealer) sender._retain_in_doubt_transfer(meta, initial, send_slot_id=12) - status.is_completed.assert_not_called() + status.is_quiesced.assert_not_called() sender._poll_in_doubt_transfers(0) assert task.resources_drained assert sender._pending_settlements[0] @@ -382,7 +384,7 @@ def test_settlement_retries_in_order_without_releasing_twice() -> None: sender._poll_in_doubt_transfers(0) assert not sender._pending_settlements[0] sender._bounce.release_send.assert_called_once_with(12) - status.is_completed.assert_called_once() + status.is_quiesced.assert_called_once() assert task.transferred_count == 1 codes = [ transfer_mod._KV_RESULT_PREFIX.unpack(call.args[0][1])[4] @@ -400,7 +402,7 @@ def test_committed_session_outcome_survives_late_settlement( direction: str, terminal_kind: str, poll_early: bool ) -> None: """Integration seam with the event-time logical-outcome commit change.""" - status = Mock(is_completed=Mock(return_value=True)) + status = Mock(is_quiesced=Mock(return_value=True)) if direction == "send": sender = Mock() sender._enforce_physical_ownership = True @@ -474,7 +476,7 @@ def test_worker_polls_retained_status_when_queue_is_idle( sender = _sender() sender._thread_local = threading.local() task = transfer_mod.KVSendTask(_chunk(), _params(), slice_id=0) - status = Mock(is_completed=Mock(side_effect=[False, True])) + status = Mock(is_quiesced=Mock(side_effect=[False, True])) assert task.begin_physical_operation(7) task.begin_backend_submission(7, object()) task.record_backend_submission(7, status) @@ -519,6 +521,6 @@ def send(_message: list[bytes]) -> None: assert task.resources_drained assert not sender._pending_settlements[0] assert dealer.send.call_count == (4 if initial_send_fails else 2) - assert status.is_completed.call_count == 2 + assert status.is_quiesced.call_count == 2 assert task.status is transfer_mod.TaskStatus.ERROR sender._shutdown = True diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..c9557005a910 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -31,6 +31,11 @@ from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.bounce.core import TransferContext +from tensorrt_llm._torch.disaggregation.native.rank_info import ( + NATIVE_OWNERSHIP_PROTOCOL_VERSION, + RankInfo, + validate_ownership_peer, +) from tensorrt_llm._torch.disaggregation.native.transfer import ( AgentResult, KVRecvTask, @@ -44,6 +49,31 @@ from tensorrt_llm.bindings import DataType, LlmRequestState from tensorrt_llm.disaggregated_params import DisaggScheduleStyle +pytestmark = pytest.mark.cpu_only + + +class _PhysicalStatusProbe: + """Keep logical failure fixed while independently advancing physical proof.""" + + def __init__(self) -> None: + self.quiesced = False + self.cancel_count = 0 + + def wait(self, timeout_ms: int | None = None) -> bool: + return False + + def is_completed(self) -> bool: + return False + + def is_quiesced(self) -> bool: + return self.quiesced + + def request_cancel(self) -> None: + self.cancel_count += 1 + + def last_status_str(self) -> str: + return "failed while scatter is active" + def _sole_piece(block_ids_per_layer_groups=None, *, tokens=0, is_last=True) -> Chunk: """A transfer's only piece; nothing in this file exercises chunking.""" @@ -172,6 +202,11 @@ def record_result( ) self._advance() + def record_failure(self, rid_slice: tuple[int, int], peer_rank: int) -> None: + assert self.context is not None and self.context.rid_slice == rid_slice + self.context.record_writer_result(peer_rank, succeeded=False) + self._advance() + def _start_checked_thread( target: Callable[[], None], @@ -749,8 +784,8 @@ def request_sender_data(endpoint: str, _payload: bytes) -> None: assert bounce.release_count == 0 assert not session.resources_drained() - # Only writer 0's REQUEST_DATA was successfully queued. The destination - # remains owned until that writer reports terminal physical evidence. + # Writer 1's send may have escaped before raising, so both writers stay + # owned until each supplies physical evidence. session.process_kv_agent_result( peer_rank=0, receiver_slice_id=0, @@ -764,6 +799,10 @@ def request_sender_data(endpoint: str, _payload: bytes) -> None: assert session.status == SessionStatus.ERROR assert not session.resources_drained() session.process_aux_agent_result(0, AgentResult.FAILED) + assert not session.resources_drained() + session.process_kv_agent_result(1, 0, True, AgentResult.FAILED) + assert not session.resources_drained() + session.process_aux_agent_result(1, AgentResult.FAILED) assert session.resources_drained() assert bounce.context is None assert bounce.scatter_count == 0 @@ -1460,6 +1499,7 @@ def test_pre_cancelled_sender_does_not_publish_from_transceiver() -> None: ) transceiver = object.__new__(KvCacheTransceiverV2) transceiver._fp4_mla_bridge_enabled = True + transceiver._enforce_physical_ownership = True transceiver._enable_pipelined_transfer = False transceiver._ever_had_send_session = False transceiver._ctx_need_tp_sync = False @@ -1542,6 +1582,7 @@ def test_bridge_rejection_releases_only_without_physical_owner( ) transceiver = SimpleNamespace( _fp4_mla_bridge_enabled=bridge_enabled, + _enforce_physical_ownership=bridge_enabled, has_retired_send_session=lambda _req: False, respond_and_send_async=lambda req: setattr( req, "state", LlmRequestState.DISAGG_TRANS_ERROR @@ -1649,6 +1690,10 @@ def test_late_request_data_to_terminal_sender_settles_aux( sender = _make_owned_sender() sender._save_peer_req_info = Mock() sender._send_failed_result_to_receiver = Mock() + sender._registrar = SimpleNamespace( + self_rank_info=_ownership_rank_info(1), + get_peer_rank_info=Mock(return_value=_ownership_rank_info(1)), + ) session = SimpleNamespace( lock=threading.Lock(), _closed=False, @@ -1872,7 +1917,9 @@ def test_legacy_aux_dispatch_remains_fail_fast() -> None: @pytest.mark.parametrize( "meta_type", [transfer_mod.WriteMetaType.KV, transfer_mod.WriteMetaType.AUX] ) -@pytest.mark.parametrize("failure_mode", ["false", "submit_exception", "wait_exception"]) +@pytest.mark.parametrize( + "failure_mode", ["false", "submit_exception", "wait_exception", "null_status"] +) def test_ambiguous_sender_result_retains_source_and_reports_in_doubt( monkeypatch: pytest.MonkeyPatch, meta_type: transfer_mod.WriteMetaType, @@ -1893,10 +1940,14 @@ def test_ambiguous_sender_result_retains_source_and_reports_in_doubt( def submit(_request): if failure_mode == "submit_exception": raise RuntimeError("submit failed") + if failure_mode == "null_status": + return None return status if failure_mode == "submit_exception": detail = "submit failed" + elif failure_mode == "null_status": + detail = "backend submission returned no physical operation handle" sender._agent = SimpleNamespace( name="nixl", submit_transfer_requests=submit, @@ -1959,12 +2010,14 @@ def submit(_request): dealer.send.assert_called_once() assert task.status == transfer_mod.TaskStatus.ERROR assert not task.resources_drained - expected_status = None if failure_mode == "submit_exception" else status + expected_status = None if failure_mode in ("submit_exception", "null_status") else status operation = task._physical_operations[peer_rank] assert operation.state is transfer_mod._PhysicalOperationState.IN_DOUBT assert operation.request is request assert operation.status is expected_status - assert sender._ownership_poisoned is not None + assert (sender._ownership_poisoned is not None) == ( + failure_mode in ("submit_exception", "null_status") + ) @pytest.mark.cpu_only @@ -2004,12 +2057,13 @@ def test_receiver_in_doubt_retains_destination_and_closes_admission( assert session.status == SessionStatus.ERROR assert not session.resources_drained() assert session.close() is False - with pytest.raises(RuntimeError, match="admission is quarantined"): - RxSession( - request_id=rid + 1, - params=DisaggregatedParams(disagg_request_id=rid + 1), - receiver=receiver, - ) + assert receiver._ownership_poisoned is None + unrelated = RxSession( + request_id=rid + 1, + params=DisaggregatedParams(disagg_request_id=rid + 1), + receiver=receiver, + ) + assert unrelated.close() # This test intentionally models evidence that never settles. session._closed = True @@ -2068,6 +2122,7 @@ def test_sender_registration_does_not_wait_for_active_transfer( ) -> None: sender = _make_owned_sender() sender._device_id, sender._registrar, sender._agent = 0, Mock(), Mock() + sender._registrar.self_rank_info = _ownership_rank_info(0) task = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=1)) assert task.begin_physical_operation(7) sender._sessions[1] = SimpleNamespace(kv_tasks=[task]) @@ -2250,7 +2305,7 @@ def test_sender_retires_only_after_backend_done() -> None: done = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=92)) done_request = Mock() - def wait() -> bool: + def wait(timeout_ms: int | None = None) -> bool: operation = done._physical_operations[7] assert operation.state is transfer_mod._PhysicalOperationState.SUBMITTED assert operation.request is done_request @@ -2453,6 +2508,7 @@ def test_fp4_mla_bridge_roots_send_and_receive_requests_before_admission() -> No ) sender = object.__new__(KvCacheTransceiverV2) sender._fp4_mla_bridge_enabled = True + sender._enforce_physical_ownership = True sender._enable_pipelined_transfer = False sender._send_reqs = {} @@ -2485,6 +2541,7 @@ def send_after_rooting(_chunk: Chunk) -> None: ) receiver = object.__new__(KvCacheTransceiverV2) receiver._fp4_mla_bridge_enabled = True + receiver._enforce_physical_ownership = True receiver._recv_sessions = {} receiver._recv_reqs = {} receiver._kv_size_rank_factor = 1 @@ -2512,6 +2569,7 @@ def receive_after_rooting(_chunk: Chunk, expected_write_bytes=None) -> None: def test_fp4_mla_bridge_rejects_requests_outside_qualified_protocol() -> None: transceiver = object.__new__(KvCacheTransceiverV2) transceiver._fp4_mla_bridge_enabled = True + transceiver._enforce_physical_ownership = True valid = SimpleNamespace( schedule_style=DisaggScheduleStyle.GENERATION_FIRST, disagg_request_id=1, @@ -2556,3 +2614,604 @@ def test_fp4_mla_bridge_rejects_requests_outside_qualified_protocol() -> None: ) transceiver.request_and_receive_async(req) assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("proof", [None, False, 1, "done", RuntimeError("query failed")]) +def test_logical_terminal_status_without_physical_proof_retains_roots(proof: object) -> None: + task = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=300)) + request = object() + status = SimpleNamespace(is_completed=Mock(return_value=True)) + if proof is not None: + status.is_quiesced = ( + Mock(side_effect=proof) if isinstance(proof, RuntimeError) else Mock(return_value=proof) + ) + assert task.begin_physical_operation(0) + task.begin_backend_submission(0, request) + task.record_backend_submission(0, status) + task.mark_physical_operation_in_doubt(0) + task.fail(RuntimeError("immutable failure")) + + assert not task.poll_in_doubt_physical_operation(0) + assert not task.resources_drained + assert task._physical_operations[0].request is request + assert task._physical_operations[0].status is status + status.is_completed.assert_not_called() + + +@pytest.mark.cpu_only +def test_failed_bounce_can_drain_without_success_or_poisoning_next_request() -> None: + sender = _make_owned_sender() + status = _PhysicalStatusProbe() + sender._agent = SimpleNamespace(submit_transfer_requests=lambda _request: status) + failed = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=301)) + request = object() + assert failed.begin_physical_operation(0) + + delivered, detail = sender._submit_transfer(failed, 0, request) + failed.fail(RuntimeError(detail)) + assert not delivered + assert sender._ownership_poisoned is None + assert status.cancel_count > 0 + assert not failed.poll_in_doubt_physical_operation(0) + assert not failed.resources_drained + + status.quiesced = True + assert not status.is_completed() + assert failed.poll_in_doubt_physical_operation(0) + assert failed.resources_drained + assert failed.status is transfer_mod.TaskStatus.ERROR + assert failed._physical_operations[0].request is None + assert failed._physical_operations[0].status is None + assert not failed.poll_in_doubt_physical_operation(0) + + independent = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=302)) + sender._agent = SimpleNamespace( + submit_transfer_requests=lambda _request: SimpleNamespace(wait=lambda timeout_ms=None: True) + ) + assert independent.begin_physical_operation(0) + assert sender._submit_transfer(independent, 0, object()) == (True, None) + assert independent.resources_drained + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("late_registration", [False, True]) +def test_cancellation_reaches_retained_or_late_status_without_ownership_lock( + late_registration: bool, +) -> None: + task = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=303)) + status = _PhysicalStatusProbe() + calls = [] + + def cancel() -> None: + assert task._physical_lock.acquire(blocking=False), "backend called under ownership lock" + task._physical_lock.release() + calls.append(True) + + status.request_cancel = cancel + assert task.begin_physical_operation(0) + task.begin_backend_submission(0, object()) + if late_registration: + task.request_physical_cancel() + assert calls == [] + task.record_backend_submission(0, status) + else: + task.record_backend_submission(0, status) + task.request_physical_cancel() + assert calls == [True] + assert not task.resources_drained + assert task.begin_physical_operation(1) + with pytest.raises(transfer_mod._TransferNotSubmittedError): + task.begin_backend_submission(1, object()) + assert task._physical_operations[1].state is transfer_mod._PhysicalOperationState.NOT_SUBMITTED + + +@pytest.mark.cpu_only +def test_two_writer_failure_cannot_reuse_kv_or_aux_until_every_physical_owner_drains() -> None: + receiver = _ReceiverProbe() + session = RxSession( + request_id=304, + params=DisaggregatedParams( + disagg_request_id=304, + schedule_style=DisaggScheduleStyle.GENERATION_FIRST, + ), + receiver=receiver, + ) + task = session.prepare_receive(_sole_piece()) + assert task is not None + task.expected_transfers = 2 + assert session.try_begin_transfer(task.slice_id, set(), writer_cohort={0, 1}) + allocation = _OneSlotAllocator(304) + session.process_kv_agent_result(0, 0, True, AgentResult.IN_DOUBT) + session.process_kv_agent_result(1, 0, True, AgentResult.IN_DOUBT) + session.process_aux_agent_result(0, AgentResult.IN_DOUBT) + session.process_aux_agent_result(1, AgentResult.FAILED) + assert session.status is SessionStatus.ERROR + allocation.apply_reuse_decision(session.close()) + assert not allocation.is_reusable + + session.process_kv_agent_result(0, 0, True, AgentResult.FAILED_QUIESCED) + allocation.apply_reuse_decision(session.close()) + assert not allocation.is_reusable + assert not task.resources_drained + session.process_kv_agent_result(1, 0, True, AgentResult.FAILED_QUIESCED) + assert task.resources_drained + allocation.apply_reuse_decision(session.close()) + assert not allocation.is_reusable + assert session._aux_physical_owner is not None + assert not session._aux_physical_owner.resources_drained + + session.process_aux_agent_result(0, AgentResult.FAILED_QUIESCED) + assert session.resources_drained() + allocation.apply_reuse_decision(session.close()) + assert allocation.is_reusable + assert allocation.release_count == 1 + assert receiver.clear_count == 1 + + +def _ownership_rank_info(version: int, name: str = "peer") -> RankInfo: + return RankInfo( + instance_name=name, + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[1], + sender_endpoints=["tcp://sender"], + self_endpoint="tcp://self", + transfer_engine_info=b"opaque-agent-descriptor", + ownership_protocol_version=version, + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("local,peer", [(0, 1), (1, 0), (1, 2), (2, 2)]) +def test_ownership_capability_mismatch_is_rejected(local: int, peer: int) -> None: + with pytest.raises(ValueError, match="ownership protocol mismatch"): + validate_ownership_peer(_ownership_rank_info(local), _ownership_rank_info(peer)) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("version", [0, NATIVE_OWNERSHIP_PROTOCOL_VERSION]) +def test_rank_info_ownership_roundtrip_preserves_legacy_schema(version: int) -> None: + import msgpack + + info = _ownership_rank_info(version) + encoded = info.to_bytes() + assert ("ownership_protocol_version" in msgpack.unpackb(encoded)) == (version != 0) + restored = RankInfo.from_bytes(encoded) + assert restored == info + validate_ownership_peer(info, restored) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("local,peer", [(0, 1), (1, 0)]) +def test_sender_rejects_capability_mismatch_before_registering_remote_agent( + monkeypatch: pytest.MonkeyPatch, local: int, peer: int +) -> None: + sender = _make_owned_sender() + sender._device_id = 0 + sender._registrar = SimpleNamespace(self_rank_info=_ownership_rank_info(local), register=Mock()) + sender._agent = SimpleNamespace(load_remote_agent=Mock()) + monkeypatch.setattr(transfer_mod.torch.cuda, "set_device", Mock()) + monkeypatch.setattr(transfer_mod.cudart, "cudaSetDevice", Mock()) + monkeypatch.setattr(transfer_mod, "CUASSERT", Mock()) + + with pytest.raises(ValueError, match="ownership protocol mismatch"): + sender._register_peer_rank( + b"sender", [b"REGISTER_RANK_INFO", _ownership_rank_info(peer).to_bytes()] + ) + sender._registrar.register.assert_not_called() + sender._agent.load_remote_agent.assert_not_called() + assert not sender._loaded_remote_agents + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("local,peer", [(0, 1), (1, 0)]) +def test_cached_peer_cannot_bypass_receive_ownership_preflight(local: int, peer: int) -> None: + receiver = _make_owned_receiver() + receiver._incompatible_peers = {} + receiver._extract_info_endpoint = Mock(return_value="tcp://cached") + receiver._should_register_peer = Mock(return_value=False) + receiver._sender_ep_instance_map = {"tcp://cached": _ownership_rank_info(peer)} + receiver._registrar = SimpleNamespace(self_rank_info=_ownership_rank_info(local)) + receiver._get_or_connect_dealer = Mock() + with pytest.raises( + (ValueError, transfer_mod.PeerIncompatibleError), match="ownership protocol mismatch" + ): + receiver._get_sender_info(DisaggregatedParams(disagg_request_id=305)) + receiver._get_or_connect_dealer.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("local,peer", [(0, 1), (1, 0)]) +def test_receive_capability_preflight_precedes_address_publication( + monkeypatch: pytest.MonkeyPatch, local: int, peer: int +) -> None: + receiver = _make_owned_receiver() + receiver._incompatible_peers = {} + receiver._extract_info_endpoint = Mock(return_value="tcp://new-peer") + receiver._should_register_peer = Mock(return_value=True) + receiver._sender_ep_instance_map = {} + receiver._registrar = SimpleNamespace(self_rank_info=_ownership_rank_info(local)) + receiver._get_or_connect_dealer = Mock() + messenger = Mock(receive=Mock(return_value=[_ownership_rank_info(peer).to_bytes()])) + monkeypatch.setattr(transfer_mod, "ZMQMessenger", Mock(return_value=messenger)) + with pytest.raises(transfer_mod.PeerIncompatibleError, match="ownership protocol mismatch"): + receiver._get_sender_info(DisaggregatedParams(disagg_request_id=306)) + messenger.send.assert_called_once_with([MessageType.REQUEST_INSTANCE_INFO]) + messenger.stop.assert_called_once() + receiver._get_or_connect_dealer.assert_not_called() + assert not receiver._sender_ep_instance_map + + +@pytest.mark.cpu_only +def test_unregistered_request_data_cannot_create_pending_physical_owner() -> None: + sender = _make_owned_sender() + sender._registrar = SimpleNamespace( + self_rank_info=_ownership_rank_info(1), + get_peer_rank_info=Mock(side_effect=KeyError("unregistered peer")), + ) + sender._save_peer_req_info = Mock() + sender._get_session = Mock() + info = transfer_mod.RecvReqInfo( + sender_req_id=18, + instance_name="unregistered", + instance_rank=2, + block_ids_per_layer_groups=[], + unique_rid=307, + ) + with pytest.raises(KeyError, match="unregistered peer"): + sender._respond_with_kv(b"receiver", [MessageType.REQUEST_DATA, info.to_bytes()]) + sender._save_peer_req_info.assert_not_called() + sender._get_session.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("pure_python", [False, True]) +@pytest.mark.parametrize( + "methods", [(), ("is_quiesced",), ("request_cancel",), ("is_quiesced", "request_cancel")] +) +def test_native_quiescence_requires_both_native_methods( + monkeypatch: pytest.MonkeyPatch, pure_python: bool, methods: tuple[str, ...] +) -> None: + import tensorrt_llm._torch.disaggregation.base.agent as agent_mod + + binding_status = SimpleNamespace(**{name: Mock() for name in methods}) + monkeypatch.setattr(agent_mod, "_cpp_binding", SimpleNamespace(TransferStatus=binding_status)) + monkeypatch.setattr(agent_mod, "_use_pure_python_transfer_agent", pure_python) + assert agent_mod.supports_native_quiescence() == (not pure_python and len(methods) == 2) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "violation", + [None, "retry", "pp", "cp", "pipeline", "sync", "layerwise", "pure_python", "timeout"], +) +def test_agent_bounce_ownership_profile_is_checked_before_publication( + monkeypatch: pytest.MonkeyPatch, violation: str | None +) -> None: + monkeypatch.setenv("TRTLLM_DISAGG_NO_RETRY", "1") + monkeypatch.delenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", raising=False) + monkeypatch.delenv("TRTLLM_DISAGG_LAYERWISE", raising=False) + monkeypatch.setattr( + transceiver_mod, "use_pure_python_transfer_agent", lambda: violation == "pure_python" + ) + mapping = SimpleNamespace( + tp_size=4, pp_size=2 if violation == "pp" else 1, cp_size=2 if violation == "cp" else 1 + ) + config = SimpleNamespace( + agent_bounce_buffer_enable=True, + kv_cache_bounce_size_mb=64, + kv_transfer_timeout_ms=0 if violation == "timeout" else 1000, + enable_pipelined_transfer=violation == "pipeline", + ) + if violation == "retry": + monkeypatch.delenv("TRTLLM_DISAGG_NO_RETRY") + if violation == "sync": + monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") + if violation == "layerwise": + monkeypatch.setenv("TRTLLM_DISAGG_LAYERWISE", "1") + if violation is None: + assert transceiver_mod._validate_agent_bounce_ownership_profile(mapping, config) + else: + with pytest.raises(ValueError, match=r"C\+\+ bounce physical ownership requires"): + transceiver_mod._validate_agent_bounce_ownership_profile(mapping, config) + config.agent_bounce_buffer_enable = False + assert not transceiver_mod._validate_agent_bounce_ownership_profile(mapping, config) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("auxiliary", [False, True]) +def test_failed_quiesced_report_retries_in_order_and_retires_once(auxiliary: bool) -> None: + sender = _make_owned_sender() + sender._bounce = Mock() + params = DisaggregatedParams(disagg_request_id=308) + task = ( + transfer_mod.AuxSendTask(params, slot=0) + if auxiliary + else transfer_mod.KVSendTask(_sole_piece(), params, slice_id=0) + ) + status = _PhysicalStatusProbe() + assert task.begin_physical_operation(0) + task.begin_backend_submission(0, object()) + task.record_backend_submission(0, status) + task.mark_physical_operation_in_doubt(0) + task.fail(RuntimeError("failed")) + meta = transfer_mod.WriteMeta( + task=task, + expected_transfers=1, + peer_name="gen0", + peer_rank=0, + peer_endpoint="receiver", + unique_rid=308, + src_ptrs=np.array([0x1000], dtype=np.int64), + dst_ptrs=np.array([0x2000], dtype=np.int64), + sizes=np.array([0x100], dtype=np.int64), + slice_id=0, + receiver_slice_id=0, + is_last_slice=True, + meta_type=transfer_mod.WriteMetaType.AUX if auxiliary else transfer_mod.WriteMetaType.KV, + ) + initial = ( + transfer_mod._make_aux_result_msg(0, 308, AgentResult.IN_DOUBT) + if auxiliary + else transfer_mod._make_kv_result_msg(0, 308, 0, True, AgentResult.IN_DOUBT) + ) + dealer = Mock() + dealer.send.side_effect = [ + RuntimeError("initial send escaped"), + None, + RuntimeError("settlement send escaped"), + None, + ] + sender._get_result_dealer = Mock(return_value=dealer) + sender._retain_in_doubt_transfer(meta, initial, send_slot_id=4) + assert not task.resources_drained + status.quiesced = True + sender._poll_in_doubt_transfers(0) + assert task.resources_drained + assert sender._pending_settlements[0] + sender._bounce.release_send.assert_called_once_with(4) + sender._poll_in_doubt_transfers(0) + sender._poll_in_doubt_transfers(0) + messages = [call.args[0] for call in dealer.send.call_args_list] + assert messages[0] == messages[1] == initial + assert messages[2] == messages[3] + if auxiliary: + assert messages[3][-1] == AgentResult.FAILED_QUIESCED.value.encode("ascii") + assert task._transfer_count == 1 + else: + result_code = transfer_mod._KV_RESULT_PREFIX.unpack(messages[3][1])[4] + assert transfer_mod._AGENT_RESULT_BY_CODE[result_code] is AgentResult.FAILED_QUIESCED + assert task.transferred_count == 1 + assert not sender._pending_settlements[0] + assert task.status is transfer_mod.TaskStatus.ERROR + sender._bounce.release_send.assert_called_once_with(4) + + +@pytest.mark.parametrize("direction", ["send", "receive"]) +@pytest.mark.parametrize("physically_drained", [False, True]) +def test_executor_free_cannot_bypass_unquiesced_session( + direction: str, physically_drained: bool +) -> None: + from tensorrt_llm._torch.disaggregation.native.retirement import RetirementWatchdog + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + request = SimpleNamespace( + py_request_id=309, + request_id=309, + py_disaggregated_params=DisaggregatedParams(disagg_request_id=309), + ) + watchdog = RetirementWatchdog(Mock()) + session = SimpleNamespace( + _enforce_physical_ownership=True, + cancel_local=Mock(), + resources_drained=Mock(return_value=physically_drained), + close=Mock(return_value=False), + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._enforce_physical_ownership = True + transceiver._send_sessions = {309: session} if direction == "send" else {} + transceiver._recv_sessions = {309: session} if direction == "receive" else {} + transceiver._transfer_worker = SimpleNamespace(_retirement_watchdog=watchdog) + transceiver._fail_unproven_transfer = Mock() + executor = SimpleNamespace( + kv_cache_transceiver=transceiver, + resource_manager=Mock(), + _prefetched_request_ids={309}, + disagg=Mock(), + ) + + with pytest.raises(RuntimeError, match="executor resource release refused"): + PyExecutor._free_request_resources(executor, request) + executor.resource_manager.free_resources.assert_not_called() + executor.disagg.forget_request.assert_not_called() + assert executor._prefetched_request_ids == {309} + assert watchdog.fatal is not None + assert watchdog.fatal.direction == direction + session.cancel_local.assert_called_once() + assert session.close.called == physically_drained + transceiver._fail_unproven_transfer.assert_called_once_with(watchdog.fatal) + + +def test_executor_free_fences_and_closes_drained_sessions_before_reuse() -> None: + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + order = [] + request = SimpleNamespace( + py_request_id=310, + request_id=310, + py_disaggregated_params=DisaggregatedParams(disagg_request_id=310), + ) + session = SimpleNamespace( + _enforce_physical_ownership=True, + cancel_local=lambda: order.append("cancel"), + resources_drained=lambda: True, + close=lambda: order.append("close") or True, + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._enforce_physical_ownership = True + transceiver._send_sessions, transceiver._recv_sessions = {}, {310: session} + transceiver._send_reqs, transceiver._recv_reqs = {}, {310: request} + executor = SimpleNamespace( + kv_cache_transceiver=transceiver, + resource_manager=SimpleNamespace(free_resources=lambda _request: order.append("free")), + _prefetched_request_ids={310}, + disagg=SimpleNamespace(forget_request=lambda _rid: order.append("forget")), + ) + PyExecutor._free_request_resources(executor, request) + assert order == ["cancel", "close", "free", "forget"] + assert not executor._prefetched_request_ids + assert not transceiver._recv_sessions + assert not transceiver._recv_reqs + + +def test_old_binding_is_rejected_before_creating_or_publishing_transfer_agent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tensorrt_llm._torch.disaggregation.base.agent as agent_mod + + worker = object.__new__(transfer_mod.TransferWorker) + worker._config = SimpleNamespace(enforce_physical_ownership=True) + create_agent = Mock() + set_device = Mock() + monkeypatch.setattr(agent_mod, "_use_pure_python_transfer_agent", False) + monkeypatch.setattr(agent_mod, "_cpp_binding", SimpleNamespace(TransferStatus=object)) + monkeypatch.setattr(transfer_mod, "_create_nixl_agent", create_agent) + monkeypatch.setattr(transfer_mod.torch.cuda, "set_device", set_device) + with pytest.raises(ValueError, match="quiescence-capable bindings"): + worker._setup_transfer_engine() + create_agent.assert_not_called() + set_device.assert_not_called() + assert not hasattr(worker, "_agent") + + +def test_generic_bounce_send_admission_exception_stays_rooted_until_drain() -> None: + rid = 311 + request = SimpleNamespace( + py_request_id=rid, + request_id=rid, + py_disaggregated_params=DisaggregatedParams( + disagg_request_id=rid, schedule_style=DisaggScheduleStyle.GENERATION_FIRST + ), + state=LlmRequestState.CONTEXT_INIT, + set_kv_cache_transfer_start=Mock(), + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._fp4_mla_bridge_enabled = False + transceiver._enforce_physical_ownership = True + transceiver._enable_pipelined_transfer = False + transceiver._send_reqs = {} + transceiver._ctx_need_tp_sync = transceiver._ctx_need_pp_sync = False + transceiver._transfer_worker = SimpleNamespace(sweep_stale_req_infos=Mock()) + transceiver._ctx_consensus = lambda rids: rids + transceiver._ctx_consensus_outcome = lambda _rids, cancelled, failed, completed, drained: ( + cancelled, + failed, + completed, + drained, + ) + drained = [False] + session = SimpleNamespace( + _enforce_physical_ownership=True, + status=SessionStatus.READY, + disagg_request_id=rid, + kv_tasks=[], + is_completed=lambda: False, + has_failed=lambda: session.status is SessionStatus.ERROR, + has_transferring_tasks=lambda: not drained[0], + resources_drained=lambda: drained[0], + wait_complete=lambda blocking=True: WaitResult.FAILED, + close=Mock(return_value=True), + ) + + def send(_chunk: Chunk) -> None: + assert transceiver._send_reqs[rid] is request + raise RuntimeError("backend admission may have escaped") + + def fail(_reason: str) -> None: + session.status = SessionStatus.ERROR + + session.send = send + session.set_exception = Mock(side_effect=fail) + transceiver._send_sessions = {rid: session} + transceiver._get_or_create_send_session = Mock(return_value=session) + transceiver._create_cache_extent = Mock(return_value=CacheExtent(name=rid, local=_sole_piece())) + transceiver._finalize_send = Mock() + transceiver.respond_and_send_async(request) + session.set_exception.assert_called_once_with( + "transfer admission failed: backend admission may have escaped" + ) + assert request.state is LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + assert transceiver.check_context_transfer_status(None) == ([], []) + assert transceiver._send_reqs[rid] is request + assert transceiver._send_sessions[rid] is session + session.close.assert_not_called() + drained[0] = True + assert transceiver.check_context_transfer_status(None) == ([], [rid]) + assert not transceiver._send_reqs and not transceiver._send_sessions + assert request.py_kv_send_session_retired is True + session.close.assert_called_once() + transceiver._finalize_send.assert_not_called() + + +def test_generic_bounce_receive_admission_exception_stays_rooted_until_drain() -> None: + rid = 312 + request = SimpleNamespace( + py_request_id=rid, + request_id=rid, + py_disaggregated_params=DisaggregatedParams( + disagg_request_id=rid, schedule_style=DisaggScheduleStyle.GENERATION_FIRST + ), + state=LlmRequestState.GENERATION_IN_PROGRESS, + set_kv_cache_transfer_start=Mock(), + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._fp4_mla_bridge_enabled = False + transceiver._enforce_physical_ownership = True + transceiver._recv_reqs, transceiver._recv_sessions = {}, {} + transceiver._kv_size_rank_factor = 1 + transceiver._gen_need_sync = False + transceiver._gen_allgather = Mock(side_effect=AssertionError("single rank must not gather")) + transceiver._mapping = SimpleNamespace(pp_size=1, world_size=1, enable_attention_dp=False) + transceiver._dist = SimpleNamespace(rank=0) + drained = [False] + session = SimpleNamespace( + _enforce_physical_ownership=True, + status=SessionStatus.READY, + _kv_tasks=[], + is_completed=lambda: False, + has_failed=lambda: session.status is SessionStatus.ERROR, + resources_drained=lambda: drained[0], + wait_complete=lambda blocking=True: WaitResult.FAILED, + close=Mock(return_value=True), + ) + admission_error = RuntimeError("destination publication may have escaped") + + def receive(_chunk: Chunk, expected_write_bytes=None) -> None: + assert transceiver._recv_reqs[rid] is request + raise admission_error + + def fail(_error: Exception) -> None: + session.status = SessionStatus.ERROR + + session.receive = receive + session.fail_admission = Mock(side_effect=fail) + transceiver._transfer_worker = SimpleNamespace(create_rx_session=Mock(return_value=session)) + transceiver._create_cache_extent = Mock(return_value=CacheExtent(name=rid, local=_sole_piece())) + transceiver._chunk_num_bytes = Mock(return_value=0) + transceiver.request_and_receive_async(request) + session.fail_admission.assert_called_with(admission_error) + assert all(call.args == (admission_error,) for call in session.fail_admission.call_args_list) + assert request.state is LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + assert transceiver.check_gen_transfer_status(None) == ([], [], []) + assert transceiver._recv_reqs[rid] is request + assert transceiver._recv_sessions[rid] is session + session.close.assert_not_called() + drained[0] = True + assert transceiver.check_gen_transfer_status(None) == ([], [rid], []) + assert not transceiver._recv_reqs and not transceiver._recv_sessions + assert request.state is LlmRequestState.DISAGG_TRANS_ERROR + session.close.assert_called_once() diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index 09eb9bfa2296..21d8bfc79d5b 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -58,8 +58,13 @@ def _sender(*, ownership: bool = True) -> transfer_mod.Sender: sender._ownership_poison_lock = threading.Lock() sender._agent = Mock() sender._registrar = Mock() + sender._registrar.self_rank_info = SimpleNamespace(ownership_protocol_version=1) sender._registrar.get_peer_rank_info.return_value = SimpleNamespace( - instance_name="gen", instance_rank=0, dp_rank=0, self_endpoint="tcp://gen:1234" + instance_name="gen", + instance_rank=0, + dp_rank=0, + self_endpoint="tcp://gen:1234", + ownership_protocol_version=1, ) sender._registrar.get_peer_overlap.return_value = SimpleNamespace(ranks=[0]) return sender @@ -182,7 +187,7 @@ def run_worker() -> None: session.kv_tasks.append(task) assert task.begin_physical_operation(0) task.begin_backend_submission(0, object()) - task.record_backend_submission(0, Mock(is_completed=Mock(return_value=True))) + task.record_backend_submission(0, Mock(is_quiesced=Mock(return_value=True))) task.mark_physical_operation_in_doubt(0) empty = np.empty(0, dtype=np.int64) meta = transfer_mod.WriteMeta(task, 1, "gen", 0, "tcp://gen:1234", 401, empty, empty, empty) @@ -309,7 +314,7 @@ def test_cancel_does_not_acknowledge_while_a_physical_write_is_unproven() -> Non assert task.begin_physical_operation(0) request = object() task.begin_backend_submission(0, request) - status = Mock(is_completed=Mock(return_value=False)) + status = Mock(is_quiesced=Mock(return_value=False)) task.record_backend_submission(0, status) task.mark_physical_operation_in_doubt(0) @@ -325,7 +330,7 @@ def test_cancel_does_not_acknowledge_while_a_physical_write_is_unproven() -> Non assert sender._get_session(401) is session assert task._physical_operations[0].request is request - status.is_completed.return_value = True + status.is_quiesced.return_value = True assert task.poll_in_doubt_physical_operation(0) _assert_session_ack(sender) assert session.close() @@ -351,7 +356,7 @@ def test_session_ack_waits_for_admitted_and_submitting_pieces(piece: str, state: if state == "admitted": task.retire_unsubmitted_physical_operation(0) else: - task.record_backend_submission(0, Mock(is_completed=Mock(return_value=True))) + task.record_backend_submission(0, Mock(is_quiesced=Mock(return_value=True))) assert task.retire_backend_done_physical_operation(0) _assert_session_ack(sender) assert session.close()