diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 01adf276f878..4a134ac4f1d5 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -999,7 +999,8 @@ class WindowBlockManager void storeNewBlock(GenerationRequest& sequence, OptionalRef llmRequest); - //! \brief Pin blocks associated with a sequence to prevent eviction. + //! \brief Pin each distinct physical block associated with a sequence once. + //! \details Blocks shared by multiple beams acquire one pin reference, not one per beam. void pinBlocks(GenerationRequest& sequence); //! \brief Release blocks of the sequence. @@ -1256,7 +1257,8 @@ class WindowBlockManager [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys); - //! \brief Unpin blocks by block ids directly + //! \brief Unpin blocks by block ids directly. + //! \details Every id must be unique and identify an outstanding pin reference. void unpinBlocksById(std::vector const& blockIds); void truncateBlocks(LlmRequest::VecTokens const& targetTokens, SizeType32 numTokensToKeep); @@ -1555,8 +1557,9 @@ class BlockManager void schedulingReleaseBlocks(LlmRequest::RequestIdType requestId); - /// @brief Pin all blocks associated with a sequence across all window managers. + /// @brief Pin each distinct physical block associated with a sequence once across all window managers. /// @param sequence The generation request whose blocks should be pinned. + /// @note Block-ID unpinning currently supports single-window managers only. void pinBlocks(GenerationRequest& sequence); void unpinBlocksById(std::vector const& blockIds); @@ -2214,6 +2217,8 @@ class BaseKVCacheManager std::vector const& blockKeys, SizeType32 windowSize) = 0; + /// @brief Consume one outstanding pin reference for each unique block id. + /// @note Callers must not pass duplicate ids or ids without an outstanding pin. virtual void unpinBlocksById(std::vector const& blockIds) = 0; /// @brief Release cached blocks for a token sequence beyond a given prefix length. diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 266cdafc8d06..336c7d526068 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -42,6 +42,7 @@ #include #include #include +#include #include namespace tc = tensorrt_llm::common; @@ -2978,11 +2979,22 @@ void BlockManager::unpinBlocksById(std::vector const& bloc void WindowBlockManager::pinBlocks(GenerationRequest& sequence) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId); + // Shared beams repeat the same BlockPtr in mAllocatedBlocksPerSeq. One + // reference is sufficient to pin that physical block. The matching unpin + // operation must likewise provide that physical block ID only once. + std::unordered_set pinnedBlocks; + pinnedBlocks.reserve(allocatedBlocks.size()); for (auto& block : allocatedBlocks) { - block->incRefCount(); + // Placeholders have no addressable pool slot or non-negative block ID, + // so they cannot produce a token accepted by unpinBlocksById. + if (!block->isPlaceholder() && pinnedBlocks.insert(block.get()).second) + { + block->incRefCount(); + } } } @@ -2993,18 +3005,36 @@ void WindowBlockManager::unpinBlocksById(std::vector const return; } + std::lock_guard lock(mLookupTree->getMutex()); + + // Validate the complete operation before changing any refcount. Raw block + // IDs do not identify which refcount contribution is a pin token, so a + // duplicate could consume a live sequence or another owner's reference. + std::unordered_set seenBlockIds; + std::vector blocksToUnpin; + seenBlockIds.reserve(blockIds.size()); + blocksToUnpin.reserve(blockIds.size()); for (auto const& blockId : blockIds) { TLLM_CHECK_WITH_INFO(blockId >= 0 && static_cast(blockId) < mAllBlocksById.size(), "Block id %d is out of range", blockId); + auto const inserted = seenBlockIds.insert(blockId).second; + TLLM_CHECK_WITH_INFO(inserted, "Duplicate block id %d in one unpin operation", blockId); auto block = mAllBlocksById[blockId]; if (block && block->getBlockId() != KVCacheBlock::kCachedBlocksRootId) { - block->decRefCount(); - if (!block->hasRefs()) - { - mEvictionPolicy->releaseBlock(block); - } + TLLM_CHECK_WITH_INFO( + block->hasRefs(), "Can't unpin block (id=%d) that is not allocated", static_cast(blockId)); + blocksToUnpin.push_back(block); + } + } + + for (auto const& block : blocksToUnpin) + { + block->decRefCount(); + if (!block->hasRefs()) + { + mEvictionPolicy->releaseBlock(block); } } } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 98c2232b4062..a9746b52f360 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -4903,16 +4903,86 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById) ASSERT_TRUE(lastBlockIdOpt.has_value()); auto const& allBlockIds = kvCacheManager.getCacheBlockIds(requestId, maxAttentionWindow)[0]; std::vector pinnedBlockIds(allBlockIds.begin(), allBlockIds.end()); + ASSERT_FALSE(pinnedBlockIds.empty()); + + // Raw block IDs are not owner-scoped pin handles. A duplicate must fail + // before it can consume the live sequence reference. + EXPECT_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), pinnedBlockIds.front()}), std::runtime_error); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest); (void) kvCacheManager.removeSequence(requestId, llmRequest); auto const freeAfterRemovePinned = kvCacheManager.getNumFreeBlocks(); EXPECT_LT(freeAfterRemovePinned, totalBlocks); - kvCacheManager.unpinBlocksById(pinnedBlockIds); + auto invalidBlockIds = pinnedBlockIds; + invalidBlockIds.push_back(totalBlocks); + EXPECT_THROW(kvCacheManager.unpinBlocksById(invalidBlockIds), std::runtime_error); + + auto const freeBlockId = [&]() -> SizeType32 + { + for (SizeType32 blockId = 0; blockId < totalBlocks; ++blockId) + { + if (std::find(pinnedBlockIds.begin(), pinnedBlockIds.end(), blockId) == pinnedBlockIds.end()) + { + return blockId; + } + } + return totalBlocks; + }(); + ASSERT_LT(freeBlockId, totalBlocks); + EXPECT_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), freeBlockId}), std::runtime_error); + EXPECT_NO_THROW(kvCacheManager.unpinBlocksById(pinnedBlockIds)); auto const freeAfterUnpin = kvCacheManager.getNumFreeBlocks(); EXPECT_EQ(freeAfterUnpin, totalBlocks); } +TEST_F(KVCacheManagerTest, PinAndUnpinSharedBeamBlocksOncePerPhysicalBlock) +{ + using namespace tensorrt_llm::batch_manager::kv_cache_manager; + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 16; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimaryPool = 8; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 8; + auto const stream = std::make_shared(); + + auto constexpr beamWidth = 2; + auto const maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool; + BlocksPerWindow const blocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, + maxAttentionWindow, maxAttentionWindow, true); + kvCacheManager.allocatePools(false); + + LlmRequest::RequestIdType requestId{0}; + auto inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7}); + tr::SamplingConfig const samplingConfig{beamWidth}; + auto llmRequest = std::make_shared(requestId, 0, inputTokens, samplingConfig, false); + kvCacheManager.addSequenceBatch( + {{{requestId, static_cast(inputTokens->size()), beamWidth}}}, {std::ref(*llmRequest)}); + + kvCacheManager.pinBlocks(requestId); + auto const& perBeamBlockIds = kvCacheManager.getCacheBlockIds(requestId, maxAttentionWindow); + std::vector perBeamPinnedBlockIds; + for (auto const& beamBlockIds : perBeamBlockIds) + { + perBeamPinnedBlockIds.insert(perBeamPinnedBlockIds.end(), beamBlockIds.begin(), beamBlockIds.end()); + } + ASSERT_FALSE(perBeamPinnedBlockIds.empty()); + std::set const uniquePinnedBlockIds(perBeamPinnedBlockIds.begin(), perBeamPinnedBlockIds.end()); + EXPECT_LT(uniquePinnedBlockIds.size(), perBeamPinnedBlockIds.size()); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest); + (void) kvCacheManager.removeSequence(requestId, llmRequest); + EXPECT_LT(kvCacheManager.getNumFreeBlocks(), kvCacheManager.getMaxNumBlocks()); + + EXPECT_NO_THROW(kvCacheManager.unpinBlocksById( + std::vector(uniquePinnedBlockIds.begin(), uniquePinnedBlockIds.end()))); + EXPECT_EQ(kvCacheManager.getNumFreeBlocks(), kvCacheManager.getMaxNumBlocks()); +} + // Regression test for NVBug 6018647: storeBlocks(pin=true) on a zero-ref block // that sits in the eviction free queue must call claimBlock() before incRefCount(). // Without the fix, unpinBlocksById inserts the block into the free queue a second diff --git a/docs/design/README.md b/docs/design/README.md new file mode 100644 index 000000000000..d5451aa02563 --- /dev/null +++ b/docs/design/README.md @@ -0,0 +1,13 @@ + + +# Design Notes + +This directory contains draft engineering design notes and implementation +plans. They are review artifacts, not part of the published Sphinx user +documentation under `docs/source`. Each note records its own status and target +scope. + +- [Python native disaggregated KV transfer ownership and lifecycle](disagg-kv-transfer-ownership/README.md) diff --git a/docs/design/disagg-kv-transfer-ownership/README.md b/docs/design/disagg-kv-transfer-ownership/README.md new file mode 100644 index 000000000000..e1d091308f0d --- /dev/null +++ b/docs/design/disagg-kv-transfer-ownership/README.md @@ -0,0 +1,1619 @@ + + +# Python Native Disaggregated KV Transfer Ownership and Lifecycle + +| Field | Value | +|---|---| +| **Owner** | Chien-Chun Hung | +| **Status** | Draft design contract | +| **Created** | 2026-07-08 | +| **Last updated** | 2026-07-21 | +| **Implementation baseline** | [PR #15618](https://github.com/NVIDIA/TensorRT-LLM/pull/15618), merged as [`3bf37d2`](https://github.com/NVIDIA/TensorRT-LLM/commit/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d) | +| **Primary target** | Python native KV cache transceiver with NIXL, direct and bounce paths | +| **Detailed plan** | [`implementation-plan.md`](implementation-plan.md) | + +## Executive Summary + +PR #15618 introduced an opt-in bounce path for the Python native KV cache +transceiver. It gathers fragmented sender KV into a contiguous arena, performs a +coalesced NIXL write, and scatters the contiguous receive region back to the +destination KV blocks. The PR also added a receive-side `TransferContext` that +waits for fan-in writers and scatter completion before releasing a bounce slot. + +That context closes an important part of the normal-path bounce lifetime, but it +is not an end-to-end transfer owner. Ownership remains split across +`LlmRequest`, `RxSession`, sender worker stack frames, bounce allocators, CUDA +events, NIXL status objects, and KV cache managers. Cancellation, timeout, +partial publication, fan-in failure, and shutdown can therefore make logical +request completion happen before every physical memory accessor is known to be +quiescent. + +This document defines the target ownership contract for the Python native +transceiver: + +- The unit of physical ownership is an **endpoint-local address unit**: a + receiver target slice or sender source chunk. Receive contexts carry the + exact per-writer ledger, and a request-level handle aggregates local contexts + for user-visible completion. +- `TransferWorker` owns a registry that strongly retains each canonical + endpoint-transfer record, its handle/gate, and its contexts. `LlmRequest` and + session objects are associated consumers, not the root owners. +- A context owns leases for every allocation and asynchronous operation that can + outlive request/session cleanup. +- Logical completion and physical retirement are independent states. +- Cancellation, timeout, session destruction, elapsed quarantine time, and an + ambiguous transport failure are not proof of quiescence. +- Direct, bounce, and mixed fan-in paths follow the same contract. +- In-doubt memory remains unavailable for reuse until terminal evidence or an + externally established remote/global quiescence fence exists. The initial + implementation chooses safety over bounded reclamation. + +The functional scope is the Python transceiver. This is not necessarily a +Python-files-only change: allocator-enforced KV leases or definitive NIXL status +may require small C++/nanobind extensions. The separate C++ +`CacheTransceiver` implementation, data path, and wire protocol are not +implementation targets. Shared `AsyncTransferManager`/`ResourceManager` +accounting and its shutdown vetoes are entered only for the lifecycle-capable +Python transceiver. Default C++ and non-disaggregated executors retain baseline +local manager teardown; this adapter neither audits nor retires C++ transport +owners. + +## Initial Containment Implementation + +This follow-up implements the first Python-runtime containment slice of the +design documented by +[PR #16347](https://github.com/NVIDIA/TensorRT-LLM/pull/16347). It is +intentionally smaller than the complete allocator- and protocol-backed +contract below. + +For the supported Python-native cohorts, this containment is always active when +`transceiver_runtime="PYTHON"` explicitly selects that transceiver; the default +C++ transceiver implementation, data path, and wire protocol are unchanged. +The default C++ and non-disaggregated shutdown paths bypass the Python ownership +drain, veto ledger, and rank-uniform votes and keep their baseline local manager +teardown. Within the Python runtime ownership is not gated by +`kv_cache_bounce_size_mb` or a second feature flag. It applies to direct +transfers and eligible bounce transfers while preserving the #15618 wire +format. This coarse runtime selection is not the gated rollout of the complete +protocol: generation identity, capability negotiation, allocator-enforced +leases, and the configuration exclusions documented below remain follow-up +work. + +Implemented in this slice: + +- a receiver registry that survives `RxSession` lookup/removal and tracks the + exact expected rank set for each `(request, slice)`; +- one cancellation/publication gate that marks a writer possibly exposed before + the ZMQ address send; +- registry-first result routing, duplicate/conflict handling, and separate + logical and physical drain states for known writer cohorts; +- direct-path cleanup gating: Python keeps the request and its destination + allocation owned until every published writer is terminal; +- an exact sender-operation ledger keyed by session, channel, slice, and + receiver identity; claim and source-owner enrollment are atomic before + descriptor construction, duplicate target requests cannot launch a second + write, and terminal frames are retained for idempotent replay; +- `SendTransferContext` retains the exact session, `LlmRequest`, operation + identity, and every descriptor request, NIXL status object, or send-bounce + slot that was returned before NIXL quiescence became ambiguous. If gather + fails before those values can be returned, the bounce allocator separately + owns the quarantined slot while the context retains the source request; + either owner vetoes sender teardown; +- `RecvBounceContext` with fixed rank-to-bounce-subrange mapping, + evidence-based slot retirement, explicit slot-settlement acknowledgement, no + time-based quarantine reuse, scatter-aware close, legacy-tail validation + against the receiver-authorized destination interval union, exact agreement + between that union and reserved bounce bytes before publication, retryable + logical-failure delivery when scatter state update fails, and lazy + construction of the interval set after slot reservation succeeds; +- transactional sender cancellation sealing and durable tombstones that cover + local and remote cancellation before or after `TxSession`, task, or request + metadata creation, including explicit no-access results for published KV and + auxiliary targets; +- serialized ZMQ socket I/O and close, bounded listener shutdown, and admission + gates that prevent launch/cancel/shutdown races; +- error cleanup and retryable, fail-closed shutdown keep request resources and + manager allocations alive until the transceiver reports physical drain; a + non-drained attempt returns `False` and retains sessions, requests, result + ingress, registrations, and mappings. Nested cancellation, worker, or session + cleanup failures also retain the transceiver ownership root for retry, and + `PyExecutor` refuses resource-manager teardown while its asynchronous transfer + manager still reports an owner. On this Python-native path, a final world + readiness vote prevents one rank from destroying managers while another rank + retains a post-connector terminal or resource-release failure. Manager + shutdown is then tracked as a + non-replayable per-manager phase, followed by a second world outcome vote; + the first ambiguous manager blocks every later dependency, and no rank + deletes its engines if any manager shutdown failed or became ambiguous; +- `PyExecutor` retains one exact request-level terminal ledger across native + and connector completion. Native success/failure, cancellation, and request + failure use first-result-wins arbitration; connector-first completion in + mixed mode retires only its own lease while native remains authoritative. + The winning response object survives enrichment and sibling retirement. The + ledger is removed only after response publication (or explicit post-loop + discard), final owner retirement, and request teardown. Because the response + sink has no acknowledgement contract, an exception after publication begins + is treated as an in-doubt side effect: the response is retained but never + replayed, preventing duplicates at the cost of possible response loss. + +This slice does **not** claim the complete ownership contract. It has no +allocator-enforced destination KV lease, transfer generation/endpoint epochs, +replay window, negotiated wire capability, allocator-issued source lease, or +structured shutdown diagnostics beyond the boolean drained result. A missing +result or ambiguous local operation therefore retains the whole Python request +and transport mappings indefinitely; destroying the local NIXL agent is not +treated as proof that a remote process can no longer submit a one-sided write. +Most current cohorts protect KV through strong `LlmRequest`/manager-allocation +retention. The V1 PP=1 partial-reuse cohort is the exception: after early +request teardown, `AsyncTransferManager` retains request-wide block-ID refcount +pins acquired by `store_blocks_for_reuse(..., pin_blocks=True)`, and the last +provider retires them through a batch-prevalidated, non-replayable +`unpinBlocksById` operation. The batch is prevalidated for range, uniqueness, +and live refcount; raw IDs do not authenticate the pin owner. An exception after +that call begins leaves the unpin outcome and request ownership in doubt rather +than replaying possible refcount decrements. +V2 retains its cache allocation through the request/cache-object ownership path +instead of that V1 pin API. Neither mechanism is an allocation-generation- +bearing source or destination lease; the common allocator lease remains +follow-up work. + +The exact writer registry is enabled only when `ctx_dp_rank` selects a +deterministic cohort. Generation-first attention-DP broadcast remains on the +pre-existing compatibility path, which assumes exactly one enumerated +candidate cohort performs the transfer. This initial slice deliberately adds no +unnegotiated wire message to close non-participants: doing so would create an +independent rolling-upgrade hang and would still lack generation-safe replay. +Consequently attention-DP broadcast and its auxiliary fan-out are not covered +by the ownership rollout until an address-free cohort selection or negotiated +`NO_REMOTE_ACCESS` protocol satisfies the normative contract below. + +Cancellation tombstones are also keyed by the existing request ID and retained +without a bounded replay window. This prevents cancel-before-create from +stranding a currently published receiver, but endpoint epochs, capacity bounds, +acknowledged retirement, and ABA protection remain follow-up work. + +Receive bounce is disabled for every multi-writer cohort in this slice. Equal +writer counts, divisible arena sizes, or uniform PP-stage layer counts do not +prove byte contributions: one local PP interval can intersect peer stages by +unequal extents, TP mappings can duplicate heads, and layer groups can have +different block sizes. Multi-writer transfer therefore uses the direct path +until the receiver can authorize byte-accurate per-writer extents and offsets +before publishing any bounce address. Single-writer bounce additionally +requires equal tokens-per-block, a complete one-to-one identity mapping over +every attention pool view on both peers, and equal corresponding physical slot +bytes. It remains direct when any part of that byte-exact proof is unavailable. + +For compatibility with the #15618 wire format, successful mode is inferred from +the optional bounce result tail: a complete tail means bounce and no tail means +direct fallback. A failed result is conservatively accounted as bounce when a +bounce target was offered because the old frame does not encode the attempted +mode. Transfer generations and explicit mode fields require negotiated protocol +hardening. + +For a non-empty legacy bounce tail, this slice validates the fixed +writer-to-bounce-source offset, positive and overflow-safe fragments, exact +aggregate bytes for that writer, containment in the receiver-authorized +destination interval union, and destination non-overlap. The legacy frame does +not carry a receiver-owned expected source-to-destination manifest per writer. +An equal-sized source-to-destination permutation within the sole enabled writer +(or across writers if multi-writer bounce is later enabled) therefore remains +undetectable. This is a residual compatibility limitation, not satisfaction of +the complete O13 fixed-manifest contract below. + +Synchronous receive timeout is intentionally fail-closed in this slice. The +timeout latches logical failure and sends cancellation, but the call cannot +return while a published target remains in doubt because there is not yet a +detached allocator lease. The allocator-backed follow-up is required to +preserve prompt timeout return semantics while safely retaining only the +affected KV allocation. + +Focused CPU unit tests cover the registry state machine, exact candidate-cohort +reconciliation, cancel-before-create replay, duplicate sender-operation +admission, multi-operation sender retention, mixed direct/bounce accounting, +socket serialization, deferred executor cleanup, and fail-closed shutdown +gates. Separate merged coverage from PR #16116 positively asserts active bounce +on its scheduled GB200 fabric-VMM cell. The CPU suite does not establish that +hardware engagement, attention-DP containment, mixed-version protocol liveness, +or an allocator-enforced KV lease; those remaining configurations stay outside +a production ownership rollout until their corresponding validation gates pass. + +## Design Decisions + +The following decisions are normative for the implementation plan. + +1. **One transfer does not have a single cross-process owner.** Sender and + receiver each have an endpoint-local context correlated by protocol identity. +2. **Physical ownership is per endpoint-local address unit.** A receive context + owns one advertised target slice; a send context owns one fixed operation or + dynamically scheduled source chunk. A request-level handle aggregates those + local contexts. +3. **The registry is the root owner.** It strongly owns one canonical + endpoint-transfer record, its handle/gate, and all contexts for a transfer + generation. Request or session removal cannot destroy that lifecycle state. +4. **KV ownership is allocator-enforced.** A strong reference to `LlmRequest` is + not a sufficient KV lease if `free_resources()` can independently recycle its + blocks. +5. **The contract covers direct and bounce paths.** Bounce being disabled or a + sender falling back to direct transfer does not bypass ownership tracking. +6. **Logical completion is separate from physical retirement.** A request may + fail promptly while its memory remains leased and unavailable for reuse. +7. **Address publication is conservative.** A writer is considered capable of + access before any message containing its address is sent. +8. **Quarantine is containment, not evidence.** Wall-clock expiry never makes an + in-doubt allocation safe to reuse. +9. **Identity is transfer-generation-safe.** Every lifecycle/control message is + matched by endpoint epochs, transfer, and a strictly increasing admission + generation allocated by the admission initiator. Endpoint-owned replay state + prevents a retired generation from recreating a transfer record. Writer + results additionally carry slice plus exact writer stream/operation identity; + no mutation is keyed only by request ID or writer count. +10. **Shutdown drains before deregistration or unmapping.** If quiescence cannot + be proven, memory remains mapped and shutdown reports a non-drained state. +11. **Wire changes are negotiated.** New transfer-generation, writer, stream, + or mode fields require capability/version negotiation or explicit + mixed-version rejection before address publication. +12. **Non-KV auxiliary transfers are a separate audit.** This project covers all + resources used to perform KV transfer, including KV descriptors and bounce + gather/scatter plans, but does not claim to own independent auxiliary-buffer + transfers. +13. **Handle enrollment closes cancellation races.** Context admission and + request-handle enrollment are one transaction. Cancellation closes the + handle to new contexts before it snapshots and cancels existing contexts. +14. **Initial binding closes the local cancellation race.** Before a canonical + handle exists, the request/session owns a one-shot admission token. + Admission atomically binds that token to the registry-owned handle, while + cancellation either aborts the pending token or forwards to the bound + handle; it cannot fall between those states. + +## Background + +### Bounce transfer introduced by PR #15618 + +KV cache pages are often physically fragmented. A direct NIXL transfer may +therefore require many descriptors. The bounce path introduced by +[PR #15618](https://github.com/NVIDIA/TensorRT-LLM/pull/15618) changes the data +path to: + +```text +source KV pages + -> CUDA gather into a contiguous sender bounce slot + -> coalesced NIXL WRITE into a contiguous receiver bounce slot + -> CUDA scatter into destination KV pages +``` + +`kv_cache_bounce_size_mb > 0` requests the feature; zero disables it. Requested +does not imply active: the Python-native factory returns `NoBounceTransport` if +the fabric-VMM arena cannot be constructed or fitted, and an active transport +still falls back per transfer when a request is ineligible or no slot is +available. Merged PR #16116 makes the bounce-configured integration test assert +the generation-side `"[kv-bounce] coalesced"` signal. Its scheduled GB200 cell +therefore provides positive active-bounce coverage; GB300 remains marker-eligible +but is not claimed here as scheduled positive coverage. Other environments may +still run only the direct fallback. + +The merged +[`bounce.TransferContext`](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/bounce/core.py#L61-L175) +tracks a receive bounce slot, writer results, scatter state, and settlement. It +does not hold source or destination KV leases, the sender bounce lease, the NIXL +operation, or the publication identity. The initial containment renames that +class to `RecvBounceContext` and keeps it as an internal receive-bounce state +object; the future general context remains separate. + +### Related work + +Status snapshot: 2026-07-21. + +| Work | Relationship to this design | +|---|---| +| [PR #15618](https://github.com/NVIDIA/TensorRT-LLM/pull/15618) | Merged Python/native bounce implementation and baseline for this follow-up. | +| [Transfer-owner review on #15618](https://github.com/NVIDIA/TensorRT-LLM/pull/15618#pullrequestreview-4612079358) and [follow-up acknowledgement](https://github.com/NVIDIA/TensorRT-LLM/pull/15618#pullrequestreview-4630057518) | Review provenance for the comprehensive ownership follow-up. | +| [PR #16116](https://github.com/NVIDIA/TensorRT-LLM/pull/16116) | Merged post-#15618 follow-up adding positive active-bounce validation and partial bounce-only lifecycle containment: it orphans/quarantines in-flight receive reservations during cancellation/session close and propagates local gather/build failures. It does not provide direct-path or allocator-enforced KV ownership, generation-safe routing, or bounded safe reclamation. | +| [PR #15780](https://github.com/NVIDIA/TensorRT-LLM/pull/15780) | Open draft proposing an independent C++ NIXL bounce implementation with a different arena and credit protocol. Useful lifecycle prior art; not a shared implementation target. | +| [PR #15139](https://github.com/NVIDIA/TensorRT-LLM/pull/15139) | Merged C++/V1 precedent for rank-consistent terminal-state consensus. This design adopts the logical/physical separation for Python native transfers. | +| [PR #15356](https://github.com/NVIDIA/TensorRT-LLM/pull/15356) | Merged bounded polling/admission work. Draining must remain non-blocking to the executor hot path. | +| [PR #15238](https://github.com/NVIDIA/TensorRT-LLM/pull/15238) | Merged C++ cancellation/quiescence prior art for gated NIXL in-flight cancellation. It is inherited conceptual precedent, not a Python-native dependency. | +| [PR #15737](https://github.com/NVIDIA/TensorRT-LLM/pull/15737) | Merged sender-liveness hardening for disaggregated KV transfer. It is operationally adjacent, not a replacement for endpoint-local ownership. | +| [PRs #15794](https://github.com/NVIDIA/TensorRT-LLM/pull/15794), [#15795](https://github.com/NVIDIA/TensorRT-LLM/pull/15795), [#15798](https://github.com/NVIDIA/TensorRT-LLM/pull/15798), and [#15799](https://github.com/NVIDIA/TensorRT-LLM/pull/15799) | Open draft C++-transceiver cancellation chain plus PyExecutor integration: buffer ownership, detached-owner/fatal cleanup, negotiation, and a generation-safe peer-protocol design. This work aligns conceptually but fills the Python-native gap. | +| [PR #15738](https://github.com/NVIDIA/TensorRT-LLM/pull/15738) | Open draft default-on policy for a qualified C++ NIXL/UCX configuration. It explicitly excludes the Python transceiver, so this work is complementary. | +| [PR #15727](https://github.com/NVIDIA/TensorRT-LLM/pull/15727) | Open Python pipelined/chunked transfer consumer. The owner must support multiple independently progressing chunks/slices and overlapping prefill/transfer lifetimes. | +| [PR #15803](https://github.com/NVIDIA/TensorRT-LLM/pull/15803) | Open draft C++-transceiver/V1 work and prior art for explicit KV transfer leases, descriptor lifetime, and constrained early release. It is not an available shared primitive. | + +The adjacent cancellation chain represented by PRs #15238, #15794, #15795, +#15798, and #15799 owns cancellation consensus and process-health policy across +a wider runtime matrix. +This note owns the physical resource-lifetime contract for the Python native +transfer path. Cancellation intent enters this design as a logical event; it +never authorizes physical release by itself. + +## Pinned-Baseline Problem Statement and Residual Gaps + +The concrete hazards in this section are anchored to the merged PR #15618 +baseline (`3bf37d2`). Each subsection states what the initial containment now +closes and what remains for the complete ownership contract. + +### Baseline and residual ownership are fragmented + +At the pinned baseline, the path has several partial owners: + +- `LlmRequest` carries request state and allocated block identifiers. +- `AsyncTransferManager` retains context-side requests during asynchronous send + work, but does not provide a symmetric allocator lease for receiver loading. +- `RxSession` owns request-facing receive state and callbacks. +- Sender worker stack frames hold the NIXL status and release sender bounce slots + in local cleanup. +- `VmmBounceTransport` owns arena allocators, registrations, CUDA streams, and + receive bounce contexts. +- KV managers remain the authority that can return blocks to their pools. + +The initial containment adds registry-owned receive state, exact sender +operation owners, request retention, and V1 pin tokens. It still has no common +allocator-generation-bearing lease, so no single object has both the +information and allocator authority to answer: **which asynchronous accessors +can still touch which allocation generation?** + +### Concrete baseline hazards and containment status + +#### Released bounce address can still be advertised + +At the pinned baseline, the +[`dispatch_task()` sequence](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/transfer.py#L1522-L1589) +reserves a receive slot, then separately finds the session, changes task state, +and publishes the address. Concurrent +[`RxSession.cancel()`](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/transfer.py#L1985-L2006) +can observe the task as not transferring and release the reservation between +those steps. Dispatch may then advertise an address that has already returned to +the allocator, creating a potential cross-request overwrite if a remote writer +uses the stale address after the slot is reused. + +**Initial-containment status:** closed for the supported deterministic writer +cohorts. Registry admission, publication gating, and cancellation are one +lifecycle transaction, and a possibly published slot remains owned until its +writers and scatter settle. + +#### First fan-in failure can outlive its session + +At the pinned baseline, the first failed writer +[`fail()`s the task immediately](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/transfer.py#L1898-L1912). +Failed-session cleanup can remove the `RxSession`, while +[`_process_kv_agent_result()`](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/transfer.py#L1684-L1710) +drops later results when that session is absent. + +For an all-bounce transfer this can strand the receive slot. For a mixed or +direct transfer, a sibling writer may still be writing destination KV. If +request cleanup recycles those blocks before that sibling is quiescent, this +creates a potential physical GPU memory use-after-release/ABA hazard, not merely +a Python object-lifetime issue. + +**Initial-containment status:** closed for supported cohorts. Results route +through a registry that outlives `RxSession`, tracks the exact writer set, and +keeps the destination owner until every possibly exposed writer is terminal. + +#### Timeout and quarantine do not prove retirement + +At the pinned PR #15618 baseline, `mark_orphaned()` exists in the bounce state +object but has no production callers through the transport interface. PR #16116 +proposes callers for cancellation/session close, but that is bounce-only +containment: direct destination KV and sender resources remain uncovered, and +quarantine still does not prove safe reclamation. Timeout, partial publication, +and other session-destruction paths therefore still need the general owner. + +The allocator's fixed-duration quarantine, if wired into production, would also +not be a transport guarantee. An elapsed grace period cannot prove that a +one-sided operation will not access the region later. + +**Initial-containment status:** timed reuse is removed. Ambiguous operations are +retained indefinitely and shutdown reports non-drained instead of reclaiming +their memory. Bounded recovery still requires generation-safe replay, +acknowledgement, and allocator-enforced leases. + +#### Shutdown order can destroy live memory + +At the pinned baseline, sender workers use an unbounded transfer wait but are +joined with a finite deadline. The enclosing +[`TransferWorker.shutdown()`](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/transfer.py#L2297-L2339) +invokes the current +[`VmmBounceTransport.close()`](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py#L382-L394) +which stops the scatter worker, deregisters arenas, and destroys their mappings. +This occurs before local +[`BaseTransferAgent.shutdown()`](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/base/agent.py#L85-L121) +is invoked. A sender worker, remote RMA, or queued scatter may still reference +those ranges. Invoking local agent shutdown first would not repair the contract: +destroying the local handle is not an external fence preventing a remote peer +from submitting a later one-sided write to an address it already received. + +**Initial-containment status:** shutdown now closes admission and refuses to +deregister or destroy mappings while any sender, receiver, result-delivery, or +bounce owner remains. A non-drained attempt retains the transceiver for retry; +an externally established remote/global quiescence fence remains future work. +Post-loop `PyExecutor.shutdown()` guarantees ownership drain, not response +delivery: it locally discards newly created transfer-completion responses +because TP/attention-DP response collectives are no longer rank-safe after a +rank-local teardown failure. + +#### Count-based fan-in is not identity-safe + +At the pinned PR #15618 baseline, the receive context records distinct peer +ranks until `len(_writer_ok) >= num_writers`. It does not retain the exact +advertised writer set or a transfer generation. An unexpected distinct writer +can count toward completion, and a stale result can collide with a later reuse +of the same `(request_id, slice_id)` key. The initial containment implemented +in this follow-up closes the first gap with an exact writer ledger; transfer +generation and endpoint-epoch identity remain future work. + +## Scope and Applicability + +### In scope + +| Axis | In-scope variants | +|---|---| +| **Transceiver runtime** | Python native `KvCacheTransceiverV2`. Here, “V2” names the transceiver, not necessarily the KV manager. | +| **Network backend** | NIXL/DEFAULT as supported by the Python transceiver. | +| **Transfer path** | Direct fragmented transfer, bounce transfer, and per-writer fallback between them. | +| **Bounce configuration** | Not requested (`kv_cache_bounce_size_mb == 0`), requested but factory-inactive (`> 0` with `NoBounceTransport` fallback), and arena-active with per-transfer bounce/direct fallback. | +| **Active-bounce hardware validation** | The merged #16116 test is GB200/GB300 MNNVL-eligible, scheduled on GB200, and positively asserts per-transfer bounce engagement there. GB300 is eligible but not claimed as scheduled coverage. Other hardware remains in direct-path scope and needs the same positive evidence before a bounce rollout. | +| **KV manager** | Python KV manager V2 and the C++-backed V1 manager exposed to Python. | +| **Direction** | Context-side send and generation-side receive. | +| **Fan-in/topology** | Single writer and every multi-writer topology supported by the Python transceiver, including TP/PP/ADP overlap or fan-in and multiple slices/chunks. | +| **Executor** | PyTorch backend and AutoDeploy when they select the Python transceiver. | +| **Resources** | KV allocations, bounce-slot leases, CUDA gather/scatter work, transfer descriptor/plan storage, NIXL operation state, result delivery state, and registration/mapping shutdown dependencies. | + +The ownership path must be active even when bounce is disabled. The bounce +configuration changes which leases exist; it does not select whether lifetime +safety applies. + +### Non-goals + +- Replacing or porting the design into the separate C++ `CacheTransceiver`, or + treating destruction of its local object as transport quiescence evidence. +- Unifying the Python bounce implementation with the independent C++ bounce + implementation in PR #15780. +- Redesigning steady-state transfer consensus, scheduler cancellation, or request admission. Shutdown-only rank-uniform ownership votes remain in scope for the Python-native path. +- Making `LlmRequest` the low-level transfer state machine. +- Creating one distributed object that owns both processes. +- Retrying or resuming a failed model request. +- Rolling back KV bytes written before a transfer fails. +- Optimizing gather/scatter kernels, descriptor coalescing, or arena sizing. +- Enabling bounce by default. +- Transparent recovery from a crashed peer without transport or process-level + quiescence evidence. +- Independent non-KV auxiliary-buffer transfer ownership. Such paths require a + separate audit against the same invariant. + +## Terminology + +| Term | Meaning | +|---|---| +| **Logical outcome** | Request-visible success, failure, or cancellation. It controls scheduling and notification, not memory reuse. | +| **Physical retirement** | Proof that no network or CUDA accessor can touch the leased allocation generation, followed by exactly-once lease release. | +| **Lease** | A token that prevents an allocator, arena, or manager from reusing or destroying a resource while the token is live. | +| **Publication** | Any attempt to send an address or descriptor to a writer that could cause the writer to access it. | +| **Admission control** | An address-free protocol record that admits, rejects, skips, or creation-closes one transfer generation. Sending it is not address publication. | +| **Generation skip** | An authenticated admission control used only when local state positively proves normal admission never became observable and no address can be published; the peer must acknowledge it before replay floors advance. | +| **Terminal evidence** | Positive evidence, defined by the backend contract, that a specific operation cannot perform later memory access. | +| **Quiescence** | Positive per-operation evidence or an externally established remote/global fence covering every relevant peer, submission path, and fabric operation. Local agent shutdown alone is not quiescence. | +| **Quarantine** | Removal from reuse while quiescence is unknown. It does not become safe through elapsed time. | +| **Transfer generation** | A strictly increasing admission sequence allocated by the admission initiator within its endpoint epoch. Every slice in that attempt and every request-level lifecycle control share it. | +| **Allocation generation** | An allocator-issued incarnation of a KV block or bounce slot. It is carried by the resource lease and is independent of the transfer generation. | +| **Endpoint replay state** | Endpoint-epoch-owned anti-replay state: a contiguous retired-generation floor plus a bounded sparse window for newer live or out-of-order-retired generations. | +| **Target slice** | The receiver address set published and physically retired as one unit. | +| **Source chunk** | A sender scheduling/lease unit. In dynamic chunking it need not be one-to-one with the receiver target slice. | +| **Creation close** | Phase 1 acknowledgement that an endpoint producer cannot create another local context for a transfer generation. It does not prove that a published address cannot be accessed. | +| **Submission fence** | Phase 3 acknowledgement that a peer cannot later submit against the covered published operation/stream addresses. | +| **Consumer** | The request/session-facing observer that receives logical completion. It is not the physical owner. | +| **Transfer admission token** | A one-shot request/session gate whose state is `PENDING`, `BOUND(handle)`, or `ABORTED`; it serializes local cancellation with initial canonical-handle admission. | + +## Normative Safety Contract + +The load-bearing invariant is: + +> For every asynchronous accessor and every memory range it may read or write, +> the range's allocation generation and registration MUST remain unchanged from +> before its address becomes observable until there is positive evidence that +> the accessor can no longer touch it. + +The following are explicitly **not** terminal evidence: + +- client cancellation; +- request or session destruction; +- scheduler timeout; +- elapsed quarantine duration; +- release of a local request handle without an abort guarantee; +- loss of the result message; +- an ambiguous `False` from a bounded wait that conflates in-progress and + terminal failure; +- local CUDA synchronization when a remote RMA may still be active. + +### Required invariants + +**O1 — Prepare each endpoint before its access boundary.** On the receiver, the +destination KV snapshot/lease, optional receive-bounce lease, exact authorized +writer cohort, fixed-operation manifest or bounded operation-stream ledger, and +receive context MUST be registered before target addresses are published. On +the sender, a preparation transaction MUST own the source KV snapshot/lease +before it resolves source addresses. The optional send-bounce lease, descriptor +storage, and send context MUST be committed before those addresses are used by +gather or NIXL. The sender-side precondition occurs after it receives the target +advertisement; it is not a cross-process prerequisite for receiver publication. + +**O2 — Mark possible access first.** Before the receiver attempts to send a +target advertisement containing an address, `publication_started()` MUST pass +through the request handle's cancellation gate and atomically set that writer +operation or authorized operation stream's exposure state to `MAY_ACCESS` and +access state to `POSSIBLE`. An advertisement-send error may revert to +`NEVER_EXPOSED`/`QUIESCED` only when the messaging layer positively guarantees +non-delivery. + +**O3 — Request lifetime is not resource lifetime.** Closing an `LlmRequest`, +`TxSession`, `RxSession`, future, or request-level transfer handle MUST NOT +directly release transfer-owned resources or bypass the context retirement +predicate. It MAY trigger retirement when that predicate is already satisfied. + +**O4 — Logical and physical state are independent.** Failure or cancellation MAY +be reported immediately, but physical retirement MUST wait for all possible +accessors. + +**O5 — Exact writer and operation accounting.** Before publication, a receive +context MUST track every authorized peer writer and use one of two negotiated +forms: + +- A fixed manifest lists every writer-operation identity, its exact expected + logical source-to-target KV mapping, and its authorized candidate + direct/bounce ranges and descriptor/scatter-segment envelope. +- A dynamic operation stream pre-authorizes one exact writer, target slice, + required logical source-to-target mapping, candidate range, stream ID, + monotonically increasing operation-sequence domain, and maximum + operation/descriptor-segment/byte budget. Each sender chunk gets a distinct + `(writer_stream_id, operation_sequence)` before local submission. An + authenticated end-of-stream record carries the final sequence high-water mark + and closes both future enrollment and future submission. The sender may emit + it only after every sequence through that mark has made an irreversible + decision: already `SUBMITTED` with owned status, or terminal + `NO_REMOTE_ACCESS`. The receiver validates every operation subrange and does + not retire the stream target until end-of-stream (or an acknowledged + submission fence) and quiescence for every submitted operation through that + high-water mark. + +The stream form supports a monolithic receiver target with a sender-determined +chunk count, as in PR #15727, without pretending the receiver knows all chunks +before publication. Actual target mode is recorded per operation after sender +selection. Duplicate advertisements and operation identities are idempotent; +contradictory reuse is rejected. Writer count alone is insufficient. Fixed +manifests have negotiated operation-count, descriptor/scatter-segment, and +authorized-byte limits, and both forms obey the per-transfer context/slice and +endpoint-global metadata limits in O19. + +Fixed-mode success requires the normalized delivered source-to-target mapping +to equal each operation's expected mapping and the manifest's complete mapping. +In-range undercoverage, gaps, duplicate segments, aggregate-byte mismatch, and +unintended within- or cross-writer overlap are invalid even when every segment +is individually authorized. Full target coverage with permuted source intervals +or wrong bounce offsets is also invalid. The initial implementation rejects +cross-writer overlap; a future explicit overlap mode would need deterministic +idempotent semantics and separate negotiation/validation. + +Dynamic-stream success likewise requires the normalized union of all operations +through the authenticated high-water mark to equal the stream's required +logical source-to-target mapping. A sequence-complete stream with in-range +undercoverage, a mapping permutation, or a coverage gap fails; lowering the +high-water mark cannot redefine the required target. + +**O6 — Registry-first result routing.** Every KV writer-operation result MUST be +routed to the transfer registry before optional session lookup or notification. +A missing consumer MUST NOT cause a physical result to be dropped. Independent +auxiliary results remain with their separate owner. + +**O7 — Allocator-enforced KV leases.** The KV manager MUST provide one atomic +`snapshot_and_lease(request, slice_spec)` operation under its allocation lock. +It validates that the request still owns each current allocation, increments +the transfer-lease count, and returns immutable descriptors containing at least +`(logical_kv_range, pool, block, allocation_generation, address, size)`. The +stable logical coordinate binds each physical source/destination interval to the +manifest's source-to-target mapping. Context construction and range +authorization use those lease tokens; copying block IDs and acquiring a lease +later is forbidden because it leaves an ABA window. `free_resources()` +atomically drops request ownership and marks leased blocks pending-free. An +allocation generation becomes reusable only when request ownership is gone and +its transfer-lease count reaches zero. Overlapping slice leases are +reference-counted, and KV-manager shutdown obeys the same rule. Holding only +block IDs or a Python request reference is insufficient. + +**O8 — Exactly-once retirement.** Duplicate, reordered, late, or contradictory +events MUST NOT double-release a lease, advance another transfer generation, or +release another allocation generation. + +**O9 — Safe uncertainty.** When quiescence is unknown, resources MUST remain +leased or quarantined. Admission may fail due to exhausted safe capacity; reuse +is forbidden. + +**O10 — Drain before teardown.** A KV or bounce registration/mapping MUST NOT be +destroyed until every relevant network accessor is quiescent (or backend-wide +quiescence established by an external remote/global fence substitutes for +missing per-operation evidence) **and** all local CUDA work that can touch it is +complete. Shutting down the receiver's local agent is not such a fence. If +shutdown also removes registrations used by auxiliary transfers, those +independent owners MUST drain as an external precondition; draining only the KV +registry is insufficient. + +**O11 — No blocking under lifecycle locks.** The only permitted nested +lifecycle-lock order is admission token, then registry, then handle, then +context; registry/remote paths never acquire a token gate. Network operations, +CUDA synchronization, allocator callbacks, and consumer callbacks MUST execute +outside all of those state locks. + +**O12 — Bounded executor polling.** Ownership draining MUST preserve the bounded +polling behavior established by PR #15356. Waiting for physical retirement MUST +NOT block the main executor loop. + +**O13 — Range authorization.** Before NIXL submission or gather/scatter launch, +every descriptor segment MUST be validated as a subset of the correct leased +allocation generation and the writer's authorized candidate range. Validation +includes device, required alignment, non-negative size, integer overflow, +aggregate byte count, and per-writer subrange boundaries. A lifetime lease does +not authorize out-of-range access. Before fixed-mode success, normalize the +actual source-to-target segment pairs and require the exact expected +per-operation and manifest mapping with no permutation, gap, duplicate, or +overlap under the negotiated policy. Bounce validation includes the expected +bounce-source offset paired with each destination interval. + +O13 is the target contract. The initial legacy-tail containment proves only the +narrower aggregate, authorized-union, and non-overlap properties described +above; exact per-writer manifest matching and permutation rejection require a +negotiated protocol extension and remain a rollout gate. + +**O14 — Endpoint-local transactional preparation.** `prepare_send()` and +`prepare_receive()` MUST be all-or-nothing construction transactions. A +preparation object owns partial KV/bounce leases until registry insertion +commits. The receiver may unwind only before its target-publication boundary. +The sender may unwind unused local resources until its own gather/NIXL launch +boundary even though the receiver already published a target. It then reports +quiescent `NO_REMOTE_ACCESS`; if that result is lost, the receiver remains in +doubt but the unused sender resources need not. After an endpoint crosses its +own boundary, failure follows normal fail-closed context retirement and never +uses construction rollback to release possibly accessible memory. + +**O15 — Independent auxiliary cleanup.** KV logical completion or cancellation +MUST NOT release a generation-first auxiliary slot or close the session that +owns it while an auxiliary RMA may still target that slot. Auxiliary transfer +ownership remains a separate audit, so it keeps an independent safe-to-clean +gate. Any configuration that cannot provide that gate is excluded from the +ownership rollout until the auxiliary path is covered. + +**O16 — Atomic handle enrollment, access gating, and closure.** A +`TransferHandle` has an `OPEN`, `SEALED`, or `ABORTED` admission/access state; +the separate logical outcome records whether an abort reason is failure, +cancellation, or shutdown. +Endpoint context insertion and handle enrollment MUST commit atomically with +respect to `cancel()` and `seal()`. Every publication, gather, NIXL, or scatter +boundary obtains the same handle gate before the context lock, verifies that +the handle is not `ABORTED` and the context's logical state still authorizes +that accessor, then marks it possible before releasing the locks. An accepted +cancellation, first logical failure, or shutdown changes the handle to +`ABORTED` and snapshots its enrolled contexts while holding that gate, then +performs context/peer callbacks outside it. If an access boundary wins the gate, +the abort observes an already-possible accessor and drains it; if the abort wins, +the boundary is rejected even before its context callback runs. A preparation +racing abort either enrolls and is included in it, or transactionally unwinds +before publication/local access. Failure/cancellation also starts the applicable +peer operation/stream close or submission-fence protocol; missing acknowledgement +keeps remote targets in doubt. + +All contexts in one handle share one `transfer_generation`. The local +producer/session is the only sealing authority. It seals only after atomically +enrolling the complete required local context set: either an immutable expected +target-slice/source-chunk manifest, or a validated final-chunk proof enrolled +with the final send context. Receiver planning may provide its target manifest +up front; incremental sender chunks cannot seal merely because the currently +known contexts completed. Logical success requires a sealed handle whose +declared contexts all succeeded. + +**O17 — Atomic, capacity-reserved pre-cancellation.** Validated remote +cancellation and lifecycle-record/context creation MUST serialize under the +registry lock. Before a peer is authorized to send lifecycle controls or +addresses for a transfer generation, admission reserves a registry-owned +`TransferHandle`/control-state slot. Cancellation either closes that live handle +or marks its zero-context record pre-cancelled; preparation under the same lock +observes the state and cannot cross an access boundary. A control for an +unnegotiated identity is a protocol violation: poison and close that endpoint +session in O(1), rejecting every later admission on it. A replacement requires +full negotiation and may retain the endpoint epoch only if its existing +`EndpointEpochState`/replay window survives; otherwise it uses a fresh epoch. +Do not create unbounded per-identity deny state for arbitrary controls. + +Capacity is backpressured before authorizing a new transfer generation, so every +valid pre-cancel has reserved durable state. A pre-cancelled or otherwise closed +handle remains until local producer sealing or an acknowledged session/endpoint +creation fence proves that no later context can be created, and until all +physical contexts retire. Age alone cannot prune it. The registry never drops a +live lifecycle record to admit new work. + +**O18 — Durable endpoint anti-replay state.** The endpoint that initiates +transfer admission MUST allocate `transfer_generation` monotonically from +endpoint-epoch state that survives session, request, and transfer-record +cleanup. On that endpoint, counter advancement, replay-window live-entry +reservation, canonical record creation, and admission-token binding commit in +one registry transaction after capacity checks; a pre-commit rejection consumes +no generation. Once committed, a generation is never rolled back. If its normal +admission control will not be sent, the initiator sends an authenticated +generation-skip record and obtains the peer acknowledgement before sending a +higher-generation admission control. A cumulative contiguous closed watermark +may compress skips, but the receiver rejects one that crosses any locally live, +published, or otherwise unclosed generation. If normal admission control was +sent or its delivery is ambiguous but no target address was published, the +initiator retries/deduplicates normal admission and rejection or abandonment +uses the authenticated creation-close/reject exchange; it cannot relabel the +generation as skipped without positive non-delivery proof. Once any target +address was published, neither skip nor creation close retires the record or +memory; full operation/stream closure and physical quiescence remain mandatory. +Both endpoint replay trackers mark an empty/fully retired generation retired +before its slot can compact. Both peers bind the admission to the negotiated +endpoint-epoch pair. +Session negotiation authenticates the initiator epoch, initial generation +floor, and maximum outstanding-generation window. The registry retains, +independently of individual transfer records, a contiguous retired-generation +floor and a bounded sparse set for newer live or out-of-order-retired +generations. Admission outside the negotiated window, at or below the retired +floor, or naming an already-retired sparse entry is rejected. Advertisement, +result, and control messages for such a generation can never recreate a +canonical record. The initiator does not send admission outside the peer's +negotiated window. For an in-window admission, the peer reserves its replay +entry before evaluating later metadata/resource capacity; a rejection marks +that entry closed and returns an authenticated close/reject acknowledgement. + +Out-of-order retirement marks the sparse entry and advances the floor only when +all preceding generations are retired or explicitly closed by the negotiated +creation-close protocol. A missing generation consumes the negotiated replay +window and eventually backpressures new admission; it is never evicted to make +progress. The sparse set compacts as the contiguous floor advances, keeping +healthy sequential churn O(maximum outstanding generations). Recreating the +endpoint owner MUST allocate a new, non-reused endpoint epoch. Messages from an +older epoch cannot bootstrap a new session or mutate the new endpoint, so its +old generation window may be discarded only after the new epoch is installed +and session admission enforces that epoch binding. Reconnecting with the same +endpoint epoch MUST resume the existing local replay state; negotiation cannot +reset the window or advance its retired floor past a locally live/unclosed +generation, and an inconsistent peer claim rejects the session. + +**O19 — Bounded metadata admission.** Configured endpoint limits and negotiated +per-transfer envelopes MUST bound canonical records, contexts/slices, fixed +writer operations, dynamic-stream operations, descriptor/scatter-plan segments, +authorized bytes, and replay entries. A fixed manifest reserves its actual +operation/authorized-byte counts and the worst-case descriptor/scatter segments +permitted by its envelope; a dynamic stream reserves its negotiated maximum +envelope. The registry atomically reserves all required worst-case metadata +credits before their first applicable boundary: record/control/replay credits +at admission before the peer is authorized, and context/manifest/stream credits +before the first address publication or sender local launch. Oversized plans or +unavailable credits fail/backpressure before that boundary, and a partial +preparation returns its credits transactionally. Live work never overshoots a +limit, and capacity reserved for an already-admitted control or stream cannot +be evicted to admit another transfer. Unused and retired credits return only at +their defined boundary. For a dynamic stream, unused envelope credits may +return only when an authenticated end-of-stream or submission-fence result both +irrevocably closes future enrollment/submission and establishes the exact final +operation set/high-water mark. A fence that only prevents future submission +cannot classify unreported prior operations as unused; those credits remain +reserved until exact ledger reconciliation or backend quiescence provides the +required evidence. Credits backing materialized operations remain until those +entries retire. + +**O20 — Atomic local admission and cancellation binding.** Before transfer +admission is scheduled, the request/session MUST create one +`TransferAdmissionToken` in the same request-lifecycle transition that makes +transfer scheduling eligible. Cancellation that wins before this transition +makes the request ineligible and prevents token creation. After creation, local +cancellation, request cleanup, and `admit_and_bind()` serialize on that token. +If cancellation wins while the token is `PENDING`, it changes the token to +`ABORTED` and later admission fails before peer authorization, publication, or +local launch. If admission wins, it creates or resolves the canonical registry +record and changes the token directly from `PENDING` to `BOUND(handle)` before +the peer is authorized; a racing or later cancellation observes that handle +and aborts its shared access gate. + +The token gate precedes the registry lock, and the registry never acquires a +request/session lock while holding a lifecycle lock. On the initiating endpoint, +a failed admission before the atomic generation/record/token commit consumes no +generation and returns all provisional credits. A committed generation whose +normal admission control is not sent is converted into an authenticated +generation skip; a higher-generation admission control cannot be sent until the +peer acknowledges it. If normal admission control was sent but no address was +published, abandonment uses the authenticated creation-close/reject exchange. +After address publication, abandonment follows the full physical-retirement +contract instead. Thus neither endpoint leaves an unretired replay gap or treats +an address-bearing transfer as empty. After binding, dropping the token or any +request-facing handle reference cannot destroy the registry-owned record, +handle, gate, or contexts. + +`begin_shutdown()` is the third participant in this arbitration. Under the +registry lock it atomically changes endpoint admission to `CLOSED` and snapshots +all already-committed records. `admit_and_bind()` that loses this race changes +its still-pending token to `ABORTED(shutdown)` and cannot create a record or +authorize the peer. A record that wins admission is visible in the shutdown +snapshot. Its later peer-authorization boundary passes through the same handle +gate: it is either marked possible before shutdown aborts the handle and is then +owned/drained, or it is rejected without sending authorization. No record or +peer authorization may appear after the shutdown snapshot unaccounted. + +**O21 — No unbudgeted healthy-path admission round trip.** Session capability +negotiation is amortized. Per-transfer generation, epoch, fixed-envelope, and +admission fields MUST piggyback on the existing address-free KV request and +target-advertisement response: the receiver commits admission before sending +that existing response, whose target advertisement also confirms acceptance. +The normal path does not add a separate admission request/ack RTT. Generation +skip, creation-close, rejection, cancellation, and shutdown-fence +acknowledgements are exceptional-path controls and may add messages. If an +implementation cannot preserve this piggybacking, its additional RTT is an +explicit design/performance change that must be budgeted and qualified before +that cohort can claim the healthy-path non-regression expectation. + +## Ownership Model + +```mermaid +flowchart TB + TW["TransferWorker"] --> TR["TransferRegistry (root owner)"] + TR --> ES["EndpointEpochState and replay window"] + TR --> ER["EndpointTransferRecord per transfer generation"] + ER --> TH["canonical TransferHandle"] + ER --> TG["shared durable access gate"] + ER --> SC["SendTransferContext per fixed operation/source chunk"] + ER --> RC["RecvTransferContext per target slice"] + RH["LlmRequest / session"] --> AT["TransferAdmissionToken"] + AT -. "BOUND consumer reference" .-> TH + SC --> TG + RC --> TG + TH --> TG + SC --> SKV["source KV lease"] + SC --> SB["optional send-bounce lease"] + SC --> GE["gather event and plan"] + SC --> NS["NIXL operation status"] + RC --> DKV["destination KV lease"] + RC --> RB["optional receive-bounce lease"] + RC --> WL["exact writer ledger"] + RC --> SE["scatter job, plan, and event"] + BT["BounceTransport"] --> BA["arenas, registrations, streams, allocators"] + SB -. "lease prevents reuse" .-> BA + RB -. "lease prevents reuse" .-> BA +``` + +### Ownership granularity + +The endpoint-local physical context key is conceptually: + +```text +(role, local_endpoint_epoch, admission_authority_epoch, + transfer_id, transfer_generation, local_context_id) +``` + +`local_context_id` is the target-slice ID on the receiver and the fixed-operation +or source-chunk ID on the sender; those IDs need not match. `role` prevents a +simultaneous local send and receive from sharing an identity. +`admission_authority_epoch` names the endpoint that allocated the monotonic +generation; the full registry identity also binds the negotiated endpoint-epoch +pair. A receive context contains ledger entries keyed by +`(peer_endpoint_epoch, peer_rank, writer_operation_or_stream_id)`. A writer that +issues multiple NIXL operations for the target uses a distinct fixed operation +ID or stream sequence for each. Lookup resolves the endpoint-local context +first, then validates the exact peer operation. One request-level +`TransferHandle` aggregates all local contexts in one transfer generation. Its +sealing proof defines the complete required local context set needed to compute +the logical request outcome. This supports fan-in, peer restart, and chunked or +pipelined transfers without serializing independent units or allowing early +completion before a later chunk enrolls. + +### Component responsibilities + +| Component | Owns | Must not own or decide | +|---|---|---| +| `TransferRegistry` | Endpoint-epoch replay state; exactly one canonical `EndpointTransferRecord` per admitted transfer generation; identity/replay lookup; shutdown drain accounting. | Request scheduling or KV allocation policy. | +| `EndpointEpochState` | Stable local endpoint epoch, monotonic generation allocation when this endpoint initiates admission, and a retired floor/bounded sparse live-or-retired window for each negotiated admission authority and endpoint-epoch pair. | Per-context resources or request outcome. | +| `EndpointTransferRecord` | Canonical handle, durable access gate, context membership, sealing/pre-cancel state, and operation/stream replay identities until its full retirement predicate holds. | KV allocator policy or remote resources. | +| `SendTransferContext` | Source KV lease, optional send-bounce lease, gather fence/plan, descriptor storage, NIXL operation, result-delivery state. | Receiver resources or request object destruction. | +| `RecvTransferContext` | Destination KV lease, optional receive-bounce lease, writer ledger, scatter fence/plan, detachable consumer callback. | Source resources or distributed request consensus. | +| `TransferHandle` | Registry-owned request-facing facade for context aggregation, sealing, logical cancellation/notification, and the shared access gate. | Physical lease release. | +| `TransferAdmissionToken` | One-shot local arbitration between request/session cancellation and initial binding to the canonical handle. | Physical resources, record retirement, or remote control routing. | +| Shared access gate | Durable admission/access serialization retained by the record and each context. | Context membership or resource release. | +| `BounceTransport` | Arena mappings, registrations, allocators, streams, and lease issuance. | End-to-end transfer outcome. | +| KV manager | Underlying KV allocation; atomic allocation-generation-bearing snapshot/lease issuance; enforcement of active transfer leases. | Network completion inference. | +| NIXL agent | Submission/progress and documented quiescence evidence. | KV block reuse. | + +The initial containment has renamed the PR #15618 bounce-only +`bounce.TransferContext` to `RecvBounceContext` and kept it an implementation +detail. It may eventually be narrowed further to `RecvBounceLease` state, but +it must not become the general owner: doing so would couple direct transfer, +request notification, KV allocation, and transport arena internals into the +bounce package. + +Before admission, the request/session holds a one-shot +`TransferAdmissionToken`; after atomic binding it exposes only a +consumer-facing reference to the canonical `TransferHandle`. Dropping either +does not drop the handle or gate. To avoid an ownership cycle, the +`EndpointTransferRecord` owns context membership and the gate, while contexts +retain only the gate, not the record/handle. The gate owns no contexts. +Retirement first removes every settled context from the record. +It then marks the generation retired in `EndpointEpochState` before dropping +the record, and only after local producer sealing or an acknowledged creation +fence proves that no later context can appear and every open stream and +pre-cancel obligation is settled. The endpoint replay floor/window, not the +deleted record, rejects later admission or advertisement replay. Queue-held +retired contexts may keep the already-closed gate alive but cannot recreate +membership or physical leases. + +### Resource retirement matrix + +In this matrix, “every writer operation” also requires every dynamic writer +stream to be closed by authenticated end-of-stream or an acknowledged +submission fence. An open stream remains a possible future accessor even when +all currently known operations are quiescent. + +| Resource | Earliest safe release | +|---|---| +| Source KV, bounce writer | Gather fence completed; NIXL subsequently reads only the send-bounce slot. | +| Source KV, direct writer | The NIXL operation is definitively quiescent. | +| Send-bounce slot | The NIXL operation reading it is definitively quiescent. | +| Receive-bounce slot | Every writer operation that may have received its address is quiescent, and scatter either completed or was conclusively suppressed. | +| Destination KV, direct writer | Every direct writer operation targeting those blocks is quiescent. | +| Destination KV, bounce writer | All bounce writers are quiescent and successful scatter completed, or scatter was conclusively suppressed because the data outcome failed. | +| Destination KV, mixed fan-in | All direct writers are quiescent, all bounce writers are quiescent, and scatter completed or was conclusively suppressed. | +| Descriptor/plan backing storage | The backend has copied it synchronously, or every asynchronous consumer is quiescent. | +| Gather/scatter CUDA event | Completion has been observed and no queued work or callback still references it. | +| Transfer context | All physical leases are released exactly once; only lightweight result-delivery/tombstone state may remain. | +| `EndpointTransferRecord` and gate | Local producer sealing or an acknowledged creation fence proves no later context can be created; every context is retired; every stream and pre-cancel obligation is settled; and its generation is atomically marked retired in endpoint replay state before record removal. | +| Endpoint replay entry | Its generation is at or below the contiguous retired floor. Sparse entries compact only when every preceding generation is retired or explicitly creation-closed; endpoint recreation uses a new epoch. | +| Arena registration and VMM mapping | All contexts that can access that arena are retired; all local CUDA work is complete; and any missing per-operation network evidence is replaced by backend-wide quiescence. | + +Resources MAY be released independently at their earliest safe boundary. For +example, bounce gather completion can release source KV while the send-bounce +lease remains active through NIXL completion. The context remains registered +until all of its responsibilities are settled. + +### Relationship to `LlmRequest` + +The request and transfer lifetimes largely overlap, so they should be associated +through a `TransferAdmissionToken` that atomically binds to the registry-owned +canonical `TransferHandle`. After binding, the request/session holds only a +consumer-facing reference. They must not be collapsed into one owner: + +- a request can become logically terminal before DMA or scatter retires; +- one request can own several independently progressing slices; +- transport result handling must survive session removal; +- retaining a request object does not necessarily prevent explicit KV + `free_resources()`; +- transport workers should not mutate scheduler-facing request state directly. + +Request cleanup does not perform a check-then-free against `TransferHandle`. +Instead, `free_resources()` atomically drops request ownership in the KV manager; +leased allocation generations become pending-free and are reclaimed +automatically when the last transfer lease is released. This avoids a race +between checking the handle and acquiring a new lease. Dropping the +consumer-facing handle reference likewise cannot destroy the canonical +endpoint-transfer record or its access gate. + +### Cancellation API contract + +The current `KvCacheTransceiver.cancel_request()` boolean combines two different +questions: whether logical cancellation was accepted and whether KV is already +safe to free. The ownership implementation separates them. Cancellation returns +a structured result with the logical outcome and an aggregate physical +disposition of `RETIRED`, `DRAINING`, or `IN_DOUBT`. `DRAINING` or `IN_DOUBT` +never delays the request-visible cancellation result. + +Local cancellation is routed through the request's `TransferAdmissionToken`. +While admission is pending, cancellation aborts that token; after binding, the +token forwards to the canonical `TransferHandle`, which carries the role, +negotiated local/admission-authority endpoint epochs, transfer ID, transfer +generation, and complete declared local context set. A bare request/transfer ID +is not a valid registry lookup key. A stale token/handle or delayed cancellation +for a prior transfer generation is idempotently rejected and cannot affect a +current context. + +Context preparation, registry insertion, and handle enrollment form one commit +transaction under the registry-to-handle lock order. Every context access +boundary uses the handle-to-context order. Accepted `cancel()` aborts the handle +and snapshots its contexts while holding the same gate, so a boundary cannot +pass between handle abort and later context callbacks. A failed enrollment +rolls back the unexposed context and its leases. After abort, no context can +newly publish an address or launch gather/NIXL/scatter work; an accessor that +won the gate first remains owned and drains normally. + +Remote cancellation follows the same rule. `CANCEL_SESSION` and any future +lifecycle-control message carry the negotiated protocol version, sender and +receiver endpoint epochs, transfer ID, and transfer generation. The registry +validates that envelope before session lookup. Under the registry lock it either +aborts the canonical handle or marks its admitted zero-context record +pre-cancelled, keyed by the full endpoint/transfer identity and never a bare +`unique_rid`. Context creation checks that same record in its commit transaction, +so a racing cancellation cannot be lost. A consumer that attaches later receives +the same aborted handle through its original token or an explicitly authorized +attach path rather than constructing a replacement. + +Pre-cancel state cannot expire by age. It retires only after local producer +sealing or the negotiated Phase 1 creation-close acknowledgement makes later +context creation impossible. Capacity was reserved when the transfer generation +was admitted; exhaustion backpressures admission before the peer is authorized +and emits configured/used/remaining capacity, oldest-age, rejection, and +creation-close retirement metrics. A prior-generation record cannot cancel a +later transfer generation. + +Once the endpoint has allocator-enforced leases, PyExecutor may immediately +drop request ownership of KV. The KV manager's `free_resources()` makes the +affected allocation generations pending-free, while the registry-owned leases +prevent reuse until retirement. +This permission is resource-specific: it does not close an `RxSession` or free a +generation-first auxiliary slot whose independent owner has not retired. +PyExecutor MUST NOT poll KV transfer task state and interpret a boolean as proof +of KV memory safety. Physical status remains available for diagnostics, +capacity accounting, and shutdown. + +The structured lifecycle result is a Python-native integration, not a required +behavior change for the C++ transceiver. PyExecutor selects a lifecycle-capable +adapter when constructing `KvCacheTransceiverV2`; other implementations keep the +existing boolean contract. The selection is explicit per transceiver instance, +not a runtime union inferred from a returned value. The shared teardown ledger +and its fail-closed checks are part of that selected Python-native path. Other +transceivers keep baseline local manager teardown; this adapter neither audits +nor retires a C++ transport owner or interprets local C++ object destruction as +physical drain. + +Phase 1 applies this behavior to the receive side only after the destination KV +lease exists. The context-side sender retains the legacy cleanup gate until the +source lease lands in Phase 2; Phase 1 therefore cannot claim complete +sender-side cancellation ownership. + +## State Model + +One enum must not conflate request outcome with memory safety. + +### Logical outcome + +```text +ACTIVE -> SUCCEEDED +ACTIVE -> FAILED +ACTIVE -> CANCELLED +``` + +Logical failure may occur on the first writer failure. Logical cancellation may +occur on client cancellation, deadline, consensus outcome, or shutdown. Neither +transition implies physical retirement. + +Logical success requires a sealed handle, every required fixed operation or +closed dynamic writer stream to succeed through its declared high-water mark, +every fixed manifest and dynamic stream to satisfy its exact expected mapping, +and all required scatter work to complete successfully. + +The first logical terminal outcome wins. Failure, cancellation, or shutdown +atomically aborts the handle gate before callbacks, so later chunks cannot +enroll or cross access boundaries. For example, cancellation already reported +to the consumer is not overwritten by a later writer failure. Later events still +update physical and process-health state for already-possible accessors. + +### Physical access, target, and data state + +Each writer operation tracks separate dimensions rather than one overloaded +terminal enum. + +**Exposure state:** + +```text +PLANNED -> NEVER_EXPOSED +PLANNED -> MAY_ACCESS +MAY_ACCESS -> NEVER_EXPOSED (only with positive non-delivery proof) +``` + +- `PLANNED`: leases exist, but publication has not started. +- `NEVER_EXPOSED`: positive proof exists that no address was observable. It is a + physically terminal exposure state. +- `MAY_ACCESS`: publication started, submission may be active, or outcome is + uncertain. + +**Access state:** `NOT_STARTED`, `POSSIBLE`, or `QUIESCED`. `QUIESCED` requires +positive per-operation evidence or a backend-wide drain that covers the +operation. A generic failure without that guarantee remains `POSSIBLE`. + +For a dynamic writer stream, individual operations have these states, while the +advertised target remains `POSSIBLE` for future operations until authenticated +end-of-stream or an acknowledged submission fence closes both enrollment and +submission. Quiescence of the currently known operations alone cannot retire an +open stream target. + +**Target mode:** `PENDING`, `DIRECT`, `BOUNCE`, `NO_REMOTE_ACCESS`, or `UNKNOWN`. +The pre-publication plan contains authorized direct and bounce candidate ranges; +the actual mode is selected later. + +**Data outcome:** `UNKNOWN`, `SUCCESS`, `FAILURE`, `ABORTED`, or `INVALID`. +Physical quiescence and data validity are independent. For example, a valid +quiescence-bearing NIXL result with malformed scatter metadata makes the data outcome +`INVALID` and suppresses scatter, but it can still prove that the network +accessor is `QUIESCED`. Backend-wide drain after a lost result can produce +`QUIESCED` with data outcome `UNKNOWN`; the logical transfer fails and no +scatter occurs. + +### Local accessor state + +Gather, each sender-side NIXL operation, and scatter are local asynchronous +accessors with their own state: + +```text +PLANNED -> MAY_ACCESS -> QUIESCED +PLANNED -> QUIESCED (positive proof that launch/queue/submission did not occur) +``` + +The preparation transaction and required leases MUST exist before resolving raw +addresses or constructing plans. The context MUST be committed before a plan or +address becomes observable to an asynchronous accessor. The accessor enters +`MAY_ACCESS` through the handle-to-context cancellation gate before CUDA launch, +queue publication, or NIXL submission. `OPEN` and `SEALED` handles permit an +already-enrolled context to cross a planned boundary only when its context state +also permits the accessor; `ABORTED` does not. +Synchronous failure before launch may move directly to `QUIESCED`; an exception +after an ambiguous launch/submission stays `MAY_ACCESS` until positive evidence +arrives. Worker death does not imply quiescence. + +### Publication contract + +Network send cannot be made atomic with local cancellation. The receiver and +sender therefore have separate ordered boundaries. + +Receiver target publication: + +1. Atomically snapshot/lease destination KV and reserve any receive-bounce + candidate. +2. Build either the exact fixed-operation manifest or the exact authorized + writer-stream ledger and candidate ranges, then transactionally insert the + receive context into the registry. +3. For each writer operation/stream, under the handle cancellation gate and then + the context lock, set exposure to `MAY_ACCESS` and access to `POSSIBLE` + before attempting to send any target address. +4. Release the lock, publish the advertisement, and record the messaging outcome + without weakening `MAY_ACCESS` unless non-delivery is proven. The + advertisement carries its target slice plus exact fixed operation ID or + writer-stream ID. Replaying the same advertisement cannot launch a second + stream/operation. + +Sender local access and submission: + +1. Validate and deduplicate the received transfer, target-slice, and fixed + operation or writer-stream identity. +2. Begin a preparation transaction. For a dynamic stream, reserve its next + operation sequence and subrange inside that transaction; neither is externally + visible yet, and both remain within the advertised stream budget. +3. Atomically snapshot/lease source KV, resolve addresses only from that lease + token, reserve any selected send-bounce slot, and build descriptor/plan + storage. Commit the operation sequence, send context, and handle enrollment + together before using an address asynchronously. A pre-boundary failure + either rolls back an unexposed reservation or records a terminal gap that + prevents successful end-of-stream validation. +4. Through the handle cancellation gate, mark gather `MAY_ACCESS` before launch. + If the sender falls back to direct, record that mode without constructing a + gather accessor. +5. Through the same gate, mark the NIXL operation `MAY_ACCESS` before submission + and retain the source or send-bounce lease through definitive quiescence as + required by the path. +6. Record and deliver the result without coupling sender resource retirement to + receiver session lifetime. For a dynamic stream, atomically enroll the final + operation before sealing the sender handle. Emit authenticated end-of-stream + only after every sequence through its high-water mark is already submitted or + terminal `NO_REMOTE_ACCESS`; the transition closes future submission before + the message is sent. A missing/duplicate sequence or later submission attempt + prevents logical success and fails closed. + +Cancellation can retire an endpoint context immediately only when each remote +operation was `NEVER_EXPOSED` or its access state is `QUIESCED` (including +positive `NO_REMOTE_ACCESS` evidence), every dynamic stream is closed to future +submission, **and** every local gather, NIXL, scatter, or queued accessor is +`PLANNED` or `QUIESCED`. Otherwise cancellation changes only the logical outcome +and starts drain; it does not release memory. In particular, a sender gather may +already be active before NIXL submission, so proof that no remote operation was +submitted is insufficient for immediate retirement. + +For partial fan-out publication, operations that never crossed the boundary +become `NEVER_EXPOSED`; operations that crossed it drain independently. The +request may fail immediately, but the context retires only after all possible +accessors do. + +### Direct, bounce, and fallback mode + +The receiver may offer a bounce target while the sender later chooses direct +fallback. Until the actual mode is known, the receiver conservatively holds both +the destination KV lease and any offered receive-bounce lease. + +The writer-operation ledger records target mode independently from exposure, +access, and data state. The mode is initially `PENDING` and normally becomes: + +- `DIRECT`: destination KV access retires when the writer operation is quiescent. +- `BOUNCE`: receive-bounce access retires when the writer operation is quiescent; successful + writers contribute a validated scatter plan. + +`NEVER_EXPOSED` is not a transfer mode. It is a physical exposure state proving +that neither candidate target was observable to that writer. + +`NO_REMOTE_ACCESS` means the address may have been exposed, but the sender has +positive evidence that it did not submit a remote operation. `UNKNOWN` means the +backend or a global drain established quiescence without reconstructing which +candidate target was used. + +A failed transfer is never scattered into usable KV. Partial direct writes need +not be rolled back; the destination KV lease remains until all writers drain, +then the failed allocation may be recycled. + +### Identity, control, and result routing + +Every wire message that can mutate lifecycle state, including cancellation, +must carry a common identity envelope with: + +- protocol version/capability; +- transfer ID; +- monotonic transfer admission generation shared by every slice in the attempt; +- sender and receiver endpoint/instance epochs, because ranks can be reused + after restart; +- lifecycle scope: transfer-wide, or the exact target slice for a slice-scoped + message. + +Every writer-scoped advertisement, result, or control also identifies the exact +writer and either a fixed writer-operation ID or a writer-stream ID. A dynamic +operation result/control additionally carries its operation sequence. The +target advertisement therefore names the ledger authorization that the sender +must echo; it never publishes an address with only a transfer/request ID. If a +writer needs multiple fixed NIXL operations, each is advertised with a distinct +preplanned ID. Replaying an identical advertisement is idempotent and launches +at most one fixed operation or stream; a conflicting replay fails closed. + +A writer result additionally identifies: + +- source chunk/slice identity when it differs from the target slice; +- actual direct/bounce mode; +- quiescence evidence and data outcome; +- validated scatter metadata when bounce succeeded. + +For a dynamic stream, an authenticated end-of-stream control/result carries the +final operation-sequence high-water mark and aggregate source-to-target +mapping/byte coverage. +End-of-stream is accepted only when every covered sequence is already submitted +or terminal `NO_REMOTE_ACCESS`; it atomically closes future enrollment and +submission but does not make submitted operations quiescent. Missing, +duplicate-conflicting, out-of-budget, overlapping, or out-of-range operations +fail the logical transfer and remain owned according to their physical state. + +The registry validates every state-mutating envelope before any request/session +lookup. Duplicate results and control messages are idempotent. Unexpected writer +operations, contradictory duplicates, wrong transfer generations, and stale +remote cancellations cannot advance any live context. Valid quiescence evidence may +advance physical access state even when result payload or scatter metadata is +invalid; that invalid data fails the logical transfer and suppresses scatter. +All ranges are checked against the registered authorization before use. + +Creating a context with an already-live identity MUST fail rather than replace +the existing context. Transfer generations MUST NOT be reused within an endpoint +epoch. Before deleting a retired canonical record, the registry atomically +marks its generation retired in the endpoint replay window. A later admission, +advertisement, result, or control at or below the contiguous retired floor, or +for an out-of-order retired sparse entry, is rejected before record creation. +Bounded diagnostic tombstones may retain richer error context, but pruning them +does not remove this compact anti-replay state or authorize resource +reclamation. Pre-cancel tombstones follow the stronger O17 retention rule and +are not age-pruned diagnostic entries. + +A result may advance access state to `QUIESCED` only when the sender has positive +evidence that its NIXL operation cannot perform later DMA. It is not enough to +report that the request was cancelled or that a bounded wait elapsed. + +### Lost results and quarantine + +In the first implementation, a receive context whose quiescence-bearing result +is lost remains in doubt until an externally established remote/global fence or +driver/fabric-level revocation explicitly prevents every relevant peer from +submitting or completing later DMA. Shutting down only the receiver's local +agent, losing Python objects, running destructors, or entering generic +process-control teardown is not evidence. The memory is not reused. This can +reduce capacity and eventually reject admission, but it does not trade memory +safety for a timer. + +Bounded in-process reclamation requires one of: + +- reliable result acknowledgement and retry; +- a receiver query that obtains per-operation quiescence evidence; +- a peer/session epoch transition with a backend guarantee that old operations + cannot access current registrations; or +- backend-wide quiescence backed by an external remote/global fence. + +Those are liveness improvements, not permission to weaken the baseline safety +contract. + +### Unknown writer cohort + +Some gen-first ADP flows broadcast discovery/request messages before the active +DP cohort is known. Before publishing memory addresses, the implementation MUST +either: + +1. select the exact fixed-operation or authorized writer-stream cohort through + an address-free handshake; + or +2. register every recipient as a potential writer operation and obtain a + quiescent `NO_REMOTE_ACCESS` acknowledgement from every non-participant. + +If neither is implemented, that topology is excluded from the first ownership +rollout even though it remains an eventual design target. An expected writer +count is not an acceptable substitute. + +### Threading and callbacks + +State changes are serialized per context. The implementation must not hold an +admission-token, registry, handle, or context lock while: + +- sending or receiving network messages; +- waiting for NIXL; +- synchronizing CUDA work; +- allocating or freeing KV/bounce memory; +- invoking request/session callbacks. + +Queue entries for gather/scatter work hold a strong context reference rather +than only raw slot IDs and callbacks. Consumer callback exceptions are recorded +but cannot interrupt physical finalization. + +## Proposed Interfaces + +Names are illustrative; behavior is normative. + +```python +class TransferAdmissionToken: + def cancel(self, reason: str) -> CancellationResult: ... + + +class TransferRegistry: + def admit_and_bind( + self, token: TransferAdmissionToken, admission: TransferAdmission + ) -> TransferHandle: ... + def prepare_send(self, plan: SendPlan) -> SendTransferContext: ... + def prepare_receive(self, plan: RecvPlan) -> RecvTransferContext: ... + def route_result(self, result: WriterOperationResult) -> None: ... + def route_control(self, control: LifecycleControlMessage) -> None: ... + def begin_shutdown(self) -> None: ... + def drain(self, deadline: float | None) -> DrainResult: ... + + +class RecvTransferContext: + def publication_started( + self, authorization: WriterAuthorizationId + ) -> PublicationToken: ... + def mark_never_published( + self, authorization: WriterAuthorizationId, proof: NonDeliveryProof + ) -> None: ... + def record_result(self, result: WriterOperationResult) -> None: ... + def record_stream_end(self, end: WriterStreamEnd) -> None: ... + def record_global_quiescence(self, evidence: QuiescenceEvidence) -> None: ... + def cancel_consumer(self, reason: str) -> None: ... + def detach_consumer(self) -> None: ... + + +class KVTransferLease: + @property + def ranges(self) -> tuple[LeasedRange, ...]: ... + + def release(self) -> None: ... + + +class KVLeaseProvider: + def snapshot_and_lease( + self, request: LlmRequest, slice_spec: SliceSpec + ) -> KVTransferLease: ... + + +class TransferHandle: + def seal(self, proof: SliceSetProof) -> None: ... + def record_failure(self, reason: str) -> LifecycleResult: ... + def cancel(self, reason: str) -> CancellationResult: ... +``` + +Required properties: + +- Methods are thread-safe and idempotent. +- `admit_and_bind()` reserves capacity, creates exactly one canonical + `EndpointTransferRecord`/handle per full identity, and moves the one-shot + token directly from `PENDING` to `BOUND(handle)` before the peer is authorized + to send addresses or lifecycle controls. An identical retry through the same + token returns the same handle. A different token for an already-live identity, + or any conflicting/retired identity, is rejected and cannot attach to or + replace canonical state; any future consumer-attach API must authorize that + role explicitly without granting replacement cancellation authority. +- `admit_and_bind()` classifies the generation against endpoint replay state + while holding the token gate and then the registry lock. It rejects an + aborted token and generations at/below the retired floor or in the sparse + retired set, returns the canonical live handle only for an exact retry through + that same bound token, and backpressures a new generation that would exceed + the bounded replay window. +- On the admission initiator, generation allocation is part of that successful + commit; pre-commit rejection does not advance the counter. A committed + generation whose admission control is not sent must be acknowledged as a + generation skip before a higher admission control is sent. If admission was + sent but no address published, peer rejection/abandonment uses authenticated + creation-close/reject on both endpoints. After address publication, those + empty-generation controls cannot retire the record or its resources. +- `begin_shutdown()` closes registry admission and snapshots committed records + under the registry lock. A losing pending token becomes shutdown-aborted; a + winning record is in the snapshot. Peer authorization is a handle-gated + boundary, so it either becomes owned before abort or is suppressed afterward. +- `prepare_send()` and `prepare_receive()` resolve the canonical record from the + full identity in their plan; callers cannot inject a replacement handle. +- Registry insertion and handle enrollment are one preparation commit with a + registry-to-handle lock order. Access boundaries use handle-to-context order. + `SEALED` rejects new enrollment but permits planned boundaries; `ABORTED` + rejects both enrollment and every later boundary. +- `SliceSetProof` is an immutable expected target-slice/source-chunk manifest or + an atomically enrolled final-chunk proof. Only the endpoint's local + producer/session may seal its handle; a timer, currently known chunk count, or + consumer completion is not sufficient. +- `WriterAuthorizationId` is a `FixedOperationId` or `WriterStreamId`, so the + publication interface marks either negotiated authorization possible. +- Fixed-operation manifests and bounded dynamic writer streams are explicit + negotiated modes. Stream end closes enrollment and submission only after + every covered sequence is submitted or terminal `NO_REMOTE_ACCESS`; it does + not imply quiescence for submitted operations. +- After transactional construction commits, resource release is private to the + context. Receiver preparation may unwind only before target publication; + sender preparation may unwind only before its local gather/NIXL launch. +- Request-facing handles cannot force physical retirement. +- `KVTransferLease` prevents allocator reuse even if request cleanup runs. +- `KVTransferLease.ranges` is an immutable allocator-produced snapshot carrying + logical KV coordinates and allocation generations; contexts do not + reconstruct it from copied block IDs or positional order. +- `BounceTransport` issues optional `SendBounceLease` and `RecvBounceLease` + objects; it remains the arena owner. +- `NoBounceTransport` issues no bounce lease, but the direct path still creates + send/receive contexts and KV leases. +- Invariant violations retain resources and emit diagnostics rather than trying + to recover through reuse. +- `CancellationResult` reports logical cancellation separately from physical + state. `DrainResult` and the transceiver-level `ShutdownResult` distinguish a + fully drained endpoint from a retryable non-drained endpoint. +- Local cancellation uses the one-shot admission token and, after binding, its + transfer-generation-bound `TransferHandle`; no lifecycle mutation looks up a + context by bare request/transfer ID. +- Pre-cancel lookup and handle/context insertion serialize in the registry. A + live unmatched pre-cancel tombstone cannot be evicted to admit new work. +- The registry retains the canonical record/handle/gate after every consumer + reference is dropped. It removes the record only after no later local context + can be created and all contexts, streams, and pre-cancel state satisfy their + retirement predicates, atomically marking the generation retired in endpoint + replay state before deletion. +- Endpoint replay state outlives transfer records and sessions within one + endpoint epoch. Out-of-order retired entries compact only by advancing a + contiguous floor; a gap backpressures admission rather than being evicted. + +### NIXL status contract + +The lifecycle owner needs at least: + +```text +IN_PROGRESS +QUIESCED_SUCCESS +QUIESCED_FAILURE +QUIESCED_ABORTED +IN_DOUBT +``` + +The current Python-facing boolean wait is insufficient for bounded drain if it +maps both timeout/in-progress and failure to `False`. The adapter may expose a +tri-state poll plus a separate quiescence guarantee, or a richer enum such as +the above. Exact names are not important; distinguishing “still possible” from +“cannot access memory” is mandatory. This classification is correctness +infrastructure required by the first retirement implementation. A richer +per-operation cancel/recovery API may remain a later liveness improvement. + +## Shutdown Contract + +When the Python-native lifecycle adapter is selected, `KvCacheTransceiverV2` is +the top-level owner of shutdown progress. Its `shutdown(deadline)` returns a +`ShutdownResult` rather than treating initiation as finalization. `DRAINED` +permits dependent teardown. `RETRYABLE_IN_DOUBT` requires the caller to retain +the transceiver, worker, registry, KV managers, and registrations and to retry +drain or escalate process health; it does not permit resource-manager teardown. +Other transceivers bypass this Python ownership drain, veto ledger, and its +rank-uniform votes and retain baseline local manager teardown. This adapter does +not derive quiescence for, audit, or retire a C++ transport owner. + +Shutdown follows this order: + +1. Stop request transfer-eligibility/token creation, then atomically close + registry admission and snapshot all committed records. A pending + `admit_and_bind()` either committed before the snapshot or is rejected with + its token shutdown-aborted. Gate and suppress any peer authorization or + address publication that did not already become possible. +2. Mark remaining consumers failed/cancelled, but keep result listeners, + transfer workers, and CUDA completion workers alive. +3. Start draining endpoint contexts while registrations and progress workers + remain available. +4. For every fixed operation or dynamic stream that may have received an + address but is not closed to future submission, establish a documented + submission fence. Its acknowledgement MUST guarantee that the peer cannot + later submit against that advertisement. A terminal `NO_REMOTE_ACCESS` + result closes its fixed operation; authenticated end-of-stream closes future + enrollment/submission for a dynamic stream, but neither substitutes for + quiescence of already-submitted work. +5. For submitted or still-in-doubt network operations, request per-operation + evidence or backend-wide quiescence through a documented external + remote/global `quiesce()`/drain fence before the deadline. Local agent + shutdown cannot satisfy this step. +6. If the deadline expires without the required operation/stream closure or + submission fence and quiescence evidence, return a non-drained result and + stop teardown at this point. +7. Synchronize remaining gather/scatter work. +8. Release transfer leases and remove retired contexts. +9. Confirm that any non-KV registration owners, including auxiliary transfer + paths, have independently drained while shared listeners, progress, and + completion workers remain alive. If this barrier misses the deadline, return + a non-drained result and stop teardown at this point. +10. Stop listeners and completion workers; then deregister memory and destroy + VMM mappings. +11. Destroy the backend agent. + +Draining only operations that are already visible to the local backend is not +sufficient. For example, a sender may pause after receiving a target address +and submit against it later. Until a terminal result or acknowledged +peer/session/endpoint fence makes that future submission impossible, the +advertised target remains live and shutdown is non-drained. + +The transport likely needs three ordered phases: fence future submission while +control paths are live, quiesce submitted work while registrations are usable, +then perform final destruction after deregistration. If the deadline expires +and either fencing or quiescence cannot be proven, shutdown returns +`RETRYABLE_IN_DOUBT` and must not deregister, unmap, reuse, or drop the affected +memory owners. The transceiver does not enter its final `_shutdown` state. The +worker remains alive and retryable, and listeners/completion workers remain +available for a later drain attempt. `PyExecutor` and executor-creation cleanup +MUST inspect this result and MUST NOT call KV-manager shutdown while any lease +is live. Destructors MUST NOT perform fallback unmapping. The caller may +escalate process health through the adjacent cancellation/poison policy, but +in-process cleanup fails closed. + +## Protocol and Rolling Upgrade Contract + +PR #15618 changed advertisement and result framing without a general protocol +negotiation layer. This follow-up must not add transfer-generation, +endpoint-epoch, stream/mode, or cancellation/control fields with another +implicit same-version assumption. + +Before any address is published, peers must either: + +1. negotiate a protocol version/capability set that includes the required + identity, result, and lifecycle-control fields; or +2. reject the mixed-version session explicitly. + +Silent downgrade is not allowed when it would remove transfer-generation-safe +identity, exact writer/stream accounting, or quiescence/mode reporting. The +negotiation approach should align with the protocol work in PRs #15798 and +#15799 where practical, without coupling this Python implementation to the C++ +transceiver internals. + +Phase 1 MUST introduce the negotiated transfer generation and both endpoint +epochs on every lifecycle-mutating advertisement, result, and control message. +The existing `unique_rid`-only identity is not a permitted compatibility mode, +regardless of an attempted non-reuse proof. Peers that do not negotiate this +identity are rejected before address publication. + +The protocol assigns one endpoint as the admission initiator. Its transfer +generation is monotonically allocated within its endpoint epoch, and the +negotiated maximum outstanding-generation window bounds peer replay state. +Admission outside that window is backpressured. Reconnecting without the +existing endpoint replay state requires a fresh endpoint epoch; a delayed +message from the old epoch cannot open or attach to the new session. + +Healthy per-transfer admission fields ride the existing KV request/target +response; they do not create another request/ack exchange before address +publication. Capability negotiation is session-scoped. Exceptional skip/close +and rejection acknowledgements remain explicit, retryable controls. + +Phase 1 negotiation also covers transfer admission, the initial replay +floor/window, fixed-manifest/context capacity envelopes, and the +creation-close/cancel acknowledgement used to retire zero-context lifecycle +state. Phase 2 adds the bounded dynamic-stream/end-of-stream contract. Phase 3 +adds the stronger submission fence for published addresses. A peer that lacks +the capability required by the selected phase is excluded before that phase's +access boundary; no control is interpreted using a weaker earlier-phase +meaning. + +## Implementation and Validation + +See [`implementation-plan.md`](implementation-plan.md) for the performance and +capacity contract, observability requirements, phased implementation plan, +validation matrix, risks, alternatives, and phase gates. diff --git a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md new file mode 100644 index 000000000000..1bdb82c8d4d4 --- /dev/null +++ b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md @@ -0,0 +1,1065 @@ + + +# Implementation Plan for Python Native Disaggregated KV Transfer Ownership + +[< Back to ownership contract](README.md) + +| Field | Value | +|---|---| +| **Owner** | Chien-Chun Hung | +| **Status** | Draft implementation plan | +| **Created** | 2026-07-08 | +| **Last updated** | 2026-07-21 | + +## Performance and Capacity Contract + +Ownership safety must not undo the performance purpose of bounce transfer. + +- Direct transfer adds no data copy and no new CUDA synchronization. +- Registry/identity lookup and fixed-operation state transitions are O(1) + average. Per-segment bounds checks are O(segments), and fixed exact-mapping + normalization is O(segments log segments) unless canonical wire ordering + permits a linear scan. Dynamic-stream interval insertion/overlap validation is + O(log(operations + segments)), and end-of-stream coverage/compaction is + O(operations + segments), all within negotiated bounds. Metadata is + O(contexts + writers + operations + segments) per transfer and bounded by + fixed-manifest or dynamic-stream envelopes plus endpoint-global capacity. +- Locks are per context or sharded; no global lock is held across NIXL or CUDA + work. +- Executor polling remains bounded and non-blocking. +- Session capability negotiation is amortized, and healthy per-transfer + admission piggybacks on the existing KV request/target response without a new + request/ack RTT. +- Source KV is released at its earliest safe per-path boundary, enabling chunked + early release rather than pinning the whole request until all slices finish. +- Bounce arena sizing remains configuration-driven and opt-in; requested, + transport-active, and per-transfer-engaged states are distinguished. + +Fail-closed retention can consume KV or bounce capacity. This is intentional +under uncertainty and must be visible operationally. It is preferable to an +invisible cross-request overwrite. Bounded reclamation is a follow-up enabled by +stronger transport/protocol evidence. + +### Expected metric impact + +This is primarily a correctness and capacity-recovery project, not a NIXL +bandwidth optimization. On a healthy transfer, the expected data path is +unchanged. The owner adds host metadata, lookup, and state transitions, but no +new direct-path CUDA copy, kernel, synchronization, or device arena. When the +requested bounce transport is active, the configured send/receive arenas +introduced by PR #15618 remain the only preallocated bounce storage. + +For the initial containment implementation specifically, no TTFT, throughput, +or NIXL bandwidth improvement should be claimed. The healthy sender adds an O(1) +operation-ledger lookup plus small state/counter transitions, exact-session +queue binding, and serialized socket access; the receiver adds registry +transitions and bounds checks. Duplicate request delivery is cheaper because it +replays or joins the existing operation instead of launching another transfer. +Lazy destination-manifest construction can reduce host allocation/CPU work when +a bounce slot is unavailable and the operation falls back to direct transfer, +but it does not change transferred bytes. The slice allocates no additional +device arena beyond the opt-in bound buffers from PR #15618. Under cancellation, +missing results, or ambiguous NIXL status, retained Python owners and bounce/KV +capacity can increase memory residence and reduce admission throughput; that is +the intended fail-closed result. The detailed metrics below are the target +instrumentation contract and are not all exported by this first slice. + +A non-empty initial legacy bounce result tail is not an exact fixed manifest. It +proves the expected bounce-source base, exact aggregate writer bytes, +receiver-authorized destination-union containment, and disjoint destinations. +It cannot reject an equal-sized source-to-destination permutation within the +sole enabled writer (or across writers if multi-writer bounce is later enabled) +because the frame does not carry the receiver's expected mapping. Exact +manifest matching and permutation rejection remain O13 follow-up gates; the +containment slice keeps every multi-writer cohort on the direct path until +byte-accurate per-writer extents can be authorized. Single-writer bounce also +falls back to direct transfer unless the complete receiver pool mapping is +provably a byte-equal bijection: both peers use equal tokens-per-block, every +attention pool view maps exactly once through `IdentityMapper`, and each mapped +physical pool has equal slot bytes. + +| Scenario | Expected impact | +|---|---| +| Healthy direct or bounce transfer | With admission piggybacked on the existing request/target response, throughput, TTFT, and transfer-task latency should remain within benchmark noise. CPU time and host memory rise slightly with O(contexts + writer operations + descriptor/scatter segments + lifecycle records + outstanding generations) metadata and state transitions. An implementation that adds a per-transfer RTT must declare and qualify that separate behavior. | +| Chunked bounce send | Unique transfer-leased source-KV byte-seconds/high-water mark can decrease because a slice may release source KV after gather completes rather than at request-wide completion; the pending-free subset shows actual admission capacity recovered after request cleanup. | +| Chunked direct send | Unique transfer-leased source-KV retention can decrease when each slice is released after its NIXL operation becomes definitively quiescent. There is no extra copy. | +| Known-terminal failure that currently strands state | Registry-first routing should recover pending-free KV and unique bounce-slot capacity as soon as the last physical accessor retires, reducing excess post-logical retention, admission loss, and fault-period goodput degradation. | +| Known-terminal direct/mixed failure that currently frees KV too early | The owner intentionally increases lease duration and possibly unique/pending-free high-water marks to the true quiescence boundary. That increase closes unsafe reuse and must not be classified as a regression. | +| Lost result or ambiguous transport failure | Exact unique bounce-slot bytes and unique/pending-free KV bytes remain unavailable, so high-water marks, admission failures, and tail latency can worsen as safe capacity is exhausted. This is an intentional safety tradeoff until bounded recovery exists. | +| Many dynamic streams reserve maximum envelopes but materialize little work | Host allocation may stay low while reserved credits reduce admissible concurrency and goodput. Reservation utilization/slack and reservation-caused rejection determine whether a configured envelope is rollout-safe. | +| Shutdown with in-flight work | Graceful shutdown may take longer or return non-drained. Unsafe deregistration/unmapping events must fall to zero. | + +The containment implementation also moves fallible local retirement preparation +before the transceiver's existing outcome-consensus barrier. It re-advertises a +provisional candidate until every rank reports cleanup-ready and adds no new +healthy-path collective round. Cancellation may be conservatively delayed while +one collective-bearing status/preparation call owns a snapshot; this affects +only overlapping control work, not the data path. + +Request-level response arbitration likewise has no expected healthy-path metric +gain. It prevents duplicate terminal responses across native, connector, and +cancellation paths. If the response sink raises after its side-effect boundary, +the implementation refuses replay because the sink exposes no acknowledgement; +that trades a possible lost response for duplicate prevention and should be +counted separately from transfer failures. + +The two extra rank-uniform votes occur only during Python-native executor +shutdown. The first +follows connector completion draining and fallible sampler/DWDP finalization, +then gates entry to resource-manager shutdown. The second publishes the outcome +of the non-replayable, per-manager shutdown phase before any engine deletion; +the phase stops at its first ambiguous manager so later dependencies remain +live. The votes do not affect +TTFT or steady-state throughput. A connector-disabled shutdown executes these +two reductions in total; a connector-enabled shutdown executes three, including +its earlier connector-drain outcome consensus. This cost applies only to an +otherwise-ready Python-native shutdown attempt and prevents rank-divergent retry +deadlocks before and after manager destruction begins. Default C++-transceiver +and non-disaggregated shutdown retain their existing local-only manager teardown. + +The main capacity metrics are: + +```text +lease_token_byte_seconds = sum(token_range_bytes * token_lifetime) +unique_transfer_leased_byte_seconds = integral(union of ranges with transfer_leases > 0) +pending_free_kv_blocked_byte_seconds = integral(bytes with request_owners == 0 and transfer_leases > 0) +reservation_utilization(kind) = materialized_credits(kind) / reserved_credits(kind) +reservation_slack(kind) = reserved_credits(kind) - materialized_credits(kind) +post_logical_retention = max(0, last_lease_release - logical_terminal_time) +early_release_lead(resource) = max(0, logical_terminal_time - lease_release_time) +``` + +Lease-token byte-seconds intentionally count overlapping tokens and measure +ownership bookkeeping. Unique transfer-leased bytes deduplicate overlapping +ranges by allocation generation. Pending-free blocked KV bytes are the subset +whose reuse is prevented solely by transfer leases; KV capacity-recovery and +admission claims use that subset rather than token sums. Token metrics may be +attributed to each operation path, but deduplicated unique and pending-free +capacity gauges are path-neutral. An optional path breakdown must use mutually +exclusive `direct_only`, `bounce_only`, and `mixed` allocation-generation +classes so its series remain additive. `post_logical_retention` +measures failure/cancellation drain without becoming negative when a source +chunk retires early. Per-resource `early_release_lead` and lease duration +measure that intended early release. Improvements should be claimed only for +the applicable scenario. A normal-path latency or bandwidth improvement is not +an expected outcome of ownership alone. + +## Observability + +At minimum, expose counters or structured logs for: + +- active send and receive contexts; +- active fixed-operation manifests and dynamic writer streams, including open + streams and their enrolled/high-water-mark operation counts; +- configured, reserved, materialized, and remaining context/slice, + fixed-manifest operation/descriptor-segment/authorized-byte, and + dynamic-stream operation/segment/byte capacity, plus reservation utilization, + unused slack, envelope/global-capacity rejection, and reservation-caused + rejection; +- contexts logically terminal but physically draining; +- in-doubt context count, bytes, age, and reason; +- pre-cancel tombstone count, oldest age, capacity rejections/backpressure, and + fence-based retirement; +- configured, used, and remaining canonical lifecycle-record/control-state + capacity; +- configured, used, and remaining endpoint replay-window capacity, retired-floor + advancement, sparse generation gaps, oldest gap age, and backpressure; +- active source/destination KV leases; +- active send/receive bounce leases; +- lease-token bytes/byte-seconds by endpoint and operation path; path-neutral + unique transfer-leased bytes/byte-seconds and pending-free KV + bytes/byte-seconds blocked solely by transfer leases, optionally split only + into mutually exclusive `direct_only`, `bounce_only`, and `mixed` classes; +- bounce requested, transport active/inactive, and direct, bounce, or fallback + writers so configured-but-inactive runs cannot claim bounce coverage; +- gather, NIXL, scatter, registry lookup, and lifecycle-lock wait latency; +- transfer admission latency and bounded capacity-rejection latency; +- healthy piggyback and exceptional skip/close/reject control-message counts and + round-trip latency, so an accidental normal-path RTT is visible; +- logical completion time, last lease release, post-logical retention, and + per-resource early-release lead; +- duplicate, unexpected, contradictory, and stale-transfer-generation results; +- partial-publication failures; +- shutdown drain duration and non-drained context count; +- rollback requested, blocked/non-drained, and completed; +- admission failures caused by safely retained capacity. + +Aggregate metrics and benchmark records use low-cardinality dimensions for +ownership mode, negotiated protocol/capability bucket, rollout cohort, +KV-manager implementation, and executor. Operation/token metrics also use path; +deduplicated capacity gauges omit it unless they use the mutually exclusive +classes above. Request/transfer identities remain in structured logs rather +than metric labels. Logs identify transfer ID, target slice, transfer +generation, writer/stream/operation identity, path, logical outcome, physical +state, and held leases without logging raw data contents. + +The existing sender `transfer_latency_ms` timer begins after bounce gather and +the receive task completes after scatter. Raw NIXL throughput can therefore hide +gather/scatter and lifecycle overhead. Benchmarks must report stage timings and +end-to-end TTFT/goodput in addition to the existing NIXL interval. + +## Implementation Plan + +The work should land as a chain of focused PRs rather than one large refactor. +Size estimates are directional and include only production code unless noted. + +### Phase 1 — Receiver safety foundation + +**Size:** Very large; approximately 18–28 production files, 2,000–3,500 +production lines, and 2,000–3,500 test lines. Metrics export or a new V1 lease +binding can increase this estimate. Phase 1 should normally be a short stack: +identity/admission/replay and fixed envelopes; receive contexts plus KV leases; +then cancellation, shutdown veto, and minimum metrics. None is independently +rollout-safe until the full Phase 1 gate passes. + +- Add `TransferRegistry`, registry-owned canonical `EndpointTransferRecord`, + durable `EndpointEpochState`, access gate/handle, and `RecvTransferContext` + above the bounce package. Requests/sessions hold only consumer references; + contexts retain the shared gate without forming a record-context ownership + cycle. +- Create a one-shot `TransferAdmissionToken` atomically with the request/session + state transition that enables transfer scheduling. Implement token-to-registry + `admit_and_bind()` so local cancellation either prevents token creation, + aborts pending admission, or observes/aborts the canonical bound handle before + any peer authorization. +- Make `begin_shutdown()` close registry admission and snapshot all committed + records in one registry transaction. A racing pending token is + shutdown-aborted if it loses; a winning record is included, and its later + peer-authorization boundary is handle-gated and either owned or suppressed. +- Assign one negotiated admission initiator. Allocate its transfer generation + monotonically within its stable endpoint epoch, committing counter advancement + only with replay-entry reservation, canonical record creation, and token + binding after capacity checks. Negotiate the authenticated initial floor and + maximum outstanding-generation window, and maintain a contiguous retired + floor plus bounded sparse live/out-of-order-retired window per negotiated + endpoint-epoch pair. Atomically mark a generation retired before deleting its + canonical record; backpressure gaps/window exhaustion rather than evicting + anti-replay state. Convert every committed generation whose normal admission + control is not sent into an authenticated generation-skip record, and block + higher-generation admission control until its acknowledgement. A cumulative + skip watermark cannot cross a locally live/published/unclosed generation. If + admission control was sent but no address published and the peer rejects, + reserve/close the peer replay entry first and retire both sides through an + authenticated creation-close/reject acknowledgement. After address + publication, require the full operation/stream closure and physical-retirement + predicate instead. +- Piggyback generation/epoch/fixed-envelope admission on the existing KV + request and target-advertisement response. Commit the receiver record before + that response; do not add a separate healthy admission acknowledgement. +- Split logical outcome, exposure, access, target, and data state. +- Key contexts by role, local endpoint epoch, and admission-authority epoch; key + receive ledger entries by peer epoch, writer rank, and fixed-operation or + writer-stream ID. +- Create the exact writer cohort, fixed-operation manifest, and range + authorization, including expected logical source-to-target segment mappings, + before publication. Dynamic sender-determined chunk streams are excluded from + Phase 1 rollout and land in Phase 2. +- Negotiate per-transfer context/slice and fixed-manifest + operation/descriptor-segment/authorized-byte limits. Reserve actual + operation/authorized-byte credits and the worst-case permitted segment + envelope atomically against configured endpoint-global capacity before + publication; reject an oversized plan or exhausted pool without partial + exposure. +- Move the publication boundary before message send. +- Put target slice and exact fixed writer-operation identity in each + advertisement. Deduplicate replay before local work so one advertisement can + launch at most one operation. +- Route every KV writer-operation result through the registry before session + lookup; keep auxiliary results on their independent ownership path. +- Route remote cancellation and every state-mutating control message through the + same transfer-generation/endpoint-epoch validation. Replace bare-`unique_rid` + pre-cancellation state with registry-atomic transfer-generation-aware + records. Reserve control-state capacity at admission, before the peer may send + lifecycle messages. An arbitrary unnegotiated control poisons/closes the + endpoint session in O(1); it does not allocate an unbounded deny entry. +- Add a minimum negotiated generation/epoch-scoped creation-close/cancel-ack + handshake in Phase 1. The peer acknowledges only after its local producer gate + prevents later context creation; retry/deduplicate the acknowledgement and use + it to retire zero-context pre-cancel state. This is not the Phase 3 submission + fence and cannot retire any published address. +- Retain unmatched pre-cancel state until producer sealing or acknowledged + creation close, enforce capacity through pre-admission backpressure, and + export configured/used/remaining capacity plus count/age/rejection metrics. +- Drain late sibling results after logical failure/session removal. +- Replace timed bounce reuse with indefinite in-doubt retention. +- Keep the initial containment's renamed `RecvBounceContext` narrow and local + to the bounce package; optionally reduce it further to `RecvBounceLease` + state rather than promoting it into the general owner. +- Add an allocator-atomic `snapshot_and_lease(request, slice_spec)` operation for + destination KV before copying block IDs, resolving pointers, or publishing + addresses. Construct the receive context and wire slice from the immutable, + allocation-generation-bearing lease snapshot; Phase 1 cannot claim + direct-path safety without it. +- Add a Python-native lifecycle cancellation adapter with a structured + logical/physical result. Select it explicitly for `KvCacheTransceiverV2` + without changing the C++ transceiver's boolean contract. Let PyExecutor drop + destination-KV ownership while the lease makes it pending-free, but preserve + independent session/auxiliary cleanup gates; retain the sender-side legacy + gate until Phase 2. Enter the shared manager drain/veto protocol only for the + lifecycle-capable Python transceiver. Default C++ and non-disaggregated paths + retain baseline local manager teardown; the adapter must not audit or retire a + C++ transport owner or treat local object destruction as quiescence evidence. +- Make receiver context insertion and `TransferHandle` enrollment atomic with + respect to handle abort/sealing. Gate every access boundary on the durable + gate before taking the context lock, abort the gate on first failure, + cancellation, or shutdown before enumerating contexts, and unwind a losing + preparation before target publication. +- Define one monotonic transfer-attempt generation shared by every slice. Let + only the receiver producer/session seal its handle, using an immutable + expected-slice manifest after all fixed contexts are atomically enrolled. +- Treat any failure lacking documented quiescence evidence as in doubt. +- Preselect the exact ADP writer cohort or exclude that topology from the first + rollout. +- Preserve a separate auxiliary-slot cleanup gate, or exclude generation-first + configurations that transfer auxiliary data from Phase 1. +- Add the minimum transceiver `ShutdownResult`, caller-retention contract, and + PyExecutor/creator/KV-manager teardown veto. If a context remains in doubt, + Phase 1 must keep its worker, registry, manager, registration, and mapping + alive; it cannot defer that safety boundary to Phase 3. +- Add negotiated transfer generation and both endpoint epochs to every + lifecycle-mutating advertisement, result, and control message. Reject peers + with the legacy `unique_rid`-only identity before address publication. +- Bind every admission to the negotiated endpoint-epoch pair. Recreating an + endpoint without its replay window requires a fresh, non-reused epoch, and + old-epoch traffic cannot bootstrap a replacement session. +- Ship minimum capacity diagnostics with the fail-closed path: in-doubt + count/bytes/oldest age/reason, lease-token and unique transfer-leased + KV/bounce bytes, pending-free KV bytes, post-logical retention, per-resource + early-release lead, and admission/fallback failures caused by retained + capacity. + +Likely files: + +- `tensorrt_llm/_torch/disaggregation/native/transfer.py` +- `tensorrt_llm/_torch/disaggregation/transceiver.py` +- `tensorrt_llm/_torch/disaggregation/native/bounce/core.py` +- `tensorrt_llm/_torch/disaggregation/native/bounce/impl.py` +- `tensorrt_llm/_torch/disaggregation/native/perf_logger.py` or the selected + serving-metrics adapter +- destination KV-manager adapter/binding and PyExecutor cancellation integration +- `tensorrt_llm/_torch/pyexecutor/py_executor.py`, + `tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`, and KV-manager teardown + adapters for the minimum non-drained shutdown veto +- a new native transfer-context/registry module +- unit and fault-injection tests + +Phase 1 is a receiver-safety milestone, not completion of the comprehensive +sender-and-receiver contract. + +### Phase 2 — Complete KV leases and sender context + +**Size:** Large; approximately 12–18 files, 1,000–1,800 production lines, and +1,200–2,200 test lines. A small C++/nanobind extension may be required. Split +source-lease/sender-context work from the dynamic-stream protocol when review or +validation scope would otherwise become too broad. + +- Complete the common source/destination `KVTransferLease` contract for both KV + managers. +- Support atomic pending-free and overlapping reference-counted slice leases. +- Acquire source leases before gather or direct submission. +- Promote the initial in-doubt `SendTransferContext` containment record into a + canonical registry-owned context, backed by allocator-issued source leases, + and track gather, NIXL, send-bounce, descriptor, and result-delivery state + independently through normal as well as ambiguous completion. +- Extend atomic `TransferHandle` enrollment to sender contexts and reject a + losing preparation before gather or NIXL launch. +- Retain the sender's canonical record/gate independently of its request-facing + handle reference, and abort it on first failure/cancellation before admitting + another chunk. +- Add a negotiated bounded dynamic-writer-stream mode for sender-determined + chunks: advertise one authorized writer/target/range/stream, assign a distinct + operation sequence per chunk, validate the operation/segment/byte budget, and + reserve its worst-case metadata envelope before publication, then close + enrollment with an authenticated end-of-stream high-water mark. +- Use an ordered interval/coverage structure so dynamic operation insertion and + overlap checks are O(log(operations + segments)), while end-of-stream + validation is one O(operations + segments) pass bounded by the negotiated + maximum. +- Track reserved versus materialized operation/segment/byte credits and unused + slack. Return only unused envelope credits when authenticated close makes + future materialization impossible and establishes the exact final operation + set, or after backend quiescence permits fail-closed record retirement; retain + credits backing materialized entries through their retirement. +- Atomically enroll the final sender context with its final-slice proof before + sealing the sender handle. Keep the receiver target leased until the stream is + closed/fenced and every operation through the high-water mark is quiescent. +- Emit end-of-stream only after every covered sequence is already submitted or + terminal `NO_REMOTE_ACCESS`; atomically close future submission first. A + failed/cancelled stream uses terminal per-operation decisions or remains open + until the applicable fence, never normal sealing as a shortcut. +- Deduplicate repeated stream advertisements and operation results, and fail + closed on missing, conflicting, out-of-range, or out-of-budget sequences. +- Integrate with or generalize `AsyncTransferManager` rather than creating a + second request-retention mechanism. +- Remove the remaining sender-side safe-to-free cancellation gate once source + leases make request cleanup pending-free. +- Permit per-chunk source release at the earliest safe boundary without forcing + request-wide pinning; PR #15803 provides C++-side prior art for the lease + shape. +- Cover direct, bounce, and mixed-mode fan-in. + +### Phase 3 — Submission fencing and safe shutdown + +**Size:** Large; approximately 8–14 production files, 700–1,300 production +lines, and 800–1,500 test lines. This phase should normally be split into a +fence-protocol PR and a shutdown-integration PR. + +- Add a negotiated peer/session/endpoint submission-fence capability and + transfer-generation/endpoint-epoch/fixed-operation-or-stream-scoped + request/ack framing for advertisements without a terminal result. An + acknowledgement closes future submission against all covered advertisements. + This strengthens the Phase 1 creation-close handshake after addresses may have + been exposed. +- Add retry/idempotency handling for lost, delayed, duplicated, contradictory, + and stale-epoch fence messages. Mixed-version peers that cannot provide the + required fence are non-drained rather than silently downgraded. +- After fencing, add backend quiesce/drain for submitted or in-doubt work and + reorder both before deregistration/unmapping. +- Extend the Phase 1 non-drained result and teardown veto with backend-wide + quiescence, full retry/drain behavior, and process-health escalation. +- Include independent auxiliary-transfer owners in the shared-registration + shutdown barrier. +- Preserve independent earliest-safe release boundaries and retain only + lightweight result-delivery/tombstone state after GPU leases retire. +- Add shutdown/quiesce-specific lifecycle metrics beyond the Phase 1 minimum. + +Likely files include `tensorrt_llm/_torch/disaggregation/transceiver.py`, +`tensorrt_llm/_torch/disaggregation/native/transfer.py`, +`tensorrt_llm/_torch/disaggregation/base/agent.py`, +`tensorrt_llm/_torch/pyexecutor/py_executor.py`, +`tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`, and KV-manager teardown +adapters in `tensorrt_llm/_torch/pyexecutor/_util.py`. + +### Phase 4 — Protocol and liveness hardening + +**Size:** Medium; approximately 300–700 production lines plus protocol tests. + +- Add result acknowledgement/retry or receiver query for bounded recovery from + message loss. +- Add a richer per-operation query/abort API beyond the minimum quiescence + classification required by Phase 1. +- Consider per-operation abort only if the backend can guarantee no later memory + access after acknowledgement. + +The original pre-implementation estimate of 4,000–7,300 production lines is +obsolete. The 2026-07-21 implementation snapshot is approximately +10,300/-1,400 +production lines, +11,000/-260 test lines, and +2,700 documentation lines. That +is too large for one effective code review. After landing the design-only +#16347 prerequisite, the implementation should be split into stacked review +units for (1) C++ V1 pin hardening, (2) native operation/session ledgers and +transport cancellation, (3) bounce containment, and (4) PyExecutor/manager +ownership and shutdown. Unit, integration, and test-list wiring travels with +the stack unit it validates. Each boundary must preserve a safe, testable +fail-closed state rather than temporarily weakening ownership. + +The C++ `CacheTransceiver` implementation, data path, and wire protocol remain +outside this scope. Cross-language work is limited to shared KV-manager or NIXL +binding contracts that the Python runtime consumes. The Python-native path may +use shared `AsyncTransferManager`/`ResourceManager` accounting to fail closed on +its own live ownership, but default C++ and non-disaggregated teardown bypass +that Python lifecycle protocol. + +## Rollout and Rollback + +Ownership selection is an endpoint-session capability decided before context +creation and before any address is published. A temporary internal deployment +gate may select eligible canaries, but `kv_cache_bounce_size_mb` is not that +gate: bounce-disabled direct transfers need the same lifetime safety. + +Rollout proceeds by configuration cohort: + +1. Land deterministic tests and allocator conformance tests before enabling any + production cohort. +2. Enable Phase 1 only where the exact writer cohort is known, both peers + negotiate the ownership protocol, the destination KV manager implements the + atomic lease contract, Phase 1 creation-close/cancel acknowledgement, + monotonic generation/replay-window contract, and fixed-manifest/context + capacity envelopes, and the minimum capacity metrics are available. + Unsupported ADP cohorts, mixed versions, and generation-first auxiliary paths + without an independent safe cleanup gate fail before address publication. + Sender-determined dynamic chunk streams remain excluded until Phase 2. +3. Canary direct and bounce paths separately for each KV manager, then expand to + fan-in, fallback, AutoDeploy, and chunked configurations as their matrix cells + pass. A bounce canary requires positive transport-active and per-transfer + bounce-engaged signals; configuration alone is not evidence. +4. Promote toward default-on only after the safety, capacity-recovery, and + quantitative performance gates pass for the applicable cells. + +Rollback never changes the owner of an active transfer. Stop new admission, +drain or fail closed all contexts negotiated under the new protocol, and only +then select another mode for newly negotiated sessions. If drain is non-terminal, +retain the owners and escalate process health; do not bypass leases to complete +rollback. A legacy path may remain available for explicitly excluded canary +cohorts during rollout, but it is not an in-place recovery mechanism for an +ownership failure. Rollback reports a blocked/non-drained result while any +context cannot retire and increments the corresponding low-cardinality cohort +metric. + +## Validation Plan + +### Deterministic unit and fault-injection tests + +- Cancel locally immediately before transfer-eligibility/token creation and + before/during `admit_and_bind()` with barriers. Cancellation prevents + scheduling, aborts the pending token with no peer authorization, or observes + and aborts the canonical bound handle; no unbound interval is permitted. + Repeat request cleanup while admission is paused. +- Race `begin_shutdown()` with transfer-eligibility/token creation, + `admit_and_bind()` before and after registry insertion, and the peer + authorization boundary. Every run must either reject/shutdown-abort the token + with no record/authorization or include the canonical record and any possible + authorization in the shutdown snapshot and drain set. +- Cancel immediately before, during, and after address publication. +- Race cancellation against receiver and sender preparation/handle enrollment + with barriers; assert that the context is either enrolled and cancelled or + fully unwound before its publication/local-launch boundary. Reject slice + creation after handle sealing. +- Close the handle, pause cancellation before it obtains a context lock, and + race publication/gather/NIXL/scatter boundary entry; assert that the closed + handle gate rejects every new boundary while previously admitted accessors + remain owned. +- Drop every request/session-facing handle reference while a context is blocked, + then race abort with an access boundary. The registry-owned record/gate must + remain live. After current contexts retire but before producer sealing or a + creation fence, reject a delayed slice and advertisement replay instead of + creating a replacement handle. +- Fully retire and delete a canonical record, reuse its former target + allocation, then deliver delayed duplicate admission and target-advertisement + messages. Endpoint replay state must reject both before record creation or + local launch. Retire generations out of order, verify sparse entries block + replay, then close the gap and verify contiguous-floor compaction. +- Fill the negotiated generation window behind one missing generation and + verify admission rejects in the same worker iteration without sleep/retry or + eviction and within the frozen wall-time ceiling. Recreate an endpoint with a + fresh epoch and prove that old-epoch traffic cannot bootstrap its session. An + arbitrary unnegotiated control must poison the current session in O(1), after + which all admission fails until fresh negotiation. Reconnect with the same + epoch and a peer-claimed floor beyond a locally live/gapped generation; + reject the inconsistent session without resetting local replay state. +- Repeatedly fail initiating admission before the generation/record/token + commit and prove that the counter/floor/window do not advance. Then inject + failures after local commit but before admission-control send, peer-side + rejection after admission control but before address publication, ambiguous + admission-control delivery, and abandonment after address publication. Lose, + delay, or duplicate the generation-skip acknowledgement and prove + higher-generation admission-control send remains blocked. Ambiguous normal + admission must retry/deduplicate and creation-close rather than becoming a + skip. Reject a contradictory cumulative watermark crossing a live or + published record. After authenticated skip or pre-publication close/reject, + both replay states advance/compact; post-publication abandonment retains the + record/resources until full physical retirement. +- Race first logical failure against next-chunk enrollment and gather/NIXL + launch; the failure must abort the shared gate before callbacks, reject new + work, and retain already-possible accessors through quiescence. +- Race cancellation with final-slice enrollment and handle sealing. Cover + missing, duplicate, reordered, and conflicting final chunks; success requires + the immutable manifest or final-slice proof to name the complete slice set. +- Writer A fails, the consumer/session closes, and writer B reports later. +- Mixed direct/bounce fan-in where one writer fails. +- Partial multi-writer publication where only a subset may have seen addresses. +- Timeout followed by late success and late failure. +- Duplicate, unexpected, contradictory, and stale-transfer-generation results. +- Duplicate and reorder target advertisements. Assert that an identical fixed + operation or writer-stream advertisement launches at most once and a + conflicting replay fails closed. +- Exercise a dynamic writer stream with out-of-order results and authenticated + end-of-stream. Cover a missing sequence, conflicting duplicate, out-of-range + subrange, operation/segment/byte budget overflow, forged/premature + end-of-stream, sequence-complete but in-range undercoverage/gaps, and + a full-coverage source-to-target permutation. Also cover an open stream whose + currently known operations are all quiescent; none may cause early target + retirement or logical success. +- Pause an enrolled sequence before NIXL submission and attempt end-of-stream. + The close must wait for an irreversible `SUBMITTED`/`NO_REMOTE_ACCESS` + decision or reject later submission. Backend drain before that decision must + not permit target unmapping. Repeat for failed and cancelled streams. +- Gather launch, descriptor construction, NIXL submission, result-send, and + scatter exceptions. +- Cancel while gather is active but before NIXL submission; assert that source + KV and any send-bounce resources remain leased until gather quiesces. +- Consumer callback exception during physical finalization. +- Attempted source KV reuse while direct NIXL read is blocked. +- Attempted destination KV reuse while direct write or scatter is blocked. +- Attempted send/receive bounce reuse while NIXL or scatter is active. +- Shutdown with blocked NIXL status, lost quiescence result, and queued/running + scatter; assert that deregistration and unmapping do not occur. +- Publish a target, pause the sender before NIXL submission, and start receiver + shutdown. Assert that a drain of currently submitted operations cannot unmap + the target; retirement requires a terminal `NO_REMOTE_ACCESS` result or an + acknowledged submission fence that prevents later use of the advertisement. +- Lose, delay, duplicate, reorder, contradict, and replay a submission-fence + acknowledgement from a stale endpoint epoch. Only the matching negotiated + acknowledgement may close future submission; mixed-version or unresolved + peers keep shutdown non-drained. +- Exercise new/new negotiation independently for Phase 1 identity, + creation-close, fixed-envelope, and replay-window capabilities; Phase 2 + dynamic-stream/end-of-stream capability; and Phase 3 submission fencing. + For each phase, reject new/legacy, missing, unknown, and downgraded capability + combinations before that phase's publication/local-launch boundary. +- Count protocol messages on healthy admission and assert generation/epoch and + fixed-envelope fields use the existing KV request/target response with no + standalone admission acknowledgement. Verify skip/close/reject controls occur + only in their exceptional scenarios. +- Publish a target, then deliver a valid Phase 1 creation-close acknowledgement. + It may close future context creation but must not retire the published target, + satisfy a submission fence, or permit shutdown while later sender submission + remains possible. Only terminal operation evidence or the negotiated Phase 3 + fence plus quiescence may release the target. +- Exactly-once release under concurrent cancel/result and reordered duplicates. +- Lease acquisition racing request `free_resources()`. +- Raw block-ID snapshot racing `free_resources()` and block reuse; assert that + only the allocator-atomic allocation-generation-bearing snapshot can create a + context. +- Overlapping slice leases on the same KV allocation generation. +- Overlap direct and bounce-path lease tokens on one allocation generation. + Per-path token series may count both, while path-neutral unique/pending-free + gauges count the range once; any `direct_only`/`bounce_only`/`mixed` breakdown + must sum exactly to the path-neutral gauge. +- Malformed, overflowing, misaligned, or out-of-range NIXL and scatter metadata. +- Fixed-manifest segments that are individually in range but truncate expected + coverage, leave a gap, duplicate/overlap within or across writers, mismatch + aggregate bytes, or permute equal-sized source-to-target intervals while + retaining full target coverage. Cover direct descriptors and wrong + bounce-source offsets in scatter metadata. All fail logical success and + suppress usable scatter while physical access still drains. Equivalent + differently segmented inputs that normalize to the exact expected mapping + pass. +- Fixed manifests and transfer context/slice sets at, below, and above their + configured operation/descriptor-segment/authorized-byte/context limits. + Concurrent admissions must atomically reserve credits without overshoot; + rejection occurs before publication/local launch, completes in the same + bounded worker iteration, and transactional unwind returns every credit. +- Open many dynamic streams with maximum envelopes but materialize only a small + fraction of their operations/segments/bytes. Verify reserved/materialized + credits, utilization/slack, reservation-caused rejection, and prompt return of + unused credits after authenticated close establishes the exact operation set, + without releasing materialized entries early. A future-submission-only fence + must retain slack until ledger reconciliation or backend quiescence. +- Synchronous pre-submit failure versus ambiguous NIXL submission failure. +- Quiesced `NO_REMOTE_ACCESS` and quiesced-with-unknown-data outcomes. +- Global NIXL quiescence while scatter remains queued or active. +- KV registry drained while an independent auxiliary transfer remains active. +- Non-drained shutdown followed by a successful retry. +- Duplicate context identity insertion. +- Retry identical admission through the same token and receive the same handle; + attempt the live identity through a different token and reject it without + attaching cancellation authority or replacing canonical state. +- Partial KV/bounce acquisition, plan-validation failure, and duplicate registry + insertion before the receiver target-publication or sender local-launch + boundary; assert endpoint-local transactional unwind with no leaked lease. + Repeat after each endpoint's boundary and assert fail-closed retention instead + of rollback. +- Simultaneous local send/receive identities, fan-in writers from different peer + epochs, and a stale peer epoch targeting a new local transfer generation. +- Keep the local endpoint epoch fixed while the remote admission authority + restarts with a fresh epoch and repeats a transfer ID/generation value. The + new full identity may be admitted, while delayed old-authority traffic is + rejected without colliding with the new record or replay window. +- Delayed cancellation through a stale handle after request ID reuse; assert + that it cannot change the current transfer generation. +- Remote cancellation arriving before context creation, followed by the matching + transfer generation; assert that the transfer-generation-aware pre-cancel + tombstone applies. +- Race remote cancellation against registry/handle/context creation; assert + that either the live handle closes or the preparation observes the tombstone, + with no access boundary in between. Fill the tombstone table and assert + that every already-admitted generation has reserved cancellation state and + new admission backpressures before overflow. Deliver a valid cancellation at + full occupancy and verify it is recorded. Lose, delay, duplicate, and replay + the Phase 1 creation-close acknowledgement; retire zero-context state only + after the matching acknowledgement or local producer sealing. +- Old-transfer-generation remote cancellation arriving after request ID reuse; assert + that neither the live context nor its session changes state. +- Logical cancellation followed immediately by request cleanup while the + physical disposition remains `DRAINING` or `IN_DOUBT`. +- Generation-first cancellation while auxiliary RMA remains active; assert that + KV may become pending-free but the session and auxiliary slot remain live. +- Gather/scatter worker death or asynchronous CUDA failure. +- Bounce send with NIXL blocked after gather; assert that the slice's source-KV + lease retires while its send-bounce lease remains live. +- Direct send with NIXL definitively quiescent and result delivery blocked; + assert that source KV retires before result delivery. +- Multi-slice send with one slice quiescent and a sibling blocked; assert that + the first slice retires without waiting for sibling or request completion. +- Non-drained transceiver shutdown retained by PyExecutor; assert that executor + and creator cleanup cannot shut down the KV manager until a retry drains. +- Request rollback with live direct and bounce contexts. Assert that admission + and mode switching remain blocked and each context keeps its negotiated owner. + Hold one context `IN_DOUBT` and require a visible blocked/non-drained rollback + without resource release; after terminal evidence drains every context, only + newly admitted sessions may negotiate the replacement mode. + +Tests must use barriers/hooks rather than probabilistic sleeps for the critical +publication and teardown races. + +### Configuration matrix + +- bounce not requested, requested but factory-inactive, and arena-active; +- positive bounce-engagement evidence on fabric-VMM cells; merged #16116 + provides a scheduled GB200 assertion, GB300 remains marker-eligible but is not + claimed as scheduled coverage, and configured-but-inactive runs count as + direct fallback; +- Python KV manager V2 and C++-backed V1; +- phase-matched new/new negotiation plus new/legacy, missing-capability, + unknown-capability, and downgrade rejection; +- single writer and multi-writer TP/PP/ADP overlap or fan-in; +- direct, bounce, and per-writer fallback; +- single slice and chunked/multi-slice transfer; +- fixed manifests and dynamic streams at small, midpoint, and negotiated-maximum + operation-count/descriptor-segment/authorized-byte budgets, including + out-of-order and end-of-stream load, plus per-transfer context/slice limits; +- high-concurrency maximum-envelope dynamic streams with low actual + operation/segment/byte materialization; +- empty, midpoint, near-capacity, and full endpoint-global + context/operation/segment/authorized-byte credit occupancy; +- empty, midpoint, near-capacity, and full lifecycle-record/pre-cancel occupancy; +- sequential and out-of-order retirement at empty, midpoint, near-capacity, and + full endpoint replay-window occupancy; +- PyTorch executor and AutoDeploy with Python transceiver; +- context-first scheduling, plus generation-first only after its independent + auxiliary-slot cleanup gate is validated; +- bounce-ineligible/non-attention layer-group fallback; +- normal completion, cancellation, timeout, peer failure, and shutdown. + +### Performance validation + +Use an A/B build or runtime gate on the same commit and compare against the +merged PR #15618 behavior. Hold model, hardware, NIXL backend, request trace, +arena size, and scheduler configuration constant. Cover direct, bounce, and +fallback paths across representative transfer sizes, concurrency, fan-in, both +KV managers, and PyTorch/AutoDeploy cells. Run warmup plus enough independent +benchmark runs to report 95% confidence intervals, using runs rather than +individual requests as the sampling unit. Freeze workload definitions and gates +before examining implementation results. Record the low-cardinality ownership +mode, protocol bucket, rollout cohort, KV manager, executor, transport-active +state, and observed writer path for every aggregate. Reclassify an intended +bounce cell that lacks positive engagement evidence as direct fallback and do +not use it to approve bounce rollout. + +Report: + +- request throughput and completed-request goodput; +- p50/p95/p99 TTFT and end-to-end transfer-task latency; +- gather, NIXL, and scatter latency/overlap, with the existing raw NIXL timer + reported separately; +- host CPU time per completed transfer and registry/lifecycle-lock contention; +- per-transfer healthy admission control message count plus exceptional + skip/close/reject RTT and retry count; +- active host metadata by context/stream/tombstone/replay-window occupancy, + retired-floor advancement and gap age, end-of-stream processing time, + reserved/materialized credits, reservation utilization/slack and caused + rejection, admission/rejection latency, per-path lease-token bytes, + path-neutral unique transfer-leased bytes, pending-free KV bytes blocked + solely by transfer leases, and each corresponding byte-seconds/high-water + mark; +- post-logical retention, per-resource early-release lead, lease duration, and + capacity-recovery time; +- admission failures and goodput during known-terminal and deliberately + ambiguous fault injection; +- shutdown drain latency and non-drained bytes. + +Initial go/no-go thresholds are scoped by scenario. + +Healthy/no-fault matrix cells: + +- throughput/goodput 95% confidence-interval lower bound is no worse than -2%; +- the upper 95% confidence bound of relative p50/p95 TTFT and end-to-end transfer + latency regression is at most 3%, and the p99 bound is at most 5%; +- the upper 95% confidence bound of host CPU time per completed transfer + regression is at most 5%; +- the direct path adds zero CUDA copies, kernels, or synchronizations; +- ownership adds no device arena beyond an active PR #15618 send/receive bounce + arena pair; +- host metadata has frozen measured per-entry budgets for contexts, writer + operations, descriptor/scatter segments, lifecycle records, and + live/sparse-retired generation entries; repeated waves, replay gaps, and floor + compaction stay within the resulting O(contexts + writer operations + + segments + lifecycle records + outstanding generations) budget; +- across predeclared small/midpoint/maximum fixed-manifest, context-set, + dynamic-stream, lifecycle-record, and replay-window sizes—including + out-of-order gaps and floor compaction—host bytes and CPU/close processing fit + the frozen per-entry complexity budget with no statistically significant + unexplained superlinear term; +- at each declared supported concurrency, maximum-envelope/low-materialization + streams have zero reservation-caused rejection; otherwise that envelope or + concurrency is excluded from the rollout cohort. Report utilization/slack and + goodput at and beyond the declared boundary; +- before the configured backpressure threshold, the upper 95% confidence bound + for p95 admission-latency regression at near-capacity lifecycle-record and + replay-window/metadata-credit occupancy is at most 5%; at capacity, rejection + performs no sleep/retry, completes in the same admission-worker iteration and + within one configured bounded-poll budget, and does not overshoot configured + record/replay/context/operation/segment/byte limits. Report p95/p99 wall time and + freeze an absolute per-cohort ceiling before implementation results are seen. + +Known-terminal fault cells: + +- after each fixed-concurrency fault wave drains, active context/lease counts and + retained bytes return to their pre-wave baseline with no monotonic growth; +- once the final required terminal/completion evidence is observed, leases are + released in the same state-machine call or the next completion-worker + iteration, as applicable; +- delayed-result tests recover all leased capacity. + +Ambiguous-result fault cells: + +- exactly the affected bytes remain retained and observable, with no unsafe + reuse, unmap, or cross-request corruption; +- fault-period goodput and latency are reported, but the healthy-path + non-regression thresholds do not apply because safe capacity exhaustion is an + expected outcome. + +If stable benchmark noise is wider than a threshold, calibrate and document a +replacement before implementation measurements are used for approval. Do not +relax a gate after seeing a regression without a separate review. + +### Acceptance criteria + +The design is implemented only when all of the following are true: + +1. No address can be published without a registered context and live leases. +2. Session/request removal cannot drop a physical quiescence result. +3. KV and bounce allocators cannot reuse leased allocation generations. +4. Every release path is idempotent and tied to quiescence evidence. +5. Direct, bounce, and mixed transfers share the same ownership path. +6. Timers alone never authorize reuse. +7. Shutdown cannot deregister or unmap memory while a context may access it. +8. New identity fields are negotiated or mixed versions are rejected before + DMA addresses are exchanged. +9. The executor remains responsive while contexts drain. +10. In-doubt capacity and shutdown state are observable. +11. Lease acquisition and `free_resources()` cannot race into block reuse. +12. Every NIXL/gather/scatter range is authorized by a live allocation-generation + lease. +13. Every local CUDA/NIXL accessor is independently tracked through quiescence. +14. Phase 1 negotiates and validates one monotonic, attempt-wide transfer + generation plus both endpoint epochs on every lifecycle-mutating wire + message; legacy identity is rejected before address publication. +15. Contexts are built only from allocator-atomic, + allocation-generation-bearing KV lease snapshots; copied block IDs are + never used to acquire a later lease. +16. Local role/epoch plus admission-authority epoch identifies the context, + while exact peer epoch/rank/operation identifies each fan-in ledger entry. +17. Logical cancellation and dropping KV request ownership do not wait for KV + physical retirement, and pending-free KV cannot be reused while its lease + remains live. +18. A non-drained shutdown result keeps the transceiver and KV managers alive and + can be retried without destructor fallback cleanup. +19. The applicable quantitative normal-path and fault-period performance gates + pass with the required stage and end-to-end metrics. +20. Unsupported topology or mixed-version cohorts are rejected before address + publication. Rollback never changes an active transfer's owner and reports a + visible blocked/non-drained state until every old-mode context retires. +21. Preparation failure unwinds every partial lease before that endpoint's + publication/launch boundary, while failure after its boundary follows + fail-closed retirement. +22. Cancellation is routed through a transfer-generation-bound handle; a bare + or reused request/transfer ID cannot mutate lifecycle state. +23. KV cleanup never releases an independently active auxiliary slot; unsupported + generation-first configurations are rejected before publication. +24. Bounce source KV retires after gather, direct source KV retires after NIXL + quiescence, and one slice never waits for unrelated sibling completion. +25. Every remote lifecycle-control message is transfer-generation/epoch + validated before session lookup; pre-cancel tombstones cannot cross transfer + generations, and an arbitrary unnegotiated control closes the endpoint + session without allocating per-identity state. +26. First failure, accepted cancellation, or shutdown atomically aborts context + admission and every later publication/local-launch boundary. Sealing is + performed only by the local producer/session with a complete immutable + manifest or atomically enrolled final-chunk proof. +27. Immediate retirement requires every remote exposure and every local + gather/NIXL/scatter accessor to be terminal; proof that no remote operation + was submitted is insufficient while a local gather remains active. +28. Shutdown fences future submission against every published address before it + treats backend quiescence as sufficient to deregister or unmap memory. +29. Every target advertisement identifies its target slice and exact fixed + operation or bounded writer stream. Duplicate advertisements are + idempotent. End-of-stream closes future submission only after every covered + sequence is submitted or terminal `NO_REMOTE_ACCESS`; dynamic streams cannot + retire until that close/fencing plus operation quiescence through the + high-water mark, and logical success additionally requires the exact + normalized required source-to-target mapping. +30. Pre-cancel lookup and context creation serialize in the registry. Capacity + is reserved at admission for every valid cancellation; unmatched state + remains until producer sealing or the Phase 1 creation-close acknowledgement + and uses pre-admission backpressure rather than unsafe eviction. +31. Bounce correctness/performance cells record positive transport-active and + per-transfer engagement evidence; `kv_cache_bounce_size_mb > 0` alone does + not count as bounce coverage. +32. The registry strongly retains exactly one canonical endpoint-transfer + record, handle, and gate after consumer cleanup and until no later local + context can be created and all physical obligations retire; it atomically + marks the generation retired before removing that record. +33. Endpoint-epoch replay state rejects delayed admission/advertisement after + record deletion, bounds out-of-order generations with a retired floor and + sparse window, backpressures gaps without eviction or admission-worker + retry/sleep, and requires a fresh non-reused endpoint epoch when that state + is recreated. +34. Fixed manifests, dynamic streams, and context/slice sets reserve negotiated + operation/segment/byte metadata credits against endpoint-global limits + before exposure; overload fails without overshoot or retry/sleep in the same + bounded admission-worker iteration, and retirement returns credits exactly + once. +35. Capacity telemetry distinguishes overlapping lease-token bytes, unique + transfer-leased allocation-generation bytes, and pending-free KV bytes + blocked solely by transfer leases; KV recovery/admission claims use the last + measure, while bounce claims use unique leased slot bytes. Deduplicated + gauges are path-neutral or use only mutually exclusive direct/bounce/mixed + classes; per-path token metrics are not summed as unique capacity. +36. Local cancellation and initial admission serialize through a one-shot token: + cancellation either prevents transfer eligibility/token creation, rejects + pending admission before peer authorization, or aborts the canonical bound + handle, with no lost-cancel interval. +37. Phase-specific new/new negotiation succeeds, while missing/downgraded/legacy + capabilities fail before the relevant access boundary; a Phase 1 creation + close never substitutes for a Phase 3 submission fence after publication. +38. Dynamic-stream telemetry distinguishes reserved from materialized + operation/segment/byte credits, exposes slack and reservation-caused + rejection, and proves declared rollout concurrency without slack-induced + rejection; authenticated close returns only credits that cannot materialize. +39. Shutdown atomically closes registry admission and snapshots committed + records; a racing pending token is shutdown-aborted or its canonical record + and any possible peer authorization are included in drain, never omitted. +40. Fixed-mode success requires the exact normalized per-operation and manifest + source-to-target KV mapping; in-range gaps, undercoverage, + duplicate/overlap, aggregate mismatch, and full-coverage source permutation + fail rather than producing corrupted KV. +41. Initiator pre-commit rejection consumes no transfer generation. Every + committed generation whose admission control is not sent is bilaterally + acknowledged as skipped before a higher admission control is sent. A + generation rejected after admission control but before address publication + is creation-closed; after publication it follows full physical retirement. + Repeated admission failure cannot create artificial peer replay gaps. +42. Session negotiation is amortized and healthy admission fields piggyback on + the existing KV request/target response; no new normal-path per-transfer RTT + is introduced, or the cohort is separately budgeted and qualified. + +## Risks and Mitigations + +| Risk | Mitigation | +|---|---| +| Lock-order deadlock across admission token, registry, session, allocator, CUDA, and callbacks | Token-to-registry-to-handle-to-context order; no blocking operations or callbacks under lifecycle locks; lock-order tests. | +| Local cancellation races initial handle binding | One-shot admission token with token-to-registry lock order; bind the canonical handle before peer authorization; deterministic race tests. | +| Shutdown races pending admission or peer authorization | Close admission and snapshot under one registry lock; shutdown-abort losing tokens; gate authorization through the bound handle; deterministic boundary races. | +| KV over-pinning or leaked leases | Exactly-once release, active-lease metrics, oldest-context diagnostics, and fault-injection coverage. | +| Capacity exhaustion from indefinite retention | Fail admission visibly; expose bytes/age/reason; add reliable result recovery or backend quiescence before bounded reclaim. | +| Treating failure as quiescence incorrectly | Track quiescence separately from data outcome and preserve an explicit in-doubt state. | +| Mixed direct/bounce accounting error | Record actual mode per exact writer and retain both candidate leases until mode is known. | +| Individually valid fixed segments leave incomplete, overlapping, or permuted KV | Manifest the exact logical source-to-target mapping, normalize delivered pairs, reject gaps/duplicates/overlap/aggregate mismatch/permutation, and test direct plus bounce pairings. | +| Dynamic stream closes early, undercovers or permutes the target, or grows without bound | Authenticate end-of-stream, require irreversible submission decisions before close, enforce sequence/operation/segment/byte budgets, validate the exact required source-to-target mapping independent of the claimed high-water mark, and retain the target while the stream is open. | +| Stale result advances a new transfer | Endpoint epochs, monotonic generation, durable retired floor/replay window, exact writer validation, and negotiated protocol. | +| Retired transfer is recreated by delayed admission | Endpoint-owned retired floor plus bounded sparse replay window; mark retirement before record deletion; reject old epochs before session creation. | +| Replay-window gap exhausts host capacity | Bound outstanding generations, expose floor/gap/window metrics, and backpressure admission without evicting anti-replay state. | +| Failed admission burns generations and fills peer replay capacity | Allocate the generation only in the successful local admission commit; acknowledge a skip before sending a higher admission control; creation-close rejection only before address publication; use full retirement afterward; repeat/lost-ack tests. | +| Admission protocol adds a healthy-path RTT | Piggyback fields on the existing KV request/target response, count normal/exceptional controls, and separately budget any cohort that cannot preserve the exchange. | +| Fixed manifests, descriptor plans, or context fan-out exhaust host metadata | Negotiate per-transfer operation/segment/byte/context envelopes, reserve endpoint-global credits before address publication, reject oversize/over-capacity plans, and test concurrent no-overshoot accounting. | +| Worst-case stream reservations suppress useful concurrency | Export reserved/materialized credits and slack, benchmark maximum-envelope/low-materialization concurrency, tune or gate supported envelopes, and return unused credits only after authenticated close. | +| Pre-cancel state exhausts host capacity | Reserve a lifecycle/control slot per admitted generation, export configured/used/remaining capacity, backpressure before admission, and retire only after producer sealing or Phase 1 creation close. | +| Request cleanup drops the access gate | Registry-own one canonical record/handle/gate through no-future-context proof and physical retirement; contexts retain only the shared gate, which owns no contexts. | +| First failure races another chunk | Atomically abort the shared gate before callbacks and reject later enrollment/access while already-possible work drains. | +| Main-loop latency regression | Preserve bounded polling; perform draining on transfer/completion workers; benchmark registry contention. | +| Shutdown hangs indefinitely | Use an explicit deadline and return non-drained status, but never convert deadline expiry into unsafe unmap/reuse. | +| Rolling-upgrade incompatibility | Capability/version negotiation or pre-DMA rejection; no silent downgrade of safety fields. | +| Creation-close acknowledgement is mistaken for a submission fence | Distinct message/capability types and predicates; after-publication tests prove Phase 1 close cannot release a target or authorize teardown. | +| Behavioral divergence between KV managers | One conformance test suite for the common lease API; implementation-specific adapters only below it. | +| Scope expands into cancellation consensus | Keep logical consensus in the adjacent cancellation design; consume its outcome through `TransferHandle`. | +| Metrics add hot-path overhead | Use O(1) counters/histograms, batch export, and benchmark instrumentation-on overhead; never scan the full registry per transfer. | +| Bounce configuration silently exercises direct fallback | Require transport-active and per-transfer path signals in tests, benchmarks, and canary telemetry. | +| A caller drops a non-drained worker or shuts down its KV manager | Structured `ShutdownResult`, strong caller retention, retry tests, and a hard teardown precondition on zero live leases. | +| Rollback bypasses an active owner | Negotiate mode before publication, stop admission, and drain existing contexts before selecting a mode for new sessions. | + +## Alternatives Considered + +### Expand the PR #15618 bounce `TransferContext` + +Rejected as the general architecture. The initial containment has since renamed +that object to `RecvBounceContext`, but it remains naturally scoped to a receive +bounce slot. Making it own direct writers, source KV, sender status, request +aggregation, and KV-manager leases would invert module boundaries and make +ownership disappear when bounce is disabled. + +### Make `LlmRequest` the owner + +Rejected. Its logical lifetime may end before physical retirement; it aggregates +multiple slices; and a live object reference does not necessarily prevent +allocator-level `free_resources()`. It remains an associated consumer through a +request-level handle. + +### Treat timeout or fixed quarantine as safe reclamation + +Rejected. Wall-clock duration is not proof that one-sided RMA has stopped. A +timer may trigger diagnostics, fail admission, or escalate process health; it +cannot authorize memory reuse. + +### Solve only the bounce path + +Rejected. Sender fallback can mix bounce and direct writers, and destination KV +is the resource at risk when a direct writer outlives logical failure. The +general owner must sit above `BounceTransport`. + +### Unify Python and C++ transceivers now + +Rejected for this project. Their implementations, flow control, and active +bounce proposals differ. They should share invariants and, where useful, +allocator/protocol primitives, but not be forced into one implementation before +the Python safety gap is closed. + +## Open Implementation Questions and Phase Gates + +These do not change the safety contract, but they determine phase sequencing. + +1. Which exact NIXL states guarantee no later DMA after success, failure, abort, + and peer loss? +2. Can the advertisement messaging layer prove non-delivery after a send error? + If not, every attempted send remains `MAY_ACCESS`. +3. Which endpoint-owned counter/state implementation provides monotonic transfer + generation allocation and the bounded replay window across session/worker + recreation, and what maximum outstanding-generation window supports expected + concurrency without excessive reservation? +4. What allocates a unique endpoint restart epoch, keeps it stable for one + endpoint lifetime, and prevents its reuse after rank/process restart? +5. Can existing V1 block pinning enforce leases through explicit + `free_resources()`, or is a new nanobind API required? +6. What is the smallest V2 KV-manager lease surface that supports per-slice early + release without request-wide over-pinning? +7. Does the NIXL wrapper need a separate `quiesce()` operation before final + shutdown so registrations can remain valid during drain? +8. Which peer/session/endpoint primitive can acknowledge a submission fence and + guarantee that an old target advertisement cannot be used afterward? +9. Is result acknowledgement/retry required for the first production rollout, + or is observable indefinite retention acceptable for all Python-native + direct and bounce transfers? +10. Which independent auxiliary-buffer paths can outlive their request/session? + They need a follow-up audit before the broader transceiver can claim complete + ownership coverage. +11. Can every claimed TP/PP/ADP topology identify its exact potential writer cohort + before address fan-out, or which cells must be gated initially? +12. Which existing serving metrics surface should export lifecycle histograms and + retained-capacity gauges without adding per-transfer registry scans? +13. How should the internal rollout gate and negotiated capability be represented + without making the safety contract depend on the bounce-size option? +14. Which `SliceSetProof` form does each endpoint use: an immutable manifest or + a final-chunk proof atomically enrolled with the final context? For PR + #15727-style scheduling, how is the final marker authenticated and tied to + the sender's operation-sequence high-water mark? +15. What conservative operation-count, descriptor-segment, and byte budget can + be authorized before opening a dynamic writer stream without constraining + valid chunk schedules? +16. Which non-GB200/GB300 environments can positively validate an active + fabric-VMM bounce arena rather than only the direct fallback? +17. Which Phase 1 control/acknowledgement closes future context creation for a + transfer generation, and how is it retried independently from the stronger + Phase 3 fence that closes future submission to published addresses? +18. What fixed-manifest operation/descriptor-segment/authorized-byte, + per-transfer context/slice, and endpoint-global metadata-credit limits cover + production fan-out without allowing one transfer to monopolize host + capacity? +19. What canonical logical KV coordinate encodes source-to-target mappings across + TP/PP/ADP, layer groups, and both KV managers so validation never depends on + physical descriptor order? + +No answer to these questions may weaken the invariant by interpreting +uncertainty as retirement. diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py b/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py index f014842d5254..740913f002ca 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py @@ -19,7 +19,15 @@ from .buffer import Buffer, SlotAllocator from .config import Config, FixedSizing, Sizing, SizingContext, config_from_size -from .core import BounceTransport, Disposition, ScatterState, TransferContext, TransferState +from .core import ( + BounceTransport, + Disposition, + GatherSourceInDoubtError, + RecvBounceContext, + ScatterState, + TransferState, + WriterState, +) from .gather_scatter import Plan from .impl import ( NoBounceTransport, @@ -31,14 +39,19 @@ scatter_write_result, ) +# Deprecated compatibility name retained for imports from the #15618 API. +TransferContext = RecvBounceContext + __all__ = [ "BounceTransport", "Buffer", "Config", "Disposition", "FixedSizing", + "GatherSourceInDoubtError", "NoBounceTransport", "Plan", + "RecvBounceContext", "ScatterState", "Sizing", "SizingContext", @@ -46,6 +59,7 @@ "TransferContext", "TransferState", "VmmBounceTransport", + "WriterState", "build_send_request", "config_from_size", "create_bounce", diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py b/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py index bb77eab32908..0d4d662368bd 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py @@ -25,6 +25,11 @@ _MIB = 1024 * 1024 +# A Buffer that was not explicitly closed has no proof that transport users +# are quiescent. Retain its VirtMem object until process exit instead of +# allowing destructor order to unmap a potentially live one-sided target. +_UNSAFE_VMM_RETENTION: list[VirtMem] = [] + def _div_up(a: int, b: int) -> int: return (a + b - 1) // b @@ -78,11 +83,23 @@ def close(self) -> None: self._vm = None # type: ignore[assignment] def __del__(self): - # A destructor must never raise, but a leaked region should be visible, so log the failure. + # Explicit owners call close() only after lifecycle drain. A destructor + # has no such evidence, so it must fail safe by intentionally leaking + # the mapping until process teardown. try: - self.close() + vm = getattr(self, "_vm", None) + if vm is None: + return + retained = globals().get("_UNSAFE_VMM_RETENTION") + if retained is not None: + retained.append(vm) + self._vm = None # type: ignore[assignment] + logger.warning( + f"[kv-bounce] retaining unclosed buffer '{getattr(self, '_name', '?')}' " + "until process exit" + ) except Exception as e: - logger.debug(f"[kv-bounce] buffer '{getattr(self, '_name', '?')}' cleanup failed: {e}") + logger.debug(f"[kv-bounce] buffer retention failed: {e}") # region starts are rounded to this for copy alignment (negligible waste) @@ -94,7 +111,15 @@ class SlotAllocator: reuses a hole freed out of order rather than skipping it. Reserve is thread-safe and blocking. The whole buffer is one registration, so a write can stripe across the network links.""" - __slots__ = ("_buf", "_cap", "_cv", "_in_use", "_quarantine", "_next_slot_id") + __slots__ = ( + "_buf", + "_cap", + "_closed", + "_cv", + "_in_use", + "_quarantine", + "_next_slot_id", + ) def __init__(self, capacity_bytes: int, phys_chunk_size: int, name: str = "kv_bounce"): if capacity_bytes <= 0: @@ -102,10 +127,11 @@ def __init__(self, capacity_bytes: int, phys_chunk_size: int, name: str = "kv_bo self._buf = Buffer(capacity_bytes, phys_chunk_size, name=name) self._cap = self._buf.size # rounded up to a chunk multiple self._in_use: Dict[int, Tuple[int, int]] = {} # each live slot maps to its start and size - # Quarantined slots not yet reusable: an orphaned writer's write may still be landing - # and cannot be aborted, so each is held out of the pool until its deadline passes. - self._quarantine: Dict[int, Tuple[int, int, float]] = {} + # Quarantined slots are never made reusable by elapsed time. A caller must provide explicit + # quiescence evidence through release_quarantined(), or close must refuse the arena. + self._quarantine: Dict[int, Tuple[int, int]] = {} self._next_slot_id = 0 + self._closed = False self._cv = threading.Condition(threading.Lock()) @property @@ -116,7 +142,7 @@ def _occupied(self): """Ranges that must not be handed out: live and quarantined, treated the same.""" for s, n in self._in_use.values(): yield s, n - for s, n, _dl in self._quarantine.values(): + for s, n in self._quarantine.values(): yield s, n def _find_free_start(self, size: int) -> Optional[int]: @@ -138,6 +164,8 @@ def reserve(self, size: int, timeout: Optional[float] = None) -> Optional[Tuple[ deadline = None if timeout is None else time.monotonic() + timeout with self._cv: while True: + if self._closed: + return None start = self._find_free_start(size) if start is not None: slot_id = self._next_slot_id @@ -156,37 +184,45 @@ def release(self, slot_id: int) -> None: self._in_use.pop(slot_id, None) self._cv.notify_all() - def quarantine(self, slot_id: int, grace_s: float) -> None: - """Hold a slot out of the free pool for the grace period instead of releasing it, because its - region may still be under an in-doubt write. An infinite grace holds it until close.""" + def quarantine(self, slot_id: int) -> None: + """Hold a slot indefinitely because an in-doubt remote writer may still access it.""" with self._cv: entry = self._in_use.pop(slot_id, None) if entry is not None: - start, size = entry - # a finite time plus infinity is infinity, so an infinite grace never expires - deadline = time.monotonic() + grace_s - self._quarantine[slot_id] = (start, size, deadline) + self._quarantine[slot_id] = entry self._cv.notify_all() - def reclaim_expired(self) -> int: - """Return quarantined slots past their deadline to the free pool and report how many. Runs - off a timer, not tied to reserve, so it makes progress even when the arena is full.""" - now = time.monotonic() + def release_quarantined(self, slot_id: int) -> bool: + """Release one quarantined slot after the caller proves remote access is quiescent.""" with self._cv: - expired = [sid for sid, (_s, _n, dl) in self._quarantine.items() if dl <= now] - for sid in expired: - del self._quarantine[sid] - if expired: + released = self._quarantine.pop(slot_id, None) is not None + if released: self._cv.notify_all() - return len(expired) + return released @property def quarantined_bytes(self) -> int: """Bytes currently held in quarantine, for observability.""" - return sum(n for _s, n, _dl in self._quarantine.values()) + with self._cv: + return sum(n for _s, n in self._quarantine.values()) + + @property + def has_outstanding(self) -> bool: + """Whether a live or quarantined slot still owns any part of the arena.""" + with self._cv: + return bool(self._in_use or self._quarantine) def reg_descs(self) -> "RegMemoryDescs": return self._buf.reg_descs() def close(self) -> None: + with self._cv: + if self._in_use or self._quarantine: + raise RuntimeError( + "[kv-bounce] cannot close an arena with live or quarantined slots" + ) + # Close admission while holding the allocator gate. A blocked or + # racing reserve must observe this state before the VMM is unmapped. + self._closed = True + self._cv.notify_all() self._buf.close() diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py index 1878fa90fd3a..4d6165d3443d 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py @@ -16,9 +16,22 @@ machine. No CUDA or NIXL imports, so it is unit-testable on CPU.""" from abc import ABC, abstractmethod -from dataclasses import dataclass, field +from bisect import bisect_right +from dataclasses import InitVar, dataclass, field from enum import Enum -from typing import Callable, Dict, List, Optional, Tuple +from heapq import merge +from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple + +_MAX_ADDRESS = (1 << 64) - 1 + + +class GatherSourceInDoubtError(RuntimeError): + """A failed gather may still access its source KV pages. + + Callers must retain the source request/KV owner until local CUDA + quiescence is established; this is not an ordinary terminal transfer + failure. + """ class TransferState(Enum): @@ -29,7 +42,7 @@ class TransferState(Enum): SCATTERING = "scattering" # all writers succeeded; scattering back into the cache COMPLETED = "completed" # scattered cleanly, slot released FAILED = "failed" # all writers done, at least one failed, slot released - QUARANTINED = "quarantined" # a writer is in doubt, slot held out of reuse + QUARANTINED = "quarantined" # logical failure; an exposed writer is still in doubt class ScatterState(Enum): @@ -39,11 +52,20 @@ class ScatterState(Enum): FAILED = "failed" +class WriterState(Enum): + """Physical-access state for one planned writer.""" + + UNEXPOSED = "unexposed" + EXPOSED = "exposed" + NO_ACCESS = "no_access" + TERMINAL = "terminal" + + class Disposition(Enum): """What to do with the slot when the transfer finishes.""" RELEASE = "release" # drained: return it to the free pool - QUARANTINE = "quarantine" # in doubt: hold it out of the pool for a grace period + QUARANTINE = "quarantine" # in doubt: hold it until explicit quiescence evidence @dataclass @@ -58,30 +80,168 @@ class Settlement: @dataclass -class TransferContext: - """Lifetime state machine for one bounced receive region. The region is released only after every - writer reports (success or failure both mean its write has drained); a writer given up on (a - cancelled or timed-out transfer, whose one-sided write cannot be aborted) makes it quarantined - instead of reused.""" +class RecvBounceContext: + """Lifetime state machine for one bounced receive region. + + ``writer_ranks`` is the exact ordered publication plan. Logical failure suppresses scatter and + prevents publication to writers that were not exposed yet, but it does not retire an exposed + writer. The slot remains owned until every possibly exposed writer reports a terminal result or + backend-wide quiescence provides equivalent evidence. + """ rid_slice: Tuple[int, int] slot_id: int base_addr: int per_writer_bytes: int - num_writers: int + writer_ranks: Tuple[int, ...] + destination_intervals: InitVar[Optional[Iterable[tuple[int, int]]]] = None on_done: Optional[Callable[[bool], None]] = None + on_logical_failure: Optional[Callable[[], None]] = None - # Whether each writer that reported succeeded, keyed by rank; presence means it reported. + # Whether each physically terminal writer succeeded, keyed by its exact planned rank. _writer_ok: Dict[int, bool] = field(default_factory=dict) - # per successful writer: where it wrote, plus the fragments to scatter back - _scatter_descs: List[tuple] = field(default_factory=list) - _orphaned: bool = False + _writer_state: Dict[int, WriterState] = field(init=False) + # Per successful bounced writer: where it wrote, plus the fragments to scatter back. + _scatter_descs: Dict[int, tuple] = field(default_factory=dict) + _scatter_destination_intervals: Tuple[tuple[int, int], ...] = field(default_factory=tuple) + _logical_failed: bool = False + _backend_quiesced: bool = False + _protocol_conflict: bool = False + _destination_intervals: Optional[Tuple[tuple[int, int], ...]] = field(init=False, default=None) scatter_state: ScatterState = ScatterState.IDLE state: TransferState = TransferState.INIT settled: bool = False + _logical_failure_notification_in_progress: bool = False + _logical_failure_notification_delivered: bool = False + + def __post_init__(self, destination_intervals: Optional[Iterable[tuple[int, int]]]) -> None: + self.writer_ranks = tuple(self.writer_ranks) + if not self.writer_ranks: + raise ValueError("RecvBounceContext requires at least one writer rank") + if any( + not isinstance(rank, int) or isinstance(rank, bool) or rank < 0 + for rank in self.writer_ranks + ): + raise ValueError( + f"RecvBounceContext writer ranks must be non-negative integers: {self.writer_ranks}" + ) + if len(set(self.writer_ranks)) != len(self.writer_ranks): + raise ValueError(f"RecvBounceContext writer ranks must be unique: {self.writer_ranks}") + self._writer_state = {rank: WriterState.UNEXPOSED for rank in self.writer_ranks} + if destination_intervals is not None: + normalized_intervals = [] + for start, size in destination_intervals: + start, size = int(start), int(size) + end = start + size + if ( + start <= 0 + or start > _MAX_ADDRESS + or size <= 0 + or size > _MAX_ADDRESS + or end > _MAX_ADDRESS + 1 + ): + raise ValueError(f"invalid trusted destination interval ({start}, {size})") + normalized_intervals.append((start, end)) + if not normalized_intervals: + raise ValueError("trusted destination intervals must not be empty") + merged_intervals = [] + for start, end in sorted(normalized_intervals): + if merged_intervals and start <= merged_intervals[-1][1]: + previous_start, previous_end = merged_intervals[-1] + merged_intervals[-1] = (previous_start, max(previous_end, end)) + else: + merged_intervals.append((start, end)) + authorized_bytes = sum(end - start for start, end in merged_intervals) + reserved_bytes = self.per_writer_bytes * len(self.writer_ranks) + if authorized_bytes != reserved_bytes: + raise ValueError( + "trusted destination interval union covers " + f"{authorized_bytes} bytes, but the bounce reservation covers " + f"{reserved_bytes} bytes" + ) + self._destination_intervals = tuple(merged_intervals) + + @staticmethod + def _non_overlapping_destination_intervals( + dst_values: Tuple[int, ...], size_values: Tuple[int, ...] + ) -> Optional[Tuple[tuple[int, int], ...]]: + intervals = [] + for dst, size in zip(dst_values, size_values, strict=True): + end = dst + size + if ( + dst <= 0 + or dst > _MAX_ADDRESS + or size <= 0 + or size > _MAX_ADDRESS + or end > _MAX_ADDRESS + 1 + ): + return None + intervals.append((dst, end)) + if not intervals: + return None + intervals.sort() + if any( + start < previous_end for (_, previous_end), (start, _) in zip(intervals, intervals[1:]) + ): + return None + return tuple(intervals) + + def _scatter_destinations_disjoint( + self, candidate_intervals: Tuple[tuple[int, int], ...] + ) -> bool: + existing_index = 0 + for start, end in candidate_intervals: + while ( + existing_index < len(self._scatter_destination_intervals) + and self._scatter_destination_intervals[existing_index][1] <= start + ): + existing_index += 1 + if ( + existing_index < len(self._scatter_destination_intervals) + and self._scatter_destination_intervals[existing_index][0] < end + ): + return False + return True + + def _scatter_destinations_authorized( + self, dst_values: Tuple[int, ...], size_values: Tuple[int, ...] + ) -> bool: + if self._destination_intervals is None: + return True + for dst, size in zip(dst_values, size_values, strict=True): + end = dst + size + if ( + dst <= 0 + or dst > _MAX_ADDRESS + or size <= 0 + or size > _MAX_ADDRESS + or end > _MAX_ADDRESS + 1 + ): + return False + interval_index = bisect_right(self._destination_intervals, (dst, _MAX_ADDRESS + 1)) - 1 + if interval_index < 0: + return False + start, interval_end = self._destination_intervals[interval_index] + if not (start <= dst and end <= interval_end): + return False + return True + + @property + def logical_failed(self) -> bool: + return self._logical_failed + + @property + def pending_exposed_writers(self) -> Tuple[int, ...]: + return tuple( + rank for rank in self.writer_ranks if self._writer_state[rank] is WriterState.EXPOSED + ) - def writer_base(self, writer_index: int) -> int: - """Where this fan-in writer writes in the region.""" + def writer_base(self, peer_rank: int) -> int: + """Return the planned sub-region base for ``peer_rank``.""" + try: + writer_index = self.writer_ranks.index(peer_rank) + except ValueError as e: + raise KeyError(f"unexpected bounce writer rank {peer_rank}") from e return self.base_addr + writer_index * self.per_writer_bytes def _writers_final(self) -> bool: @@ -89,14 +249,121 @@ def _writers_final(self) -> bool: TransferState.SCATTERING, TransferState.COMPLETED, TransferState.FAILED, - TransferState.QUARANTINED, ) - def _all_writers_reported(self) -> bool: - return len(self._writer_ok) >= self.num_writers + def _all_writers_terminal(self) -> bool: + return all( + state in (WriterState.NO_ACCESS, WriterState.TERMINAL) + for state in self._writer_state.values() + ) def _all_writers_succeeded(self) -> bool: - return self._all_writers_reported() and all(self._writer_ok.values()) + return self._all_writers_terminal() and all( + self._writer_ok.get(rank, False) for rank in self.writer_ranks + ) + + def set_on_done(self, on_done: Optional[Callable[[bool], None]]) -> None: + if on_done is not None: + self.on_done = on_done + + def set_on_logical_failure(self, on_logical_failure: Optional[Callable[[], None]]) -> None: + """Install the notification used when launched local CUDA work fails. + + This notification is deliberately separate from ``on_done``. The + latter acknowledges physical slot settlement, which is not safe after + a scatter error without a positive CUDA quiescence fence. + """ + if on_logical_failure is not None: + self.on_logical_failure = on_logical_failure + + def begin_logical_failure_notification(self) -> Optional[Callable[[], None]]: + """Claim the one retryable logical-failure notification, if ready.""" + if ( + self.scatter_state is not ScatterState.FAILED + or self.on_logical_failure is None + or self._logical_failure_notification_delivered + or self._logical_failure_notification_in_progress + ): + return None + self._logical_failure_notification_in_progress = True + return self.on_logical_failure + + def finish_logical_failure_notification(self, delivered: bool) -> None: + """Commit or reopen a claimed logical-failure notification.""" + if not self._logical_failure_notification_in_progress: + return + self._logical_failure_notification_in_progress = False + if delivered: + self._logical_failure_notification_delivered = True + + def mark_writer_exposed(self, peer_rank: int) -> bool: + """Enter the publication boundary for a planned writer. + + False means publication must be suppressed because the writer is unknown, the context is + logically failed, or its operation is already terminal. Repeated calls for an already + exposed writer are idempotent. + """ + writer_state = self._writer_state.get(peer_rank) + if writer_state is None or self._writers_final() or self._logical_failed: + return False + if writer_state is WriterState.UNEXPOSED: + self._writer_state[peer_rank] = WriterState.EXPOSED + if self.state is TransferState.INIT: + self.state = TransferState.ACTIVE + return True + return writer_state is WriterState.EXPOSED + + def mark_writer_no_access(self, peer_rank: int, *, succeeded: bool = True) -> bool: + """Record positive evidence that ``peer_rank`` cannot access the receive slot.""" + writer_state = self._writer_state.get(peer_rank) + if writer_state is None or self._writers_final(): + return False + if writer_state in (WriterState.NO_ACCESS, WriterState.TERMINAL): + return False + self._writer_state[peer_rank] = WriterState.NO_ACCESS + self._writer_ok[peer_rank] = succeeded + if not succeeded: + self.mark_logical_failure() + return True + + def mark_logical_failure(self) -> None: + """Fail logically and close every not-yet-exposed writer. + + Writers already exposed remain in doubt and keep the slot owned. Because this transition + gates future calls to :meth:`mark_writer_exposed`, changing an unexposed writer to + ``NO_ACCESS`` is safe and atomic when called under the transport lock. + """ + if self.settled: + return + self._logical_failed = True + self._scatter_descs.clear() + self._scatter_destination_intervals = () + for rank, writer_state in self._writer_state.items(): + if writer_state is WriterState.UNEXPOSED: + self._writer_state[rank] = WriterState.NO_ACCESS + self._writer_ok[rank] = False + if self.pending_exposed_writers and self.state is not TransferState.SCATTERING: + self.state = TransferState.QUARANTINED + + def mark_protocol_conflict(self) -> None: + """Retain the slot until backend-wide quiescence resolves ambiguity.""" + if self.settled: + return + self._protocol_conflict = True + self.mark_logical_failure() + + def mark_backend_quiesced(self) -> None: + """Retire every remote accessor after backend-wide quiescence is proven.""" + if self.settled: + return + self._backend_quiesced = True + self.mark_logical_failure() + for rank, writer_state in self._writer_state.items(): + if writer_state not in (WriterState.NO_ACCESS, WriterState.TERMINAL): + self._writer_state[rank] = WriterState.NO_ACCESS + self._writer_ok[rank] = False + if self.state is TransferState.QUARANTINED: + self.state = TransferState.ACTIVE # Mutations: call only while holding the transport's reservation lock. def record_writer_result( @@ -107,45 +374,110 @@ def record_writer_result( src_base: Optional[int] = None, dst_ptrs=None, sizes=None, - ) -> None: + ) -> bool: """Record one writer's terminal report. Repeat or late reports are ignored, so duplicate or out-of-order messages are harmless.""" # A duplicate or late report can still arrive after the writer set is final, because # scatter runs on another thread while the region stays live. For example a retransmitted # notification, or a stray failure that would flip a good transfer to failed; drop it. - if self._writers_final() or peer_rank in self._writer_ok: - return + writer_state = self._writer_state.get(peer_rank) + if ( + writer_state is None + or self._writers_final() + or writer_state in (WriterState.NO_ACCESS, WriterState.TERMINAL) + ): + return False + + # A result itself proves that the writer crossed its publication boundary. This implicit + # transition preserves compatibility while callers are migrated to mark exposure before + # sending an address. + if writer_state is WriterState.UNEXPOSED: + self._writer_state[peer_rank] = WriterState.EXPOSED + + scatter_intervals = None + tail_absent = dst_ptrs is None and sizes is None and src_base is None + if succeeded and not tail_absent: + valid_tail = ( + dst_ptrs is not None + and sizes is not None + and int(dst_ptrs.size) == int(sizes.size) + and int(dst_ptrs.size) > 0 + ) + expected_src = self.writer_base(peer_rank) + valid_tail = valid_tail and src_base is not None and int(src_base) == expected_src + if valid_tail: + dst_values = tuple(int(ptr) for ptr in dst_ptrs) + size_values = tuple(int(size) for size in sizes) + scatter_intervals = self._non_overlapping_destination_intervals( + dst_values, size_values + ) + # ``src_base`` names this writer's fixed bounce sub-region and + # scatter consumes the fragments contiguously from that base. + # Requiring the exact contribution prevents a success tail + # from silently truncating the authorized source mapping while + # still allowing any equivalent fragment segmentation. + valid_tail = ( + scatter_intervals is not None and sum(size_values) == self.per_writer_bytes + ) + valid_tail = valid_tail and self._scatter_destinations_authorized( + dst_values, size_values + ) + valid_tail = valid_tail and self._scatter_destinations_disjoint(scatter_intervals) + if not valid_tail: + succeeded = False + + self._writer_state[peer_rank] = WriterState.TERMINAL self._writer_ok[peer_rank] = succeeded - if succeeded and dst_ptrs is not None and int(dst_ptrs.size) > 0: - self._scatter_descs.append( - (src_base if src_base is not None else self.base_addr, dst_ptrs, sizes) + if ( + succeeded + and not self._logical_failed + and dst_ptrs is not None + and int(dst_ptrs.size) > 0 + ): + assert scatter_intervals is not None + self._scatter_descs[peer_rank] = (int(src_base), dst_ptrs, sizes) + self._scatter_destination_intervals = tuple( + merge(self._scatter_destination_intervals, scatter_intervals) ) + elif not succeeded: + self.mark_logical_failure() if self.state is TransferState.INIT: self.state = TransferState.ACTIVE - - def mark_orphaned(self) -> None: - """Give up on writers that never reported (cancel, timeout, shutdown): a one-sided write - cannot be aborted, so the region is quarantined on settle. No-op once scattering or done.""" - if self._writers_final(): - return - self._orphaned = True + return True def begin_scatter(self) -> None: self.state = TransferState.SCATTERING self.scatter_state = ScatterState.QUEUED def sorted_scatter_descs(self) -> List[tuple]: - """Success fragments ordered by source, for a deterministic scatter.""" - return sorted(self._scatter_descs, key=lambda t: t[0]) + """Success fragments ordered by the exact writer publication plan.""" + return [ + self._scatter_descs[rank] for rank in self.writer_ranks if rank in self._scatter_descs + ] def finish_scatter(self, ok: bool) -> None: + """Record the local scatter outcome. + + A CUDA error is not positive evidence that queued memory accesses have + stopped. ``FAILED`` therefore remains owned indefinitely rather than + making the slot or destination KV observable as drained. + """ self.scatter_state = ScatterState.DONE if ok else ScatterState.FAILED + if not ok: + # Logical failure is known even though the slot and destination + # remain physically owned until local CUDA quiescence is proven. + self.mark_logical_failure() + + def suppress_scatter(self) -> None: + """Fail work that was queued but provably never launched on CUDA.""" + self.mark_logical_failure() + self.scatter_state = ScatterState.DONE def ready_to_scatter(self) -> bool: """All writers succeeded, there is data to scatter, and scatter has not started.""" return ( self.state is TransferState.ACTIVE - and not self._orphaned + and not self._logical_failed and self._all_writers_succeeded() and bool(self._scatter_descs) ) @@ -153,24 +485,27 @@ def ready_to_scatter(self) -> bool: def ready_to_settle(self) -> bool: if self.settled: return False - if self._orphaned: - return True # in doubt: settle now and quarantine - if not self._all_writers_reported(): - return False # a writer has not reported yet - if self._all_writers_succeeded() and self._scatter_descs: - return self.scatter_state in (ScatterState.DONE, ScatterState.FAILED) + if self._protocol_conflict and not self._backend_quiesced: + return False + if not self._all_writers_terminal(): + return False # an exposed writer may still access the slot + if self.scatter_state in (ScatterState.QUEUED, ScatterState.FAILED): + return False # queued or failed CUDA work has no positive local fence + if not self._logical_failed and self._all_writers_succeeded() and self._scatter_descs: + return self.scatter_state is ScatterState.DONE return True # nothing to scatter, or a failure among them: either way drained def settle(self) -> Optional[Settlement]: """Finalize the transfer once and return the decision, or None if already done. Every termination path calls this; only the first call has any effect.""" - if self.settled: + if self.settled or not self.ready_to_settle(): return None self.settled = True - if self._orphaned: - self.state = TransferState.QUARANTINED - return Settlement(self.slot_id, Disposition.QUARANTINE, False, self.on_done) - success = self._all_writers_succeeded() and self.scatter_state is not ScatterState.FAILED + success = ( + not self._logical_failed + and self._all_writers_succeeded() + and self.scatter_state is not ScatterState.FAILED + ) self.state = TransferState.COMPLETED if success else TransferState.FAILED return Settlement(self.slot_id, Disposition.RELEASE, success, self.on_done) @@ -192,12 +527,28 @@ def release_send(self, slot_id) -> None: """Release a send region after its write completes.""" @abstractmethod - def reserve(self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = None) -> bool: - """Reserve a region and record its address for the senders. False falls back to per-fragment.""" + def quarantine_send(self, slot_id) -> None: + """Permanently retain a send region whose NIXL access remains ambiguous.""" @abstractmethod - def writer_base(self, rid_slice, writer_index: int) -> Optional[int]: - """Where the given fan-in writer writes in the region.""" + def reserve( + self, + recv_req, + writer_ranks: Sequence[int] = (), + *, + timeout: Optional[float] = None, + destination_intervals: Optional[Iterable[tuple[int, int]]] = None, + destination_intervals_factory: Optional[Callable[[], Iterable[tuple[int, int]]]] = None, + ) -> bool: + """Reserve for one exact writer-rank sequence, or return False for direct fallback. + + ``destination_intervals_factory`` is evaluated only after allocator admission succeeds, so + an expensive trusted destination manifest is not built for a transfer that falls back. + """ + + @abstractmethod + def writer_base(self, rid_slice, peer_rank: int) -> Optional[int]: + """Where the given exact fan-in writer rank writes in the region.""" @abstractmethod def is_bounced(self, rid_slice) -> bool: @@ -211,16 +562,58 @@ def release_idle_reservation(self, rid_slice) -> None: def orphan_reservation(self, rid_slice) -> None: """Give up on an in-flight reservation (cancel/timeout/lost result); quarantine, don't leak.""" + @abstractmethod + def mark_writer_exposed(self, rid_slice, peer_rank: int) -> bool: + """Mark the writer possibly able to access the slot before publishing its address.""" + + @abstractmethod + def record_no_access( + self, rid_slice, peer_rank: int, *, succeeded: bool = True, on_done=None + ) -> None: + """Record positive evidence that a planned writer cannot access the slot.""" + + @abstractmethod + def mark_logical_failure(self, rid_slice, on_done=None) -> None: + """Suppress future publication and drain already-exposed writers.""" + + @abstractmethod + def mark_protocol_conflict(self, rid_slice, on_done=None) -> None: + """Retain an ambiguous receive slot until backend-wide quiescence.""" + + @abstractmethod + def mark_backend_quiesced(self, rid_slice=None, on_done=None) -> None: + """Record backend-wide evidence and report when the bounce resource settles.""" + + @abstractmethod + def retry_settlements(self, scope=None) -> bool: + """Retry durable settlement acknowledgements in scope. + + ``scope`` is ``None`` for the whole transport, a request ID for every + slice of that request, or an exact ``(request_id, slice_id)`` key. + Return whether the selected scope has no pending acknowledgement. + """ + @abstractmethod def record_result( - self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None + self, + rid_slice, + peer_rank, + dst_ptrs=None, + sizes=None, + src_base=None, + on_done=None, + on_logical_failure=None, ) -> None: - """Handle a writer's success; scatter and finalize once all writers reported.""" + """Handle a writer's success; scatter and finalize once all writers reported. + + ``on_logical_failure`` reports a launched local scatter failure without + claiming that the bounce slot has physically settled. + """ @abstractmethod - def record_failure(self, rid_slice, peer_rank) -> None: + def record_failure(self, rid_slice, peer_rank, on_done=None) -> None: """Handle a writer's failure; free the region once all writers reported.""" @abstractmethod def close(self) -> None: - """Stop the scatter worker, deregister memory, free the arenas.""" + """Stop the scatter worker, deregister its memory, and free the arenas.""" diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py index 21b81b250ff5..405ecfb47ae2 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py @@ -116,23 +116,40 @@ def _uniform_nelem(sizes: np.ndarray): # reusable pinned staging for the metadata copy, one per stream, so the copy is a true async transfer _meta_lock = threading.Lock() -_meta_buffers = {} # each stream maps to its pinned host buffer, device buffer, and capacity +_meta_buffers = {} # each (device, stream) maps to pinned/device buffers and capacity + + +def _meta_buffer_key(stream_handle: int, dev) -> tuple[int, int]: + """Return a cache key that cannot alias an equal stream handle on another device.""" + device_index = dev.index + if device_index is None: + raise ValueError("CUDA metadata buffers require an explicit device index") + return int(device_index), int(stream_handle) def _get_meta_buffers(stream_handle: int, need: int, dev): """Return the pinned host and device buffers for the stream, large enough for the request, growing them under a lock.""" + key = _meta_buffer_key(stream_handle, dev) with _meta_lock: - ent = _meta_buffers.get(stream_handle) + ent = _meta_buffers.get(key) if ent is None or ent[2] < need: new_cap = max(need, (ent[2] * 2 if ent else 0), 4096) pinned = torch.empty(new_cap, dtype=torch.int64, pin_memory=prefer_pinned()) devt = torch.empty(new_cap, dtype=torch.int64, device=dev) ent = (pinned, devt, new_cap) - _meta_buffers[stream_handle] = ent + _meta_buffers[key] = ent return ent[0], ent[1] +def release_meta_buffers(stream_handle: int, device_id: int) -> None: + """Release metadata staging owned by a quiesced stream.""" + with _meta_lock: + buffers = _meta_buffers.pop((int(device_id), int(stream_handle)), None) + # Tensor destruction may enter the CUDA allocator; do it outside the cache lock. + del buffers + + def _launch_batched_copy( dst_addrs: np.ndarray, src_addrs: np.ndarray, sizes: np.ndarray, stream ) -> bool: diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index a9597b7cbb2a..f8517dcbcaab 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -18,6 +18,8 @@ import queue import threading +from collections.abc import Iterable, Sequence +from dataclasses import dataclass from typing import Callable, Dict, List, Optional import numpy as np @@ -38,15 +40,41 @@ from .buffer import SlotAllocator from .config import SizingContext, fit_within_free -from .core import BounceTransport, Disposition, Settlement, TransferContext -from .gather_scatter import Plan, gather_contiguous, scatter_contiguous +from .core import ( + BounceTransport, + Disposition, + GatherSourceInDoubtError, + RecvBounceContext, + ScatterState, + Settlement, +) +from .gather_scatter import Plan, gather_contiguous, release_meta_buffers, scatter_contiguous RidSlice = tuple # the request id and slice id a region serves _MIB = 1024 * 1024 -_SCATTER_POLL_S = 0.5 # how often the scatter worker wakes to re-check the stop flag and reclaim _RESERVE_TIMEOUT_S = 0.2 # max wait for a bounce region before falling back to per-fragment _CLOSE_JOIN_S = 2.0 # max wait for the scatter thread to drain on close -_QUARANTINE_GRACE_S = 60.0 # how long an orphaned region is held out of reuse +_WORKER_START_TIMEOUT_S = 2.0 + +# A partially constructed transport whose rollback could not finish must stay +# reachable until its registrations, streams, and mappings can be retried. +_INCOMPLETE_TRANSPORTS: set[object] = set() + + +class IncompleteBounceInitializationError(RuntimeError): + """Construction failed and at least one transport resource remains owned.""" + + def __init__(self, owner: "VmmBounceTransport", cause: Exception): + super().__init__(f"{cause} (bounce initialization rollback remains incomplete)") + self.owner = owner + + +@dataclass +class _PendingSettlement: + settlement: Settlement + physical_committed: bool = False + physical_in_progress: bool = False + callback_in_progress: bool = False class VmmBounceTransport(BounceTransport): @@ -95,7 +123,6 @@ def __init__( phys_chunk_size: int, block_bytes_per_group: List[int], min_blocks: int = 96, - quarantine_grace_s: float = _QUARANTINE_GRACE_S, name: str = "kv_bounce", ): self._agent = agent @@ -105,62 +132,298 @@ def __init__( # Below this many blocks, skip bounce: coalescing only pays off for long context (the default # is roughly twelve thousand tokens; a heuristic, and tunable). self._min_blocks = min_blocks - # how long an orphaned region is held out of reuse; must outlast the worst in-flight write - self._quarantine_grace_s = quarantine_grace_s - - # one registered region each for sending and receiving - self._send_alloc = SlotAllocator(capacity_bytes, phys_chunk_size, name=f"{name}_send") - self._recv_alloc = SlotAllocator(capacity_bytes, phys_chunk_size, name=f"{name}_recv") - self._reg_descs = [self._send_alloc.reg_descs(), self._recv_alloc.reg_descs()] - for d in self._reg_descs: - self._agent.register_memory(d) - - self._send_stream = self._new_stream() - self._send_stream_lock = threading.Lock() - + self._send_alloc = None + self._recv_alloc = None + self._send_alloc_closed = False + self._recv_alloc_closed = False + self._reg_descs = [] + # Only descriptors still registered with the agent live here. Successful + # teardown removes each entry immediately so a retry never deregisters it twice. + self._registered_descs = [] + # CUDA events remain owned until destruction succeeds. A failed destroy + # is retried by close rather than losing the handle. + self._pending_events = [] + self._send_stream = None + self._send_stream_healthy = True + self._scatter_stream = None + self._scatter_stream_healthy = True + self._scatter_thread = None + self._scatter_q = None + self._scatter_ready = None + self._scatter_start_error = None + self._scatter_worker_error = None + # gather_scatter metadata is reused per stream and remains live until + # its async H2D copy completes, so one owner holds this through launch + # and event completion. + self._send_stream_lock = threading.RLock() self._init_recv_state() - self._start_scatter_worker(name) + + try: + # One registered region each for sending and receiving. Construction + # is transactional: every later failure rolls back the resources + # that were successfully created before it. + self._send_alloc = SlotAllocator(capacity_bytes, phys_chunk_size, name=f"{name}_send") + self._recv_alloc = SlotAllocator(capacity_bytes, phys_chunk_size, name=f"{name}_recv") + self._reg_descs = [self._send_alloc.reg_descs(), self._recv_alloc.reg_descs()] + for desc in self._reg_descs: + # Registration can fail after partially mutating the backend. + # Track the attempt first so rollback conservatively requires + # successful deregistration before unmapping its arena. + self._registered_descs.append(desc) + self._agent.register_memory(desc) + + self._send_stream = self._new_stream() + self._start_scatter_worker(name) + except Exception as e: + if not self._rollback_initialization(): + _INCOMPLETE_TRANSPORTS.add(self) + raise IncompleteBounceInitializationError(self, e) from e + raise logger.info( f"[kv-bounce] Transport: send+recv regions of " f"{self._send_alloc.capacity / _MIB:.1f}MiB each" ) + def _stop_scatter_worker(self) -> None: + """Stop an initialized scatter worker once; safe to retry after a timeout.""" + thread = getattr(self, "_scatter_thread", None) + if thread is None: + return + if thread.is_alive(): + scatter_q = getattr(self, "_scatter_q", None) + if scatter_q is not None: + scatter_q.put(None) + thread.join(timeout=_CLOSE_JOIN_S) + if thread.is_alive(): + raise RuntimeError("[kv-bounce] scatter worker did not drain during close") + self._scatter_thread = None + + def _destroy_stream(self, attr_name: str) -> None: + """Evict quiesced metadata and destroy one CUDA stream. + + Callers establish that the stream has no outstanding work before this + point. Evict the cache before releasing the CUDA handle so another + thread cannot create a new stream with the recycled handle and then + have its metadata removed by this teardown. + """ + stream = getattr(self, attr_name, None) + if stream is None: + return + release_meta_buffers(stream, self._device_id) + CUASSERT(cudart.cudaStreamDestroy(stream)) + setattr(self, attr_name, None) + + def _destroy_streams(self) -> list[Exception]: + errors = [] + # The scatter worker must be stopped before its stream; send ownership + # checks in close() prove no gather can still use the send stream. + for attr_name in ("_scatter_stream", "_send_stream"): + try: + self._destroy_stream(attr_name) + except Exception as e: + errors.append(e) + return errors + + def _destroy_event(self, event) -> None: + CUASSERT(cudart.cudaEventDestroy(event)) + self._pending_events.remove(event) + + def _destroy_pending_events(self) -> list[Exception]: + errors = [] + for event in list(getattr(self, "_pending_events", ())): + try: + self._destroy_event(event) + except Exception as e: + errors.append(e) + return errors + + def _deregister_registered_descriptors(self) -> list[Exception]: + errors = [] + remaining = [] + for desc in getattr(self, "_registered_descs", ()): + try: + self._agent.deregister_memory(desc) + except Exception as e: + errors.append(e) + remaining.append(desc) + self._registered_descs = remaining + return errors + + def _close_allocators(self) -> list[Exception]: + errors = [] + for attr_name, closed_attr in ( + ("_send_alloc", "_send_alloc_closed"), + ("_recv_alloc", "_recv_alloc_closed"), + ): + allocator = getattr(self, attr_name, None) + if allocator is None or getattr(self, closed_attr, False): + continue + try: + allocator.close() + except Exception as e: + errors.append(e) + else: + setattr(self, closed_attr, True) + return errors + + def _rollback_initialization(self) -> bool: + """Retry fail-safe rollback; return whether every resource retired.""" + self._accepting_reservations = False + worker_stopped = True + try: + self._stop_scatter_worker() + except Exception as e: + worker_stopped = False + logger.error(f"[kv-bounce] constructor worker rollback failed: {e}") + + stream_errors = self._destroy_streams() if worker_stopped else [] + for error in stream_errors: + logger.error(f"[kv-bounce] constructor stream rollback failed: {error}") + event_errors = ( + self._destroy_pending_events() if worker_stopped and not stream_errors else [] + ) + for error in event_errors: + logger.error(f"[kv-bounce] constructor event rollback failed: {error}") + deregistration_errors = self._deregister_registered_descriptors() + for error in deregistration_errors: + logger.error(f"[kv-bounce] constructor deregistration rollback failed: {error}") + + # Never unmap an allocator while a registration, worker, or CUDA stream + # may still refer to it. Buffer.__del__ intentionally retains such VMM. + streams_destroyed = all( + getattr(self, attr_name, None) is None + for attr_name in ("_scatter_stream", "_send_stream") + ) + if ( + worker_stopped + and streams_destroyed + and not self._pending_events + and not self._registered_descs + ): + for error in self._close_allocators(): + logger.error(f"[kv-bounce] constructor allocator rollback failed: {error}") + allocators_closed = all( + getattr(self, allocator_attr, None) is None or getattr(self, closed_attr, False) + for allocator_attr, closed_attr in ( + ("_send_alloc", "_send_alloc_closed"), + ("_recv_alloc", "_recv_alloc_closed"), + ) + ) + complete = ( + worker_stopped + and streams_destroyed + and not self._pending_events + and not self._registered_descs + and allocators_closed + ) + if complete: + self._closed = True + _INCOMPLETE_TRANSPORTS.discard(self) + return complete + + def retry_initialization_rollback(self) -> None: + """Retry incomplete constructor cleanup without losing the owning object.""" + if not self._rollback_initialization(): + raise RuntimeError("[kv-bounce] initialization rollback is still incomplete") + def _init_recv_state(self) -> None: # Live per-transfer state, guarded by a leaf lock: mutate and decide under it, then release # it before any CUDA sync, allocator call, or callback. - self._reserved_map: Dict[RidSlice, TransferContext] = {} + self._reserved_map: Dict[RidSlice, RecvBounceContext] = {} self._reserved_map_lock = threading.Lock() + self._accepting_reservations = True + self._pending_reservations = 0 + self._pending_send_reservations = 0 + self._pending_settlements: Dict[RidSlice, _PendingSettlement] = {} + self._closed = False + self._close_lock = threading.Lock() def _start_scatter_worker(self, name: str) -> None: # Scatter runs on its own thread: it ends in a blocking sync, so keeping it off the # completion handler lets that handler keep draining other transfers. self._scatter_q: "queue.Queue" = queue.Queue() self._scatter_stream = self._new_stream() - self._stop = threading.Event() + self._scatter_ready = threading.Event() + self._scatter_start_error = None self._scatter_thread = threading.Thread( target=self._scatter_loop, name=f"{name}-scatter", daemon=True ) self._scatter_thread.start() + if not self._scatter_ready.wait(timeout=_WORKER_START_TIMEOUT_S): + raise RuntimeError("[kv-bounce] scatter worker initialization timed out") + if self._scatter_start_error is not None: + raise RuntimeError( + "[kv-bounce] scatter worker failed to initialize" + ) from self._scatter_start_error def _new_stream(self): return CUASSERT(cudart.cudaStreamCreate())[0] - def _gather_blocking(self, src_addr: int, write_meta, total: int) -> None: - """Gather the scattered fragments into the send region and block until done. The whole gather - runs under the stream lock so a second sender thread can't overwrite the shared staging buffer - mid-copy and corrupt this region; only the fast gather serializes, the writes stay parallel.""" + def _launch_gather(self, src_addr: int, write_meta, total: int): + """Launch the gather into the send region and return an event to wait on.""" plan = Plan(write_meta.src_ptrs, write_meta.dst_ptrs, write_meta.sizes, total) with self._send_stream_lock: - gather_contiguous( - src_addr, plan.src_ptrs, plan.sizes, plan.offsets, stream=self._send_stream - ) - event = CUASSERT(cudart.cudaEventCreate())[0] + event = None try: + gather_contiguous( + src_addr, plan.src_ptrs, plan.sizes, plan.offsets, stream=self._send_stream + ) + event = CUASSERT(cudart.cudaEventCreate())[0] + self._pending_events.append(event) CUASSERT(cudart.cudaEventRecord(event, self._send_stream)) + except Exception: + if event is not None: + try: + self._destroy_event(event) + except Exception as cleanup_error: + logger.error( + f"[kv-bounce] failed to destroy gather event after launch error: " + f"{cleanup_error}" + ) + raise + return event + + def _wait_gather(self, event) -> None: + if event is not None: + try: CUASSERT(cudart.cudaEventSynchronize(event)) - finally: - CUASSERT(cudart.cudaEventDestroy(event)) + except Exception: + try: + self._destroy_event(event) + except Exception as cleanup_error: + logger.error( + f"[kv-bounce] failed to destroy gather event after wait error: " + f"{cleanup_error}" + ) + raise + self._destroy_event(event) + + def _rollback_send_slot(self, slot_id: int) -> bool: + """Retire a failed send slot; return whether source access is proven quiescent.""" + try: + with self._send_stream_lock: + CUASSERT(cudart.cudaStreamSynchronize(self._send_stream)) + except Exception as e: + # A failed CUDA fence leaves the gather's physical-access state + # ambiguous. Never hand the region to another transfer. + self._send_stream_healthy = False + try: + self._send_alloc.quarantine(slot_id) + except Exception as quarantine_error: + # Poison the stream before this fallible bookkeeping step. The + # caller must still retain the transfer/source owner even when + # the allocator cannot move the slot into its quarantine map. + logger.critical( + f"[kv-bounce] failed to quarantine send slot {slot_id} after " + f"stream fence failure; arena ownership remains in doubt: " + f"{quarantine_error}" + ) + logger.error(f"[kv-bounce] retained send slot {slot_id} after stream fence failed: {e}") + return False + else: + self._send_alloc.release(slot_id) + return True def _make_write(self, src_addr: int, write_meta, total: int): # one coalesced descriptor spanning the whole region @@ -176,11 +439,18 @@ def _make_write(self, src_addr: int, write_meta, total: int): ) return TransferRequest(TransferOp.WRITE, src, dst, write_meta.peer_name, None) - def _reserve_and_gather(self, write_meta, *, timeout): - """Reserve a send slot and gather into it, or None on send-region backpressure. Eligibility - was already decided by the receiver, so the sender only falls back under backpressure.""" + def _reserve_send_slot(self, write_meta, *, timeout): + """Reserve a send slot, or return None on backpressure/unhealthy transport.""" total = int(write_meta.sizes.sum()) - res = self._send_alloc.reserve(total, timeout=timeout) + with self._reserved_map_lock: + if not self._accepting_reservations or not self._send_stream_healthy: + return None + self._pending_send_reservations += 1 + try: + res = self._send_alloc.reserve(total, timeout=timeout) + finally: + with self._reserved_map_lock: + self._pending_send_reservations -= 1 if res is None: logger.debug( f"[kv-bounce] in-place: no send region space for {total // _MIB}MiB within {timeout}s " @@ -188,26 +458,60 @@ def _reserve_and_gather(self, write_meta, *, timeout): ) return None slot_id, src_addr = res - try: - self._gather_blocking(src_addr, write_meta, total) - except Exception: - self._send_alloc.release(slot_id) # free the slot if the gather raises - raise return slot_id, src_addr, total def build_request(self, write_meta): - """Gather into a send slot and build the coalesced write, or None on backpressure. The gather - blocks (and frees the slot on failure) inside _reserve_and_gather.""" - gathered = self._reserve_and_gather(write_meta, timeout=_RESERVE_TIMEOUT_S) - if gathered is None: # backpressure: fall back + """Gather into a send slot and build the coalesced write, or None on backpressure. Frees the + slot if the gather raises.""" + reserved = self._reserve_send_slot(write_meta, timeout=_RESERVE_TIMEOUT_S) + if reserved is None: # backpressure or unhealthy stream: fall back return None - slot_id, src_addr, total = gathered - return self._make_write(src_addr, write_meta, total), slot_id + slot_id, src_addr, total = reserved + with self._send_stream_lock: + # Another gather may have poisoned the stream while this caller + # reserved outside the stream-critical section. + if not self._send_stream_healthy: + self._send_alloc.release(slot_id) + return None + try: + event = self._launch_gather(src_addr, write_meta, total) + self._wait_gather(event) + request = self._make_write(src_addr, write_meta, total) + except Exception as e: + if not self._rollback_send_slot(slot_id): + raise GatherSourceInDoubtError( + f"gather for send slot {slot_id} failed without a positive CUDA fence" + ) from e + raise + return request, slot_id def release_send(self, slot_id) -> None: """Release a send region after its write has completed.""" self._send_alloc.release(slot_id) + def quarantine_send(self, slot_id) -> None: + """Retain a send slot until an external NIXL quiescence mechanism exists.""" + self._send_alloc.quarantine(slot_id) + + @staticmethod + def _normalize_writer_ranks(writer_ranks: Sequence[int]) -> tuple[int, ...]: + """Return an exact ordered writer plan, or empty for safe direct fallback. + + Integer counts are deliberately rejected: synthesizing ``range(count)`` + associates ownership with ranks that may not be the actual writers. + """ + if isinstance(writer_ranks, (bool, int)): + return () + try: + ranks = tuple(writer_ranks) + except TypeError: + return () + if any(not isinstance(rank, int) or isinstance(rank, bool) or rank < 0 for rank in ranks): + return () + if len(set(ranks)) != len(ranks): + return () + return ranks + @staticmethod def _skip_bounce(reason: str, *, warn_key: Optional[str] = None) -> bool: """Log why a transfer falls back to the per-fragment path and return False, so the guards @@ -217,140 +521,304 @@ def _skip_bounce(reason: str, *, warn_key: Optional[str] = None) -> bool: return False def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + writer_ranks: Sequence[int] = (), + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + destination_intervals: Optional[Iterable[tuple[int, int]]] = None, + destination_intervals_factory: Optional[Callable[[], Iterable[tuple[int, int]]]] = None, ) -> bool: - """Reserve a region and create its state, recording the address for the senders. Returns - False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total - must divide across the writers.""" - nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) + """Reserve a region and create its state, recording the address for one sender. + + Multi-writer bounce is intentionally rejected until the caller can + provide receiver-owned, byte-accurate offsets and extents for every + writer. Divisibility and uniform slot sizes do not constrain the size + of the NIXL descriptor that a remote writer will submit. + """ + ranks = self._normalize_writer_ranks(writer_ranks) + if not ranks: + return self._skip_bounce("writer plan is empty, invalid, or contains duplicate ranks") + if destination_intervals is not None and destination_intervals_factory is not None: + return self._skip_bounce( + "trusted destination intervals and their lazy factory are mutually exclusive" + ) + num_writers = len(ranks) + if num_writers > 1: + return self._skip_bounce( + "multi-writer bounce requires byte-accurate receiver-owned writer extents", + warn_key="kv-bounce-multiwriter-without-extents", + ) + with self._reserved_map_lock: + if not self._accepting_reservations or not self._scatter_stream_healthy: + return self._skip_bounce("transport is closing") + valid_block_counts = [ + int(np.count_nonzero(np.asarray(block_ids, dtype=np.int64) >= 0)) + for block_ids in recv_req.block_ids_per_layer_groups + ] + nblocks = sum(valid_block_counts) if nblocks < self._min_blocks: return self._skip_bounce(f"{nblocks} blocks < min {self._min_blocks} (too small)") total = 0 - for g, block_ids in enumerate(recv_req.block_ids_per_layer_groups): + for g, valid_blocks in enumerate(valid_block_counts): if g >= len(self._block_bytes_per_group): return self._skip_bounce(f"layer group {g} has no known slot size (e.g. mamba)") - total += int(block_ids.size) * self._block_bytes_per_group[g] + total += valid_blocks * self._block_bytes_per_group[g] if total <= 0: return self._skip_bounce(f"computed transfer size {total} <= 0") - if num_writers > 1 and total % num_writers != 0: - return self._skip_bounce( - f"fan-in {total}B across {num_writers} senders is not an even split " - f"({total % num_writers}B remainder); head-mismatch explosion NOT mitigated", - warn_key="kv-bounce-uneven-fanin", - ) - if num_writers > 1: - # Fan-in gives each writer an equal share of the region, which only matches where it - # writes when all writers send the same bytes. Equal layer count guarantees that only - # when the per-block sizes match, so require that here, else fall back. - present_slot_bytes = { - self._block_bytes_per_group[g] - for g, block_ids in enumerate(recv_req.block_ids_per_layer_groups) - if int(block_ids.size) > 0 - } - if len(present_slot_bytes) > 1: - return self._skip_bounce( - f"fan-in across {num_writers} senders with non-uniform layer-group slot bytes " - f"{sorted(present_slot_bytes)}; the equal split would overrun a sub-region", - warn_key="kv-bounce-heterogeneous-fanin", - ) if total > self._recv_alloc.capacity: # too big to ever fit, unlike transient backpressure return self._skip_bounce( f"transfer {total // _MIB}MiB exceeds the {self._recv_alloc.capacity // _MIB}MiB bounce " f"region; raise the bounce arena size to re-enable coalescing", warn_key="kv-bounce-oversize", ) - res = self._recv_alloc.reserve(total, timeout=timeout) - if res is None: - return self._skip_bounce( - f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)" - ) - slot_id, addr = res - recv_req.bounce_dst_base = addr with self._reserved_map_lock: - ctx = TransferContext( - rid_slice=(recv_req.unique_rid, recv_req.slice_id), - slot_id=slot_id, - base_addr=addr, - per_writer_bytes=total // num_writers, - num_writers=num_writers, + if not self._accepting_reservations or not self._scatter_stream_healthy: + return self._skip_bounce("transport is closing") + self._pending_reservations += 1 + try: + res = self._recv_alloc.reserve(total, timeout=timeout) + if res is None: + return self._skip_bounce( + f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)" + ) + slot_id, addr = res + try: + if destination_intervals_factory is not None: + destination_intervals = destination_intervals_factory() + if destination_intervals is None: + raise ValueError( + "trusted destination interval factory returned no intervals" + ) + ctx = RecvBounceContext( + rid_slice=(recv_req.unique_rid, recv_req.slice_id), + slot_id=slot_id, + base_addr=addr, + per_writer_bytes=total // num_writers, + writer_ranks=ranks, + destination_intervals=destination_intervals, + ) + except (TypeError, ValueError) as e: + logger.error(f"[kv-bounce] invalid receive ownership plan: {e}") + self._recv_alloc.release(slot_id) + return self._skip_bounce("receive ownership plan is invalid") + except Exception: + self._recv_alloc.release(slot_id) + raise + + duplicate = False + with self._reserved_map_lock: + if ( + not self._accepting_reservations + or ctx.rid_slice in self._reserved_map + or ctx.rid_slice in self._pending_settlements + ): + duplicate = True + else: + self._reserved_map[ctx.rid_slice] = ctx + if duplicate: + self._recv_alloc.release(slot_id) + return self._skip_bounce("transport is closing or transfer context already exists") + recv_req.bounce_dst_base = addr + # Positive marker: every fallback guard passed, so this transfer + # provably uses the coalesced bounce path. + logger.info_once( + f"[kv-bounce] coalesced {nblocks} blocks / {total // _MIB}MiB into one region " + f"across {num_writers} writer(s)", + key="kv-bounce-coalesced", ) - self._reserved_map[ctx.rid_slice] = ctx # inactive until the first writer reports - # Positive marker: all fall-back guards above passed, so this transfer provably takes the - # coalesced-bounce WRITE path. Logged once (per process) so an e2e test can assert that - # bounce actually engaged instead of silently falling back to the per-fragment path. - logger.info_once( - f"[kv-bounce] coalesced {nblocks} blocks / {total // _MIB}MiB into one region " - f"across {num_writers} writer(s)", - key="kv-bounce-coalesced", - ) - return True + return True + finally: + with self._reserved_map_lock: + self._pending_reservations -= 1 - def writer_base(self, rid_slice: RidSlice, writer_index: int) -> Optional[int]: - """Where the given fan-in writer writes in the region.""" + def writer_base(self, rid_slice: RidSlice, peer_rank: int) -> Optional[int]: + """Where the exact planned fan-in writer writes in the region.""" with self._reserved_map_lock: ctx = self._reserved_map.get(rid_slice) - return None if ctx is None else ctx.writer_base(writer_index) + if ctx is None: + return None + try: + return ctx.writer_base(peer_rank) + except KeyError: + return None def is_bounced(self, rid_slice: RidSlice) -> bool: with self._reserved_map_lock: - return rid_slice in self._reserved_map + return rid_slice in self._reserved_map or rid_slice in self._pending_settlements def release_idle_reservation(self, rid_slice: RidSlice) -> None: - """Immediately release a reservation cancelled before any address went out; no write can be - in flight. Idempotent. Drained transfers finalize through the result path instead.""" - with self._reserved_map_lock: - ctx = self._reserved_map.pop(rid_slice, None) - if ctx is not None: - self._recv_alloc.release(ctx.slot_id) + """Compatibility cancellation entry point. + + A genuinely idle context releases immediately. If any writer was exposed, this only closes + future publication and retains the slot until those writers become terminal. + """ + self.mark_logical_failure(rid_slice) def orphan_reservation(self, rid_slice: RidSlice) -> None: - """Give up on a reservation whose write may still be in flight (cancel/timeout/lost result). - The write can't be aborted, so quarantine the region (reclaimed later) rather than releasing - or leaking it. Idempotent; a no-op once the transfer has settled.""" - self._apply(rid_slice, lambda ctx: ctx.mark_orphaned()) - - def _apply(self, rid_slice: RidSlice, mutate: Callable[[TransferContext], None]) -> None: - """Mutate the state under the lock, then do what it asks (scatter or settle) with the lock - released, never holding it across a CUDA sync, a queue put, or a callback. No-op if the - region is already gone.""" - scatter: Optional[tuple] = None - settlement: Optional[Settlement] = None + """Retain an ambiguous reservation until backend-wide quiescence.""" + self.mark_protocol_conflict(rid_slice) + + def _apply(self, rid_slice: RidSlice, mutate: Callable[[RecvBounceContext], None]) -> None: + """Mutate and enqueue under the lock, then settle without holding it across allocator or + callback work. No-op if the region is already gone.""" + retry_settlement = False with self._reserved_map_lock: ctx = self._reserved_map.get(rid_slice) if ctx is None: - return - mutate(ctx) - if ctx.ready_to_scatter(): - ctx.begin_scatter() - scatter = (ctx, ctx.sorted_scatter_descs()) - elif ctx.ready_to_settle(): - settlement = ctx.settle() - if settlement is not None: - self._reserved_map.pop(rid_slice, None) - if scatter is not None: - self._enqueue_scatter(*scatter) - if settlement is not None: - self._commit(settlement) - - def _enqueue_scatter(self, ctx: TransferContext, descs: List[tuple]) -> None: + retry_settlement = rid_slice in self._pending_settlements + else: + mutate(ctx) + if ctx.ready_to_scatter(): + ctx.begin_scatter() + if self._scatter_stream_healthy: + # Queue insertion is inside the publication/close lock. This guarantees + # close's FIFO poison pill cannot overtake an accepted scatter. + self._enqueue_scatter(ctx, ctx.sorted_scatter_descs()) + else: + # This context never reached CUDA because a prior job + # poisoned the shared stream. + ctx.suppress_scatter() + if ctx.ready_to_settle(): + settlement = ctx.settle() + if settlement is not None: + self._reserved_map.pop(rid_slice, None) + self._pending_settlements[rid_slice] = _PendingSettlement(settlement) + retry_settlement = True + if retry_settlement: + self._commit_pending_settlement(rid_slice) + self._retry_logical_failure_notification(rid_slice) + + def _retry_logical_failure_notification(self, rid_slice: RidSlice) -> bool: + """Deliver logical scatter failure without releasing its physical lease.""" + with self._reserved_map_lock: + ctx = self._reserved_map.get(rid_slice) + if ctx is None: + return True + callback = ctx.begin_logical_failure_notification() + if callback is None: + return not ( + ctx.scatter_state is ScatterState.FAILED + and ctx.on_logical_failure is not None + and not ctx._logical_failure_notification_delivered + ) + try: + callback() + except Exception as e: + with self._reserved_map_lock: + current = self._reserved_map.get(rid_slice) + if current is ctx: + current.finish_logical_failure_notification(False) + logger.error( + f"[kv-bounce] logical scatter-failure notification failed " + f"(slot={ctx.slot_id}); retaining it for retry: {e}" + ) + return False + with self._reserved_map_lock: + current = self._reserved_map.get(rid_slice) + if current is ctx: + current.finish_logical_failure_notification(True) + return True + + def _enqueue_scatter(self, ctx: RecvBounceContext, descs: List[tuple]) -> None: """Hand the per-writer fragments to the worker. Each is scattered from its own source, so a writer that fell back to the in-place path cannot shift where the others are read from.""" self._scatter_q.put((ctx, descs)) - def _commit(self, settlement: Settlement) -> None: - """Carry out the decision: release or quarantine the slot, then fire the callback once. No - lock is held.""" - if settlement.disposition is Disposition.QUARANTINE: - self._recv_alloc.quarantine(settlement.slot_id, self._quarantine_grace_s) - else: - self._recv_alloc.release(settlement.slot_id) - if settlement.on_done is not None: + def _commit_pending_settlement(self, rid_slice: RidSlice) -> bool: + """Commit physical release once and durably retry its ownership acknowledgement.""" + with self._reserved_map_lock: + pending = self._pending_settlements.get(rid_slice) + if pending is None: + return True + if pending.physical_in_progress or pending.callback_in_progress: + return False + settlement = pending.settlement + commit_physical = not pending.physical_committed + if commit_physical: + pending.physical_in_progress = True + + if commit_physical: try: - settlement.on_done(settlement.success) - except Exception as e: # never let the callback strand the arena + if settlement.disposition is Disposition.QUARANTINE: + self._recv_alloc.quarantine(settlement.slot_id) + else: + self._recv_alloc.release(settlement.slot_id) + except Exception as e: + with self._reserved_map_lock: + pending.physical_in_progress = False + logger.error( + f"[kv-bounce] physical settlement failed " + f"(slot={settlement.slot_id}, disposition={settlement.disposition.name}); " + f"retaining it for retry: {e}" + ) + return False + with self._reserved_map_lock: + pending.physical_in_progress = False + pending.physical_committed = True + + callback = settlement.on_done + if callback is not None: + with self._reserved_map_lock: + current = self._pending_settlements.get(rid_slice) + if current is not pending or pending.callback_in_progress: + return current is None + pending.callback_in_progress = True + try: + callback(settlement.success) + except Exception as e: + with self._reserved_map_lock: + pending.callback_in_progress = False logger.error( - f"[kv-bounce] completion callback failed (slot={settlement.slot_id}): {e}" + f"[kv-bounce] completion acknowledgement failed " + f"(slot={settlement.slot_id}); retaining it for retry: {e}" ) + return False + else: + # Acknowledge and remove atomically so no concurrent retry can + # invoke an already successful callback a second time. + with self._reserved_map_lock: + pending.callback_in_progress = False + if self._pending_settlements.get(rid_slice) is pending: + self._pending_settlements.pop(rid_slice, None) + return True + + with self._reserved_map_lock: + if self._pending_settlements.get(rid_slice) is pending: + self._pending_settlements.pop(rid_slice, None) + return True + + @staticmethod + def _settlement_in_scope(rid_slice: RidSlice, scope) -> bool: + """Whether an exact key belongs to a global, request, or exact-key retry.""" + if scope is None: + return True + if isinstance(scope, tuple): + return rid_slice == scope + return rid_slice[0] == scope + + def _retry_pending_settlements(self, scope=None) -> bool: + with self._reserved_map_lock: + keys = tuple( + key for key in self._pending_settlements if self._settlement_in_scope(key, scope) + ) + for key in keys: + self._commit_pending_settlement(key) + with self._reserved_map_lock: + return not any( + self._settlement_in_scope(key, scope) for key in self._pending_settlements + ) + + def retry_settlements(self, scope=None) -> bool: + """Retry pending acknowledgements globally, for one request RID, or for one exact key.""" + with self._reserved_map_lock: + keys = tuple(key for key in self._reserved_map if self._settlement_in_scope(key, scope)) + notifications_delivered = True + for key in keys: + notifications_delivered = ( + self._retry_logical_failure_notification(key) and notifications_delivered + ) + return notifications_delivered and self._retry_pending_settlements(scope) def record_result( self, @@ -360,66 +828,281 @@ def record_result( sizes=None, src_base=None, on_done: Optional[Callable[[bool], None]] = None, + on_logical_failure: Optional[Callable[[], None]] = None, ) -> None: """A writer reported success. The completion callback fires only after the scatter lands, so the reader never sees completion before the cache is in place.""" - def mut(ctx: TransferContext) -> None: - if on_done is not None: - ctx.on_done = on_done - ctx.record_writer_result( + def mut(ctx: RecvBounceContext) -> None: + accepted = ctx.record_writer_result( peer_rank, succeeded=True, src_base=src_base, dst_ptrs=dst_ptrs, sizes=sizes ) + if accepted: + ctx.set_on_done(on_done) + ctx.set_on_logical_failure(on_logical_failure) self._apply(rid_slice, mut) - def record_failure(self, rid_slice: RidSlice, peer_rank: int) -> None: + def record_failure( + self, + rid_slice: RidSlice, + peer_rank: int, + on_done: Optional[Callable[[bool], None]] = None, + ) -> None: """A writer reported failure (it has drained). The region is freed only once every writer has reported, not here.""" - self._apply(rid_slice, lambda ctx: ctx.record_writer_result(peer_rank, succeeded=False)) - def _scatter_loop(self): - CUASSERT(cudart.cudaSetDevice(self._device_id)) - while not self._stop.is_set(): + def mut(ctx: RecvBounceContext) -> None: + if ctx.record_writer_result(peer_rank, succeeded=False): + ctx.set_on_done(on_done) + + self._apply(rid_slice, mut) + + def mark_writer_exposed(self, rid_slice: RidSlice, peer_rank: int) -> bool: + """Atomically retain the slot for ``peer_rank`` before its address is published.""" + with self._reserved_map_lock: + ctx = self._reserved_map.get(rid_slice) + return False if ctx is None else ctx.mark_writer_exposed(peer_rank) + + def record_no_access( + self, + rid_slice: RidSlice, + peer_rank: int, + *, + succeeded: bool = True, + on_done: Optional[Callable[[bool], None]] = None, + ) -> None: + """Settle a planned writer using proof that it cannot access the bounce slot.""" + + def mut(ctx: RecvBounceContext) -> None: + if ctx.mark_writer_no_access(peer_rank, succeeded=succeeded): + ctx.set_on_done(on_done) + + self._apply(rid_slice, mut) + + def mark_logical_failure( + self, + rid_slice: RidSlice, + on_done: Optional[Callable[[bool], None]] = None, + ) -> None: + """Suppress scatter/publication while retaining any possibly exposed writer.""" + + def mut(ctx: RecvBounceContext) -> None: + ctx.set_on_done(on_done) + ctx.mark_logical_failure() + + self._apply(rid_slice, mut) + + def mark_protocol_conflict( + self, + rid_slice: RidSlice, + on_done: Optional[Callable[[bool], None]] = None, + ) -> None: + """Suppress work and retain the slot until backend-wide quiescence.""" + + def mut(ctx: RecvBounceContext) -> None: + ctx.set_on_done(on_done) + ctx.mark_protocol_conflict() + + self._apply(rid_slice, mut) + + def mark_backend_quiesced( + self, + rid_slice: Optional[RidSlice] = None, + on_done: Optional[Callable[[bool], None]] = None, + ) -> None: + """Use backend-wide quiescence as terminal evidence for in-doubt remote access.""" + + def mut(ctx: RecvBounceContext) -> None: + ctx.set_on_done(on_done) + ctx.mark_backend_quiesced() + + if rid_slice is not None: + self._apply(rid_slice, mut) + return + with self._reserved_map_lock: + keys = tuple(self._reserved_map) + for key in keys: + self._apply(key, mut) + + def _suppress_queued_scatters(self) -> None: + """Fail queue entries that provably never launched after the stream was poisoned.""" + while True: try: - item = self._scatter_q.get(timeout=_SCATTER_POLL_S) + item = self._scatter_q.get_nowait() except queue.Empty: - # idle: reclaim quarantine past its grace period, independent of any reserve call - self._recv_alloc.reclaim_expired() - continue - if item is None: - break # poison pill from close: wake and exit - ctx, descs = item - ok = True + return try: - # Scatter each writer's fragments from its own source, never one global offset, so a - # missing or fallback writer cannot shift where the others are read from. - for src_base, dst_ptrs, sizes in descs: - p = Plan(dst_ptrs, dst_ptrs, sizes, int(sizes.sum())) - scatter_contiguous( - src_base, p.dst_ptrs, p.sizes, p.offsets, stream=self._scatter_stream - ) - CUASSERT(cudart.cudaStreamSynchronize(self._scatter_stream)) - except Exception as e: - # a scatter failure must not kill the worker nor be reported as success - ok = False - logger.error(f"[kv-bounce] scatter failed (slot={ctx.slot_id}): {e}") - # record the outcome and settle; completion fires only after the sync above - self._apply(ctx.rid_slice, lambda c, ok=ok: c.finish_scatter(ok)) + if item is None: + continue + ctx, _descs = item + self._apply(ctx.rid_slice, lambda current: current.suppress_scatter()) + finally: + self._scatter_q.task_done() - def close(self) -> None: - self._stop.set() - # poison pill: wake the worker now instead of waiting out its poll - self._scatter_q.put(None) - if self._scatter_thread.is_alive(): - self._scatter_thread.join(timeout=_CLOSE_JOIN_S) - for d in self._reg_descs: + def _fail_scatter_worker(self, error: Exception) -> None: + """Fail closed if an unexpected error escapes normal per-scatter handling.""" + with self._reserved_map_lock: + self._scatter_stream_healthy = False + self._accepting_reservations = False + self._scatter_worker_error = error + logger.error( + f"[kv-bounce] scatter worker stopped unexpectedly; " + f"disabling receive-bounce admission: {error}" + ) + try: + # Entries still in the queue never launched, so they can settle as + # failures. The entry that escaped remains owned because its local + # CUDA-access state is ambiguous. + self._suppress_queued_scatters() + except Exception as cleanup_error: + logger.error( + f"[kv-bounce] failed to suppress unlaunched scatters after worker failure; " + f"retaining them for teardown retry: {cleanup_error}" + ) + + def _fail_current_scatter(self, ctx: RecvBounceContext) -> None: + """Latch logical failure for the dequeued item without claiming a fence. + + This is the last-resort boundary for an exception that escapes the + normal per-scatter path, including a state-update failure after CUDA + work was launched. The context remains physically owned because the + exception site does not itself prove local stream quiescence. + """ + with self._reserved_map_lock: + current = self._reserved_map.get(ctx.rid_slice) + if current is ctx: + current.finish_scatter(False) + self._retry_logical_failure_notification(ctx.rid_slice) + + def _scatter_loop(self) -> None: + """Run the worker with a fail-closed boundary around unexpected errors.""" + try: + self._run_scatter_loop() + except Exception as e: + self._fail_scatter_worker(e) + + def _run_scatter_loop(self) -> None: + try: + CUASSERT(cudart.cudaSetDevice(self._device_id)) + except Exception as e: + self._scatter_start_error = e + self._scatter_ready.set() + return + self._scatter_ready.set() + while True: try: - self._agent.deregister_memory(d) - except Exception: - pass - self._send_alloc.close() - self._recv_alloc.close() + item = self._scatter_q.get() + if item is None: + return # FIFO poison pill: every earlier scatter has finished + ctx, descs = item + try: + ok = True + try: + # Scatter each writer's fragments from its own source, never one global + # offset, so a missing or fallback writer cannot shift where the others + # are read from. + for src_base, dst_ptrs, sizes in descs: + p = Plan(dst_ptrs, dst_ptrs, sizes, int(sizes.sum())) + scatter_contiguous( + src_base, + p.dst_ptrs, + p.sizes, + p.offsets, + stream=self._scatter_stream, + ) + # The metadata staging buffer is shared per stream, so it cannot be + # refilled for another writer until this writer's async metadata copy + # has completed. + CUASSERT(cudart.cudaStreamSynchronize(self._scatter_stream)) + except Exception as e: + # A scatter failure must not kill the worker nor be reported as success. + ok = False + logger.error(f"[kv-bounce] scatter failed (slot={ctx.slot_id}): {e}") + + # Success settles after the positive fence above. Failure is retained and + # poisons the shared stream: a CUDA error does not prove queued accesses + # stopped, so no later job may use it. + def finish(c, ok=ok): + if not ok: + self._scatter_stream_healthy = False + c.finish_scatter(ok) + + self._apply(ctx.rid_slice, finish) + if not ok: + self._suppress_queued_scatters() + return + except Exception: + self._fail_current_scatter(ctx) + raise + finally: + self._scatter_q.task_done() + + def close(self) -> None: + """Drain the worker and destroy the arenas, but never underneath a live lease.""" + with self._close_lock: + self._close_locked() + + def _close_locked(self) -> None: + """Retryable close body, serialized so resources are retired at most once.""" + # Close reservation admission before inspecting drain state. If this + # attempt finds live work and returns retryable failure, no new receive + # or send lease may race the later retry. + with self._reserved_map_lock: + if self._closed: + return + self._accepting_reservations = False + if not self._retry_pending_settlements(): + raise RuntimeError("[kv-bounce] cannot close with unacknowledged settlements") + with self._reserved_map_lock: + non_scatter_contexts = sum( + ctx.scatter_state is not ScatterState.QUEUED for ctx in self._reserved_map.values() + ) + if non_scatter_contexts: + raise RuntimeError( + f"[kv-bounce] cannot close with {non_scatter_contexts} live receive contexts" + ) + if self._pending_reservations or self._pending_send_reservations: + raise RuntimeError("[kv-bounce] cannot close while a reservation is pending") + if self._send_alloc.has_outstanding: + raise RuntimeError("[kv-bounce] cannot close with outstanding arena slots") + if not self._reserved_map and self._recv_alloc.has_outstanding: + raise RuntimeError("[kv-bounce] cannot close with outstanding arena slots") + + # FIFO poison pill makes the worker complete every item already accepted before exiting. + self._stop_scatter_worker() + with self._reserved_map_lock: + if self._reserved_map or self._recv_alloc.has_outstanding: + raise RuntimeError("[kv-bounce] scatter drain left outstanding receive contexts") + if not self._retry_pending_settlements(): + raise RuntimeError("[kv-bounce] cannot close with unacknowledged settlements") + + stream_errors = self._destroy_streams() + if stream_errors: + raise RuntimeError( + f"[kv-bounce] failed to destroy {len(stream_errors)} CUDA stream(s)" + ) from stream_errors[0] + + event_errors = self._destroy_pending_events() + if event_errors: + raise RuntimeError( + f"[kv-bounce] failed to destroy {len(event_errors)} CUDA event(s)" + ) from event_errors[0] + + deregistration_errors = self._deregister_registered_descriptors() + if deregistration_errors: + raise RuntimeError( + f"[kv-bounce] failed to deregister " + f"{len(deregistration_errors)} memory descriptor(s)" + ) from deregistration_errors[0] + + allocator_errors = self._close_allocators() + if allocator_errors: + raise RuntimeError( + f"[kv-bounce] failed to close {len(allocator_errors)} arena allocator(s)" + ) from allocator_errors[0] + with self._reserved_map_lock: + self._closed = True class NoBounceTransport(BounceTransport): @@ -435,12 +1118,21 @@ def build_request(self, write_meta): def release_send(self, slot_id) -> None: pass + def quarantine_send(self, slot_id) -> None: + pass + def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + writer_ranks: Sequence[int] = (), + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + destination_intervals: Optional[Iterable[tuple[int, int]]] = None, + destination_intervals_factory: Optional[Callable[[], Iterable[tuple[int, int]]]] = None, ) -> bool: return False - def writer_base(self, rid_slice, writer_index: int): + def writer_base(self, rid_slice, peer_rank: int): return None def is_bounced(self, rid_slice) -> bool: @@ -452,18 +1144,59 @@ def release_idle_reservation(self, rid_slice) -> None: def orphan_reservation(self, rid_slice) -> None: pass + def mark_writer_exposed(self, rid_slice, peer_rank: int) -> bool: + return False + + def record_no_access( + self, rid_slice, peer_rank: int, *, succeeded: bool = True, on_done=None + ) -> None: + pass + + def mark_logical_failure(self, rid_slice, on_done=None) -> None: + pass + + def mark_protocol_conflict(self, rid_slice, on_done=None) -> None: + pass + + def mark_backend_quiesced(self, rid_slice=None, on_done=None) -> None: + pass + + def retry_settlements(self, scope=None) -> bool: + return True + def record_result( - self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None + self, + rid_slice, + peer_rank, + dst_ptrs=None, + sizes=None, + src_base=None, + on_done=None, + on_logical_failure=None, ): pass - def record_failure(self, rid_slice, peer_rank) -> None: + def record_failure(self, rid_slice, peer_rank, on_done=None) -> None: pass def close(self) -> None: pass +class CleanupPendingBounceTransport(NoBounceTransport): + """Direct-path transport that retains a partially constructed VMM owner.""" + + def __init__(self, owner: VmmBounceTransport): + self._owner = owner + + def close(self) -> None: + owner = self._owner + if owner is None: + return + owner.retry_initialization_rollback() + self._owner = None + + def create_bounce(agent, cfg, *, device_id: int, page_table) -> BounceTransport: """Build the real transport from the config, or the disabled one when bounce is off, it cannot fit, or the fabric allocation races.""" @@ -474,9 +1207,13 @@ def create_bounce(agent, cfg, *, device_id: int, page_table) -> BounceTransport: agent, cfg, device_id=device_id, block_bytes_per_group=block_bytes_per_group(page_table) ) return transport if transport is not None else NoBounceTransport() - except ( - Exception - ) as e: # rare race: memory taken between the free-memory query and the allocation + except IncompleteBounceInitializationError as e: + logger.error( + f"[kv-bounce] disabled with incomplete initialization cleanup: {e}; " + "retaining cleanup owner" + ) + return CleanupPendingBounceTransport(e.owner) + except Exception as e: # rare race: memory taken between the query and allocation logger.warning(f"[kv-bounce] disabled (alloc failed: {e}); using in-place path") return NoBounceTransport() @@ -526,8 +1263,14 @@ def decode_result_tail(message): def block_bytes_per_group(page_table) -> list: - """Byte size of one cache block for each leading attention layer group, stopping at the first - non-attention group.""" + """Byte size of one cache block for each leading attention layer group. + + A layer group can expose multiple pool views (for example, the ordinary KV pool plus an indexer + pool). The sender gathers fragments from every mapped view, so the receive reservation must + cover every pool view. This intentionally mirrors the sender's per-view descriptor loop: even + repeated views of one physical pool contribute repeated gathered extents. Stop at the first + non-attention group because bounce does not size Mamba state. + """ from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool @@ -536,5 +1279,10 @@ def block_bytes_per_group(page_table) -> list: for lg_idx, lg in enumerate(page_table.layer_groups): if not isinstance(lg, AttentionLayerGroup): break - out.append(int(get_physical_pool(page_table, lg_idx, 0).slot_bytes)) + out.append( + sum( + int(get_physical_pool(page_table, lg_idx, int(pool_view.pool_idx)).slot_bytes) + for pool_view in lg.pool_views + ) + ) return out diff --git a/tensorrt_llm/_torch/disaggregation/native/messenger.py b/tensorrt_llm/_torch/disaggregation/native/messenger.py index ceb6aa626ed9..83fa491a7c27 100644 --- a/tensorrt_llm/_torch/disaggregation/native/messenger.py +++ b/tensorrt_llm/_torch/disaggregation/native/messenger.py @@ -1,5 +1,21 @@ +# 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 time from abc import ABC, abstractmethod -from threading import Event, Lock, Thread +from threading import Event, Lock, Thread, current_thread from typing import Callable, Optional import zmq @@ -7,6 +23,9 @@ from tensorrt_llm import logger from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip +_CONTROL_STOP = b"STOP" +_CONTROL_WAKE = b"WAKE" + class MessengerInterface(ABC): """ @@ -87,8 +106,20 @@ def __init__(self, mode: str, endpoint: Optional[str] = None) -> None: self._mode = mode self._socket = self._context.socket(self.SOCKET_MODES[mode]) self._endpoint: Optional[str] = None + # Serialize state transitions separately from socket I/O. stop() must + # release _lock before joining the listener, because a listener + # callback may be entering send(). _socket_io_lock then fences every + # main-socket send/receive from close(). self._lock = Lock() + self._stop_lock = Lock() + self._socket_io_lock = Lock() + self._control_send_lock = Lock() + self._io_waiters_lock = Lock() + self._io_waiters = 0 + self._io_waiters_drained = Event() + self._io_waiters_drained.set() self._closed = False + self._closing = False self._stop_event = Event() self._listener_thread: Optional[Thread] = None self._initialize_control_sockets() @@ -117,14 +148,63 @@ def _initialize_control_sockets(self) -> None: def start(self) -> None: pass + def _signal_listener(self, command: bytes) -> None: + """Wake a blocking listener poll without ever blocking the caller.""" + try: + with self._control_send_lock: + listener = self._listener_thread + with self._lock: + stopping = self._closing or self._closed + if ( + listener is None + or listener is current_thread() + or not listener.is_alive() + or (command == _CONTROL_WAKE and stopping) + ): + return + self._internal_socket.send(command, flags=zmq.DONTWAIT) + except zmq.Again: + # A queued control frame already wakes the poller. The listener + # also observes _stop_event directly, so another frame is not + # required for correctness. + pass + + def _acquire_socket_io(self) -> None: + """Wake the listener, yield its blocking poll, and acquire main-socket ownership.""" + with self._io_waiters_lock: + self._io_waiters += 1 + self._io_waiters_drained.clear() + try: + self._signal_listener(_CONTROL_WAKE) + self._socket_io_lock.acquire() + finally: + with self._io_waiters_lock: + self._io_waiters -= 1 + if self._io_waiters == 0: + self._io_waiters_drained.set() + def send(self, messages: list[bytes], recipient: Optional[bytes] = None) -> None: - if recipient: - self._socket.send_multipart([recipient] + messages) - else: - self._socket.send_multipart(messages) + self._acquire_socket_io() + try: + with self._lock: + if self._closing or self._closed: + raise RuntimeError("ZMQMessenger is stopping or closed") + if recipient: + self._socket.send_multipart([recipient] + messages) + else: + self._socket.send_multipart(messages) + finally: + self._socket_io_lock.release() def receive(self) -> list[bytes]: - return self._socket.recv_multipart() + self._acquire_socket_io() + try: + with self._lock: + if self._closing or self._closed: + raise RuntimeError("ZMQMessenger is stopping or closed") + return self._socket.recv_multipart() + finally: + self._socket_io_lock.release() def start_listener( self, @@ -134,8 +214,6 @@ def start_listener( assert self._mode in ["ROUTER", "REP"], ( "Listener can only be started in ROUTER or REP modes" ) - if self._listener_thread and self._listener_thread.is_alive(): - raise RuntimeError("Listener already running") def handle_listener_exceptions( exception: Exception, on_error: Optional[Callable[[Exception], None]] @@ -148,16 +226,37 @@ def handle_listener_exceptions( def listener() -> None: poller = zmq.Poller() - poller.register(self._socket, zmq.POLLIN) - poller.register(self._control_socket, zmq.POLLIN) + with self._socket_io_lock: + poller.register(self._socket, zmq.POLLIN) + poller.register(self._control_socket, zmq.POLLIN) while not self._stop_event.is_set(): - events = dict(poller.poll(timeout=100)) try: - if self._control_socket in events: + messages = None + control_message = None + with self._socket_io_lock: + events = dict(poller.poll(timeout=100)) + if self._control_socket in events: + control_message = self._control_socket.recv() + elif self._socket in events: + with self._lock: + stopping = self._closing or self._closed + if not stopping: + messages = self._socket.recv_multipart() + if control_message == _CONTROL_STOP: self._stop_event.set() - elif self._socket in events: - messages = self.receive() + elif control_message == _CONTROL_WAKE: + # Do not immediately reacquire the socket ahead of the + # send/receive thread that woke this poll. + while not self._stop_event.is_set() and not self._io_waiters_drained.wait( + timeout=0.1 + ): + pass + elif control_message is not None: + logger.warning( + f"Ignoring unknown messenger control frame {control_message!r}" + ) + elif messages is not None: persist = on_message(messages) if persist is False: self._stop_event.set() @@ -170,42 +269,95 @@ def listener() -> None: self._stop_event.set() - logger.info(f"Starting Messenger listener thread for {self._endpoint}") - self._listener_thread = Thread(target=listener, daemon=True) - self._listener_thread.start() + # Serialize listener creation with stop() so a thread cannot start + # against sockets that teardown has already decided to close. + with self._stop_lock: + with self._lock: + if self._closing or self._closed: + raise RuntimeError("ZMQMessenger is stopping or closed") + if self._listener_thread and self._listener_thread.is_alive(): + raise RuntimeError("Listener already running") + logger.info(f"Starting Messenger listener thread for {self._endpoint}") + self._listener_thread = Thread(target=listener, daemon=True) + self._listener_thread.start() + + def stop(self, timeout: float = 5) -> None: + def _close_socket(name: str, socket: zmq.Socket) -> Optional[str]: + if socket.closed: + return None - def stop(self, timeout: int = 5) -> None: - def _close_socket(socket: zmq.Socket) -> None: try: - if not socket.closed: - socket.setsockopt(zmq.LINGER, 0) - socket.close() + socket.setsockopt(zmq.LINGER, 0) except Exception as e: - logger.error(f"Error closing socket: {e}") + # Do not close with an unknown linger policy. Leaving the + # socket open makes the failed teardown explicitly retryable. + return f"{name} socket linger configuration failed: {e}" - with self._lock: - if self._closed: - return - self._closed = True + try: + socket.close() + except Exception as e: + return f"{name} socket close failed: {e}" + if not socket.closed: + return f"{name} socket remained open after close" + return None + + # A separate stop lock keeps retryable stop attempts single-threaded + # without holding the state lock across listener.join(). + with self._stop_lock: + timeout = max(float(timeout), 0.0) + deadline = time.monotonic() + timeout + with self._lock: + if self._closed: + return + self._closing = True logger.debug("Stopping ZMQMessenger...") self._stop_event.set() - self._internal_socket.send(b"STOP") - if self._listener_thread: - self._internal_socket.send(b"STOP") - self._listener_thread.join(timeout) - if self._listener_thread.is_alive(): - logger.warning("Listener thread did not terminate within timeout") - - _close_socket(self._socket) - _close_socket(self._internal_socket) - _close_socket(self._control_socket) + listener = self._listener_thread + if listener and listener.is_alive() and listener is not current_thread(): + self._signal_listener(_CONTROL_STOP) + listener.join(max(0.0, deadline - time.monotonic())) + if listener.is_alive(): + raise RuntimeError( + "ZMQMessenger listener thread did not terminate within timeout" + ) + + # A send/receive admitted before _closing was set may still be in + # progress. Taking the same I/O lock waits for it to finish and + # prevents any later operation from racing socket close. + close_errors = [] + if not self._socket_io_lock.acquire(timeout=max(0.0, deadline - time.monotonic())): + raise RuntimeError("ZMQMessenger main socket I/O did not quiesce within timeout") + try: + error = _close_socket("main", self._socket) + if error is not None: + close_errors.append(error) + finally: + self._socket_io_lock.release() + # Fence a waiter that observed the listener before join but had + # not yet sent its wake frame. It rechecks liveness under this + # lock, so no control-socket operation can race close. + with self._control_send_lock: + for name, socket in ( + ("internal", self._internal_socket), + ("control", self._control_socket), + ): + error = _close_socket(name, socket) + if error is not None: + close_errors.append(error) + + if close_errors: + raise RuntimeError("Failed to stop ZMQMessenger: " + "; ".join(close_errors)) try: if not self._context.closed: self._context.term() except Exception as e: - logger.error(f"Error terminating ZMQ context: {e}") + raise RuntimeError( + f"Failed to stop ZMQMessenger: context termination failed: {e}" + ) from e + with self._lock: + self._closed = True @property def endpoint(self) -> str: diff --git a/tensorrt_llm/_torch/disaggregation/native/receive_lifecycle.py b/tensorrt_llm/_torch/disaggregation/native/receive_lifecycle.py new file mode 100644 index 000000000000..6581140fbc26 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/native/receive_lifecycle.py @@ -0,0 +1,915 @@ +# 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. +"""Receiver-side ownership state for native disaggregated KV transfers. + +This module deliberately depends only on the Python standard library. It owns +the concurrency-sensitive lifecycle decisions, while ``transfer.py`` and the +bounce implementation execute the returned actions. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from enum import Enum +from typing import Iterable, TypeAlias + +TransferKey: TypeAlias = tuple[int, int] + + +class WriterMode(Enum): + """The path a writer used to access receiver memory.""" + + PENDING = "pending" + DIRECT = "direct" + BOUNCE = "bounce" + NO_REMOTE_ACCESS = "no_remote_access" + UNKNOWN = "unknown" + + +class ExposureState(Enum): + """Whether a writer could have observed a receiver target address.""" + + UNEXPOSED = "unexposed" + POSSIBLY_EXPOSED = "possibly_exposed" + PUBLISHED = "published" + NEVER_EXPOSED = "never_exposed" + + +class WriterResult(Enum): + """A writer's terminal data result.""" + + SUCCESS = "success" + FAILED = "failed" + + +class LogicalState(Enum): + """Outcome visible to the request/session consumer.""" + + PENDING = "pending" + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + + +class PhysicalState(Enum): + """Whether memory may still be touched by a transfer accessor.""" + + ACTIVE = "active" + DRAINING = "draining" + IN_DOUBT = "in_doubt" + DRAINED = "drained" + + +class BounceState(Enum): + """Receiver bounce-slot scatter lifecycle.""" + + NONE = "none" + IDLE = "idle" + SCATTERING = "scattering" + DONE = "done" + FAILED = "failed" + SUPPRESSED = "suppressed" + + +class LifecycleAction(Enum): + """A side effect that the transfer/bounce adapters must execute.""" + + NOTIFY_SUCCESS = "notify_success" + NOTIFY_FAILURE = "notify_failure" + NOTIFY_CANCELLED = "notify_cancelled" + START_BOUNCE_SCATTER = "start_bounce_scatter" + RELEASE_BOUNCE = "release_bounce" + CONTEXT_DRAINED = "context_drained" + + +@dataclass(frozen=True) +class LifecycleUpdate: + """Result of one atomic lifecycle transition.""" + + key: TransferKey + accepted: bool + logical_state: LogicalState + physical_state: PhysicalState + writers_terminal: bool + bounce_pending: bool + actions: tuple[LifecycleAction, ...] = () + publication_allowed: bool = False + duplicate: bool = False + conflict: bool = False + reason: str | None = None + + +@dataclass(frozen=True) +class WriterSnapshot: + """Immutable diagnostic view of one expected writer.""" + + rank: int + exposure: ExposureState + mode: WriterMode + result: WriterResult | None + conflict: bool + backend_quiesced: bool + + +@dataclass(frozen=True) +class ContextSnapshot: + """Immutable diagnostic view of a receive transfer context.""" + + key: TransferKey + logical_state: LogicalState + physical_state: PhysicalState + bounce_state: BounceState + writers_terminal: bool + bounce_pending: bool + bounce_settled: bool + consumer_attached: bool + protocol_conflict: bool + expected_writers: frozenset[int] + writers: tuple[WriterSnapshot, ...] + + +@dataclass(frozen=True) +class RegistrySnapshot: + """Immutable diagnostic view of the complete registry.""" + + accepting: bool + backend_quiesced: bool + contexts: tuple[ContextSnapshot, ...] + + +@dataclass +class WriterRecord: + """Mutable state for one member of the exact expected-writer set.""" + + rank: int + exposure: ExposureState = ExposureState.UNEXPOSED + mode: WriterMode = WriterMode.PENDING + result: WriterResult | None = None + conflict: bool = False + backend_quiesced: bool = False + + @property + def is_quiescent(self) -> bool: + """Return whether this writer can no longer access receiver memory.""" + + return self.backend_quiesced or (self.result is not None and not self.conflict) + + def snapshot(self) -> WriterSnapshot: + """Return an immutable diagnostic view.""" + + return WriterSnapshot( + rank=self.rank, + exposure=self.exposure, + mode=self.mode, + result=self.result, + conflict=self.conflict, + backend_quiesced=self.backend_quiesced, + ) + + +@dataclass +class RecvTransferContext: + """Canonical receiver-side state for one request slice.""" + + key: TransferKey + expected_writers: frozenset[int] + has_bounce_slot: bool + consumer_attached: bool = True + writers: dict[int, WriterRecord] = field(init=False) + logical_state: LogicalState = field(default=LogicalState.PENDING, init=False) + physical_state: PhysicalState = field(default=PhysicalState.ACTIVE, init=False) + bounce_state: BounceState = field(init=False) + bounce_settled: bool = field(default=False, init=False) + protocol_conflict: bool = field(default=False, init=False) + conflict_requires_backend_quiescence: bool = field(default=False, init=False) + in_doubt: bool = field(default=False, init=False) + logical_reason: str | None = field(default=None, init=False) + _notified_logical_state: LogicalState | None = field(default=None, init=False) + _bounce_release_emitted: bool = field(default=False, init=False) + _drained_emitted: bool = field(default=False, init=False) + + def __post_init__(self) -> None: + self.writers = {rank: WriterRecord(rank) for rank in self.expected_writers} + self.bounce_state = BounceState.IDLE if self.has_bounce_slot else BounceState.NONE + + def snapshot(self) -> ContextSnapshot: + """Return an immutable diagnostic view.""" + + return ContextSnapshot( + key=self.key, + logical_state=self.logical_state, + physical_state=self.physical_state, + bounce_state=self.bounce_state, + writers_terminal=self.writers_terminal, + bounce_pending=self.bounce_pending, + bounce_settled=self.bounce_settled, + consumer_attached=self.consumer_attached, + protocol_conflict=self.protocol_conflict, + expected_writers=self.expected_writers, + writers=tuple(self.writers[rank].snapshot() for rank in sorted(self.writers)), + ) + + @property + def writers_terminal(self) -> bool: + """Return whether every expected writer is physically quiescent.""" + + return all(writer.is_quiescent for writer in self.writers.values()) + + @property + def bounce_pending(self) -> bool: + """Return whether a requested bounce scatter has not settled.""" + + return self.bounce_state is BounceState.SCATTERING + + +class RecvTransferRegistry: + """Thread-safe root owner for receiver transfer lifecycle state. + + Address publication and cancellation arbitrate under the same lock. + ``begin_publication`` changes the writer to ``POSSIBLY_EXPOSED`` before + returning permission to send, so cancellation can never mistake an + authorized publication for an idle reservation. + """ + + def __init__(self) -> None: + self._contexts: dict[TransferKey, RecvTransferContext] = {} + self._request_keys: dict[int, set[TransferKey]] = {} + self._accepting = True + self._backend_quiesced = False + self._lock = threading.RLock() + + def prepare( + self, + key: TransferKey, + expected_writers: Iterable[int], + *, + has_bounce_slot: bool, + consumer_attached: bool = True, + ) -> LifecycleUpdate: + """Create the canonical context before any target can be published.""" + + writers = frozenset(expected_writers) + if not writers: + raise ValueError("expected_writers must not be empty") + if any(not isinstance(rank, int) or isinstance(rank, bool) or rank < 0 for rank in writers): + raise ValueError("expected_writers must contain non-negative integer ranks") + + with self._lock: + if not self._accepting: + return self._rejected_update(key, "registry admission is closed") + existing = self._contexts.get(key) + if existing is not None: + same = ( + existing.expected_writers == writers + and existing.has_bounce_slot == has_bounce_slot + ) + return self._update( + existing, + accepted=False, + duplicate=same, + conflict=not same, + reason=("context already prepared" if same else "conflicting context identity"), + ) + context = RecvTransferContext( + key=key, + expected_writers=writers, + has_bounce_slot=has_bounce_slot, + consumer_attached=consumer_attached, + ) + self._contexts[key] = context + self._request_keys.setdefault(key[0], set()).add(key) + return self._update(context) + + def begin_publication(self, key: TransferKey, writer_rank: int) -> LifecycleUpdate: + """Atomically authorize at most one target publication for a writer.""" + + with self._lock: + context = self._contexts.get(key) + if context is None: + return self._rejected_update(key, "unknown transfer context") + writer = context.writers.get(writer_rank) + if writer is None: + return self._unexpected_local_writer(context, writer_rank) + if context.logical_state is not LogicalState.PENDING: + return self._update(context, accepted=False, reason="publication gate is closed") + if writer.exposure is not ExposureState.UNEXPOSED: + return self._update( + context, + accepted=False, + duplicate=True, + reason="writer publication already decided", + ) + writer.exposure = ExposureState.POSSIBLY_EXPOSED + return self._update(context, publication_allowed=True) + + def mark_published(self, key: TransferKey, writer_rank: int) -> LifecycleUpdate: + """Record that an authorized publication completed successfully.""" + + with self._lock: + context, writer, error = self._lookup_writer(key, writer_rank) + if error is not None: + return error + assert context is not None and writer is not None + if writer.exposure is ExposureState.PUBLISHED: + return self._update(context, accepted=False, duplicate=True) + if writer.exposure is not ExposureState.POSSIBLY_EXPOSED: + return self._publication_conflict(context, writer, "publication was not authorized") + writer.exposure = ExposureState.PUBLISHED + return self._update(context) + + def mark_publication_ambiguous(self, key: TransferKey, writer_rank: int) -> LifecycleUpdate: + """Keep an authorized writer possibly exposed after an ambiguous send.""" + + with self._lock: + context, writer, error = self._lookup_writer(key, writer_rank) + if error is not None: + return error + assert context is not None and writer is not None + if writer.exposure is ExposureState.POSSIBLY_EXPOSED: + return self._update(context, accepted=False, duplicate=True) + if writer.exposure is ExposureState.PUBLISHED: + return self._update(context, accepted=False, duplicate=True) + return self._publication_conflict(context, writer, "publication was not authorized") + + def mark_never_published(self, key: TransferKey, writer_rank: int) -> LifecycleUpdate: + """Record definitive proof that a writer never received a target.""" + + with self._lock: + context, writer, error = self._lookup_writer(key, writer_rank) + if error is not None: + return error + assert context is not None and writer is not None + if writer.exposure is ExposureState.NEVER_EXPOSED: + return self._update(context, accepted=False, duplicate=True) + if writer.exposure is ExposureState.PUBLISHED: + return self._publication_conflict( + context, writer, "published target cannot be undone" + ) + if writer.result is not None: + return self._publication_conflict( + context, writer, "writer already reported a result" + ) + writer.exposure = ExposureState.NEVER_EXPOSED + writer.mode = WriterMode.NO_REMOTE_ACCESS + writer.result = WriterResult.FAILED + return self._evaluate(context) + + def record_result( + self, + key: TransferKey, + writer_rank: int, + result: WriterResult, + mode: WriterMode, + ) -> LifecycleUpdate: + """Record one exact writer's terminal result and return newly safe actions.""" + + if mode is WriterMode.PENDING: + raise ValueError("a terminal result cannot use WriterMode.PENDING") + if result is WriterResult.SUCCESS and mode not in (WriterMode.DIRECT, WriterMode.BOUNCE): + raise ValueError("a successful result must identify DIRECT or BOUNCE mode") + if mode is WriterMode.NO_REMOTE_ACCESS and result is not WriterResult.FAILED: + raise ValueError("NO_REMOTE_ACCESS cannot report successful data delivery") + + with self._lock: + context = self._contexts.get(key) + if context is None: + return self._rejected_update(key, "unknown transfer context") + writer = context.writers.get(writer_rank) + if writer is None: + return self._unexpected_result_writer(context, writer_rank) + if writer.result is not None: + return self._dedupe_or_conflict_result(context, writer, result, mode) + exposed = writer.exposure in ( + ExposureState.POSSIBLY_EXPOSED, + ExposureState.PUBLISHED, + ) + allowed_modes = ( + (WriterMode.DIRECT, WriterMode.BOUNCE) + if context.has_bounce_slot + else (WriterMode.DIRECT,) + ) + if exposed and mode not in allowed_modes: + writer.mode = mode + writer.result = result + return self._writer_result_conflict( + context, + writer, + f"writer {writer_rank} reported {mode.value}, expected one of " + f"{tuple(item.value for item in allowed_modes)}", + ) + if not exposed and mode is not WriterMode.NO_REMOTE_ACCESS: + writer.mode = mode + writer.result = result + return self._writer_result_conflict( + context, + writer, + f"writer {writer_rank} reported access without publication", + ) + writer.mode = mode + writer.result = result + return self._evaluate(context) + + def record_protocol_conflict( + self, key: TransferKey, writer_rank: int, reason: str + ) -> LifecycleUpdate: + """Fail closed on a writer frame that is not terminal evidence. + + Receiving such a frame proves that the claimed identity may belong to + an exposed or stale operation. Without generation-safe replay, neither + an expected-but-unpublished writer nor an unexpected writer permits the + no-publication fast path. Backend-wide quiescence is required before + the target can be retired. + """ + + if not reason: + raise ValueError("protocol conflict reason must not be empty") + with self._lock: + context = self._contexts.get(key) + if context is None: + return self._rejected_update(key, "unknown transfer context") + writer = context.writers.get(writer_rank) + if writer is not None: + writer.conflict = True + if writer.exposure is ExposureState.UNEXPOSED: + writer.exposure = ExposureState.POSSIBLY_EXPOSED + if writer.mode is WriterMode.PENDING: + writer.mode = WriterMode.UNKNOWN + context.protocol_conflict = True + context.conflict_requires_backend_quiescence = True + context.in_doubt = True + return self._fail_and_evaluate(context, reason, conflict=True) + + def finish_bounce_scatter(self, key: TransferKey, succeeded: bool) -> LifecycleUpdate: + """Acknowledge that bounce scatter/settlement and slot release finished.""" + + with self._lock: + context = self._contexts.get(key) + if context is None: + return self._rejected_update(key, "unknown transfer context") + if context.bounce_settled: + same = (context.bounce_state is not BounceState.FAILED) == succeeded + return self._update( + context, + accepted=False, + duplicate=same, + conflict=not same, + reason=(None if same else "conflicting scatter result"), + ) + if not context.has_bounce_slot: + return self._update( + context, + accepted=False, + conflict=True, + reason="context has no bounce resource", + ) + if not context.writers_terminal: + return self._update( + context, + accepted=False, + conflict=True, + reason="bounce resource settled before its accessors drained", + ) + context.bounce_settled = True + context._bounce_release_emitted = True + context.bounce_state = BounceState.DONE if succeeded else BounceState.FAILED + if not succeeded: + return self._fail_and_evaluate(context, "bounce scatter failed") + return self._evaluate(context) + + def fail_context(self, key: TransferKey, reason: str) -> LifecycleUpdate: + """Latch logical failure after local preparation or publication fails. + + Unexposed writers are closed immediately. Any writer that crossed the + publication boundary remains owned until it reports terminal or the + backend supplies global quiescence evidence. + """ + + if not reason: + raise ValueError("failure reason must not be empty") + with self._lock: + context = self._contexts.get(key) + if context is None: + return self._rejected_update(key, "unknown transfer context") + if context.logical_state is not LogicalState.PENDING: + return self._update( + context, + accepted=False, + duplicate=True, + reason="logical outcome is already terminal", + ) + context.logical_state = LogicalState.FAILED + context.logical_reason = reason + context.in_doubt = any( + writer.exposure in (ExposureState.POSSIBLY_EXPOSED, ExposureState.PUBLISHED) + and not writer.is_quiescent + for writer in context.writers.values() + ) + return self._evaluate(context, reason=reason) + + def cancel_request(self, unique_rid: int) -> tuple[LifecycleUpdate, ...]: + """Logically cancel every slice while preserving exposed resources.""" + + with self._lock: + return tuple( + self._cancel_context(context, "request cancelled") + for context in self._contexts_for_request(unique_rid) + ) + + def timeout_request(self, unique_rid: int) -> tuple[LifecycleUpdate, ...]: + """Logically fail every slice while preserving unknown physical access.""" + + with self._lock: + updates = [] + for context in self._contexts_for_request(unique_rid): + if context.logical_state is LogicalState.PENDING: + context.logical_state = LogicalState.FAILED + context.logical_reason = "request timed out" + context.in_doubt = True + updates.append(self._evaluate(context)) + return tuple(updates) + + def detach_consumer(self, unique_rid: int) -> tuple[LifecycleUpdate, ...]: + """Detach request/session callbacks without changing physical ownership.""" + + with self._lock: + updates = [] + for context in self._contexts_for_request(unique_rid): + context.consumer_attached = False + updates.append(self._update(context)) + return tuple(updates) + + def begin_shutdown(self) -> tuple[LifecycleUpdate, ...]: + """Close admission and logically cancel every active context.""" + + with self._lock: + self._accepting = False + return tuple( + self._cancel_context(context, "receiver shutdown") + for context in sorted(self._contexts.values(), key=lambda item: item.key) + if context.physical_state is not PhysicalState.DRAINED + ) + + def mark_backend_quiesced(self) -> tuple[LifecycleUpdate, ...]: + """Apply backend-wide proof that no remote accessor can remain active.""" + + with self._lock: + self._backend_quiesced = True + updates = [] + for context in sorted(self._contexts.values(), key=lambda item: item.key): + if context.physical_state is PhysicalState.DRAINED: + continue + missing_data = False + for writer in context.writers.values(): + writer.backend_quiesced = True + if writer.result is None: + missing_data = True + context.conflict_requires_backend_quiescence = False + if missing_data and context.logical_state is LogicalState.PENDING: + context.logical_state = LogicalState.FAILED + context.logical_reason = "backend quiesced before all writer results" + updates.append(self._evaluate(context)) + return tuple(updates) + + def context_snapshot(self, key: TransferKey) -> ContextSnapshot | None: + """Return an immutable view of a context, if it exists.""" + + with self._lock: + context = self._contexts.get(key) + return None if context is None else context.snapshot() + + def target_mode(self, key: TransferKey) -> WriterMode | None: + """Return the target mode offered to this context without allocating a snapshot.""" + + with self._lock: + context = self._contexts.get(key) + if context is None: + return None + return WriterMode.BOUNCE if context.has_bounce_slot else WriterMode.DIRECT + + def snapshot(self) -> RegistrySnapshot: + """Return an immutable view of registry admission and all contexts.""" + + with self._lock: + return RegistrySnapshot( + accepting=self._accepting, + backend_quiesced=self._backend_quiesced, + contexts=tuple( + context.snapshot() + for context in sorted(self._contexts.values(), key=lambda item: item.key) + ), + ) + + def is_request_drained(self, unique_rid: int) -> bool: + """Return whether all contexts for a request are physically drained.""" + + with self._lock: + return all( + context.physical_state is PhysicalState.DRAINED + for context in self._contexts_for_request(unique_rid) + ) + + def is_drained(self) -> bool: + """Return whether every registered context is physically drained.""" + + with self._lock: + return all( + context.physical_state is PhysicalState.DRAINED + for context in self._contexts.values() + ) + + def retire_request(self, unique_rid: int) -> tuple[TransferKey, ...]: + """Remove drained request contexts while retaining any unsafe context.""" + + with self._lock: + retired = tuple( + context.key + for context in self._contexts_for_request(unique_rid) + if context.physical_state is PhysicalState.DRAINED + ) + for key in retired: + del self._contexts[key] + self._request_keys[unique_rid].discard(key) + if not self._request_keys.get(unique_rid): + self._request_keys.pop(unique_rid, None) + return retired + + def retire_request_if_drained(self, unique_rid: int) -> bool: + """Atomically retire all request contexts, or none if any is still active.""" + + with self._lock: + contexts = self._contexts_for_request(unique_rid) + if any(context.physical_state is not PhysicalState.DRAINED for context in contexts): + return False + for context in contexts: + del self._contexts[context.key] + self._request_keys.pop(unique_rid, None) + return True + + def _contexts_for_request(self, unique_rid: int) -> list[RecvTransferContext]: + keys = self._request_keys.get(unique_rid, ()) + return [self._contexts[key] for key in sorted(keys) if key in self._contexts] + + def _lookup_writer( + self, key: TransferKey, writer_rank: int + ) -> tuple[RecvTransferContext | None, WriterRecord | None, LifecycleUpdate | None]: + context = self._contexts.get(key) + if context is None: + return None, None, self._rejected_update(key, "unknown transfer context") + writer = context.writers.get(writer_rank) + if writer is None: + return context, None, self._unexpected_local_writer(context, writer_rank) + return context, writer, None + + def _cancel_context(self, context: RecvTransferContext, reason: str) -> LifecycleUpdate: + if context.logical_state is LogicalState.PENDING: + context.logical_state = LogicalState.CANCELLED + context.logical_reason = reason + context.in_doubt = True + return self._evaluate(context) + + def _close_unexposed_writers(self, context: RecvTransferContext) -> None: + if context.logical_state is LogicalState.PENDING: + return + for writer in context.writers.values(): + if writer.exposure is ExposureState.UNEXPOSED and writer.result is None: + writer.exposure = ExposureState.NEVER_EXPOSED + writer.mode = WriterMode.NO_REMOTE_ACCESS + writer.result = WriterResult.FAILED + + def _evaluate( + self, + context: RecvTransferContext, + *, + accepted: bool = True, + duplicate: bool = False, + conflict: bool = False, + reason: str | None = None, + ) -> LifecycleUpdate: + actions: list[LifecycleAction] = [] + + if context.logical_state is LogicalState.PENDING: + results = [writer.result for writer in context.writers.values()] + if any(result is WriterResult.FAILED for result in results): + context.logical_state = LogicalState.FAILED + context.logical_reason = reason or "writer failed" + elif all(result is WriterResult.SUCCESS for result in results): + has_bounce_result = any( + writer.mode is WriterMode.BOUNCE for writer in context.writers.values() + ) + if has_bounce_result: + if context.bounce_state is BounceState.IDLE: + context.bounce_state = BounceState.SCATTERING + actions.append(LifecycleAction.START_BOUNCE_SCATTER) + elif context.bounce_state is BounceState.DONE: + context.logical_state = LogicalState.SUCCEEDED + context.logical_reason = "all writers and bounce scatter succeeded" + elif context.bounce_state is BounceState.FAILED: + context.logical_state = LogicalState.FAILED + context.logical_reason = "bounce scatter failed" + else: + context.logical_state = LogicalState.SUCCEEDED + context.logical_reason = "all writers succeeded" + + self._close_unexposed_writers(context) + + if context.logical_state is not LogicalState.PENDING: + actions.extend(self._notification_actions(context)) + + all_quiescent = context.writers_terminal + scatter_active = context.bounce_pending + if context.logical_state is not LogicalState.PENDING and context.has_bounce_slot: + if context.bounce_state is BounceState.IDLE: + context.bounce_state = BounceState.SUPPRESSED + + ready_for_resource_settlement = ( + all_quiescent + and not scatter_active + and not context.conflict_requires_backend_quiescence + ) + if ready_for_resource_settlement and context.has_bounce_slot and not context.bounce_settled: + context.physical_state = PhysicalState.DRAINING + if not context._bounce_release_emitted: + context._bounce_release_emitted = True + actions.append(LifecycleAction.RELEASE_BOUNCE) + elif ready_for_resource_settlement: + context.physical_state = PhysicalState.DRAINED + if not context._drained_emitted: + context._drained_emitted = True + actions.append(LifecycleAction.CONTEXT_DRAINED) + elif context.in_doubt or context.conflict_requires_backend_quiescence: + context.physical_state = PhysicalState.IN_DOUBT + elif context.logical_state is not LogicalState.PENDING or scatter_active: + context.physical_state = PhysicalState.DRAINING + else: + context.physical_state = PhysicalState.ACTIVE + + return self._update( + context, + accepted=accepted, + duplicate=duplicate, + conflict=conflict, + actions=tuple(actions), + reason=reason, + ) + + def _notification_actions(self, context: RecvTransferContext) -> tuple[LifecycleAction, ...]: + if not context.consumer_attached: + return () + notified = context._notified_logical_state + # A protocol contradiction discovered while the context is still live + # overrides an earlier optimistic success. Emit one corrective failure + # so aggregate/session state cannot remain successful while ownership is + # physically in doubt. + if context.logical_state is LogicalState.FAILED and notified is not LogicalState.FAILED: + context._notified_logical_state = LogicalState.FAILED + return (LifecycleAction.NOTIFY_FAILURE,) + if notified is not None: + return () + context._notified_logical_state = context.logical_state + if context.logical_state is LogicalState.SUCCEEDED: + return (LifecycleAction.NOTIFY_SUCCESS,) + if context.logical_state is LogicalState.CANCELLED: + return (LifecycleAction.NOTIFY_CANCELLED,) + return (LifecycleAction.NOTIFY_FAILURE,) + + def _fail_and_evaluate( + self, context: RecvTransferContext, reason: str, *, conflict: bool = False + ) -> LifecycleUpdate: + if context.logical_state is LogicalState.PENDING or ( + conflict and context.logical_state is LogicalState.SUCCEEDED + ): + context.logical_state = LogicalState.FAILED + context.logical_reason = reason + return self._evaluate(context, accepted=False, conflict=conflict, reason=reason) + + def _dedupe_or_conflict_result( + self, + context: RecvTransferContext, + writer: WriterRecord, + result: WriterResult, + mode: WriterMode, + ) -> LifecycleUpdate: + if writer.result is result: + if writer.mode is mode or mode is WriterMode.UNKNOWN: + return self._update(context, accepted=False, duplicate=True) + if writer.mode is WriterMode.UNKNOWN: + writer.mode = mode + return self._update(context, accepted=False, duplicate=True) + writer.conflict = True + context.protocol_conflict = True + context.conflict_requires_backend_quiescence = True + context.in_doubt = True + return self._fail_and_evaluate( + context, + f"contradictory result from writer {writer.rank}", + conflict=True, + ) + + def _unexpected_result_writer( + self, context: RecvTransferContext, writer_rank: int + ) -> LifecycleUpdate: + context.protocol_conflict = True + context.conflict_requires_backend_quiescence = True + context.in_doubt = True + return self._fail_and_evaluate( + context, + f"result from unexpected writer {writer_rank}", + conflict=True, + ) + + def _writer_result_conflict( + self, context: RecvTransferContext, writer: WriterRecord, reason: str + ) -> LifecycleUpdate: + writer.conflict = True + context.protocol_conflict = True + context.conflict_requires_backend_quiescence = True + context.in_doubt = True + return self._fail_and_evaluate(context, reason, conflict=True) + + def _unexpected_local_writer( + self, context: RecvTransferContext, writer_rank: int + ) -> LifecycleUpdate: + context.protocol_conflict = True + return self._fail_and_evaluate( + context, + f"unexpected writer {writer_rank}", + conflict=True, + ) + + def _publication_conflict( + self, context: RecvTransferContext, writer: WriterRecord, reason: str + ) -> LifecycleUpdate: + writer.conflict = True + context.protocol_conflict = True + context.conflict_requires_backend_quiescence = True + context.in_doubt = True + return self._fail_and_evaluate(context, reason, conflict=True) + + @staticmethod + def _update( + context: RecvTransferContext, + *, + accepted: bool = True, + actions: tuple[LifecycleAction, ...] = (), + publication_allowed: bool = False, + duplicate: bool = False, + conflict: bool = False, + reason: str | None = None, + ) -> LifecycleUpdate: + return LifecycleUpdate( + key=context.key, + accepted=accepted, + logical_state=context.logical_state, + physical_state=context.physical_state, + writers_terminal=context.writers_terminal, + bounce_pending=context.bounce_pending, + actions=actions, + publication_allowed=publication_allowed, + duplicate=duplicate, + conflict=conflict, + reason=reason, + ) + + @staticmethod + def _rejected_update(key: TransferKey, reason: str) -> LifecycleUpdate: + return LifecycleUpdate( + key=key, + accepted=False, + logical_state=LogicalState.PENDING, + physical_state=PhysicalState.ACTIVE, + writers_terminal=False, + bounce_pending=False, + reason=reason, + ) + + +__all__ = [ + "BounceState", + "ContextSnapshot", + "ExposureState", + "LifecycleAction", + "LifecycleUpdate", + "LogicalState", + "PhysicalState", + "RecvTransferContext", + "RecvTransferRegistry", + "RegistrySnapshot", + "TransferKey", + "WriterMode", + "WriterRecord", + "WriterResult", + "WriterSnapshot", +] diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 228a321ec23f..e63e7f707ce6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.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. + from __future__ import annotations import os @@ -6,6 +21,7 @@ import threading import time import weakref +from contextlib import nullcontext from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, List, Optional, Union @@ -45,9 +61,18 @@ from tensorrt_llm._torch.disaggregation.native.peer import 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.receive_lifecycle import ( + LifecycleAction, + LifecycleUpdate, + PhysicalState, + RecvTransferRegistry, + WriterMode, + WriterResult, +) from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -63,6 +88,12 @@ # Number of worker threads for KV transfer queues (default: 1) KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "1")) +# Standalone TransferWorker users do not have KvCacheTransceiverV2's owner map. +# Retain a worker whose shutdown cannot prove drain so its manager, listeners, +# registrations, and endpoint-owned sessions cannot be destructed underneath +# a late one-sided write. +_NON_DRAINED_TRANSFER_WORKERS: set["TransferWorker"] = set() + @dataclass class RecvReqInfo: @@ -120,6 +151,37 @@ class WriteMetaType(Enum): AUX = "AUX" +class SendOperationState(Enum): + PENDING = "PENDING" + IN_DOUBT = "IN_DOUBT" + TERMINAL = "TERMINAL" + + +SendOperationKey = tuple[str, int] + + +@dataclass +class SendOperationRecord: + """One exact sender operation admitted for a task and receiver rank.""" + + state: SendOperationState + terminal_message: Optional[tuple[bytes, ...]] = None + result_delivered: bool = False + + +UnboundTerminalKey = tuple[WriteMetaType, str, int, Optional[int]] +PreSessionTerminalKey = tuple[int, WriteMetaType, str, int, Optional[int]] + + +@dataclass +class UnboundTerminalResult: + """A no-access result created before its channel task exists.""" + + req_info: RecvReqInfo + message: tuple[bytes, ...] + result_delivered: bool = False + + @dataclass class WriteMeta: task: Union[SendTaskBase, "KVRecvTask"] @@ -136,6 +198,11 @@ class WriteMeta: is_last_slice: bool = False meta_type: WriteMetaType = WriteMetaType.KV bounce_dst_base: Optional[int] = None + # Bind queued work to the exact session that authorized it. Looking a + # session up again by request ID can attach stale work to a replacement. + session: Optional["TxSession"] = None + source_access_enrolled: bool = False + operation_key: Optional[SendOperationKey] = None class MessageType: @@ -160,6 +227,38 @@ class AgentResult(Enum): FAILED = "FAILED" +class TransferSourceInDoubtError(RuntimeError): + """A transfer-agent exception did not prove source access quiescent.""" + + +@dataclass +class SendTransferContext: + """Own one sender operation and every object needed to keep it safe. + + A context is normally lexical to a worker invocation. If submit/wait + raises without terminal evidence, Sender retains the context so the exact + session/source owner, transfer request/status handle, and bounce lease + cannot be destroyed independently. + """ + + write_meta: WriteMeta + session: "TxSession" + source_owner: object | None = None + transfer_request: object | None = None + transfer_status: object | None = None + bounce_slot_id: Optional[int] = None + result_delivered: bool = False + + +@dataclass(frozen=True) +class _PendingLifecycleDelivery: + """A registry transition retained until its session consumer acknowledges it.""" + + update: LifecycleUpdate + peer_rank: Optional[int] + succeeded: bool + + # KV_AGENT_RESULT prefix in one struct frame (was 5 ascii frames serialized/parsed under the # GIL per slice per writer): instance_rank, unique_rid, slice_id, is_last, status. The optional # bounce tail follows at message[2:]. @@ -190,7 +289,7 @@ def _make_kv_result_msg( class SendTaskBase: - def __init__(self, params: DisaggregatedParams): + def __init__(self, params: DisaggregatedParams, session: Optional["TxSession"] = None): self.status = TaskStatus.INIT self._event = threading.Event() self._exception: Optional[Exception] = None @@ -198,15 +297,150 @@ def __init__(self, params: DisaggregatedParams): self._params = params self._unique_rid: Optional[int] = params.disagg_request_id self._perf_timer = PerfTimer() if perf_log_manager.enabled else None + self._session = session + # Count every queued, active, or in-doubt operation independently. A + # task-level ERROR is only a logical result; sibling operations may + # still be reading the same source allocation. + self._source_access_count = 0 + # The task itself scopes the channel (KV slice or session-level AUX). + # A receiver instance/rank may therefore authorize at most one physical + # operation for this task. Terminal frames are cached for idempotent + # replay instead of launching a second write. + self._send_operations: dict[SendOperationKey, SendOperationRecord] = {} + + def begin_source_access(self) -> None: + with self.lock: + self._source_access_count += 1 + + def finish_source_access(self) -> None: + with self.lock: + if self._source_access_count <= 0: + raise RuntimeError("source access ownership underflow") + self._source_access_count -= 1 + + @property + def source_access_active(self) -> bool: + with self.lock: + return self._source_access_count > 0 + + def admit_operation( + self, + key: SendOperationKey, + *, + allow_source_access: bool, + no_access_message: tuple[bytes, ...], + ) -> tuple[bool, SendOperationState, Optional[tuple[bytes, ...]], bool]: + """Admit one exact physical operation or return its existing state. + + The caller holds the owning session's lifecycle lock. The task lock + makes terminal updates from a worker atomic with duplicate admission. + """ + with self.lock: + existing = self._send_operations.get(key) + if existing is not None: + return ( + False, + existing.state, + existing.terminal_message, + existing.result_delivered, + ) + if allow_source_access: + self._send_operations[key] = SendOperationRecord(SendOperationState.PENDING) + self._source_access_count += 1 + return True, SendOperationState.PENDING, None, False + self._send_operations[key] = SendOperationRecord( + SendOperationState.TERMINAL, no_access_message + ) + return True, SendOperationState.TERMINAL, no_access_message, False + + def operation_snapshot( + self, key: SendOperationKey + ) -> Optional[tuple[SendOperationState, Optional[tuple[bytes, ...]], bool]]: + with self.lock: + record = self._send_operations.get(key) + if record is None: + return None + return record.state, record.terminal_message, record.result_delivered + + def mark_operation_in_doubt(self, key: Optional[SendOperationKey]) -> None: + if key is None: + return + with self.lock: + record = self._send_operations.get(key) + if record is None: + raise RuntimeError(f"sender operation {key} was not admitted") + if record.state is not SendOperationState.TERMINAL: + record.state = SendOperationState.IN_DOUBT + + def cache_terminal_message( + self, key: Optional[SendOperationKey], message: list[bytes] | tuple[bytes, ...] + ) -> None: + if key is None: + return + cached = tuple(message) + with self.lock: + record = self._send_operations.get(key) + if record is None: + raise RuntimeError(f"sender operation {key} was not admitted") + if record.state is SendOperationState.TERMINAL: + if record.terminal_message != cached: + raise RuntimeError(f"conflicting terminal result for sender operation {key}") + return + record.state = SendOperationState.TERMINAL + record.terminal_message = cached + + def mark_operation_result_delivered(self, key: Optional[SendOperationKey]) -> None: + if key is None: + return + with self.lock: + record = self._send_operations.get(key) + if record is None or record.state is not SendOperationState.TERMINAL: + raise RuntimeError(f"sender operation {key} has no terminal result to deliver") + record.result_delivered = True + + def pending_terminal_messages( + self, + ) -> list[tuple[SendOperationKey, tuple[bytes, ...]]]: + with self.lock: + return [ + (key, record.terminal_message) + for key, record in self._send_operations.items() + if record.state is SendOperationState.TERMINAL + and record.terminal_message is not None + and not record.result_delivered + ] + + @property + def has_pending_result_delivery(self) -> bool: + with self.lock: + return any( + record.state is SendOperationState.TERMINAL + and record.terminal_message is not None + and not record.result_delivered + for record in self._send_operations.values() + ) + + def mark_transferring(self) -> bool: + """Start physical work without reviving a terminal task.""" + with self.lock: + if self.status is not TaskStatus.INIT: + return self.status is TaskStatus.TRANSFERRING + self.status = TaskStatus.TRANSFERRING + return True def fail(self, exc: Exception) -> None: - self._exception = exc - self.status = TaskStatus.ERROR - self._event.set() + with self.lock: + self._exception = exc + self.status = TaskStatus.ERROR + self._event.set() - def complete(self) -> None: - self.status = TaskStatus.TRANSFERRED - self._event.set() + def complete(self) -> bool: + with self.lock: + if self.status is TaskStatus.ERROR: + return False + self.status = TaskStatus.TRANSFERRED + self._event.set() + return True def wait(self, timeout: Optional[float] = None) -> bool: """Block until terminal state. Returns True if done, False on timeout.""" @@ -230,8 +464,13 @@ def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int class AuxSendTask(SendTaskBase): - def __init__(self, params: DisaggregatedParams, slot: Optional[int]): - super().__init__(params) + def __init__( + self, + params: DisaggregatedParams, + slot: Optional[int], + session: Optional["TxSession"] = None, + ): + super().__init__(params, session) self._slot = slot self._transfer_count = 0 @@ -244,8 +483,9 @@ def __init__( slice_id: int, prompt_len: Optional[int] = None, beam_width: int = 1, + session: Optional["TxSession"] = None, ): - super().__init__(params) + super().__init__(params, session) self.slice_id = slice_id self.transferred_count = 0 self._slice = kv_slice @@ -259,6 +499,7 @@ class Sender(SenderBase): # RecvReqInfo that never gets consumed. Entries older than this # are evicted during periodic sweeps. _STALE_REQ_INFO_TTL_S = 120.0 + _TERMINAL_RESULT_RETRY_INTERVAL_S = 0.05 def __init__( self, @@ -266,6 +507,7 @@ def __init__( agent: BaseTransferAgent, bounce=None, ): + self._shutdown_attempt_lock = threading.Lock() self._registrar = peer_registrar self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent @@ -274,12 +516,36 @@ def __init__( self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() self._messenger = ZMQMessenger(mode="ROUTER") - self._dealers = {} # used by listener thread only (single-threaded path) + # Control sends can originate from the listener or executor thread. + # ZMQMessenger serializes each socket; this lock additionally makes the + # endpoint-to-socket cache's check/create transition atomic. + self._dealers = {} + self._dealers_lock = threading.Lock() self._thread_local = threading.local() # per-thread DEALER cache for worker threads - self._sessions = {} # unique_rid -> TxSession + # Close direct TxSession admission before shutdown can append queue + # sentinels. Listener callbacks that were already admitted may finish + # enqueueing before messenger.stop() joins them. + self._operation_admission_lock = threading.RLock() + self._dealer_admission_closed = False + # Worker-local sockets can only be detached by their owning thread. + # Failed closes move here so Sender remains their retry owner. + self._failed_thread_dealers: list[ZMQMessenger] = [] + self._failed_thread_dealers_lock = threading.Lock() + # Sender is the lifecycle root for standalone TransferWorker users. + # Keep sessions strong until explicit, drain-checked clear_session(). + self._sessions: dict[int, TxSession] = {} self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids self._pre_cancelled_rids: set[int] = set() + self._pre_session_terminal_results: dict[PreSessionTerminalKey, UnboundTerminalResult] = {} + self._pre_session_terminal_results_lock = threading.Lock() + self._pre_session_terminal_retry_lock = threading.Lock() + self._next_pre_session_terminal_retry_at = 0.0 + self._in_doubt_transfers: list[SendTransferContext] = [] + self._in_doubt_transfers_lock = threading.Lock() self._shutdown = False + self._listener_stopped = False + self._shutdown_complete = False + self._shutdown_sentinels_sent = False self._instance_rank = self._registrar.self_rank_info.instance_rank # Guards concurrent add() from the listener thread. self._loaded_remote_agents: set[str] = set() @@ -305,23 +571,35 @@ def __init__( def endpoint(self): return self._messenger.endpoint - def _add_req_info(self, unique_rid: int, instance_rank: int, req_info: RecvReqInfo): + def _add_req_info(self, unique_rid: int, req_info: RecvReqInfo): with self._peer_requests_lock: if unique_rid not in self._peer_requests: self._peer_requests[unique_rid] = {} self._peer_requests_timestamps[unique_rid] = time.monotonic() - self._peer_requests[unique_rid][instance_rank] = req_info + key = (req_info.instance_rank, req_info.slice_id) + self._peer_requests[unique_rid][key] = req_info - def _is_req_ready(self, unique_rid: int, expected_count: int) -> bool: + def _is_req_ready( + self, + unique_rid: int, + expected_count: int, + slice_id: Optional[int] = None, + ) -> bool: with self._peer_requests_lock: requests = self._peer_requests.get(unique_rid) if not requests: return False - return len(requests) == expected_count + ranks = { + info.instance_rank + for info in requests.values() + if slice_id is None or info.slice_id == slice_id + } + return len(ranks) == expected_count def _get_req_info(self, unique_rid: Optional[int]) -> Optional[dict]: with self._peer_requests_lock: - return self._peer_requests.get(unique_rid) + requests = self._peer_requests.get(unique_rid) + return None if requests is None else dict(requests) def _get_first_req_info(self, unique_rid: Optional[int]) -> Optional[RecvReqInfo]: with self._peer_requests_lock: @@ -338,10 +616,14 @@ def _remove_req_info(self, unique_rid: int): def sweep_stale_req_infos(self): """Evict RecvReqInfo entries that have no matching TxSession and exceed the TTL. - Called opportunistically from the listener thread when a new REQUEST_DATA - arrives. With gen-first ADP broadcast, non-assigned DP ranks accumulate - entries that are never consumed; this sweep prevents unbounded growth. + Called from the scheduler's sender progress path. With gen-first ADP + broadcast, non-assigned DP ranks accumulate entries that are never + consumed; this sweep prevents unbounded growth. """ + # The scheduler already calls this method as its sender-side progress + # hook. Reuse that cadence to retry pre-session terminal decisions so a + # transient control-send failure does not wait until process shutdown. + self._retry_pre_session_terminal_results() now = time.monotonic() with self._peer_requests_lock: stale_rids = [ @@ -353,7 +635,11 @@ def sweep_stale_req_infos(self): return for rid in stale_rids: with self._sessions_lock, self._peer_requests_lock: - if rid not in self._sessions and rid in self._peer_requests: + if ( + rid not in self._sessions + and rid not in self._pre_cancelled_rids + and rid in self._peer_requests + ): self._peer_requests.pop(rid, None) self._peer_requests_timestamps.pop(rid, None) logger.debug(f"Swept stale RecvReqInfo for rid={rid}") @@ -361,43 +647,75 @@ def sweep_stale_req_infos(self): def setup_session(self, tx_session: "TxSession"): unique_rid = tx_session.disagg_request_id pre_cancel = False - with self._sessions_lock: - self._sessions[unique_rid] = weakref.ref(tx_session) - if unique_rid in self._pre_cancelled_rids: - pre_cancel = True - self._pre_cancelled_rids.discard(unique_rid) - if pre_cancel: - tx_session.cancel() - return - - req_info = self._get_first_req_info(unique_rid) - - if req_info: - peer_ri = self._registrar.get_peer_rank_info( - req_info.instance_name, req_info.instance_rank - ) - expected_count = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) - if self._is_req_ready(unique_rid, expected_count): - with tx_session.lock: - tx_session.receiver_ready = True + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + published = False + try: + with self._sessions_lock: + if self._shutdown: + raise RuntimeError("Sender is shutting down; new sessions are not accepted") + existing = self._sessions.get(unique_rid) + if existing is not None and existing is not tx_session: + raise RuntimeError(f"TxSession {unique_rid} is already registered") + self._sessions[unique_rid] = tx_session + published = True + if unique_rid in self._pre_cancelled_rids: + pre_cancel = True + if pre_cancel: + tx_session.cancel() + return + + req_info = self._get_first_req_info(unique_rid) + + if req_info: + peer_ri = self._registrar.get_peer_rank_info( + req_info.instance_name, req_info.instance_rank + ) + expected_count = len( + self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks + ) + if self._is_req_ready(unique_rid, expected_count, req_info.slice_id): + with tx_session.lock: + tx_session.receiver_ready = True + except Exception: + if published: + with self._sessions_lock: + if self._sessions.get(unique_rid) is tx_session: + self._sessions.pop(unique_rid, None) + raise return def _get_session(self, unique_rid: Optional[int]) -> Optional["TxSession"]: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: - return None - session = session_ref() - if session is None: - logger.warning(f"TxSession {unique_rid} has been garbage collected") - return None - return session + return self._sessions.get(unique_rid) def _enqueue(self, write_meta: WriteMeta): # Route by (unique_rid, peer_rank) so that: # - Same peer's slices stay ordered on one thread (is_last_slice correctness) # - Different peers can run on different threads (better load balancing) - thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads - self._send_task_queues[thread_idx].put(write_meta) + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_shutdown_sentinels_sent", False): + raise RuntimeError("Sender queues are closed for shutdown") + thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads + self._send_task_queues[thread_idx].put(write_meta) + + def _enqueue_owned(self, write_meta: WriteMeta) -> None: + """Enroll source ownership before queued work can outlive its caller.""" + enrolled_here = not write_meta.source_access_enrolled + if enrolled_here: + write_meta.task.begin_source_access() + write_meta.source_access_enrolled = True + try: + self._enqueue(write_meta) + except Exception: + if enrolled_here: + write_meta.task.finish_source_access() + write_meta.source_access_enrolled = False + raise + + def _retain_in_doubt_transfer(self, context: SendTransferContext) -> None: + with self._in_doubt_transfers_lock: + self._in_doubt_transfers.append(context) def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger: """Get or create a per-thread DEALER socket via threading.local(). @@ -413,6 +731,24 @@ def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) return dealers[endpoint] + def _close_thread_dealers(self, thread_idx: int, dealers: dict[str, ZMQMessenger]) -> None: + """Detach worker-local sockets while preserving failed closes.""" + for endpoint, dealer in dealers.items(): + try: + dealer.stop() + except Exception as e: + failed_lock = getattr(self, "_failed_thread_dealers_lock", None) + if failed_lock is not None: + with failed_lock: + self._failed_thread_dealers.append(dealer) + else: + self._failed_thread_dealers.append(dealer) + logger.warning( + f"_process_task_queue[{thread_idx}]: failed to stop dealer " + f"for endpoint {endpoint}; retaining it for retry: {e}" + ) + dealers.clear() + def _process_task_queue(self, thread_idx: int): device_id = self._device_id torch.cuda.set_device(device_id) @@ -437,22 +773,28 @@ def _process_task_queue(self, thread_idx: int): f"_process_task_queue[{thread_idx}]: unhandled exception for " f"unique_rid={write_meta.unique_rid}: {e}" ) - write_meta.task.fail(e) + from .bounce import GatherSourceInDoubtError + + if isinstance(e, (GatherSourceInDoubtError, TransferSourceInDoubtError)): + # Keep the per-operation source owner active. The + # retained SendTransferContext owns the exact session, + # request, and every handle returned before ambiguity. + # A gather failure can occur before the bounce slot is + # returned; in that case the bounce allocator owns its + # quarantined slot independently. + logger.critical( + f"_process_task_queue[{thread_idx}]: retaining source owner for " + f"unique_rid={write_meta.unique_rid} after ambiguous source access" + ) + else: + write_meta.task.fail(e) finally: # Clean up this thread's DEALER sockets. threading.local storage # is only accessible from the owning thread, so shutdown must # happen here rather than in Sender.shutdown(). dealers = getattr(self._thread_local, "dealers", None) if dealers: - for endpoint, dealer in dealers.items(): - try: - dealer.stop() - except Exception as e: - logger.warning( - f"_process_task_queue[{thread_idx}]: failed to stop dealer " - f"for endpoint {endpoint}: {e}" - ) - dealers.clear() + self._close_thread_dealers(thread_idx, dealers) @staticmethod @nvtx_range("_make_agent_request") @@ -493,211 +835,397 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque TransferOp.WRITE, src_memory_descs, dst_memory_descs, write_meta.peer_name, None ) - @nvtx_range("_deliver_kv_to_agent") - def _deliver_kv_to_agent(self, write_meta: WriteMeta): - assert write_meta.src_ptrs.size == write_meta.dst_ptrs.size == write_meta.sizes.size, ( - f"WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" + def _failed_write_meta_message(self, write_meta: WriteMeta) -> list[bytes]: + if write_meta.meta_type is WriteMetaType.AUX: + return [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(write_meta.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ] + return _make_kv_result_msg( + self._instance_rank, + write_meta.unique_rid, + write_meta.slice_id if write_meta.slice_id is not None else 0, + write_meta.is_last_slice, + AgentResult.FAILED, ) - with self._sessions_lock: - session = self._get_session(write_meta.unique_rid) - if session is None: - msg = ( - f"_deliver_kv_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" + def _send_worker_operation_message( + self, + task: SendTaskBase, + write_meta: WriteMeta, + message: list[bytes], + ) -> bool: + """Cache then submit a terminal frame from the owning worker thread.""" + task.cache_terminal_message(write_meta.operation_key, message) + try: + self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send(message) + except Exception as e: + logger.error( + f"sender result delivery remains pending for rid={write_meta.unique_rid} " + f"writer={write_meta.peer_rank}: {e}" ) - logger.error(msg) - write_meta.task.fail(RuntimeError(msg)) - return - assert write_meta.slice_id is not None - task = session.kv_tasks[write_meta.slice_id] - timer = task._perf_timer - if timer: - timer.record_push_end(write_meta.peer_rank) - # Hold session.lock to serialize the INIT→TRANSFERRING transition with - # cancel(): prevents cancel_request() from freeing KV pages while a - # worker is about to write into them. - with session.lock: - status = session.status - if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - should_abort = True - else: - task.status = TaskStatus.TRANSFERRING - should_abort = False + return False + task.mark_operation_result_delivered(write_meta.operation_key) + return True - if should_abort: - logger.warning( - f"_deliver_kv_to_agent: session {write_meta.unique_rid} already " - f"in {status.value} state; sending FAILED to receiver" - ) - # Task may have been enqueued after cancel() already iterated kv_tasks, - # so its future was never set by cancel(). Set it here as a fallback. - task.fail( - RuntimeError(f"session {write_meta.unique_rid} {status.value}, transfer aborted") - ) - self._get_or_connect_dealer(write_meta.peer_endpoint).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, + @nvtx_range("_deliver_kv_to_agent") + def _deliver_kv_to_agent(self, write_meta: WriteMeta): + task = write_meta.task + if not isinstance(task, KVSendTask): + raise RuntimeError("queued KV work has an invalid task type") + if not write_meta.source_access_enrolled: + task.begin_source_access() + write_meta.source_access_enrolled = True + + from .bounce import GatherSourceInDoubtError, build_send_request, encode_result_tail + + session = write_meta.session + context = None + source_terminal = False + timer = task._perf_timer + try: + if not (write_meta.src_ptrs.size == write_meta.dst_ptrs.size == write_meta.sizes.size): + raise ValueError( + f"WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" ) - ) - return - - from .bounce import build_send_request, encode_result_tail - - agent_result = AgentResult.SUCCESS - send_slot_id = None - if write_meta.src_ptrs.size > 0: - try: - request, send_slot_id = build_send_request( - self._bounce, - write_meta, - lambda: Sender._make_agent_request(write_meta, device_id=self._device_id), + if session is None: + # Compatibility for standalone callers that construct + # WriteMeta directly. Normal queued work is bound to an exact + # session before admission. + with self._sessions_lock: + session = self._get_session(write_meta.unique_rid) + if session is None: + raise RuntimeError( + f"_deliver_kv_to_agent: TxSession {write_meta.unique_rid} " + "not found or already GC'd" ) - except Exception as e: - # Don't let a gather fault escape: without a result the receiver would hang and its - # region leak. Tell the receiver it failed and fail the local task instead. - logger.error( - f"_deliver_kv_to_agent: failed to build the KV send request for " - f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" + if write_meta.slice_id is None: + raise RuntimeError("queued KV work has no slice ID") + if ( + write_meta.slice_id >= len(session.kv_tasks) + or session.kv_tasks[write_meta.slice_id] is not task + ): + raise RuntimeError( + f"queued KV work no longer belongs to its session: " + f"unique_rid={write_meta.unique_rid} slice={write_meta.slice_id}" ) - task.fail(RuntimeError(f"build_send_request failed: {e}")) - self._get_or_connect_dealer(write_meta.peer_endpoint).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, + + context = SendTransferContext( + write_meta=write_meta, + session=session, + source_owner=session.source_owner, + ) + if timer: + timer.record_push_end(write_meta.peer_rank) + # Hold session.lock to serialize the INIT→TRANSFERRING transition + # with cancel(). Source ownership was enrolled before queueing. + with session.lock: + status = session.status + if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + should_abort = True + else: + should_abort = not task.mark_transferring() + + if should_abort: + logger.warning( + f"_deliver_kv_to_agent: session {write_meta.unique_rid} already " + f"in {status.value} state; sending FAILED to receiver" + ) + task.fail( + RuntimeError( + f"session {write_meta.unique_rid} {status.value}, transfer aborted" ) ) + # The source decision is terminal before result I/O. A socket + # error must not strand the source lease. + source_terminal = True + self._send_worker_operation_message( + task, write_meta, self._failed_write_meta_message(write_meta) + ) return - if timer: - timer.record_transfer_start(write_meta.peer_rank) - try: - status = self._agent.submit_transfer_requests(request) - if not status.wait(): - agent_result = AgentResult.FAILED - last_status = getattr(status, "last_status_str", lambda: "")() - agent_name = getattr(self._agent, "name", "") - detail = ( - f"KV transfer agent failed: " - f"unique_rid={write_meta.unique_rid} " - f"slice={write_meta.slice_id} " - f"peer_rank={write_meta.peer_rank} " - f"peer_endpoint={write_meta.peer_endpoint} " - f"op={getattr(request, 'op', '?')} " - f"remote={getattr(request, 'remote_name', '?')} " - f"src_size={int(write_meta.src_ptrs.size)} " - f"dst_size={int(write_meta.dst_ptrs.size)} " - f"nixl_status={last_status} agent={agent_name}" + + agent_result = AgentResult.SUCCESS + send_slot_id = None + if write_meta.src_ptrs.size > 0: + try: + request, send_slot_id = build_send_request( + self._bounce, + write_meta, + lambda: Sender._make_agent_request(write_meta, device_id=self._device_id), ) - logger.error(detail) - task.fail(RuntimeError(detail)) - finally: - if send_slot_id is not None: - self._bounce.release_send(send_slot_id) - if timer: - timer.record_transfer_end(write_meta.peer_rank) + context.transfer_request = request + context.bounce_slot_id = send_slot_id + except GatherSourceInDoubtError: + task.mark_operation_in_doubt(write_meta.operation_key) + self._retain_in_doubt_transfer(context) + raise + except Exception as e: + # Gather rollback positively fenced local source access, + # and no network request was submitted. + source_terminal = True + task.fail(e) + self._send_worker_operation_message( + task, write_meta, self._failed_write_meta_message(write_meta) + ) + return + transfer_terminal = False + try: + if timer: + timer.record_transfer_start(write_meta.peer_rank) + transfer_status = self._agent.submit_transfer_requests(request) + context.transfer_status = transfer_status + transfer_succeeded = transfer_status.wait() + transfer_terminal = True + source_terminal = True + if not transfer_succeeded: + agent_result = AgentResult.FAILED + last_status = getattr( + transfer_status, "last_status_str", lambda: "" + )() + agent_name = getattr(self._agent, "name", "") + detail = ( + f"KV transfer agent failed: " + f"unique_rid={write_meta.unique_rid} " + f"slice={write_meta.slice_id} " + f"peer_rank={write_meta.peer_rank} " + f"peer_endpoint={write_meta.peer_endpoint} " + f"op={getattr(request, 'op', '?')} " + f"remote={getattr(request, 'remote_name', '?')} " + f"src_size={int(write_meta.src_ptrs.size)} " + f"dst_size={int(write_meta.dst_ptrs.size)} " + f"nixl_status={last_status} agent={agent_name}" + ) + logger.error(detail) + task.fail(RuntimeError(detail)) + except Exception as e: + task.mark_operation_in_doubt(write_meta.operation_key) + # Retain the complete context first. Bounce quarantine is + # itself fallible and must never be able to lose the source + # owner or the backend status handle. + self._retain_in_doubt_transfer(context) + if send_slot_id is not None: + try: + self._bounce.quarantine_send(send_slot_id) + except Exception as quarantine_error: + logger.critical( + f"failed to quarantine send slot {send_slot_id} for " + f"request {write_meta.unique_rid}; retained transfer " + f"context remains authoritative: {quarantine_error}" + ) + raise TransferSourceInDoubtError( + f"KV transfer-agent operation is in doubt for " + f"request {write_meta.unique_rid} slice={write_meta.slice_id} " + f"writer={write_meta.peer_rank}" + ) from e + finally: + if transfer_terminal and send_slot_id is not None: + self._bounce.release_send(send_slot_id) + context.bounce_slot_id = None + else: + source_terminal = True + if timer: + timer.record_transfer_end(write_meta.peer_rank) - ## TODO: just last slice need to send task state? - tail = ( - encode_result_tail(write_meta) - if send_slot_id is not None and agent_result == AgentResult.SUCCESS - else None - ) - result_msg = _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - write_meta.is_last_slice, - agent_result, - tail=tail, - ) - self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send(result_msg) + tail = ( + encode_result_tail(write_meta) + if send_slot_id is not None and agent_result == AgentResult.SUCCESS + else None + ) + result_msg = _make_kv_result_msg( + self._instance_rank, + write_meta.unique_rid, + write_meta.slice_id, + write_meta.is_last_slice, + agent_result, + tail=tail, + ) + delivered = self._send_worker_operation_message(task, write_meta, result_msg) + context.result_delivered = delivered - if timer: - timer.record_task_end(write_meta.peer_rank) - ri = self._registrar.self_rank_info - task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) + if timer: + timer.record_task_end(write_meta.peer_rank) + ri = self._registrar.self_rank_info + task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) - with task.lock: - task.transferred_count += 1 - count = task.transferred_count + with task.lock: + task.transferred_count += 1 + count = task.transferred_count - if count > write_meta.expected_transfers: - session.set_exception( - f"KV slice {write_meta.slice_id} received more than {write_meta.expected_transfers} transfers" - ) - elif count == write_meta.expected_transfers: - if task.is_done: - task.status = TaskStatus.ERROR + if count > write_meta.expected_transfers: session.set_exception( - f"KV slice {write_meta.slice_id} task already resolved on completion" + f"KV slice {write_meta.slice_id} received more than " + f"{write_meta.expected_transfers} transfers" ) - else: - task.complete() + elif count == write_meta.expected_transfers: + if not task.complete(): + session.set_exception( + f"KV slice {write_meta.slice_id} task failed before all writers completed" + ) - logger.debug( - f"deliver_kv_to_agent completed: unique_rid={write_meta.unique_rid}, " - f"slice_id={write_meta.slice_id}, agent_result={agent_result}" - ) + logger.debug( + f"deliver_kv_to_agent completed: unique_rid={write_meta.unique_rid}, " + f"slice_id={write_meta.slice_id}, agent_result={agent_result}" + ) + except (GatherSourceInDoubtError, TransferSourceInDoubtError): + raise + except Exception as e: + # Every fenced failure is terminal for local source access. Cache + # the result before attempting I/O so cancellation/REQUEST_DATA can + # replay it later. + source_terminal = True + task.fail(e) + snapshot = ( + task.operation_snapshot(write_meta.operation_key) + if write_meta.operation_key is not None + else None + ) + if snapshot is None or snapshot[0] is not SendOperationState.TERMINAL: + self._send_worker_operation_message( + task, write_meta, self._failed_write_meta_message(write_meta) + ) + raise + finally: + if source_terminal: + task.finish_source_access() + write_meta.source_access_enrolled = False @nvtx_range("_deliver_aux_to_agent") def _deliver_aux_to_agent(self, write_meta: WriteMeta): - session = self._get_session(write_meta.unique_rid) - if session is None: - msg = f"_deliver_aux_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" - logger.error(msg) - write_meta.task.fail(RuntimeError(msg)) - return - aux_task = session.aux_task - assert aux_task is not None, f"aux_task is None for session {write_meta.unique_rid}" + aux_task = write_meta.task + if not isinstance(aux_task, AuxSendTask): + raise RuntimeError("queued AUX work has an invalid task type") + if not write_meta.source_access_enrolled: + aux_task.begin_source_access() + write_meta.source_access_enrolled = True + + session = write_meta.session + context = None + source_terminal = False timer = aux_task._perf_timer - if timer: - timer.record_push_end(write_meta.peer_rank) + try: + if not (write_meta.src_ptrs.size == write_meta.dst_ptrs.size == write_meta.sizes.size): + raise ValueError( + f"AUX WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" + ) + if session is None: + with self._sessions_lock: + session = self._get_session(write_meta.unique_rid) + if session is None: + raise RuntimeError( + f"_deliver_aux_to_agent: TxSession {write_meta.unique_rid} " + "not found or already GC'd" + ) + if session.aux_task is not aux_task: + raise RuntimeError( + f"queued AUX work no longer belongs to its session: " + f"unique_rid={write_meta.unique_rid}" + ) - agent_result = AgentResult.SUCCESS - if write_meta.src_ptrs.size > 0: - request = Sender._make_agent_request(write_meta, device_id=self._device_id) - if timer: - timer.record_transfer_start(write_meta.peer_rank) - if not self._agent.submit_transfer_requests(request).wait(): - agent_result = AgentResult.FAILED - session.set_exception("aux transfer agent request failed") + context = SendTransferContext( + write_meta=write_meta, + session=session, + source_owner=session.source_owner, + ) if timer: - timer.record_transfer_end(write_meta.peer_rank) + timer.record_push_end(write_meta.peer_rank) + + with session.lock: + status = session.status + if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + should_abort = True + else: + should_abort = not aux_task.mark_transferring() + + if should_abort: + aux_task.fail( + RuntimeError(f"session {write_meta.unique_rid} {status.value}, AUX aborted") + ) + source_terminal = True + self._send_worker_operation_message( + aux_task, write_meta, self._failed_write_meta_message(write_meta) + ) + return + + agent_result = AgentResult.SUCCESS + if write_meta.src_ptrs.size > 0: + try: + request = Sender._make_agent_request(write_meta, device_id=self._device_id) + context.transfer_request = request + except Exception as e: + source_terminal = True + aux_task.fail(e) + self._send_worker_operation_message( + aux_task, write_meta, self._failed_write_meta_message(write_meta) + ) + return + try: + if timer: + timer.record_transfer_start(write_meta.peer_rank) + transfer_status = self._agent.submit_transfer_requests(request) + context.transfer_status = transfer_status + transfer_succeeded = transfer_status.wait() + source_terminal = True + except Exception as e: + aux_task.mark_operation_in_doubt(write_meta.operation_key) + self._retain_in_doubt_transfer(context) + raise TransferSourceInDoubtError( + f"AUX transfer-agent operation is in doubt for " + f"request {write_meta.unique_rid} writer={write_meta.peer_rank}" + ) from e + if not transfer_succeeded: + agent_result = AgentResult.FAILED + session.set_exception("aux transfer agent request failed") + if timer: + timer.record_transfer_end(write_meta.peer_rank) + else: + source_terminal = True - self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send( - [ + result_msg = [ MessageType.AUX_AGENT_RESULT, str(self._instance_rank).encode("ascii"), str(write_meta.unique_rid).encode("ascii"), agent_result.value.encode("ascii"), ] - ) + delivered = self._send_worker_operation_message(aux_task, write_meta, result_msg) + context.result_delivered = delivered - if timer: - timer.record_task_end(write_meta.peer_rank) - ri = self._registrar.self_rank_info - aux_task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) - - with aux_task.lock: - aux_task._transfer_count += 1 - count = aux_task._transfer_count - - if count == write_meta.expected_transfers: - if aux_task.is_done: - aux_task.status = TaskStatus.ERROR - session.set_exception("aux task already resolved on completion") - else: - aux_task.complete() - elif count > write_meta.expected_transfers: - session.set_exception( - f"aux task received more than {write_meta.expected_transfers} transfers" + if timer: + timer.record_task_end(write_meta.peer_rank) + ri = self._registrar.self_rank_info + aux_task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) + + with aux_task.lock: + aux_task._transfer_count += 1 + count = aux_task._transfer_count + + if count == write_meta.expected_transfers: + if not aux_task.complete(): + session.set_exception("aux task failed before all writers completed") + elif count > write_meta.expected_transfers: + session.set_exception( + f"aux task received more than {write_meta.expected_transfers} transfers" + ) + except TransferSourceInDoubtError: + raise + except Exception as e: + source_terminal = True + aux_task.fail(e) + snapshot = ( + aux_task.operation_snapshot(write_meta.operation_key) + if write_meta.operation_key is not None + else None ) + if snapshot is None or snapshot[0] is not SendOperationState.TERMINAL: + self._send_worker_operation_message( + aux_task, write_meta, self._failed_write_meta_message(write_meta) + ) + raise + finally: + if source_terminal: + aux_task.finish_source_access() + write_meta.source_access_enrolled = False @staticmethod def _align_kv_blocks( @@ -892,6 +1420,7 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write slice_id=task.slice_id, is_last_slice=task._slice.is_last_slice, bounce_dst_base=req_info.bounce_dst_base, + session=task._session, ) def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> WriteMeta: @@ -933,27 +1462,312 @@ def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> Wri peer_endpoint=peer_ri.self_endpoint, unique_rid=task._unique_rid, meta_type=WriteMetaType.AUX, + session=task._session, ) - def dispatch_task( - self, - task: KVSendTask | AuxSendTask, - req_info_snapshot: Optional[dict] = None, - ): - # req_info_snapshot may be pre-fetched under session.lock by the caller to keep the - # critical section small. When not provided, we fetch it here (legacy / standalone path). - if req_info_snapshot is None: - req_info_snapshot = dict(self._get_req_info(task._unique_rid) or {}) - for info in req_info_snapshot.values(): - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - if isinstance(task, KVSendTask): - trans_meta = self._build_kv_write_meta(task, info) - else: - trans_meta = self._build_aux_write_meta(task, info) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) + @staticmethod + def _operation_key(req_info: RecvReqInfo) -> SendOperationKey: + return req_info.instance_name, req_info.instance_rank + + def _failed_operation_message( + self, task: KVSendTask | AuxSendTask, req_info: RecvReqInfo + ) -> tuple[bytes, ...]: + if isinstance(task, KVSendTask): + return tuple( + _make_kv_result_msg( + self._instance_rank, + req_info.unique_rid, + task.slice_id, + task._slice.is_last_slice, + AgentResult.FAILED, + ) + ) + return ( + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(req_info.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ) + + def _send_operation_message( + self, req_info: RecvReqInfo, message: tuple[bytes, ...] | list[bytes] + ) -> bool: + """Best-effort delivery for a cached terminal result. + + The result remains in the operation ledger after an I/O failure, so an + identical request/cancel replay can retry it without launching work. + """ + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_dealer_admission_closed", False): + return False + try: + peer_ri = self._registrar.get_peer_rank_info( + req_info.instance_name, req_info.instance_rank + ) + self._get_or_connect_dealer(peer_ri.self_endpoint).send(list(message)) + return True + except Exception as e: + logger.warning( + f"failed to deliver cached sender result for rid={req_info.unique_rid} " + f"receiver={req_info.instance_name}/{req_info.instance_rank}: {e}" + ) + return False + + def _send_unbound_terminal_result( + self, + session: "TxSession", + req_info: RecvReqInfo, + meta_type: WriteMetaType, + message: tuple[bytes, ...], + ) -> None: + """Cache and send a terminal decision made before a task exists.""" + key = session.cache_unbound_terminal_result(meta_type, req_info, message) + if self._send_operation_message(req_info, message): + session.mark_unbound_terminal_result_delivered(key) + + @staticmethod + def _pre_session_terminal_key( + meta_type: WriteMetaType, req_info: RecvReqInfo + ) -> PreSessionTerminalKey: + slice_id = req_info.slice_id if meta_type is WriteMetaType.KV else None + return ( + req_info.unique_rid, + meta_type, + req_info.instance_name, + req_info.instance_rank, + slice_id, + ) + + def _send_pre_session_terminal_result( + self, + req_info: RecvReqInfo, + meta_type: WriteMetaType, + message: tuple[bytes, ...], + ) -> None: + """Retain a no-access result even when no TxSession exists yet.""" + key = self._pre_session_terminal_key(meta_type, req_info) + results_lock = self._pre_session_terminal_results_lock + with results_lock: + record = self._pre_session_terminal_results.get(key) + if record is None: + record = UnboundTerminalResult(req_info=req_info, message=message) + self._pre_session_terminal_results[key] = record + elif record.message != message: + raise RuntimeError(f"conflicting pre-session terminal result for {key}") + if self._send_operation_message(req_info, message): + with results_lock: + current = self._pre_session_terminal_results.get(key) + if current is record: + current.result_delivered = True + + def _retry_pre_session_terminal_results(self, *, force: bool = False) -> None: + results_lock = getattr(self, "_pre_session_terminal_results_lock", None) + if results_lock is None: + return + retry_lock = getattr(self, "_pre_session_terminal_retry_lock", None) + if retry_lock is None or not retry_lock.acquire(blocking=False): + return + try: + now = time.monotonic() + next_retry_at = getattr(self, "_next_pre_session_terminal_retry_at", 0.0) + if not force and now < next_retry_at: + return + self._next_pre_session_terminal_retry_at = now + self._TERMINAL_RESULT_RETRY_INTERVAL_S + with results_lock: + pending = [ + (key, record) + for key, record in self._pre_session_terminal_results.items() + if not record.result_delivered + ] + for key, record in pending: + if self._send_operation_message(record.req_info, record.message): + with results_lock: + current = self._pre_session_terminal_results.get(key) + if current is record: + current.result_delivered = True + finally: + retry_lock.release() + + def _has_pending_pre_session_terminal_results(self) -> bool: + results_lock = getattr(self, "_pre_session_terminal_results_lock", None) + if results_lock is None: + return False + with results_lock: + return any( + not record.result_delivered + for record in self._pre_session_terminal_results.values() + ) + + def retry_terminal_results(self, session: "TxSession", *, force: bool = False) -> None: + """Retry each pending terminal frame at most once per polling interval.""" + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_dealer_admission_closed", False): + return + retry_lock = getattr(session, "_terminal_retry_lock", None) + if retry_lock is None or not retry_lock.acquire(blocking=False): + return + try: + now = time.monotonic() + if not force and now < session._next_terminal_retry_at: + return + session._next_terminal_retry_at = now + self._TERMINAL_RESULT_RETRY_INTERVAL_S + with session.lock: + tasks: list[SendTaskBase] = list(session.kv_tasks) + if session.aux_task is not None: + tasks.append(session.aux_task) + unbound = [ + (key, record.req_info, record.message) + for key, record in session._unbound_terminal_results.items() + if not record.result_delivered + ] + + req_infos = list((self._get_req_info(session.disagg_request_id) or {}).values()) + info_by_operation = { + self._operation_key(req_info): req_info for req_info in req_infos + } + for task in tasks: + for key, message in task.pending_terminal_messages(): + req_info = info_by_operation.get(key) + if req_info is None: + continue + if self._send_operation_message(req_info, message): + task.mark_operation_result_delivered(key) + for key, req_info, message in unbound: + if self._send_operation_message(req_info, message): + session.mark_unbound_terminal_result_delivered(key) + finally: + retry_lock.release() + + def _dispatch_operation( + self, + task: KVSendTask | AuxSendTask, + req_info: RecvReqInfo, + *, + force_no_access: bool = False, + ) -> Optional[Exception]: + """Claim and dispatch one exact task/receiver operation at most once.""" + session = task._session + if session is None: + # Compatibility for standalone callers. Production tasks are bound + # to an exact TxSession and use the operation ledger below. + try: + trans_meta = ( + self._build_kv_write_meta(task, req_info) + if isinstance(task, KVSendTask) + else self._build_aux_write_meta(task, req_info) + ) + self._enqueue_owned(trans_meta) + return None + except Exception as e: + task.fail(e) + self._send_operation_message( + req_info, self._failed_operation_message(task, req_info) + ) + return e + + key = self._operation_key(req_info) + failed_message = self._failed_operation_message(task, req_info) + with session.lock: + if isinstance(task, KVSendTask): + belongs_to_session = ( + 0 <= task.slice_id < len(session.kv_tasks) + and session.kv_tasks[task.slice_id] is task + ) + else: + belongs_to_session = session.aux_task is task + if not belongs_to_session: + return RuntimeError( + f"sender task no longer belongs to session {session.disagg_request_id}" + ) + allow_source_access = ( + not force_no_access + and getattr(session, "_accepting_operations", True) + and not getattr(session, "_closed", False) + and session.status not in (SessionStatus.ERROR, SessionStatus.CANCELLED) + ) + newly_recorded, state, cached_message, _result_delivered = task.admit_operation( + key, + allow_source_access=allow_source_access, + no_access_message=failed_message, + ) + + if not newly_recorded: + # A duplicate REQUEST_DATA is also an acknowledgement retry. Replay + # the exact cached decision even if an earlier send returned + # successfully: DEALER send completion is not remote receipt. + if ( + state is SendOperationState.TERMINAL + and cached_message is not None + and self._send_operation_message(req_info, cached_message) + ): + task.mark_operation_result_delivered(key) + return None + if state is SendOperationState.TERMINAL: + assert cached_message is not None + if self._send_operation_message(req_info, cached_message): + task.mark_operation_result_delivered(key) + return None + + trans_meta = None + try: + if task._perf_timer is not None: + task._perf_timer.record_task_start(req_info.instance_rank) + trans_meta = ( + self._build_kv_write_meta(task, req_info) + if isinstance(task, KVSendTask) + else self._build_aux_write_meta(task, req_info) + ) + trans_meta.operation_key = key + trans_meta.source_access_enrolled = True + if task._perf_timer is not None: + task._perf_timer.record_push_start(trans_meta.peer_rank) + self._enqueue_owned(trans_meta) + return None + except Exception as e: + # Descriptor construction and queue admission have no asynchronous + # source accessor. Settle this exact writer and continue dispatching + # the rest of the cohort. + task.fail(e) + try: + task.cache_terminal_message(key, failed_message) + finally: + task.finish_source_access() + if trans_meta is not None: + trans_meta.source_access_enrolled = False + if self._send_operation_message(req_info, failed_message): + task.mark_operation_result_delivered(key) + return e + + def dispatch_task( + self, + task: KVSendTask | AuxSendTask, + req_info_snapshot: Optional[dict] = None, + ): + # req_info_snapshot may be pre-fetched under session.lock by the caller to keep the + # critical section small. When not provided, we fetch it here (legacy / standalone path). + if req_info_snapshot is None: + req_info_snapshot = dict(self._get_req_info(task._unique_rid) or {}) + infos = list(req_info_snapshot.values()) + if isinstance(task, KVSendTask): + infos = [ + info for info in infos if info.slice_id is None or info.slice_id == task.slice_id + ] + else: + # AUX is session-level. Multiple KV slices can carry the same AUX + # target, so dispatch at most once per receiver rank. + by_rank = {} + for info in infos: + by_rank.setdefault(info.instance_rank, info) + infos = list(by_rank.values()) + errors = [] + for info in infos: + error = self._dispatch_operation(task, info) + if error is not None: + errors.append(error) + if errors: + raise errors[0] def _start_listener(self): def handle_message(messages: list[bytes]): @@ -1005,51 +1819,203 @@ def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): ) def _handle_cancel_session(self, message: list[bytes]): - unique_rid = int(message[1]) - session = None - with self._sessions_lock: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_dealer_admission_closed", False): + return + unique_rid = int(message[1]) + with self._sessions_lock: + # Keep the tombstone durable. A cancel can arrive before both the + # TxSession and REQUEST_DATA; either late arrival must still close + # without accessing the published destination. self._pre_cancelled_rids.add(unique_rid) - else: - session = session_ref() + session = self._sessions.get(unique_rid) if session is None: - self._pre_cancelled_rids.add(unique_rid) - if session is not None: - session.cancel() + with self._peer_requests_lock: + req_infos = list(self._peer_requests.get(unique_rid, {}).values()) + else: + req_infos = [] + if session is not None: + session.cancel() + return + aux_receivers: set[SendOperationKey] = set() + for info in req_infos: + self._send_failed_result_to_receiver(info) + operation_key = self._operation_key(info) + if info.aux_slot is not None and operation_key not in aux_receivers: + self._send_aux_failed_result_to_receiver(info) + aux_receivers.add(operation_key) + + def cancel_session(self, session: "TxSession") -> None: + """Durably seal a local sender session and settle every unused channel. + + This is invoked for both local and remote cancellation. Installing the + tombstone before result/control I/O makes a later REQUEST_DATA replay + the same no-access decision after the TxSession has been removed. + """ + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_shutdown_complete", False) or getattr( + self, "_dealer_admission_closed", False + ): + return + unique_rid = session.disagg_request_id + with self._sessions_lock: + self._pre_cancelled_rids.add(unique_rid) + with self._peer_requests_lock: + req_infos = list(self._peer_requests.get(unique_rid, {}).values()) + + aux_receivers: set[SendOperationKey] = set() + for info in req_infos: + settle_aux = self._operation_key(info) not in aux_receivers + self._settle_cancelled_info(session, info, settle_aux=settle_aux) + if info.aux_slot is not None: + aux_receivers.add(self._operation_key(info)) + self.send_cancel_to_receivers(unique_rid) + + def _settle_cancelled_info( + self, + session: "TxSession", + info: RecvReqInfo, + *, + settle_aux: bool, + force_no_access: bool = False, + ) -> None: + """Replay terminal state or emit no-access for one advertised target.""" + with session.lock: + kv_tasks = [ + task + for task in session.kv_tasks + if info.slice_id is None or task.slice_id == info.slice_id + ] + aux_task = session.aux_task + + if kv_tasks: + for task in kv_tasks: + error = self._dispatch_operation(task, info, force_no_access=force_no_access) + if error is not None: + logger.error( + f"failed to settle cancelled KV operation for rid={info.unique_rid}: " + f"{error}" + ) + else: + slice_id = info.slice_id if info.slice_id is not None else 0 + self._send_unbound_terminal_result( + session, + info, + WriteMetaType.KV, + tuple( + _make_kv_result_msg( + self._instance_rank, + info.unique_rid, + slice_id, + True, + AgentResult.FAILED, + ) + ), + ) + + if settle_aux and info.aux_slot is not None: + if aux_task is None: + self._send_unbound_terminal_result( + session, + info, + WriteMetaType.AUX, + ( + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(info.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ), + ) + else: + error = self._dispatch_operation(aux_task, info, force_no_access=force_no_access) + if error is not None: + logger.error( + f"failed to settle cancelled AUX operation for rid={info.unique_rid}: " + f"{error}" + ) @nvtx_range("_respond_with_kv") def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_dealer_admission_closed", False): + return + self._respond_with_kv_admitted(_send_id, message) + + def _respond_with_kv_admitted(self, _send_id: bytes, message: list[bytes]): # _sessions_lock prevents a race between session lookup and req_info save. # session.lock serializes _enqueue calls from both paths. info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) with self._sessions_lock: + shutting_down = getattr(self, "_shutdown", False) + if shutting_down: + # A request admitted after shutdown began is a durable no-access + # decision. Keep the tombstone even if no local session exists. + self._pre_cancelled_rids.add(info.unique_rid) + pre_cancelled = info.unique_rid in self._pre_cancelled_rids session = self._get_session(info.unique_rid) - if session is None: + if shutting_down: + try: + self._save_peer_req_info(info) + except Exception as e: + # Readiness computation is irrelevant after admission has + # closed. Preserve the address-bearing request so the + # durable no-access result can still be retried. + self._add_req_info(info.unique_rid, info) + logger.warning( + f"Sender shutdown could not compute readiness for late " + f"request {info.unique_rid}: {e}" + ) + else: self._save_peer_req_info(info) - return - with session.lock: - self._save_peer_req_info(info) - tasks = list(session.kv_tasks) - # No tasks: no worker will send KV_AGENT_RESULT FAILED to the receiver. - # Send it directly to unblock the receiver's TRANSFERRING task future; - # CANCEL_SESSION alone would leave it stuck indefinitely. - if not tasks and session.status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + if session is None: + if pre_cancelled: self._send_failed_result_to_receiver(info) - return + if info.aux_slot is not None: + self._send_aux_failed_result_to_receiver(info) + return + if shutting_down: + self._settle_cancelled_info( + session, + info, + settle_aux=True, + force_no_access=True, + ) + return + if pre_cancelled: + # setup_session() normally performs this transition. Handle the + # constructor/listener race directly so no operation can slip in + # between tombstone observation and session cancellation. + session.cancel() + return + with session.lock: + tasks = [ + task + for task in session.kv_tasks + if info.slice_id is None or task.slice_id == info.slice_id + ] + terminal = session.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + aux_task = session.aux_task + if terminal: + self._settle_cancelled_info(session, info, settle_aux=True) + return for task in tasks: - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - trans_meta = self._build_kv_write_meta(task, info) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) + error = self._dispatch_operation(task, info) + if error is not None: + logger.error(f"failed to dispatch KV operation for rid={info.unique_rid}: {error}") + if aux_task is not None and info.aux_slot is not None: + error = self._dispatch_operation(aux_task, info) + if error is not None: + logger.error(f"failed to dispatch AUX operation for rid={info.unique_rid}: {error}") def _send_failed_result_to_receiver(self, info: RecvReqInfo): - try: - peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) - slice_id = info.slice_id if info.slice_id is not None else 0 - self._get_or_connect_dealer(peer_ri.self_endpoint).send( + slice_id = info.slice_id if info.slice_id is not None else 0 + self._send_pre_session_terminal_result( + info, + WriteMetaType.KV, + tuple( _make_kv_result_msg( self._instance_rank, info.unique_rid, @@ -1057,25 +2023,45 @@ def _send_failed_result_to_receiver(self, info: RecvReqInfo): True, # is_last_slice AgentResult.FAILED, ) - ) - except Exception as e: - logger.warning( - f"_respond_with_kv: failed to abort receiver for rid={info.unique_rid}: {e}" - ) + ), + ) + + def _send_aux_failed_result_to_receiver(self, info: RecvReqInfo) -> None: + """Report a terminal AUX failure when cancellation prevented access.""" + self._send_pre_session_terminal_result( + info, + WriteMetaType.AUX, + ( + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(info.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ), + ) def _get_or_connect_dealer(self, endpoint: Optional[str]): - if endpoint is None: - raise ValueError("Sender: peer endpoint is None; peer may not have registered yet") - if endpoint not in self._dealers: - self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) - return self._dealers[endpoint] + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_dealer_admission_closed", False): + raise RuntimeError("Sender control dealers are closed for shutdown") + if endpoint is None: + raise ValueError("Sender: peer endpoint is None; peer may not have registered yet") + lock = getattr(self, "_dealers_lock", None) + if lock is None: + if endpoint not in self._dealers: + self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) + return self._dealers[endpoint] + with lock: + if endpoint not in self._dealers: + self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) + return self._dealers[endpoint] def _save_peer_req_info(self, peer_transfer_req_info: RecvReqInfo): req_info = peer_transfer_req_info - self._add_req_info(req_info.unique_rid, req_info.instance_rank, req_info) + self._add_req_info(req_info.unique_rid, req_info) peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) expected_transfers = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) - if self._is_req_ready(req_info.unique_rid, expected_transfers): + if self._is_req_ready(req_info.unique_rid, expected_transfers, req_info.slice_id): session = self._get_session(req_info.unique_rid) if session is not None and not session.receiver_ready: session.receiver_ready = True @@ -1089,63 +2075,251 @@ def has_all_peer_req_infos(self, unique_rid: int) -> bool: def _has_all_peer_req_infos(self, req_info: RecvReqInfo) -> bool: peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) expected_transfers = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) - return self._is_req_ready(req_info.unique_rid, expected_transfers) + return self._is_req_ready(req_info.unique_rid, expected_transfers, req_info.slice_id) def clear_session(self, unique_rid: int): - with self._sessions_lock: - if unique_rid in self._sessions: - del self._sessions[unique_rid] - self._remove_req_info(unique_rid) + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + with self._sessions_lock: + if unique_rid in self._sessions: + del self._sessions[unique_rid] + self._remove_req_info(unique_rid) def send_cancel_to_receivers(self, unique_rid: int) -> None: """Notify all receivers involved in this session to cancel.""" - # Snapshot under the lock to avoid RuntimeError if the listener thread - # mutates the dict via _add_req_info() concurrently. - with self._peer_requests_lock: - req_info_map = self._peer_requests.get(unique_rid) - req_infos = list(req_info_map.values()) if req_info_map else [] - for req_info in req_infos: + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self, "_dealer_admission_closed", False): + return + # Snapshot under the lock to avoid RuntimeError if the listener thread + # mutates the dict via _add_req_info() concurrently. + with self._peer_requests_lock: + req_info_map = self._peer_requests.get(unique_rid) + req_infos = list(req_info_map.values()) if req_info_map else [] + for req_info in req_infos: + try: + peer_ri = self._registrar.get_peer_rank_info( + req_info.instance_name, req_info.instance_rank + ) + self._get_or_connect_dealer(peer_ri.self_endpoint).send( + [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] + ) + except Exception as e: + logger.warning(f"send_cancel_to_receivers: failed for rid={unique_rid}: {e}") + + def shutdown(self) -> bool: + """Stop local send work without invalidating resources under live workers.""" + attempt_lock = getattr(self, "_shutdown_attempt_lock", None) + if attempt_lock is None: + attempt_lock = threading.Lock() + self._shutdown_attempt_lock = attempt_lock + with attempt_lock: + return self._shutdown_locked() + + def _shutdown_locked(self) -> bool: + if getattr(self, "_shutdown_complete", False): + return True + cancel_failed_session_ids: set[int] = set() + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + # TxSession creation/send/cancel shares this gate, so shutdown seals + # every registered session before listener and worker teardown. + self._shutdown = True + sessions_lock = getattr(self, "_sessions_lock", None) + if sessions_lock is None: + sessions = [] + else: + with sessions_lock: + sessions = list(self._sessions.values()) + for session in sessions: + try: + session.cancel() + except Exception as e: + cancel_failed_session_ids.add(id(session)) + logger.error( + f"Sender.shutdown: failed to seal TxSession " + f"{session.disagg_request_id}: {e}" + ) + if not getattr(self, "_listener_stopped", False): + # Quiesce listener before invalidate to avoid set/map mutation + # races. Do not advance teardown if stop fails; retry must call it + # again rather than treating the listener as quiescent. try: - peer_ri = self._registrar.get_peer_rank_info( - req_info.instance_name, req_info.instance_rank - ) - self._get_or_connect_dealer(peer_ri.self_endpoint).send( - [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] - ) + self._messenger.stop() except Exception as e: - logger.warning(f"send_cancel_to_receivers: failed for rid={unique_rid}: {e}") - - def shutdown(self): - if self._shutdown: - return - self._shutdown = True - - # Quiesce listener before invalidate to avoid set/map mutation races. - self._messenger.stop() - - for q in self._send_task_queues: - q.put(None) + logger.error(f"Sender.shutdown: listener stop failed: {e}") + return False + self._listener_stopped = True + if not getattr(self, "_shutdown_sentinels_sent", False): + admission_lock = getattr(self, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + for q in self._send_task_queues: + q.put(None) + self._shutdown_sentinels_sent = True for t in self._worker_threads: t.join(timeout=5) + workers_alive = any(t.is_alive() for t in self._worker_threads) + progress_failed = bool(cancel_failed_session_ids) or workers_alive + if workers_alive: + logger.error( + "Sender.shutdown: worker threads are still active; retaining remote agents" + ) - # Snapshot under lock as defense in depth. + # Worker-local DEALER sockets are detached by their owning threads. + # A failed close is transferred to this Sender and remains here until + # a retry positively closes it. + failed_thread_lock = getattr(self, "_failed_thread_dealers_lock", None) + if failed_thread_lock is None: + failed_thread_dealers = list(getattr(self, "_failed_thread_dealers", ())) + else: + with failed_thread_lock: + failed_thread_dealers = list(self._failed_thread_dealers) + thread_dealer_failed = False + for dealer in failed_thread_dealers: + try: + dealer.stop() + if failed_thread_lock is None: + self._failed_thread_dealers = [ + item for item in self._failed_thread_dealers if item is not dealer + ] + else: + with failed_thread_lock: + self._failed_thread_dealers = [ + item for item in self._failed_thread_dealers if item is not dealer + ] + except Exception as e: + thread_dealer_failed = True + logger.warning(f"Failed to retry worker-local dealer shutdown: {e}") + progress_failed = progress_failed or thread_dealer_failed + + in_doubt_lock = getattr(self, "_in_doubt_transfers_lock", None) + in_doubt_count = 0 + if in_doubt_lock is not None: + with in_doubt_lock: + in_doubt_count = len(self._in_doubt_transfers) + if in_doubt_count: + logger.critical( + f"Sender.shutdown: retaining agent and source owners for " + f"{in_doubt_count} in-doubt transfer(s)" + ) + progress_failed = True + with admission_lock if admission_lock is not None else nullcontext(): + sessions_lock = getattr(self, "_sessions_lock", None) + if sessions_lock is None: + sessions = [] + else: + with sessions_lock: + sessions = list(self._sessions.values()) + retry_failed_session_ids: set[int] = set() + for session in sessions: + try: + self.retry_terminal_results(session, force=True) + except Exception as e: + retry_failed_session_ids.add(id(session)) + progress_failed = True + logger.error( + f"Sender.shutdown: failed to retry terminal results for " + f"TxSession {session.disagg_request_id}: {e}" + ) + try: + self._retry_pre_session_terminal_results(force=True) + except Exception as e: + progress_failed = True + logger.error(f"Sender.shutdown: failed to retry pre-session terminal results: {e}") + active_session_ids: set[int] = set() + for session in sessions: + try: + active = session.has_transferring_tasks() + except Exception as e: + active = True + logger.error( + f"Sender.shutdown: failed to inspect TxSession " + f"{session.disagg_request_id}: {e}" + ) + if active: + active_session_ids.add(id(session)) + progress_failed = True + logger.critical( + f"Sender.shutdown: retaining agent for active source request " + f"{session.disagg_request_id}" + ) + pre_session_results_pending = self._has_pending_pre_session_terminal_results() + if pre_session_results_pending: + progress_failed = True + logger.critical( + "Sender.shutdown: retaining control dealers for pending pre-session results" + ) + for session in sessions: + session_id = id(session) + if session_id in ( + cancel_failed_session_ids | retry_failed_session_ids | active_session_ids + ): + continue + try: + session.close() + except Exception as e: + progress_failed = True + logger.error( + f"Sender.shutdown: failed to retire TxSession " + f"{session.disagg_request_id}: {e}" + ) + continue + # TxSession.close() normally unregisters itself. Keep shutdown + # aggregation correct for compatible session implementations + # that only release their own resources. + self.clear_session(session.disagg_request_id) + if not progress_failed: + # No later listener callback or session mutation can create + # another control socket/result obligation after this gate. + self._dealer_admission_closed = True + + if progress_failed: + return False + + invalidation_failed = False with self._loaded_remote_agents_lock: loaded_agents = list(self._loaded_remote_agents) - self._loaded_remote_agents.clear() # Invalidate all loaded remote agents to release fabric/POSIX FD resources. for agent_name in loaded_agents: try: self._agent.invalidate_remote_agent(agent_name) + with self._loaded_remote_agents_lock: + self._loaded_remote_agents.discard(agent_name) except Exception as e: + invalidation_failed = True logger.warning( f"Failed to invalidate remote agent '{agent_name}' during shutdown: {e}" ) - for dealer in self._dealers.values(): + if invalidation_failed: + return False + dealer_failed = False + dealers_lock = getattr(self, "_dealers_lock", None) + if dealers_lock is None: + dealer_items = list(self._dealers.items()) + else: + with dealers_lock: + dealer_items = list(self._dealers.items()) + for endpoint, dealer in dealer_items: try: dealer.stop() + if dealers_lock is None: + if self._dealers.get(endpoint) is dealer: + self._dealers.pop(endpoint, None) + else: + dealer_failed = True + else: + with dealers_lock: + if self._dealers.get(endpoint) is dealer: + self._dealers.pop(endpoint, None) + else: + dealer_failed = True except Exception as e: + dealer_failed = True logger.warning(f"Failed to stop dealer during Sender shutdown: {e}") - self._dealers.clear() + if dealer_failed: + return False + self._shutdown_complete = True + return True def __del__(self): try: @@ -1156,8 +2330,19 @@ def __del__(self): def __enter__(self): return self - def __exit__(self, _exc_type, _exc_val, _exc_tb): - self.shutdown() + def __exit__(self, exc_type, _exc_val, _exc_tb): + try: + drained = self.shutdown() + except Exception as e: + if exc_type is None: + raise + logger.error(f"Sender context cleanup failed while propagating an exception: {e}") + return False + if not drained: + if exc_type is None: + raise RuntimeError("Sender context exited before transfer resources drained") + logger.error("Sender context retained resources while propagating an exception") + return False class TxSession(TxSessionBase): @@ -1170,28 +2355,56 @@ def __init__( timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, beam_width: int = 1, + source_owner: Optional[LlmRequest] = None, ): - super().__init__( - sender, - SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), - ) - self._timeout_s = timeout_s - self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST - self._sender: Sender # narrow base class type for Pylance - self.request_id = request_id + # Keep the allocator-level request owner alive until every source + # operation is terminal. Enroll it before fallible initialization, + # then clear the parameter so propagated tracebacks are ownership- + # neutral. + self._source_owner = source_owner + source_owner = None self._aux_buffer = aux_buffer - self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None - self.receiver_ready: bool = False - self.kv_tasks = [] - self.aux_task = None - self.lock = threading.Lock() - - self._exception: Optional[Exception] = None - self._closed = False - self._terminal_status: Optional[SessionStatus] = None - # Must be last: makes session visible to listener thread, - # so all attributes above must be initialized first. - self._sender.setup_session(self) + self.aux_slot = None + try: + super().__init__( + sender, + SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), + ) + self._timeout_s = timeout_s + self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST + self._sender: Sender # narrow base class type for Pylance + self.request_id = request_id + self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None + # AUX contents become immutable once an AuxSendTask is published. + # This prevents fill_slot() from racing a worker/NIXL read of the + # same slot. + self._aux_packed = False + self.receiver_ready: bool = False + self.kv_tasks = [] + self.aux_task = None + self.lock = threading.Lock() + self._close_lock = threading.Lock() + + self._exception: Optional[Exception] = None + self._closed = False + self._accepting_operations = True + self._terminal_status: Optional[SessionStatus] = None + # Cancellation can precede task creation. Keep those no-access + # frames in the session lifecycle so a transient send failure + # cannot be discarded by close(). + self._unbound_terminal_results: dict[UnboundTerminalKey, UnboundTerminalResult] = {} + self._terminal_retry_lock = threading.Lock() + self._next_terminal_retry_at = 0.0 + # Must be last: makes session visible to listener thread, so all + # attributes above must be initialized first. + self._sender.setup_session(self) + except Exception: + self._source_owner = None + if self._aux_buffer is not None and self.aux_slot is not None: + self._aux_buffer.free_slot(self.aux_slot) + self.aux_slot = None + self._closed = True + raise @property def disagg_request_id(self) -> int: @@ -1205,6 +2418,48 @@ def disagg_request_id(self) -> int: return params.ctx_request_id return self.request_id + @property + def source_owner(self) -> Optional[LlmRequest]: + return self._source_owner + + @staticmethod + def _unbound_terminal_key( + meta_type: WriteMetaType, req_info: RecvReqInfo + ) -> UnboundTerminalKey: + slice_id = req_info.slice_id if meta_type is WriteMetaType.KV else None + return meta_type, req_info.instance_name, req_info.instance_rank, slice_id + + def cache_unbound_terminal_result( + self, + meta_type: WriteMetaType, + req_info: RecvReqInfo, + message: tuple[bytes, ...], + ) -> UnboundTerminalKey: + key = self._unbound_terminal_key(meta_type, req_info) + with self.lock: + record = self._unbound_terminal_results.get(key) + if record is None: + self._unbound_terminal_results[key] = UnboundTerminalResult( + req_info=req_info, + message=message, + ) + elif record.message != message: + raise RuntimeError(f"conflicting unbound terminal result for {key}") + return key + + def mark_unbound_terminal_result_delivered(self, key: UnboundTerminalKey) -> None: + with self.lock: + record = self._unbound_terminal_results.get(key) + if record is None: + raise RuntimeError(f"unbound terminal result {key} was not recorded") + record.result_delivered = True + + def _retry_terminal_results(self) -> None: + sender = getattr(self, "_sender", None) + retry = getattr(sender, "retry_terminal_results", None) + if callable(retry): + retry(self) + @property def status(self) -> SessionStatus: if self._terminal_status is not None: @@ -1221,39 +2476,91 @@ def status(self) -> SessionStatus: return SessionStatus.READY if self.receiver_ready else SessionStatus.INIT def send(self, slice: KVSlice) -> None: - with self.lock: - params = self._base_args.params - slice_id = len(self.kv_tasks) - task = KVSendTask( - slice, - params, - slice_id, - prompt_len=self._base_args.prompt_len, - beam_width=self._base_args.beam_width, - ) - task._unique_rid = self.disagg_request_id - self.kv_tasks.append(task) - req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) - self._sender.dispatch_task(task, req_info_snapshot) + admission_lock = getattr(self._sender, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self._sender, "_shutdown", False): + raise RuntimeError("Sender is shutting down; new KV operations are not accepted") + with self.lock: + if not self._accepting_operations or self._closed: + raise RuntimeError( + f"TxSession {self.disagg_request_id} is closed to new KV operations" + ) + params = self._base_args.params + slice_id = len(self.kv_tasks) + task = KVSendTask( + slice, + params, + slice_id, + prompt_len=self._base_args.prompt_len, + beam_width=self._base_args.beam_width, + session=self, + ) + task._unique_rid = self.disagg_request_id + self.kv_tasks.append(task) + req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) + self._sender.dispatch_task(task, req_info_snapshot) def send_aux(self) -> AuxSendTask: - with self.lock: - params = self._base_args.params - task = AuxSendTask(params, self.aux_slot) - task._unique_rid = self.disagg_request_id - self.aux_task = task - req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) - self._sender.dispatch_task(task, req_info_snapshot) - return task + admission_lock = getattr(self._sender, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + if getattr(self._sender, "_shutdown", False): + raise RuntimeError("Sender is shutting down; new AUX operations are not accepted") + with self.lock: + if not self._accepting_operations or self._closed: + raise RuntimeError( + f"TxSession {self.disagg_request_id} is closed to new AUX operations" + ) + if self.aux_task is not None: + raise RuntimeError( + f"TxSession {self.disagg_request_id} already has an AUX operation" + ) + if not getattr(self, "_aux_packed", False): + raise RuntimeError( + f"TxSession {self.disagg_request_id} AUX data must be packed before send" + ) + params = self._base_args.params + task = AuxSendTask(params, self.aux_slot, session=self) + task._unique_rid = self.disagg_request_id + self.aux_task = task + req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) + self._sender.dispatch_task(task, req_info_snapshot) + return task def pack_aux(self, request: LlmRequest) -> None: """Fill the aux buffer slot with token data from the given request.""" - assert self._aux_buffer is not None, "No aux_buffer set for this session" - assert self.aux_slot is not None, "No aux_slot set for this session" - self._aux_buffer.fill_slot(self.aux_slot, request) + admission_lock = getattr(self._sender, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + with self.lock: + if ( + not self._accepting_operations + or self._closed + or self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + or getattr(self._sender, "_shutdown", False) + ): + raise RuntimeError( + f"TxSession {self.disagg_request_id} is closed to AUX packing" + ) + if self.aux_task is not None: + raise RuntimeError( + f"TxSession {self.disagg_request_id} AUX data is frozen after send" + ) + if getattr(self, "_aux_packed", False): + raise RuntimeError( + f"TxSession {self.disagg_request_id} AUX data is already packed" + ) + aux_buffer = self._aux_buffer + aux_slot = self.aux_slot + assert aux_buffer is not None, "No aux_buffer set for this session" + assert aux_slot is not None, "No aux_slot set for this session" + aux_buffer.fill_slot(aux_slot, request) + with self.lock: + self._aux_packed = True def is_completed(self) -> bool: """Non-blocking check: has the transfer completed successfully?""" + self._retry_terminal_results() + if self.has_transferring_tasks(): + return False status = self.status if self._need_aux: return status == SessionStatus.FULLY_TRANSFERRED @@ -1261,6 +2568,11 @@ def is_completed(self) -> bool: def has_failed(self) -> bool: """Non-blocking check: has the transfer failed or been cancelled?""" + self._retry_terminal_results() + # A logical failure is not safe request-cleanup evidence while a local + # gather/NIXL worker may still read source KV pages. + if self.has_transferring_tasks(): + return False if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED): return True if any(task.status == TaskStatus.ERROR for task in self.kv_tasks): @@ -1270,30 +2582,53 @@ def has_failed(self) -> bool: def cancel(self) -> None: """Cancel the session and notify the remote receiver. - Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). - Only INIT tasks have their events signalled immediately. + Safe to call multiple times. Active physical work keeps running, but + every unfinished task latches logical failure immediately. Source + access and pending-result ledgers remain authoritative for retirement. The lock serializes with _deliver_kv_to_agent() so has_transferring_tasks() is accurate the moment this returns. """ - with self.lock: - if self._terminal_status == SessionStatus.CANCELLED: - return - self._terminal_status = SessionStatus.CANCELLED - exc = RuntimeError(f"TxSession {self.disagg_request_id} cancelled") - for task in self.kv_tasks: - if task.status == TaskStatus.INIT: - task.fail(exc) - if self.aux_task is not None and self.aux_task.status == TaskStatus.INIT: - self.aux_task.fail(exc) - # Send outside the lock to avoid holding it during I/O. - self._sender.send_cancel_to_receivers(self.disagg_request_id) + admission_lock = getattr(self._sender, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + self._latch_terminal_failure( + SessionStatus.CANCELLED, + RuntimeError(f"TxSession {self.disagg_request_id} cancelled"), + record_exception=False, + ) + # Tombstone, settle unused channels, and perform I/O outside the + # session lock. Repeated cancel retries cached terminal decisions. + self._sender.cancel_session(self) def has_transferring_tasks(self) -> bool: - """True if any KV task is currently mid-write (TRANSFERRING). + """True while source access or terminal-result delivery is outstanding. cancel_request() must return False while this is True. """ - return any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks) + bound_pending = any( + getattr(t, "source_access_active", False) + or t.status == TaskStatus.TRANSFERRING + or getattr(t, "has_pending_result_delivery", False) + for t in self.kv_tasks + ) or ( + self.aux_task is not None + and ( + getattr(self.aux_task, "source_access_active", False) + or self.aux_task.status == TaskStatus.TRANSFERRING + or getattr(self.aux_task, "has_pending_result_delivery", False) + ) + ) + unbound_results = getattr(self, "_unbound_terminal_results", {}) + lock = getattr(self, "lock", None) + if lock is None: + unbound_pending = any( + not record.result_delivered for record in unbound_results.values() + ) + else: + with lock: + unbound_pending = any( + not record.result_delivered for record in unbound_results.values() + ) + return bound_pending or unbound_pending def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: """Poll or block until KV (and optionally aux) transfer finishes. @@ -1302,11 +2637,15 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: With blocking=False: polls non-blockingly; returns None if any KV task or aux is not yet done. """ - if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + self._retry_terminal_results() + terminal_failure = self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + if terminal_failure and not self.has_transferring_tasks(): return WaitResult.FAILED if not self.kv_tasks: return None if not blocking: + if self.has_transferring_tasks(): + return None has_pending = False for task in self.kv_tasks: if task.status == TaskStatus.ERROR: @@ -1322,47 +2661,116 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: return WaitResult.FAILED if self.aux_task.status != TaskStatus.TRANSFERRED: return None - return WaitResult.COMPLETED + return WaitResult.FAILED if terminal_failure else WaitResult.COMPLETED for task in self.kv_tasks: if not task.wait(timeout=self._timeout_s): return WaitResult.TIMEOUT if task.status == TaskStatus.ERROR: - return WaitResult.FAILED + return None if self.has_transferring_tasks() else WaitResult.FAILED if self._need_aux and self.aux_task is not None: if not self.aux_task.wait(timeout=self._timeout_s): return WaitResult.TIMEOUT if self.aux_task.status == TaskStatus.ERROR: - return WaitResult.FAILED - return WaitResult.COMPLETED + return None if self.has_transferring_tasks() else WaitResult.FAILED + if self.has_transferring_tasks(): + return None + return WaitResult.FAILED if terminal_failure else WaitResult.COMPLETED - def set_exception(self, reason: str = ""): - msg = f"TxSession {self.disagg_request_id} exception" - if reason: - msg += f": {reason}" + def _latch_terminal_failure( + self, + status: SessionStatus, + exception: Exception, + *, + record_exception: bool, + ) -> bool: + """Atomically publish the first terminal failure and fail open tasks.""" + if status not in (SessionStatus.ERROR, SessionStatus.CANCELLED): + raise ValueError(f"invalid terminal failure status {status}") with self.lock: - self._exception = RuntimeError(msg) - self._terminal_status = SessionStatus.ERROR + self._accepting_operations = False + if self._terminal_status is not None: + return False + self._terminal_status = status + if record_exception: + self._exception = exception for task in self.kv_tasks: if not task.is_done: - task.fail(self._exception) + task.fail(exception) if self.aux_task is not None and not self.aux_task.is_done: - self.aux_task.fail(self._exception) + self.aux_task.fail(exception) + return True + + def set_exception(self, reason: str = ""): + msg = f"TxSession {self.disagg_request_id} exception" + if reason: + msg += f": {reason}" + sender = getattr(self, "_sender", None) + admission_lock = getattr(sender, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + self._latch_terminal_failure( + SessionStatus.ERROR, + RuntimeError(msg), + record_exception=True, + ) + # ERROR seals admission just like cancellation. Install the durable + # tombstone outside the session lock so late REQUEST_DATA is settled. + if sender is not None: + sender.cancel_session(self) @property def exception(self) -> Optional[Exception]: return self._exception def close(self): + # Constructor rollback marks the partially initialized object closed. + # Check that terminal marker before touching locks or retry ledgers that + # may not have been created yet; __del__ can run after any init failure. + if getattr(self, "_closed", False): + return + close_lock = getattr(self, "_close_lock", None) + if close_lock is None: + close_lock = threading.Lock() + self._close_lock = close_lock + sender = getattr(self, "_sender", None) + admission_lock = getattr(sender, "_operation_admission_lock", None) + with admission_lock if admission_lock is not None else nullcontext(): + with close_lock: + self._close_locked() + + def _close_locked(self): if getattr(self, "_closed", False): return - self._closed = True + self._retry_terminal_results() + with self.lock: + if getattr(self, "_closed", False): + return + self._accepting_operations = False + source_tasks = list(self.kv_tasks) + if self.aux_task is not None: + source_tasks.append(self.aux_task) + if any( + getattr(task, "source_access_active", False) + or task.status in (TaskStatus.INIT, TaskStatus.TRANSFERRING) + or getattr(task, "has_pending_result_delivery", False) + for task in source_tasks + ) or any( + not record.result_delivered + for record in getattr(self, "_unbound_terminal_results", {}).values() + ): + raise RuntimeError( + f"cannot close TxSession {self.disagg_request_id}: " + "source work is pending or active" + ) if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None # Unregister from Sender; keep fields alive for in-flight worker threads. if self._sender is not None: self._sender.clear_session(self.disagg_request_id) + with self.lock: + self._source_owner = None + self._closed = True def __enter__(self): return self @@ -1391,6 +2799,24 @@ def __init__( self.status = TaskStatus.INIT self.expected_transfers = 0 self.last_slice_count = 0 + # The exact-writer registry is enabled only when ctx_dp_rank selects one + # deterministic sender cohort. ADP broadcast remains on the legacy path + # until the protocol selects its writer cohort before address publication. + self.lifecycle_managed = False + # Legacy ADP broadcast cannot name the selected writer cohort before + # publication. Keep a conservative endpoint-local drain record so a + # first failure cannot release destination memory while sibling writers + # may still be active. + self.publication_started = False + self._legacy_results: dict[int, AgentResult] = {} + self._legacy_result_conflict = False + self._legacy_backend_quiesced = False + # One exact set for deterministic ctx_dp_rank, or one candidate set per + # DP group for ADP broadcast. A legacy target is drained only when the + # observed terminal identities equal one complete candidate cohort. + self._valid_writer_cohorts: tuple[frozenset[int], ...] = () + self._exposed_writer_ranks: set[int] = set() + self._publication_closed = False self._unique_rid = unique_rid self._kv_slice = kv_slice @@ -1416,6 +2842,79 @@ def wait(self, timeout: Optional[float] = None) -> bool: def is_done(self) -> bool: return self._event.is_set() + @property + def legacy_resources_drained(self) -> bool: + """Whether an unregistered ADP receive can no longer be accessed. + + This is a compatibility path, not the exact-writer ownership contract. + It preserves the existing ADP assumption that exactly one advertised + candidate cohort performs the transfer. Full attention-DP containment + requires a negotiated writer-selection protocol before publication. + """ + if self.lifecycle_managed or not self.publication_started or self._legacy_backend_quiesced: + return True + return ( + self._publication_closed + and not self._legacy_result_conflict + and len(self.complete_writer_cohorts(self._legacy_results)) == 1 + ) + + def set_valid_writer_cohorts(self, cohorts) -> None: + """Install the exact writer identities that may form one transfer.""" + normalized: list[frozenset[int]] = [] + for cohort in cohorts: + ranks = frozenset(int(rank) for rank in cohort) + if not ranks: + raise ValueError("writer cohort must not be empty") + if any(rank < 0 for rank in ranks): + raise ValueError("writer cohort ranks must be non-negative") + if ranks not in normalized: + normalized.append(ranks) + if not normalized: + raise ValueError("at least one writer cohort is required") + self._valid_writer_cohorts = tuple(normalized) + self.expected_transfers = len(normalized[0]) + + def writer_identities_valid(self, peer_rank: int, existing_ranks) -> bool: + """Whether ``peer_rank`` was one of the broadcast address recipients.""" + del existing_ranks + return int(peer_rank) in self._exposed_writer_ranks + + def is_complete_writer_cohort(self, ranks) -> bool: + """Whether results identify exactly one candidate cohort.""" + if not self._publication_closed: + return False + identities = frozenset(ranks) + return any(cohort == identities for cohort in self._valid_writer_cohorts) + + def complete_writer_cohorts(self, ranks) -> tuple[frozenset[int], ...]: + """Return the candidate cohort exactly named by the result identities.""" + if not self._publication_closed: + return () + identities = frozenset(ranks) + return tuple(cohort for cohort in self._valid_writer_cohorts if cohort == identities) + + def mark_writer_exposed(self, peer_rank: int) -> None: + """Record that this request may have carried KV and AUX addresses.""" + self.publication_started = True + self._exposed_writer_ranks.add(int(peer_rank)) + + def close_publication(self) -> None: + """Prove that dispatch will advertise no additional writer targets.""" + self._publication_closed = True + + def record_legacy_result(self, peer_rank: int, status: AgentResult) -> bool: + """Record a legacy terminal result; return whether it is a new result.""" + previous = self._legacy_results.get(peer_rank) + if previous is not None: + if previous is not status: + self._legacy_result_conflict = True + return False + self._legacy_results[peer_rank] = status + if not self.writer_identities_valid(peer_rank, self._legacy_results): + self._legacy_result_conflict = True + return True + def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int): if self._perf_timer is None: return @@ -1440,12 +2939,32 @@ def __init__( self._agent = agent self._bounce = bounce self._dealers = {} + self._dealers_lock = threading.Lock() + self._dealer_admission_open = True + self._shutdown_attempt_lock = threading.Lock() self._sender_ep_instance_map = {} self._messenger = ZMQMessenger(mode="ROUTER") - self._sessions = {} # unique_rid -> RxSession + # Strong endpoint ownership is intentional. A receive session carrying + # an AUX/legacy target must not disappear while a remote writer may + # still access it; clear_session removes it only after drain proof. + self._sessions: dict[int, RxSession] = {} self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids self._pre_cancelled_rids: set[int] = set() + self._recv_registry = RecvTransferRegistry() + # Bounce settlement retries must replay the transition that performed + # physical finalization. Calling finish_bounce_scatter() again only + # returns a duplicate update because its one-shot actions have already + # been consumed by the registry. + self._bounce_lifecycle_delivery_lock = threading.Lock() + self._pending_bounce_lifecycle_deliveries: dict[ + tuple[int, int], _PendingLifecycleDelivery + ] = {} + self._pending_bounce_logical_failure_deliveries: dict[ + tuple[int, int], _PendingLifecycleDelivery + ] = {} + self._shutdown_started = False + self._listener_stopped = False self._shutdown = False self._start_listener() @@ -1455,43 +2974,303 @@ def __init__( def endpoint(self): return self._messenger.endpoint - def shutdown(self): + @property + def transfers_drained(self) -> bool: + """Whether all managed, legacy, and AUX receive targets are retired.""" + try: + if not self._bounce.retry_settlements(): + return False + except Exception as e: + logger.error(f"failed to retry receiver bounce settlements: {e}") + return False + if not self._recv_registry.is_drained(): + return False + with self._sessions_lock: + entries = list(self._sessions.values()) + for entry in entries: + session = entry() if isinstance(entry, weakref.ReferenceType) else entry + if session is None: + # A lost legacy/AUX owner is not positive quiescence evidence. + return False + if session.has_untracked_receive_activity(): + return False + return True + + def begin_shutdown(self) -> None: + """Close receive admission while keeping result routing alive. + + TransferWorker calls this before quiescing the backend. Published + writers continue to own their targets until either their terminal + result arrives or backend-wide quiescence is proven. + """ + errors: list[tuple[str, Exception]] = [] + with self._sessions_lock: + # Session enrollment and the shutdown snapshot share this lock, so + # no ADP compatibility session can appear after admission closes. + self._shutdown_started = True + entries = list(self._sessions.values()) + for entry in entries: + session = entry() if isinstance(entry, weakref.ReferenceType) else entry + if session is not None: + try: + seal_admission = getattr(session, "seal_receive_admission", None) + if seal_admission is not None: + seal_admission() + session.cancel() + except Exception as e: + errors.append((f"session {session.disagg_request_id}", e)) + logger.error( + f"Receiver.begin_shutdown: cancellation failed for " + f"request {session.disagg_request_id}: {e}" + ) + try: + updates = self._recv_registry.begin_shutdown() + except Exception as e: + updates = () + errors.append(("receive registry", e)) + logger.error(f"Receiver.begin_shutdown: registry cancellation failed: {e}") + for update in updates: + try: + if update.physical_state is not PhysicalState.DRAINED: + self._bounce.mark_logical_failure( + update.key, + on_done=lambda succeeded, key=update.key: self._finish_bounce( + key, succeeded + ), + ) + except Exception as e: + errors.append((f"receive context {update.key}", e)) + logger.error( + f"Receiver.begin_shutdown: physical cancellation failed for " + f"request {update.key[0]} slice={update.key[1]}: {e}" + ) + try: + self._handle_lifecycle_update(update) + except Exception as e: + errors.append((f"receive notification {update.key}", e)) + logger.error( + f"Receiver.begin_shutdown: lifecycle notification failed for " + f"request {update.key[0]} slice={update.key[1]}: {e}" + ) + if errors: + detail = "; ".join(f"{owner}: {error}" for owner, error in errors) + raise RuntimeError( + f"Receiver.begin_shutdown encountered {len(errors)} error(s): {detail}" + ) from errors[0][1] + + def mark_backend_quiesced(self) -> None: + """Apply an externally proven remote/global writer fence. + + A local ``BaseTransferAgent.shutdown()`` is explicitly insufficient: + callers must first establish that peer processes and the fabric can no + longer submit one-sided writes to any advertised target. + """ + with self._sessions_lock: + sessions = tuple(self._sessions.values()) + for entry in sessions: + session = entry() if isinstance(entry, weakref.ReferenceType) else entry + if session is not None: + session.mark_backend_quiesced() + # Mark registry writers quiescent first. Bounce settlement then calls + # back into that registry, after which request cleanup may proceed. + for update in self._recv_registry.mark_backend_quiesced(): + if LifecycleAction.RELEASE_BOUNCE in update.actions: + self._bounce.mark_backend_quiesced( + update.key, + on_done=lambda succeeded, key=update.key: self._finish_bounce(key, succeeded), + ) + self._handle_lifecycle_update(update) + + def shutdown(self) -> bool: + attempt_lock = getattr(self, "_shutdown_attempt_lock", None) + if attempt_lock is None: + attempt_lock = threading.Lock() + self._shutdown_attempt_lock = attempt_lock + with attempt_lock: + return self._shutdown_serialized() + + def _shutdown_serialized(self) -> bool: if getattr(self, "_shutdown", False): - return - self._shutdown = True - for dealer in self._dealers.values(): + return True + self.begin_shutdown() + if not self.transfers_drained: + return False + # Admission is sealed and every admitted dispatch has exited. Close + # dealer admission and take the final snapshot in one transaction so a + # late control path cannot create an unowned socket after this point. + dealers_lock = getattr(self, "_dealers_lock", None) + if dealers_lock is None: + self._dealer_admission_open = False + dealers = list(self._dealers.items()) + else: + with dealers_lock: + self._dealer_admission_open = False + dealers = list(self._dealers.items()) + # No writer remains, so closing result ingress cannot strand ownership. + if not getattr(self, "_listener_stopped", False): + try: + self._messenger.stop() + except Exception as e: + logger.error(f"Receiver.shutdown: listener stop failed: {e}") + return False + self._listener_stopped = True + failed = False + for endpoint, dealer in dealers: try: dealer.stop() + if dealers_lock is None: + current = self._dealers.get(endpoint) + if current is dealer: + self._dealers.pop(endpoint, None) + elif current is not None: + failed = True + else: + with dealers_lock: + current = self._dealers.get(endpoint) + if current is dealer: + self._dealers.pop(endpoint, None) + elif current is not None: + failed = True except Exception as e: + failed = True logger.warning(f"Failed to stop dealer during Receiver shutdown: {e}") - self._dealers.clear() - self._messenger.stop() + if failed: + return False + self._shutdown = True + return True def clear_session(self, unique_rid: int): with self._sessions_lock: + entry = self._sessions.get(unique_rid) + session = entry() if isinstance(entry, weakref.ReferenceType) else entry + # resources_drained() can settle bounce work and route a callback back + # through Receiver._get_session(). Never call it while holding the + # session-map lock. + if session is not None and not session.resources_drained(): + raise RuntimeError( + f"cannot close RxSession {unique_rid}: " + "legacy or auxiliary receive resources are not drained" + ) + with self._sessions_lock: + if self._sessions.get(unique_rid) is not entry: + raise RuntimeError( + f"cannot close RxSession {unique_rid}: session ownership changed" + ) + if not self._recv_registry.retire_request_if_drained(unique_rid): + raise RuntimeError( + f"cannot close RxSession {unique_rid}: " + "receive transfer resources are not drained" + ) self._sessions.pop(unique_rid, None) def setup_session(self, rx_session: RxSessionBase): + unique_rid = rx_session.disagg_request_id pre_cancel = False with self._sessions_lock: - self._sessions[rx_session.disagg_request_id] = weakref.ref(rx_session) - if rx_session.disagg_request_id in self._pre_cancelled_rids: - pre_cancel = True - self._pre_cancelled_rids.discard(rx_session.disagg_request_id) + if getattr(self, "_shutdown_started", False): + raise RuntimeError("Receiver is shutting down; new sessions are not accepted") + existing = self._sessions.get(unique_rid) + if existing is not None and existing is not rx_session: + raise RuntimeError(f"RxSession {unique_rid} is already registered") + self._sessions[unique_rid] = rx_session + # Without generation-safe replay, cancellation tombstones are + # durable. Applying one to this session must not make a later + # replacement with the same request ID admissible. + pre_cancel = unique_rid in self._pre_cancelled_rids if pre_cancel: - rx_session.cancel() + try: + rx_session.cancel() + except Exception: + # Session enrollment and pre-cancellation are one transaction. + # Constructor rollback can now release pre-publication local + # resources, while the durable tombstone protects a retry. + with self._sessions_lock: + if self._sessions.get(unique_rid) is rx_session: + self._sessions.pop(unique_rid, None) + raise + + def record_cancel_tombstone(self, unique_rid: int) -> None: + """Durably prevent a cancelled request ID from being re-admitted.""" + with self._sessions_lock: + self._pre_cancelled_rids.add(unique_rid) def _get_session(self, unique_rid: Optional[int]) -> Optional["RxSession"]: with self._sessions_lock: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: + entry = self._sessions.get(unique_rid) + if entry is None: return None - session = session_ref() + session = entry() if isinstance(entry, weakref.ReferenceType) else entry if session is None: logger.warning(f"RxSession {unique_rid} has been garbage collected") return None return session + def _handle_lifecycle_update( + self, update: LifecycleUpdate, peer_rank: Optional[int] = None + ) -> None: + """Optionally notify the session after the registry owns the transition.""" + session = self._get_session(update.key[0]) + if session is not None: + session.process_lifecycle_update(update, peer_rank=peer_rank) + + def _finish_bounce( + self, key: tuple[int, int], succeeded: bool, peer_rank: Optional[int] = None + ) -> None: + """Finalize bounce ownership and durably deliver its lifecycle update. + + The bounce transport retries this callback when consumer delivery + raises. The registry transition itself is one-shot, so retain the + original update and replay it until ``process_lifecycle_update`` + returns successfully. + """ + with self._bounce_lifecycle_delivery_lock: + delivery = self._pending_bounce_lifecycle_deliveries.get(key) + if delivery is None: + delivery = _PendingLifecycleDelivery( + update=self._recv_registry.finish_bounce_scatter(key, succeeded), + peer_rank=peer_rank, + succeeded=succeeded, + ) + self._pending_bounce_lifecycle_deliveries[key] = delivery + elif delivery.succeeded != succeeded: + raise RuntimeError(f"conflicting bounce settlement retry for {key}") + + self._handle_lifecycle_update(delivery.update, peer_rank=delivery.peer_rank) + + with self._bounce_lifecycle_delivery_lock: + if self._pending_bounce_lifecycle_deliveries.get(key) is delivery: + self._pending_bounce_lifecycle_deliveries.pop(key, None) + + def _fail_bounce_logically(self, key: tuple[int, int]) -> None: + """Report local scatter failure without claiming physical settlement. + + The transport retains the bounce slot and destination owner because a + CUDA error is not a positive stream fence. Keep the one-shot registry + update replayable until its optional session consumer accepts it. + """ + with self._bounce_lifecycle_delivery_lock: + deliveries = getattr(self, "_pending_bounce_logical_failure_deliveries", None) + if deliveries is None: + deliveries = {} + self._pending_bounce_logical_failure_deliveries = deliveries + delivery = deliveries.get(key) + if delivery is None: + delivery = _PendingLifecycleDelivery( + update=self._recv_registry.fail_context( + key, + "bounce scatter failed without a positive local CUDA fence", + ), + peer_rank=None, + succeeded=False, + ) + deliveries[key] = delivery + + self._handle_lifecycle_update(delivery.update) + + with self._bounce_lifecycle_delivery_lock: + if deliveries.get(key) is delivery: + deliveries.pop(key, None) + def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: self_ri = self._registrar.self_rank_info assert task._params.ctx_request_id is not None, ( @@ -1512,25 +3291,141 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: ) @staticmethod - def _fanin_bounce_safe(overlap, peer_ri) -> bool: - """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. - The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV (should_send_kv) yet still - count in expected_transfers, so the live writers overflow their slots; - * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; - for PP fan-in, every overlapping stage must hold the same number of layers - (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage - list isn't available (shorter than the fan-in degree), be conservative and fall back. - Otherwise fall back to the per-fragment path (correct, just not coalesced). - Equal layer count means equal bytes only when the per-block sizes match; reserve() rejects - the mismatched case, so this only needs the count to split evenly.""" - if overlap.duplicate_head_factor != 1: + def _fanin_bounce_safe(overlap) -> bool: + """Whether equal-size fan-in bounce is safe for this overlap. + + Exact writer ranks do not change the equal-split requirement. TP head + duplication is unsafe in either direction. PP fan-in is always direct: + uniform peer PP-stage layer counts do not prove that this receiver's + local layer interval intersects each writer by an equal number of bytes. + """ + return ( + overlap.overlap_pp_size == 1 + and overlap.duplicate_head_factor == 1 + and overlap.peer_duplicate_head_factor == 1 + ) + + def _single_writer_bounce_exact(self, peer_info: RankInfo) -> bool: + """Whether the receiver can prove that one writer maps every reserved byte. + + ``IdentityMapper`` proves positional mapping, not physical extent + equality. Bounce admission additionally requires identical block + geometry and a byte-equal bijection over every attention pool view so + the sender's coalesced write cannot exceed the receiver-owned slot. + """ + self_ri = self._registrar.self_rank_info + local_layers = getattr(self_ri, "layer_num_per_pp", None) + peer_layers = getattr(peer_info, "layer_num_per_pp", None) + if local_layers is not None and peer_layers is not None: + if sum(local_layers) != sum(peer_layers): + return False + + get_pool_mapping = getattr(self._registrar, "get_pool_mapping", None) + get_kv_map = getattr(self._registrar, "get_kv_map", None) + page_table = getattr(self_ri, "page_table", None) + peer_page_table = getattr(peer_info, "page_table", None) + # If topology metadata is absent, the exact byte contribution is + # unknowable. Fail closed to the direct descriptor path. + if ( + not callable(get_pool_mapping) + or not callable(get_kv_map) + or page_table is None + or peer_page_table is None + or page_table.tokens_per_block != peer_page_table.tokens_per_block + ): return False - if overlap.overlap_pp_size > 1: - lpp = getattr(peer_ri, "layer_num_per_pp", None) - if not lpp or len(lpp) < overlap.overlap_pp_size or len(set(lpp)) != 1: + + from ..resource.utils import get_physical_pool + from .mixers.attention.peer import IdentityMapper + + try: + mapping = get_pool_mapping(peer_info) + expected_local_pool_keys = { + (layer_group_id, pool_view_id) + for layer_group_id, layer_group in enumerate(page_table.layer_groups) + if isinstance(layer_group, AttentionLayerGroup) + for pool_view_id, _pool_view in enumerate(layer_group.pool_views) + } + expected_peer_pool_keys = { + (layer_group_id, pool_view_id) + for layer_group_id, layer_group in enumerate(peer_page_table.layer_groups) + if isinstance(layer_group, AttentionLayerGroup) + for pool_view_id, _pool_view in enumerate(layer_group.pool_views) + } + mapped_peer_pool_keys = tuple(mapping.values()) + if ( + set(mapping) != expected_local_pool_keys + or len(mapped_peer_pool_keys) != len(set(mapped_peer_pool_keys)) + or set(mapped_peer_pool_keys) != expected_peer_pool_keys + ): return False - return True + + for self_key, peer_key in mapping.items(): + if not isinstance(get_kv_map(peer_info, self_key, peer_key), IdentityMapper): + return False + self_lg, self_pool_view_id = self_key + peer_lg, peer_pool_view_id = peer_key + self_pool_view = page_table.layer_groups[self_lg].pool_views[self_pool_view_id] + peer_pool_view = peer_page_table.layer_groups[peer_lg].pool_views[peer_pool_view_id] + self_pool = get_physical_pool(page_table, self_lg, self_pool_view.pool_idx) + peer_pool = get_physical_pool(peer_page_table, peer_lg, peer_pool_view.pool_idx) + if self_pool.slot_bytes != peer_pool.slot_bytes: + return False + return True + except Exception as e: + logger.warning_once( + f"Receive bounce disabled because the exact single-writer byte mapping " + f"could not be proven: {e}", + key="kv-bounce-single-writer-plan-unproven", + ) + return False + + def _destination_intervals(self, task: KVRecvTask) -> set[tuple[int, int]]: + """Build trusted allocation ranges for bounce-tail validation. + + The sender may choose a mapped subset of these blocks, but it may not + direct local scatter outside the destination request's current KV and + Mamba slots. This manifest is derived locally before publication and + never from the result frame. + """ + extractor = self._registrar.self_extractor + page_table = extractor.page_table + intervals: set[tuple[int, int]] = set() + for layer_group_id, block_ids in enumerate(task._kv_slice.block_ids_per_layer_groups): + if layer_group_id >= len(page_table.layer_groups): + raise ValueError( + f"destination layer group {layer_group_id} is absent from the page table" + ) + layer_group = page_table.layer_groups[layer_group_id] + if not isinstance(layer_group, AttentionLayerGroup): + continue + for pool_view_idx, _pool_view in enumerate(layer_group.pool_views): + region = extractor.extract( + np.asarray(block_ids, dtype=np.int64), + layer_group_id=layer_group_id, + pool_idx=pool_view_idx, + ) + size = int(region.memory.bytes_per_region) + intervals.update((int(ptr), size) for ptr in region.memory.ptrs) + + mamba_slot = task._kv_slice.mamba_state_index + if mamba_slot is not None: + for layer_group in page_table.layer_groups: + if not isinstance(layer_group, MambaLayerGroup): + continue + for pool in (layer_group.conv_states, layer_group.ssm_states): + if pool is None: + continue + slot = int(mamba_slot) + if slot < 0 or slot >= pool.num_slots: + raise ValueError( + f"destination Mamba slot {slot} is outside [0, {pool.num_slots})" + ) + intervals.add( + (int(pool.base_address + slot * pool.slot_bytes), int(pool.slot_bytes)) + ) + + return intervals def dispatch_task(self, task: KVRecvTask): params = task._params @@ -1544,61 +3439,251 @@ def dispatch_task(self, task: KVRecvTask): if sender_dp_rank is not None: # Normal path: ctx_dp_rank is known, send to overlapping ranks. peer_overlap = self._registrar.get_peer_overlap(peer_infos, sender_dp_rank) + valid_writer_cohorts = (tuple(peer_overlap.ranks),) else: # Gen-first with ADP: ctx_dp_rank unknown — broadcast REQUEST_DATA # to ALL ctx sender ranks so every DP group receives it. # get_peer_overlap returns ranks for one DP group (topology is # symmetric), so use dp_rank=0 as representative. dp_size = peer_infos.dp_size - dp0_overlap = self._registrar.get_peer_overlap(peer_infos, 0) + dp_overlaps = [ + self._registrar.get_peer_overlap(peer_infos, dp) for dp in range(dp_size) + ] + dp0_overlap = dp_overlaps[0] # Union of overlapping ranks across all DP groups for broadcast (deduplicated) - all_ranks_set: set[int] = set(dp0_overlap.ranks) - for dp in range(1, dp_size): - all_ranks_set.update(self._registrar.get_peer_overlap(peer_infos, dp).ranks) - all_ranks = list(all_ranks_set) + all_ranks_set: set[int] = set() + for overlap in dp_overlaps: + all_ranks_set.update(overlap.ranks) + all_ranks = sorted(all_ranks_set) logger.debug( f"Receiver.dispatch_task: ADP broadcast path, dp_size={dp_size}, " f"all_ranks={all_ranks}" ) peer_overlap = type(dp0_overlap)(ranks=all_ranks) + valid_writer_cohorts = tuple(tuple(overlap.ranks) for overlap in dp_overlaps) # In gen-first ADP broadcast, peer_overlap contains the union of all DP # groups, but expected_transfers should reflect per-DP-group count since # only one DP group will actually process the context request. - if sender_dp_rank is not None: - task.expected_transfers = len(peer_overlap.ranks) - else: - task.expected_transfers = len(dp0_overlap.ranks) - # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: - # _fanin_bounce_safe() (TP-by-head / even-PP), and never under ADP broadcast (sender_dp_rank - # None), where the real writer count exceeds expected_transfers and would overflow the slot. - topo_overlap = peer_overlap if sender_dp_rank is not None else dp0_overlap - allow_bounce = task.expected_transfers == 1 or ( - sender_dp_rank is not None and self._fanin_bounce_safe(topo_overlap, peer_infos) + task.set_valid_writer_cohorts(valid_writer_cohorts) + # The transport currently accepts only one writer, and only when the + # receiver can prove that its complete slot layout maps byte-for-byte. + # Fan-in remains direct until byte-accurate per-writer extents are + # authorized before any address is published. ADP broadcast also + # remains direct because it does not select the exact writer cohort. + allow_bounce = ( + self._bounce.enabled + and sender_dp_rank is not None + # The bound-buffer layout currently covers attention pool views; + # Mamba state bytes remain on the direct descriptor path. + and task._kv_slice.mamba_state_index is None + and peer_overlap.overlap_pp_size == 1 + and task.expected_transfers == 1 + and self._single_writer_bounce_exact(peer_infos) ) - bounced = allow_bounce and self._bounce.reserve(receiver_req, task.expected_transfers) session = self._get_session(task._unique_rid) if session is None: raise RuntimeError( f"dispatch_task: RxSession {task._unique_rid} not found; " "session may have been closed before dispatch" ) - session.mark_transferring(task.slice_id) - # Cache sender endpoints so cancel() can send CANCEL_SESSION to them. - session._sender_endpoints.update( - peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks + + writer_ranks = tuple(int(rank) for rank in peer_overlap.ranks) + sender_endpoints = {peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks} + key = (receiver_req.unique_rid, receiver_req.slice_id) + # Allocator waits are deliberately outside the publication/cancel gate. + # If cancellation wins while reserve blocks, the reservation is still + # unexposed and is released after the gate is acquired. + bounced = allow_bounce and self._bounce.reserve( + receiver_req, + writer_ranks, + destination_intervals_factory=lambda: self._destination_intervals(task), ) + + # This lock is the publication gate shared with RxSession.cancel(). + with session.lock: + session._sender_endpoints.update(sender_endpoints) + if session._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + dispatch_cancelled = True + if bounced: + self._bounce.release_idle_reservation(key) + else: + dispatch_cancelled = False + if sender_dp_rank is not None: + try: + update = self._recv_registry.prepare( + key, + writer_ranks, + has_bounce_slot=bounced, + ) + except Exception: + if bounced: + self._bounce.release_idle_reservation(key) + raise + if not update.accepted: + if bounced: + self._bounce.release_idle_reservation(key) + task.fail( + RuntimeError( + f"receive lifecycle admission failed for {key}: {update.reason}" + ) + ) + if session._terminal_status is None: + session._terminal_status = SessionStatus.ERROR + task.close_publication() + session._close_aux_publication_locked() + return + task.lifecycle_managed = True + task.status = TaskStatus.TRANSFERRING + + if dispatch_cancelled: + # Cancellation may have happened before sender endpoints were + # known. No target was reserved or published, so forwarding the + # cancel now is sufficient and cleanup remains immediately safe. + if task.status is not TaskStatus.ERROR: + task.fail(RuntimeError(f"RxSession {receiver_req.unique_rid} cancelled")) + session.close_task_publication(task) + self.send_cancel_to_senders(receiver_req.unique_rid, sender_endpoints) + return + # Fan-in: each sender gets its own sub-region base (writers must not overwrite); else serialize once. fanin_bounce = bounced and task.expected_transfers > 1 - key = (receiver_req.unique_rid, receiver_req.slice_id) receiver_req_bytes = None if fanin_bounce else receiver_req.to_bytes() - for i, rank in enumerate(peer_overlap.ranks): + + for rank in writer_ranks: + if task.lifecycle_managed: + # Mark POSSIBLY_EXPOSED in both ownership layers while holding + # the same gate as cancellation. Once this returns, even a ZMQ + # send exception is ambiguous and the target must be retained. + with session.lock: + aux_publication_blocked = not session._aux_writer_exposure_allowed_locked() + if aux_publication_blocked: + publication = self._recv_registry.fail_context( + key, "session AUX publication gate is already closed" + ) + bounce_allowed = False + else: + publication = self._recv_registry.begin_publication(key, rank) + if publication.publication_allowed and bounced: + bounce_allowed = self._bounce.mark_writer_exposed(key, rank) + else: + bounce_allowed = publication.publication_allowed + if publication.publication_allowed and bounce_allowed: + task.mark_writer_exposed(rank) + if not session._mark_aux_writer_exposed_locked(rank): + aux_publication_blocked = True + publication = self._recv_registry.fail_context( + key, "session AUX publication gate closed during exposure" + ) + if aux_publication_blocked: + try: + if bounced: + self._bounce.mark_logical_failure( + key, + on_done=lambda succeeded, key=key, rank=rank: self._finish_bounce( + key, succeeded, rank + ), + ) + finally: + try: + self._handle_lifecycle_update(publication, peer_rank=rank) + finally: + session.close_task_publication(task) + self.send_cancel_to_senders(receiver_req.unique_rid, sender_endpoints) + return + if not publication.publication_allowed: + continue + if not bounce_allowed: + # No address was sent. Settle both ownership layers with + # the same explicit no-access proof and let the bounce + # callback acknowledge physical slot release. + update = self._recv_registry.mark_never_published(key, rank) + self._bounce.record_no_access( + key, + rank, + succeeded=False, + on_done=lambda succeeded, key=key, rank=rank: self._finish_bounce( + key, succeeded, rank + ), + ) + self._handle_lifecycle_update(update, peer_rank=rank) + session.close_task_publication(task) + return + else: + # ADP broadcast has no exact selected-writer context. The + # session lock is still the publication/cancellation gate: + # once this flag is set, cleanup must await the complete + # legacy terminal-result cohort. + with session.lock: + if ( + session._terminal_status + in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ) + or not session._aux_writer_exposure_allowed_locked() + ): + publication_allowed = False + else: + task.mark_writer_exposed(rank) + publication_allowed = session._mark_aux_writer_exposed_locked(rank) + if not publication_allowed: + if task.status is not TaskStatus.ERROR: + task.fail( + RuntimeError( + f"AUX publication gate closed for request " + f"{receiver_req.unique_rid} slice={receiver_req.slice_id}" + ) + ) + session.close_task_publication(task) + self.send_cancel_to_senders(receiver_req.unique_rid, sender_endpoints) + return + if task._perf_timer is not None: task._perf_timer.record_task_start(rank) if fanin_bounce: - receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) + receiver_req.bounce_dst_base = self._bounce.writer_base(key, rank) receiver_req_bytes = receiver_req.to_bytes() - self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) + try: + self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) + except Exception as e: + if not task.lifecycle_managed: + session.close_task_publication(task) + raise + self._recv_registry.mark_publication_ambiguous(key, rank) + update = self._recv_registry.fail_context( + key, f"target publication to writer {rank} failed: {e}" + ) + if bounced: + self._bounce.mark_logical_failure( + key, + on_done=lambda succeeded, key=key, rank=rank: self._finish_bounce( + key, succeeded, rank + ), + ) + self._handle_lifecycle_update(update, peer_rank=rank) + logger.error( + f"Receiver failed to publish target for request {receiver_req.unique_rid} " + f"slice={receiver_req.slice_id} writer={rank}: {e}" + ) + # Keep the request in the transceiver maps so its destination + # blocks cannot be reused while this send remains ambiguous. + session.close_task_publication(task) + return + if task.lifecycle_managed: + publication = self._recv_registry.mark_published(key, rank) + if publication.conflict: + if bounced: + self._bounce.mark_protocol_conflict( + key, + on_done=lambda succeeded, key=key, rank=rank: self._finish_bounce( + key, succeeded, rank + ), + ) + self._handle_lifecycle_update(publication, peer_rank=rank) + session.close_task_publication(task) + return + session.close_task_publication(task) return @staticmethod @@ -1615,12 +3700,18 @@ def _should_register_peer(self, params: DisaggregatedParams) -> bool: def _get_or_connect_dealer(self, endpoint: Optional[str]): if endpoint is None: raise ValueError("Receiver: peer endpoint is None; peer may not have registered yet") - if endpoint not in self._dealers: - self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) - return self._dealers[endpoint] + with self._dealers_lock: + if not getattr(self, "_dealer_admission_open", True): + raise RuntimeError("Receiver is shutting down; dealer admission is closed") + if endpoint not in self._dealers: + self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) + return self._dealers[endpoint] def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: info_endpoint = self._extract_info_endpoint(params) + with self._dealers_lock: + if not getattr(self, "_dealer_admission_open", True): + raise RuntimeError("Receiver is shutting down; sender discovery is closed") if self._should_register_peer(params): logger.info(f"Registering peer in first request to endpoint '{info_endpoint}'") messenger = ZMQMessenger(mode="DEALER", endpoint=info_endpoint) @@ -1684,13 +3775,13 @@ def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) session = None with self._sessions_lock: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: - self._pre_cancelled_rids.add(unique_rid) - else: - session = session_ref() - if session is None: - self._pre_cancelled_rids.add(unique_rid) + # Retain the tombstone even when a live session consumes this + # cancellation. Request IDs are not generation-safe, so deleting it + # would let a replacement session escape a delayed replay. + self._pre_cancelled_rids.add(unique_rid) + entry = self._sessions.get(unique_rid) + if entry is not None: + session = entry() if isinstance(entry, weakref.ReferenceType) else entry if session is not None: session.cancel() @@ -1700,23 +3791,174 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" ) return - peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code = ( - _KV_RESULT_PREFIX.unpack(message[1]) - ) - from .bounce import decode_result_tail + peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code = ( + _KV_RESULT_PREFIX.unpack(message[1]) + ) + if sender_slice_id < 0: + logger.error( + f"KV result references invalid negative slice {sender_slice_id} " + f"for request {unique_rid}; dropping status" + ) + return + from .bounce import decode_result_tail + + key = (unique_rid, sender_slice_id) + target_mode = self._recv_registry.target_mode(key) + if target_mode is not None: + try: + agent_result = _AGENT_RESULT_BY_CODE[status_code] + except KeyError: + # A newer sender may define a non-terminal state. Without + # negotiated enum semantics this frame is not quiescence + # evidence, so fail logically but retain the writer target. + update = self._recv_registry.record_protocol_conflict( + key, + peer_rank, + f"unknown writer result code {status_code} from rank {peer_rank}", + ) + if target_mode is WriterMode.BOUNCE: + self._bounce.mark_protocol_conflict( + key, + on_done=lambda succeeded, key=key, peer_rank=peer_rank: self._finish_bounce( + key, succeeded, peer_rank + ), + ) + logger.error( + f"writer {peer_rank} returned unknown result code {status_code} " + f"for request {unique_rid} slice={sender_slice_id}" + ) + self._handle_lifecycle_update(update, peer_rank=peer_rank) + return + + try: + dst_ptrs, sizes, src_base = decode_result_tail(message) + except Exception as e: + logger.error( + f"malformed KV result tail for request {unique_rid} " + f"slice={sender_slice_id} writer={peer_rank}: {e}" + ) + dst_ptrs, sizes, src_base = None, None, None + agent_result = AgentResult.FAILED + + expected_bounce = target_mode is WriterMode.BOUNCE + tail_present = len(message) > 2 + complete_tail = ( + len(message) == 5 + and dst_ptrs is not None + and sizes is not None + and src_base is not None + ) + malformed_bounce_success = ( + expected_bounce + and agent_result is AgentResult.SUCCESS + and tail_present + and not complete_tail + ) + if malformed_bounce_success: + logger.error( + f"bounce writer {peer_rank} returned a malformed result tail " + f"for request {unique_rid} slice={sender_slice_id}" + ) + agent_result = AgentResult.FAILED + + if expected_bounce: + # Sender bounce allocation is per writer and may fall back to + # the direct descriptor path. A complete success tail proves + # BOUNCE; no-tail success is the backward-compatible DIRECT + # fallback. Failure is conservatively BOUNCE because it may + # have accessed the advertised slot. + mode = ( + WriterMode.BOUNCE + if complete_tail or agent_result is AgentResult.FAILED + else WriterMode.DIRECT + ) + else: + # Any tail on a direct-only target is a mode conflict. + mode = WriterMode.BOUNCE if tail_present else WriterMode.DIRECT + writer_result = ( + WriterResult.SUCCESS if agent_result is AgentResult.SUCCESS else WriterResult.FAILED + ) + update = self._recv_registry.record_result(key, peer_rank, writer_result, mode) + if not update.accepted: + if update.conflict: + if expected_bounce: + self._bounce.mark_protocol_conflict( + key, + on_done=lambda succeeded, key=key, peer_rank=peer_rank: ( + self._finish_bounce(key, succeeded, peer_rank) + ), + ) + logger.error( + f"receive lifecycle conflict for request {unique_rid} " + f"slice={sender_slice_id} writer={peer_rank}: {update.reason}" + ) + self._handle_lifecycle_update(update, peer_rank=peer_rank) + return + + if expected_bounce: + + def on_done( + succeeded: bool, + key: tuple[int, int] = key, + peer_rank: int = peer_rank, + ) -> None: + self._finish_bounce(key, succeeded, peer_rank) + + if mode is WriterMode.DIRECT: + self._bounce.record_no_access( + key, + peer_rank, + succeeded=agent_result is AgentResult.SUCCESS, + on_done=on_done, + ) + elif agent_result is AgentResult.SUCCESS: + self._bounce.record_result( + key, + peer_rank, + dst_ptrs, + sizes, + src_base, + on_done, + on_logical_failure=lambda key=key: self._fail_bounce_logically(key), + ) + else: + self._bounce.record_failure(key, peer_rank, on_done=on_done) + + # This optional notification intentionally happens after the + # registry and bounce transport have consumed the terminal result. + # A closed/GC'd RxSession can no longer strand physical ownership. + self._handle_lifecycle_update(update, peer_rank=peer_rank) + return - dst_ptrs, sizes, src_base = decode_result_tail(message) session = self._get_session(unique_rid) if session is None: logger.warning( f"_process_kv_agent_result: session {unique_rid} not found (already closed?), dropping status" ) return + try: + agent_result = _AGENT_RESULT_BY_CODE[status_code] + except KeyError: + session.process_kv_protocol_conflict( + peer_rank, + sender_slice_id, + f"unknown writer result code {status_code}", + ) + return + try: + dst_ptrs, sizes, src_base = decode_result_tail(message) + except Exception as e: + session.process_kv_protocol_conflict( + peer_rank, + sender_slice_id, + f"malformed writer result tail: {e}", + ) + return session.process_kv_agent_result( peer_rank, sender_slice_id, is_last_slice, - _AGENT_RESULT_BY_CODE[status_code], + agent_result, dst_ptrs=dst_ptrs, sizes=sizes, src_base=src_base, @@ -1732,7 +3974,17 @@ def _process_aux_agent_result(self, _send_id: bytes, message: list[bytes]): f"_process_aux_agent_result: session {unique_rid} not found (already closed?), dropping status" ) return - session.process_aux_agent_result(peer_rank, AgentResult(status)) + try: + agent_result = AgentResult(status) + except ValueError: + # An unknown value is not terminal evidence. Fail the request + # logically but keep the AUX lease until this exact writer later + # reports a known terminal result or a global fence retires it. + session.process_aux_protocol_conflict( + peer_rank, f"unknown AUX writer result {status!r}" + ) + return + session.process_aux_agent_result(peer_rank, agent_result) def _request_sender_data(self, endpoint: str, receiver_info_bytes: bytes): # receiver_info serialized once and reused for every peer rank (block-table msgpack isn't free at fan-out). @@ -1749,8 +4001,19 @@ def __del__(self): def __enter__(self): return self - def __exit__(self, _exc_type, _exc_val, _exc_tb): - self.shutdown() + def __exit__(self, exc_type, _exc_val, _exc_tb): + try: + drained = self.shutdown() + except Exception as e: + if exc_type is None: + raise + logger.error(f"Receiver context cleanup failed while propagating an exception: {e}") + return False + if not drained: + if exc_type is None: + raise RuntimeError("Receiver context exited before transfer resources drained") + logger.error("Receiver context retained resources while propagating an exception") + return False class RxSession(RxSessionBase): @@ -1763,26 +4026,65 @@ def __init__( timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, beam_width: int = 1, + destination_owner: object | None = None, ): - super().__init__( - receiver, - SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), - ) - self._timeout_s = timeout_s - self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST - self._receiver: Receiver # narrow base class type for Pylance - self.request_id = request_id + # Retain the request that owns its KV allocation until every accessor + # and scatter is drained. Enroll it before fallible initialization, + # then clear the parameter so propagated tracebacks are ownership- + # neutral. + self._destination_owner = destination_owner + destination_owner = None self._aux_buffer = aux_buffer - self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None - self._exception: Optional[Exception] = None - self._closed = False - self._terminal_status: Optional[SessionStatus] = None - self._kv_tasks: list[KVRecvTask] = [] - self._aux_count = 0 - self._aux_status: TaskStatus = TaskStatus.INIT - self._sender_endpoints: set[str] = set() - self.lock = threading.Lock() - self._receiver.setup_session(self) + self.aux_slot = None + try: + super().__init__( + receiver, + SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), + ) + self._timeout_s = timeout_s + self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST + self._receiver: Receiver # narrow base class type for Pylance + self.request_id = request_id + # Receive dispatch owns the address-publication order for every + # slice in this session. Keep it separate from ``lock`` so + # cancellation can still close the publication gate while a + # bounce reservation blocks. + self._receive_lock = threading.Lock() + self._close_lock = threading.Lock() + self.lock = threading.Lock() + self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None + self._exception: Optional[Exception] = None + self._closed = False + self._accepting_receives = True + self._active_receive_dispatches = 0 + self._last_slice_admitted = False + self._receiver_retired = False + self._terminal_status: Optional[SessionStatus] = None + self._kv_tasks: list[KVRecvTask] = [] + self._aux_results: dict[int, AgentResult] = {} + # AUX storage is session-scoped and its address can be repeated in + # every KV-slice publication. Keep one exposure ledger across + # slices; a later suppressed publication cannot revoke an earlier + # exposure. + self._aux_exposed_writer_ranks: set[int] = set() + self._aux_publication_closed = not self._need_aux + self._aux_result_conflict = False + # Allocating the session-scoped AUX slot does not create a remote + # accessor. It becomes an active ownership lease only when its + # address is actually published to a writer. + self._aux_drained = True + self._aux_status: TaskStatus = TaskStatus.INIT + self._sender_endpoints: set[str] = set() + self._selected_writer_cohort: Optional[frozenset[int]] = None + self._receiver.setup_session(self) + except Exception: + self._destination_owner = None + if self._aux_buffer is not None and self.aux_slot is not None: + self._aux_buffer.free_slot(self.aux_slot) + self.aux_slot = None + self._accepting_receives = False + self._closed = True + raise @property def disagg_request_id(self) -> int: @@ -1816,18 +4118,342 @@ def mark_transferring(self, slice_id: int): with self.lock: self._kv_tasks[slice_id].status = TaskStatus.TRANSFERRING - def receive(self, slice: KVSlice) -> None: - params = self._base_args.params - slice_id = len(self._kv_tasks) - task = KVRecvTask( - self.disagg_request_id, - slice, - slice_id, - params, - aux_slot=self.aux_slot, + def _mark_adp_cohort_conflict_locked( + self, task: KVRecvTask, *, channel: str, identities: frozenset[int] + ) -> None: + """Fail one compatibility channel without using ambiguity as drain proof.""" + detail = ( + f"ADP {channel.upper()} writer identities {sorted(identities)} conflict with " + f"the selected/candidate cohorts for request {self.disagg_request_id} " + f"slice={task.slice_id}" + ) + exc = RuntimeError(detail) + if channel == "kv": + task._legacy_result_conflict = True + task.fail(exc) + if self._need_aux: + # Cohort selection is session-wide. Evidence that a different + # cohort handled KV also invalidates an earlier AUX drain based + # on that selection, because every publication carried the AUX + # address for the same candidate writers. + self._aux_result_conflict = True + self._aux_drained = False + self._aux_status = TaskStatus.ERROR + self._exception = exc + elif channel == "aux": + self._aux_result_conflict = True + self._aux_drained = False + self._aux_status = TaskStatus.ERROR + self._exception = exc + # Conversely, an AUX result outside the selected cohort disproves + # session-wide cohort selection. Keep every legacy KV target from + # that compatibility session fail-closed. + for kv_task in self._kv_tasks: + if kv_task.lifecycle_managed or not kv_task.publication_started: + continue + kv_task._legacy_result_conflict = True + kv_task.fail(exc) + else: + raise ValueError(f"unknown ADP result channel {channel!r}") + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + logger.error(detail) + + def _select_adp_cohort_locked( + self, + task: KVRecvTask, + results: dict[int, AgentResult], + *, + channel: str, + ) -> None: + """Select one cohort only while all observations remain compatible.""" + identities = frozenset(results) + if not identities: + return + selected = self._selected_writer_cohort + if selected is not None: + if not identities <= selected: + self._mark_adp_cohort_conflict_locked(task, channel=channel, identities=identities) + return + + compatible = tuple(cohort for cohort in task._valid_writer_cohorts if identities <= cohort) + complete = tuple(cohort for cohort in compatible if cohort == identities) + if len(complete) == 1 and len(compatible) == 1: + self._selected_writer_cohort = complete[0] + return + if not compatible: + self._mark_adp_cohort_conflict_locked(task, channel=channel, identities=identities) + + def _legacy_channel_succeeded_locked( + self, task: KVRecvTask, results: dict[int, AgentResult] + ) -> bool: + """Validate success from every member of the selected legacy cohort.""" + selected = self._selected_writer_cohort + if selected is None: + return False + return all(results.get(rank) is AgentResult.SUCCESS for rank in selected) + + def _advance_legacy_kv_locked(self, task: KVRecvTask) -> None: + if task.lifecycle_managed or not task.legacy_resources_drained or task.is_done: + return + if not task._legacy_result_conflict and self._legacy_channel_succeeded_locked( + task, task._legacy_results + ): + task.complete() + return + detail = ( + f"ADP KV writer outcome did not match the selected cohort for " + f"request {self.disagg_request_id} slice={task.slice_id}" + ) + task.fail(RuntimeError(detail)) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + + def _advance_aux_locked(self, task: KVRecvTask) -> None: + if not self._need_aux or self._aux_drained: + return + selected = self._selected_writer_cohort + # ADP publishes to candidate cohorts but, by the compatibility + # contract, exactly one cohort performs the transfer. Once selected, + # only that cohort can be a physical accessor. Before selection, every + # actually published rank remains a candidate accessor. + # A protocol conflict invalidates cohort selection as physical + # quiescence evidence. Keep every identity that may have observed the + # slot in the drain ledger until it reports terminal or a backend-wide + # fence retires it. + if self._aux_result_conflict or selected is None: + required = self._aux_exposed_writer_ranks + else: + required = selected + all_terminal = self._aux_publication_closed and required <= self._aux_results.keys() + if not all_terminal: + return + self._aux_drained = True + if ( + selected is not None + and not self._aux_result_conflict + and self._legacy_channel_succeeded_locked(task, self._aux_results) + ): + self._aux_status = TaskStatus.TRANSFERRED + return + self._aux_status = TaskStatus.ERROR + self._exception = RuntimeError( + f"ADP AUX writer outcome did not match the selected cohort for " + f"request {self.disagg_request_id}" ) - self._kv_tasks.append(task) - self._receiver.dispatch_task(task) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + + def _aux_writer_exposure_allowed_locked(self) -> bool: + """Whether another publication may carry the session AUX address.""" + return not (self._need_aux and self.aux_slot is not None and self._aux_publication_closed) + + def _retain_aux_writer_exposure_locked(self, peer_rank: int) -> None: + """Retain one possible AUX accessor until terminal drain evidence.""" + if not self._need_aux: + return + self._aux_exposed_writer_ranks.add(int(peer_rank)) + self._aux_drained = False + + def _mark_aux_writer_exposed_locked(self, peer_rank: int) -> bool: + """Record one possible AUX accessor, rejecting exposure after closure.""" + if not self._aux_writer_exposure_allowed_locked(): + return False + if self._need_aux and self.aux_slot is not None: + self._retain_aux_writer_exposure_locked(peer_rank) + return True + + def _close_aux_publication_locked(self) -> None: + """Prove no later KV slice can expose the session-level AUX slot.""" + if not self._need_aux or self._aux_publication_closed: + return + self._aux_publication_closed = True + if not self._kv_tasks: + # With no remote evidence, closing before the first KV task proves + # that the slot was never exposed. A protocol-conflict frame may + # itself be evidence of an accessor, however, and only a terminal + # result or backend-wide fence can retire that ledger entry. + self._aux_drained = not self._aux_exposed_writer_ranks + self._aux_status = TaskStatus.ERROR + return + cohort_task = self._kv_tasks[0] + self._select_adp_cohort_locked(cohort_task, self._aux_results, channel="aux") + self._advance_aux_locked(cohort_task) + + def close_task_publication(self, task: KVRecvTask) -> None: + """Close a dispatch plan and consume results that raced its final send.""" + with self.lock: + task.close_publication() + self._select_adp_cohort_locked(task, task._legacy_results, channel="kv") + self._advance_legacy_kv_locked(task) + + if self._need_aux and self._kv_tasks: + if task._kv_slice.is_last_slice or self._terminal_status in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ): + self._close_aux_publication_locked() + cohort_task = self._kv_tasks[0] + self._select_adp_cohort_locked(cohort_task, self._aux_results, channel="aux") + self._advance_aux_locked(cohort_task) + + def _validate_receive_admission_locked(self) -> None: + """Reject receive work that cannot enter the serialized dispatch.""" + if not self._accepting_receives: + raise RuntimeError( + f"RxSession {self.disagg_request_id} is closing; " + "new receive slices are not accepted" + ) + if self._last_slice_admitted: + raise RuntimeError( + f"RxSession {self.disagg_request_id} already admitted its last slice" + ) + + def receive(self, slice: KVSlice) -> None: + # Serializing through dispatch is required because AUX storage and its + # publication ledger are session-wide. A later last slice must not + # close/free the slot while an earlier slice can still expose it. + # Check the admission gate before waiting for the dispatch lock so a + # receive started after shutdown is sealed cannot block behind an + # already admitted dispatch. Repeat the check after serialization to + # close the race with cancellation or close. + with self.lock: + self._validate_receive_admission_locked() + with self._receive_lock: + with self.lock: + self._validate_receive_admission_locked() + params = self._base_args.params + slice_id = len(self._kv_tasks) + task = KVRecvTask( + self.disagg_request_id, + slice, + slice_id, + params, + aux_slot=self.aux_slot, + ) + self._kv_tasks.append(task) + self._active_receive_dispatches += 1 + if slice.is_last_slice: + self._last_slice_admitted = True + try: + self._receiver.dispatch_task(task) + except Exception as dispatch_error: + # Admission and dispatch form one transaction for standalone + # callers too. Cancellation settles any context prepared + # before the exception, while preserving every possibly + # exposed target until terminal evidence arrives. + cleanup_errors: list[Exception] = [] + try: + self.cancel() + except Exception as cleanup_error: + cleanup_errors.append(cleanup_error) + with self.lock: + task.close_publication() + task.fail(dispatch_error) + self._exception = dispatch_error + self._terminal_status = SessionStatus.ERROR + self._close_aux_publication_locked() + if cleanup_errors: + logger.error( + f"RxSession {self.disagg_request_id} retained ownership after dispatch " + f"cleanup failed: {cleanup_errors[0]}" + ) + raise + finally: + with self.lock: + self._active_receive_dispatches -= 1 + + def process_lifecycle_update( + self, update: LifecycleUpdate, *, peer_rank: Optional[int] = None + ) -> None: + """Apply only the consumer-visible part of a registry transition. + + The registry and bounce transport have already consumed the physical + transition before this method runs. Consequently losing this session + cannot lose writer accounting, scatter completion, or slot ownership. + """ + slice_id = update.key[1] + with self.lock: + if not 0 <= slice_id < len(self._kv_tasks): + logger.error( + f"receive lifecycle update references missing slice {slice_id} " + f"for request {self.disagg_request_id}" + ) + return + task = self._kv_tasks[slice_id] + actions = set(update.actions) + if LifecycleAction.NOTIFY_CANCELLED in actions: + if task.status is not TaskStatus.ERROR: + task.fail(RuntimeError(f"RxSession {self.disagg_request_id} cancelled")) + self._terminal_status = SessionStatus.CANCELLED + self._close_aux_publication_locked() + return + if LifecycleAction.NOTIFY_FAILURE in actions: + detail = ( + f"KV receive lifecycle failed for request {self.disagg_request_id} " + f"slice={slice_id}" + ) + if update.reason: + detail += f": {update.reason}" + task.fail(RuntimeError(detail)) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + self._close_aux_publication_locked() + logger.error(detail) + return + if LifecycleAction.NOTIFY_SUCCESS not in actions: + return + if task.status is TaskStatus.ERROR: + return + + if peer_rank is not None: + try: + if task._perf_timer is not None: + task._perf_timer.record_task_end(peer_rank) + ri = self._receiver._registrar.self_rank_info + task.print_perf_info(peer_rank, ri.instance_name, ri.instance_rank) + except Exception as e: + logger.warning( + f"KV transfer perf logging failed for request {self.disagg_request_id} " + f"slice={slice_id}: {e}" + ) + task.complete() + logger.debug( + f"KV transfer complete for request {self.disagg_request_id} slice={slice_id}" + ) + + def process_kv_protocol_conflict( + self, peer_rank: int, sender_slice_id: int, reason: str + ) -> None: + """Fail logically without treating an unrecognized frame as terminal.""" + with self.lock: + if not 0 <= sender_slice_id < len(self._kv_tasks): + logger.error( + f"KV protocol conflict references missing slice {sender_slice_id} " + f"for request {self.disagg_request_id}: {reason}" + ) + return + task = self._kv_tasks[sender_slice_id] + # A frame for this identity is itself evidence that the peer may + # have observed a target; never use the no-publication fast path. + task.mark_writer_exposed(peer_rank) + task._legacy_result_conflict = True + detail = ( + f"KV protocol conflict for request {self.disagg_request_id} " + f"slice={sender_slice_id} writer={peer_rank}: {reason}" + ) + exc = RuntimeError(detail) + task.fail(exc) + self._exception = exc + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + # This frame is remote evidence, not a new local publication. + # Preserve the possible accessor even if the publication gate has + # already closed so AUX drain remains fail-closed. + if self._need_aux and self.aux_slot is not None: + self._retain_aux_writer_exposure_locked(peer_rank) + self._close_aux_publication_locked() + logger.error(detail) def process_kv_agent_result( self, @@ -1839,96 +4465,51 @@ def process_kv_agent_result( sizes=None, src_base=None, ): + del is_last_slice, dst_ptrs, sizes, src_base with self.lock: - assert sender_slice_id < len(self._kv_tasks), ( - f"Receiver got slice_id={sender_slice_id} from sender but only has " - f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}. " - f"Sender/receiver slice count mismatch." - ) + if not 0 <= sender_slice_id < len(self._kv_tasks): + logger.error( + f"Receiver got invalid slice_id={sender_slice_id} from sender but only has " + f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}; " + "dropping status" + ) + return task = self._kv_tasks[sender_slice_id] - if status == AgentResult.SUCCESS: - from .bounce import scatter_write_result - - on_done = None - if is_last_slice: - task.last_slice_count += 1 - if task.last_slice_count == task.expected_transfers: - # Completing message: defer task.complete()+perf until the scatter has actually - # landed. scatter_write_result fires this inline for the non-bounced path, or on - # the scatter worker (after cudaStreamSynchronize) for the bounced path, so the - # gen consumer never observes completion before the KV is scattered into place. - request_id = self.request_id - ri = self._receiver._registrar.self_rank_info - instance_name, instance_rank = ri.instance_name, ri.instance_rank - - def on_done( - success, - task=task, - peer_rank=peer_rank, - sender_slice_id=sender_slice_id, - request_id=request_id, - instance_name=instance_name, - instance_rank=instance_rank, - ): - # Runs on the scatter worker thread for the bounced path. Touches only this - # task's own status/_event/_perf_timer (no RxSession.lock, no shared session - # state), so it is lock-free. complete() sets status before _event, keeping - # wait_complete's status-first poll correct. - if not success: - task.fail( - RuntimeError( - f"KV bounce scatter failed for request {request_id} " - f"slice={sender_slice_id}" - ) - ) - return - if task.status == TaskStatus.ERROR: - return # a concurrent FAILED writer already failed it; don't un-fail - try: - if task._perf_timer is not None: - task._perf_timer.record_task_end(peer_rank) - task.print_perf_info(peer_rank, instance_name, instance_rank) - except Exception as e: # perf is best-effort; never block completion - logger.warning( - f"KV transfer perf logging failed for request {request_id} " - f"slice={sender_slice_id}: {e}" - ) - task.complete() - logger.debug( - f"KV transfer complete for request {request_id} " - f"slice={sender_slice_id}" - ) - - scatter_write_result( - self._receiver._bounce, - (self.disagg_request_id, task.slice_id), - peer_rank, - dst_ptrs, - sizes, - src_base, - on_done, + if task.lifecycle_managed: + # Managed results are consumed by Receiver's registry before + # optional session notification. Reaching the compatibility + # handler means the context was already retired; a late + # duplicate must not be counted as a legacy writer. + logger.debug( + f"Ignoring late managed KV result for request {self.disagg_request_id} " + f"slice={sender_slice_id} writer={peer_rank}" ) - elif status == AgentResult.FAILED: + return + is_new = task.record_legacy_result(peer_rank, status) + if task._legacy_result_conflict: detail = ( - f"KV transfer failed for request {self.request_id} slice={sender_slice_id} " - f"peer_rank={peer_rank} is_last_slice={is_last_slice} " - f"(reported by remote agent; see sender-side log for nixl_status)" - ) - logger.error(detail) - # Drain-before-release: record this writer FAILED; the owner frees the shared region - # only once every fan-in writer is terminal (freeing now could race a sibling's RMA). - self._receiver._bounce.record_failure( - (self.disagg_request_id, task.slice_id), peer_rank + f"conflicting legacy KV result identities for request {self.request_id} " + f"slice={sender_slice_id} writer={peer_rank}" ) task.fail(RuntimeError(detail)) - if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR + if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR - else: + self._close_aux_publication_locked() + logger.error(detail) + return + if not is_new: + return + if status not in (AgentResult.SUCCESS, AgentResult.FAILED): raise ValueError( f"Session {self.request_id} received unknown task status: {status.value}" ) + self._select_adp_cohort_locked(task, task._legacy_results, channel="kv") + if task._legacy_result_conflict: + self._close_aux_publication_locked() + return + self._advance_legacy_kv_locked(task) - def process_aux_agent_result(self, _peer_rank: int, status: AgentResult): + def process_aux_agent_result(self, peer_rank: int, status: AgentResult): # Aux is session-level (not per-slice); expected_transfers is identical # across all kv_tasks, so any task provides the right count. with self.lock: @@ -1938,28 +4519,66 @@ def process_aux_agent_result(self, _peer_rank: int, status: AgentResult): ) return task = self._kv_tasks[0] - if status == AgentResult.SUCCESS: - self._aux_count += 1 - - if self._aux_count == task.expected_transfers: - self._aux_status = TaskStatus.TRANSFERRED - elif self._aux_count > task.expected_transfers: - self._aux_status = TaskStatus.ERROR + # A terminal response proves this writer can no longer access the + # aux slot. Count identities, not success-only responses, so one + # failure cannot release the slot while siblings remain active. + previous = self._aux_results.get(peer_rank) + if previous is not None: + if previous is not status: + self._aux_result_conflict = True self._exception = RuntimeError( - f"Session {self.request_id} received too many aux transfers" + f"Session {self.request_id} received conflicting aux results " + f"from writer {peer_rank}" ) + self._aux_status = TaskStatus.ERROR if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR + self._close_aux_publication_locked() logger.error(str(self._exception)) - elif status == AgentResult.FAILED: - self._aux_status = TaskStatus.ERROR - self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") + return + identity_valid = peer_rank in self._aux_exposed_writer_ranks + if not identity_valid: + # The frame itself is evidence that this identity participated + # in some publication generation. Retain it in the physical + # ledger even while latching the protocol conflict. + self._retain_aux_writer_exposure_locked(peer_rank) + self._aux_results[peer_rank] = status + if not identity_valid: + self._aux_result_conflict = True + self._exception = RuntimeError( + f"Session {self.request_id} received an unexpected " + f"aux writer identity {peer_rank}" + ) if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR - else: + self._close_aux_publication_locked() + + if status not in (AgentResult.SUCCESS, AgentResult.FAILED): raise ValueError( f"Session {self.request_id} received unknown aux send status: {status}" ) + if self._aux_result_conflict: + self._aux_status = TaskStatus.ERROR + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + logger.error(str(self._exception)) + self._select_adp_cohort_locked(task, self._aux_results, channel="aux") + self._advance_aux_locked(task) + + def process_aux_protocol_conflict(self, peer_rank: int, reason: str) -> None: + """Latch an AUX protocol failure without treating it as quiescence.""" + with self.lock: + self._retain_aux_writer_exposure_locked(peer_rank) + self._aux_result_conflict = True + self._aux_status = TaskStatus.ERROR + self._exception = RuntimeError( + f"AUX protocol conflict for request {self.disagg_request_id} " + f"writer={peer_rank}: {reason}" + ) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + self._close_aux_publication_locked() + logger.error(str(self._exception)) @property def exception(self) -> Optional[Exception]: @@ -1984,58 +4603,204 @@ def unpack_aux(self, request: LlmRequest) -> None: }, } + def mark_backend_quiesced(self) -> None: + """Retire compatibility targets after an external remote/global fence.""" + with self.lock: + exc = RuntimeError( + f"RxSession {self.disagg_request_id} ended without complete terminal results " + "before remote/global quiescence" + ) + for task in self._kv_tasks: + if task.lifecycle_managed or task.legacy_resources_drained: + continue + task._legacy_backend_quiesced = True + task.fail(exc) + if self._need_aux and not self._aux_drained: + self._aux_drained = True + self._aux_status = TaskStatus.ERROR + self._exception = exc + if any(task.status is TaskStatus.ERROR for task in self._kv_tasks) or ( + self._need_aux and self._aux_status is TaskStatus.ERROR + ): + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + + def has_untracked_receive_activity(self) -> bool: + """Whether legacy KV or AUX memory can still have remote accessors.""" + with self.lock: + if getattr(self, "_active_receive_dispatches", 0): + return True + if any( + not task.lifecycle_managed and not task.legacy_resources_drained + for task in self._kv_tasks + ): + return True + return self._need_aux and not self._aux_drained + + def seal_receive_admission(self) -> None: + """Prevent queued or future receive calls from entering dispatch.""" + with self.lock: + self._accepting_receives = False + + def resources_drained(self) -> bool: + """Whether request cleanup can release every receive-side target.""" + try: + if not self._receiver._bounce.retry_settlements(self.disagg_request_id): + return False + except Exception as e: + logger.error( + f"failed to retry bounce settlement for request {self.disagg_request_id}: {e}" + ) + return False + if not self._receiver._recv_registry.is_request_drained(self.disagg_request_id): + return False + return not self.has_untracked_receive_activity() + def is_completed(self) -> bool: """Non-blocking check: has the transfer completed successfully?""" + if not self.resources_drained(): + return False status = self.status if self._need_aux: return status == SessionStatus.FULLY_TRANSFERRED return status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) def has_failed(self) -> bool: - """Non-blocking check: has the transfer failed or been cancelled?""" + """Whether failure is safe to expose to request cleanup. + + Logical failure is latched immediately, but without an allocator KV + lease this method must remain false until every registry-owned writer + and bounce scatter is physically drained. + """ + if not self.resources_drained(): + return False return self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) def cancel(self) -> None: """Cancel the session and notify the remote sender. - Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). - Only INIT tasks have their events signalled immediately. - The lock serializes with process_kv_agent_result() / process_aux_agent_result(). + Safe to call multiple times. The lock is also the address-publication + gate, so cancellation either closes an unexposed writer or observes it + as possibly exposed and retains its resources until terminal evidence. """ + errors: list[tuple[str, Exception]] = [] with self.lock: - if self._terminal_status == SessionStatus.CANCELLED: - return + self._accepting_receives = False + try: + self._receiver.record_cancel_tombstone(self.disagg_request_id) + except Exception as e: + errors.append(("receiver cancellation tombstone", e)) self._terminal_status = SessionStatus.CANCELLED exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") + try: + updates = self._receiver._recv_registry.cancel_request(self.disagg_request_id) + except Exception as e: + updates = () + errors.append(("receive registry", e)) for task in self._kv_tasks: - rid_slice = (self.disagg_request_id, task.slice_id) - if task.status == TaskStatus.INIT: - # INIT = reserved but no write in flight, so freeing its bounce reservation here - # is safe. - self._receiver._bounce.release_idle_reservation(rid_slice) - task.fail(exc) - elif task.status == TaskStatus.TRANSFERRING: - # A write may still be mid-flight, so quarantine the region rather than freeing - # it; this keeps a cancelled transfer from leaking. No-op when bounce is off. - self._receiver._bounce.orphan_reservation(rid_slice) + if task.lifecycle_managed: + if task.status is not TaskStatus.TRANSFERRED: + task.fail(exc) + else: + # Legacy ADP broadcast has no exact writer context. Bounce + # is disabled for that path, but keep the compatibility + # cleanup for standalone legacy users. Retry release even + # after logical failure because the first attempt may have + # raised before retiring an idle reservation. + try: + self._receiver._bounce.release_idle_reservation( + (self.disagg_request_id, task.slice_id) + ) + except Exception as e: + errors.append((f"legacy slice {task.slice_id}", e)) + if task.status is not TaskStatus.TRANSFERRED: + task.fail(exc) + self._close_aux_publication_locked() + sender_endpoints = set(self._sender_endpoints) + for update in updates: + if update.physical_state is PhysicalState.DRAINED: + continue + try: + self._receiver._bounce.mark_logical_failure( + update.key, + on_done=lambda succeeded, key=update.key: self._receiver._finish_bounce( + key, succeeded + ), + ) + except Exception as e: + errors.append((f"receive context {update.key}", e)) # Send outside the lock to avoid holding it during I/O. - self._receiver.send_cancel_to_senders(self.disagg_request_id, self._sender_endpoints) + try: + self._receiver.send_cancel_to_senders(self.disagg_request_id, sender_endpoints) + except Exception as e: + errors.append(("sender cancellation", e)) + if errors: + detail = "; ".join(f"{owner}: {error}" for owner, error in errors) + raise RuntimeError( + f"RxSession {self.disagg_request_id} cancellation encountered " + f"{len(errors)} error(s): {detail}" + ) from errors[0][1] def has_transferring_tasks(self) -> bool: - """True if any KV task is currently mid-write (TRANSFERRING). + """True while request cleanup could race a receive-side accessor. cancel_request() must return False while this is True. """ - return any(t.status == TaskStatus.TRANSFERRING for t in self._kv_tasks) + return not self.resources_drained() def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: """Poll or block until transfer completes. With blocking=False (default): polls non-blockingly; returns None if any KV task or aux is not yet done — caller should re-poll next cycle. - With blocking=True: waits up to _timeout_s for each task. - Returns WaitResult.COMPLETED on full success, WaitResult.FAILED on error/timeout. + With blocking=True: the configured timeout latches failure and asks the + sender to cancel, but this call remains fail-closed until published + targets drain. A future allocator-backed KV lease will allow logical + timeout to return without retaining the whole request. """ + if not self.resources_drained(): + if not blocking: + return None + deadline = None if self._timeout_s is None else time.monotonic() + self._timeout_s + timeout_latched = False + while not self.resources_drained(): + if deadline is not None and time.monotonic() >= deadline and not timeout_latched: + timeout_latched = True + with self.lock: + # Timeout is a local cancellation decision. Seal + # admission and retain the request-id tombstone before + # sending cancellation so a delayed result or replay + # cannot attach to a replacement session. + self._accepting_receives = False + self._receiver.record_cancel_tombstone(self.disagg_request_id) + timeout_updates = self._receiver._recv_registry.timeout_request( + self.disagg_request_id + ) + exc = TimeoutError( + f"RxSession {self.disagg_request_id} timed out while receive " + "resources were still active" + ) + for task in self._kv_tasks: + if task.status is not TaskStatus.TRANSFERRED: + task.fail(exc) + if self._need_aux and not self._aux_drained: + self._aux_status = TaskStatus.ERROR + self._exception = exc + self._close_aux_publication_locked() + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + sender_endpoints = set(self._sender_endpoints) + for update in timeout_updates: + self._receiver._bounce.mark_logical_failure( + update.key, + on_done=lambda succeeded, key=update.key: self._receiver._finish_bounce( + key, succeeded + ), + ) + self._receiver.send_cancel_to_senders(self.disagg_request_id, sender_endpoints) + logger.warning(str(exc)) + time.sleep(0.001) + if not blocking: # Use task.status instead of task.wait(timeout=0): task.complete() # sets status before event, so a GIL switch between the two steps @@ -2069,17 +4834,42 @@ def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: def close(self): if getattr(self, "_closed", False): return - self._closed = True - if self._aux_buffer is not None and self.aux_slot is not None: - self._aux_buffer.free_slot(self.aux_slot) - self.aux_slot = None - # Unregister from Receiver; keep fields alive for in-flight listener messages. - if self._receiver is not None: - # Reclaim any bounce region still live at teardown (closed mid-transfer) so it isn't - # leaked; a no-op for finished or non-bounce transfers. - for task in self._kv_tasks: - self._receiver._bounce.orphan_reservation((self.disagg_request_id, task.slice_id)) - self._receiver.clear_session(self.disagg_request_id) + with self._close_lock: + if self._closed: + return + # Close admission before waiting on dispatch so a queued receive + # cannot overtake this close attempt when the active dispatch + # releases ``_receive_lock``. + with self.lock: + self._accepting_receives = False + # Wait for an admitted receive to finish dispatch before evaluating + # drain state. This gives retirement a stable publication frontier + # and also serializes all AUX release. + with self._receive_lock: + with self.lock: + publication_pending = any( + not task._publication_closed for task in self._kv_tasks + ) + terminal_gate_closed = self._terminal_status in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ) + if publication_pending and not terminal_gate_closed: + raise RuntimeError( + f"cannot close RxSession {self.disagg_request_id}: " + "target publication is still pending" + ) + # clear_session is the fail-closed retirement gate. It raises + # rather than dropping the only Python-level destination owner + # while a direct writer is still active. + if self._receiver is not None and not self._receiver_retired: + self._receiver.clear_session(self.disagg_request_id) + self._receiver_retired = True + if self._aux_buffer is not None and self.aux_slot is not None: + self._aux_buffer.free_slot(self.aux_slot) + self.aux_slot = None + self._destination_owner = None + self._closed = True def __enter__(self): return self @@ -2098,6 +4888,7 @@ class RankInfoServer: def __init__(self, rank_info: RankInfo, addr: Optional[str] = None, port: Optional[int] = None): self._rank_info = rank_info self._shutdown = False # must be set before _start_listener() so __del__ is safe + self._shutdown_started = False if addr is None and port is None: endpoint = f"tcp://{get_local_ip()}:*" else: @@ -2109,12 +4900,18 @@ def __init__(self, rank_info: RankInfo, addr: Optional[str] = None, port: Option def endpoint(self) -> str: return self._messenger.endpoint - def shutdown(self): + def shutdown(self) -> bool: if self._shutdown: - return - self._shutdown = True + return True + self._shutdown_started = True logger.debug("RankInfoServer.shutdown() called") - self._messenger.stop() + try: + self._messenger.stop() + except Exception as e: + logger.error(f"RankInfoServer.shutdown: listener stop failed: {e}") + return False + self._shutdown = True + return True def _start_listener(self): def handle_message(messages: list[bytes]) -> bool: @@ -2146,8 +4943,21 @@ def __del__(self): def __enter__(self): return self - def __exit__(self, _exc_type, _exc_val, _exc_tb): - self.shutdown() + def __exit__(self, exc_type, _exc_val, _exc_tb): + try: + stopped = self.shutdown() + except Exception as e: + if exc_type is None: + raise + logger.error( + f"RankInfoServer context cleanup failed while propagating an exception: {e}" + ) + return False + if not stopped: + if exc_type is None: + raise RuntimeError("RankInfoServer context exited before its listener stopped") + logger.error("RankInfoServer retained its listener while propagating an exception") + return False def _create_nixl_agent(name: str) -> NixlTransferAgent: @@ -2187,6 +4997,10 @@ class TransferWorkerConfig: class TransferWorker: def __init__(self, config: TransferWorkerConfig): + self._shutdown_attempt_lock = threading.Lock() + self._shutdown_started = False + self._shutdown = False + self._session_admission_lock = threading.Lock() self._config = config kvm = config.kv_cache_manager self._aux_buffer = _make_aux_buffer( @@ -2207,30 +5021,38 @@ def populate_instance_and_rank_info(self, endpoints: list[str], layer_num_per_pp self._rank_info.layer_num_per_pp = layer_num_per_pp def create_tx_session(self, request: LlmRequest) -> TxSession: - params = request.py_disaggregated_params - assert params is not None - return TxSession( - request_id=request.py_request_id, - params=params, - sender=self._sender, - aux_buffer=self._aux_buffer, - timeout_s=self._config.tx_timeout_s, - prompt_len=request.prompt_len, - beam_width=request.py_beam_width, - ) + with self._session_admission_lock: + if self._shutdown_started: + raise RuntimeError("TransferWorker is shutting down; new sessions are not accepted") + params = request.py_disaggregated_params + assert params is not None + return TxSession( + request_id=request.py_request_id, + params=params, + sender=self._sender, + aux_buffer=self._aux_buffer, + timeout_s=self._config.tx_timeout_s, + prompt_len=request.prompt_len, + beam_width=request.py_beam_width, + source_owner=request, + ) def create_rx_session(self, request: LlmRequest) -> RxSession: - params = request.py_disaggregated_params - assert params is not None - return RxSession( - request_id=request.py_request_id, - params=params, - receiver=self._receiver, - aux_buffer=self._aux_buffer, - timeout_s=self._config.rx_timeout_s, - prompt_len=request.prompt_len, - beam_width=request.py_beam_width, - ) + with self._session_admission_lock: + if self._shutdown_started: + raise RuntimeError("TransferWorker is shutting down; new sessions are not accepted") + params = request.py_disaggregated_params + assert params is not None + return RxSession( + request_id=request.py_request_id, + params=params, + receiver=self._receiver, + aux_buffer=self._aux_buffer, + timeout_s=self._config.rx_timeout_s, + prompt_len=request.prompt_len, + beam_width=request.py_beam_width, + destination_owner=request, + ) def has_all_peer_req_infos_for_send(self, unique_rid: int) -> bool: return self._sender.has_all_peer_req_infos(unique_rid) @@ -2314,49 +5136,108 @@ def page_table(self): assert self._rank_info is not None return self._rank_info.page_table - def shutdown(self): + def _retain_for_shutdown_retry(self) -> bool: + _NON_DRAINED_TRANSFER_WORKERS.add(self) + return False + + def shutdown(self) -> bool: + """Try to tear down transport resources in ownership-safe order. + + A local agent shutdown is not proof that remote processes can no + longer submit one-sided writes. If receive results have not provided + terminal evidence, this method returns ``False`` and intentionally + keeps listeners, registrations, and mappings alive for a later retry. + """ + attempt_lock = getattr(self, "_shutdown_attempt_lock", None) + if attempt_lock is None: + attempt_lock = threading.Lock() + self._shutdown_attempt_lock = attempt_lock + with attempt_lock: + return self._shutdown_serialized() + + def _shutdown_serialized(self) -> bool: if getattr(self, "_shutdown", False): - return - self._shutdown = True - # Use getattr guards: __init__ may have failed partway, leaving some - # attributes unset. Without them, __del__ -> shutdown() raises - # AttributeError and ZMQ resources from already-created sub-objects - # are never cleaned up. + _NON_DRAINED_TRANSFER_WORKERS.discard(self) + return True + # Install the strong retry owner before the first fallible nested + # teardown. This also covers exceptions reached from __del__/__exit__. + _NON_DRAINED_TRANSFER_WORKERS.add(self) + try: + return self._shutdown_impl() + except Exception as e: + logger.error(f"TransferWorker.shutdown: retaining resources after failure: {e}") + return False + + def _shutdown_impl(self) -> bool: + if getattr(self, "_agent_shutdown_failed", False): + return self._retain_for_shutdown_retry() + admission_lock = getattr(self, "_session_admission_lock", None) + if admission_lock is None: + self._shutdown_started = True + else: + with admission_lock: + self._shutdown_started = True + + # Use getattr guards because __init__ may have failed partway. rank_info_server = getattr(self, "_rank_info_server", None) - if rank_info_server is not None: - rank_info_server.shutdown() - sender = getattr(self, "_sender", None) - if sender is not None: - sender.shutdown() + if rank_info_server is not None and not rank_info_server.shutdown(): + return self._retain_for_shutdown_retry() + receiver = getattr(self, "_receiver", None) if receiver is not None: - receiver.shutdown() - # Close the bounce transport (stops its scatter thread, deregisters its own descriptors and - # frees the VMM buffers) once the receiver listener is stopped so no new scatter is enqueued. - # No-op for the non-bounced path (NoBounceTransport.close); without this the daemon thread + fabric - # buffers leak until process exit. + # Close admission while preserving result ingress. + receiver.begin_shutdown() + + sender = getattr(self, "_sender", None) + if sender is not None and not sender.shutdown(): + return self._retain_for_shutdown_retry() + + # Sender work is locally terminal, but remote writers are independent. + # Keep the receiver and every mapping live until their exact terminal + # results (including AUX/legacy paths) have drained. + if receiver is not None and not receiver.transfers_drained: + logger.warning( + "TransferWorker.shutdown: receive transfers remain active; " + "retaining listeners and registered memory" + ) + return self._retain_for_shutdown_retry() + if receiver is not None and not receiver.shutdown(): + return self._retain_for_shutdown_retry() + bounce = getattr(self, "_bounce", None) if bounce is not None: try: bounce.close() except Exception as e: - logger.warning(f"TransferWorker.shutdown: bounce close failed: {e}") - # Deregister NIXL memory before agent.shutdown so pinned GPU memory is released - # (e.g. when the KV cache manager is recreated after profiling). + logger.error(f"TransferWorker.shutdown: retaining unsafe/live bounce mapping: {e}") + return self._retain_for_shutdown_retry() + agent = getattr(self, "_agent", None) if agent is not None: registered = getattr(self, "_registered_mem", []) - while registered: - desc = registered.pop(0) + for desc in list(registered): try: agent.deregister_memory(desc) except Exception as e: - logger.warning(f"TransferWorker.shutdown: deregister_memory failed: {e}") + logger.error( + f"TransferWorker.shutdown: retaining registered memory after " + f"deregister failure: {e}" + ) + return self._retain_for_shutdown_retry() + registered.remove(desc) try: agent.shutdown() except Exception as e: - logger.warning(f"TransferWorker.shutdown: agent.shutdown error: {e}") + # Some bindings clear their local handle before propagating an + # error. Latch this outcome so a later no-op cannot be mistaken + # for successful quiescence/cleanup. + self._agent_shutdown_failed = True + logger.error(f"TransferWorker.shutdown: agent shutdown failed: {e}") + return self._retain_for_shutdown_retry() self._agent = None + self._shutdown = True + _NON_DRAINED_TRANSFER_WORKERS.discard(self) + return True def __del__(self): try: @@ -2367,5 +5248,18 @@ def __del__(self): def __enter__(self): return self - def __exit__(self, _exc_type, _exc_val, _exc_tb): - self.shutdown() + def __exit__(self, exc_type, _exc_val, _exc_tb): + try: + drained = self.shutdown() + except Exception as e: + if exc_type is None: + raise + logger.error( + f"TransferWorker context cleanup failed while propagating an exception: {e}" + ) + return False + if not drained: + if exc_type is None: + raise RuntimeError("TransferWorker context exited before resources drained") + logger.error("TransferWorker retained resources while propagating an exception") + return False diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 4fe7fa28e50f..5e52de97db55 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -1,6 +1,24 @@ +# 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 threading import time import uuid from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass from itertools import chain from typing import Any, Callable, Dict, List, Optional, cast @@ -39,6 +57,49 @@ from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping +# A non-drained transceiver is deliberately retained even if its caller drops +# the last reference. It owns the LlmRequest/session objects that keep receive +# targets from being reclaimed before terminal transport evidence arrives. +_NON_DRAINED_TRANSCEIVERS: set["KvCacheTransceiverV2"] = set() + + +@dataclass +class _ContextRetirementProgress: + """One rank's durable preparation of a context-side terminal candidate.""" + + rid: int + session: TxSessionBase + request: LlmRequest + local_outcome: str + mark_complete: bool + outcome: Optional[str] = None + report_result: bool = True + session_closed: bool = False + maps_validated: bool = False + state_updated: bool = False + request_map_retired: bool = False + session_map_retired: bool = False + + +@dataclass +class _GenerationRetirementProgress: + """One rank's durable preparation of a generation-side terminal candidate.""" + + rid: int + session: RxSessionBase + request: LlmRequest + local_outcome: str + outcome: Optional[str] = None + report_result: bool = True + kv_size_set: bool = False + aux_applied: bool = False + history_checked: bool = False + session_closed: bool = False + maps_validated: bool = False + state_updated: bool = False + request_map_retired: bool = False + session_map_retired: bool = False + def _find_consensus_request_ids(request_ids_all_ranks, sync_size): frequency_map = defaultdict(int) @@ -54,6 +115,10 @@ def _find_consensus_request_ids(request_ids_all_ranks, sync_size): class KvCacheTransceiverV2(KvCacheTransceiver): + @property + def requires_physical_drain_before_request_release(self) -> bool: + return True + def __init__( self, mapping: Mapping, @@ -61,6 +126,16 @@ def __init__( kv_cache_manager: KVCacheManager, cache_transceiver_config: CacheTransceiverConfig, ): + self._session_admission_lock = threading.RLock() + # Serialize every collective-bearing status/preparation call. Shutdown + # does not wait on this lock: it closes admission and asks active + # sessions to cancel, then returns non-drained until the active call + # leaves the collective stream. + self._status_call_lock = threading.Lock() + self._active_status_calls = 0 + self._retirement_fault: Optional[BaseException] = None + self._shutdown_started = False + self._shutdown = False self._dist: Distributed = dist self._kv_cache_manager = kv_cache_manager self._mapping = mapping @@ -100,6 +175,16 @@ def __init__( self._send_reqs = {} self._recv_reqs = {} self._wait_reqs = {} + # Exact wait owners whose cancellation raced a collective-bearing + # prepare call. Promotion must ignore them until cancel_request retries + # after that call exits and retires the wait entry. + self._cancelled_wait_reqs: Dict[int, LlmRequest] = {} + # Terminal candidates are enrolled provisionally before the existing + # outcome consensus. Keep their exact local owners and phase progress + # until this rank commits the distributed decision, including across + # retryable local preparation failures. + self._context_retirements: Dict[int, _ContextRetirementProgress] = {} + self._generation_retirements: Dict[int, _GenerationRetirementProgress] = {} self._page_table = self._transfer_worker.page_table # _slice_num_bytes() is this rank's KV shard, so scale by tp_size to get the request total (kv_cache_size), # except under attention DP where the local count already is the total. @@ -147,25 +232,207 @@ def _exchange_rank_info(self): logger.info(f"layer_num_per_pp: {layer_num_per_pp}") logger.info(f"self._context_info_endpoint: {self._context_info_endpoint}") - def shutdown(self): - if getattr(self, "_shutdown", False): - return - self._shutdown = True - for session in list(self._send_sessions.values()): - session.close() - for session in list(self._recv_sessions.values()): - session.close() - self._send_sessions.clear() - self._send_reqs.clear() - self._recv_sessions.clear() - self._recv_reqs.clear() - self._transfer_worker.shutdown() + def _get_session_admission_lock(self): + """Return the launch/shutdown gate, including partial-init cleanup.""" + admission_lock = getattr(self, "_session_admission_lock", None) + if admission_lock is None: + admission_lock = threading.RLock() + self._session_admission_lock = admission_lock + return admission_lock + + def _get_status_call_lock(self): + """Return the collective-stream lock, including object.__new__ fixtures.""" + status_lock = getattr(self, "_status_call_lock", None) + if status_lock is None: + status_lock = threading.Lock() + self._status_call_lock = status_lock + return status_lock + + @contextmanager + def _status_call(self): + """Serialize one complete collective-bearing transceiver call. + + Once admitted, a call finishes its collective sequence even if a + concurrent shutdown closes admission. Shutdown may initiate session + cancellation, but it cannot drain retirement ledgers until the active + call leaves this scope. + """ + status_lock = self._get_status_call_lock() + status_lock.acquire() + admitted = False + try: + with self._get_session_admission_lock(): + fault = getattr(self, "_retirement_fault", None) + if fault is not None: + raise RuntimeError( + "KV cache transceiver retirement invariant failed" + ) from fault + if not getattr(self, "_shutdown_started", False): + self._active_status_calls = getattr(self, "_active_status_calls", 0) + 1 + admitted = True + yield admitted + finally: + if admitted: + with self._get_session_admission_lock(): + self._active_status_calls -= 1 + status_lock.release() + + def _latch_retirement_fault(self, error: BaseException) -> None: + """Fail stop after an asserted-infallible retirement commit fails.""" + with self._get_session_admission_lock(): + if getattr(self, "_retirement_fault", None) is None: + self._retirement_fault = error + self._shutdown_started = True + _NON_DRAINED_TRANSCEIVERS.add(self) + + @staticmethod + def _retirement_sessions(retirements: dict) -> set[int]: + return {id(progress.session) for progress in retirements.values()} + + def _cancel_sessions_for_shutdown(self) -> set[int]: + """Initiate cancellation without retiring any request/session owner.""" + context_retirement_sessions = self._retirement_sessions(self._get_context_retirements()) + generation_retirement_sessions = self._retirement_sessions( + self._get_generation_retirements() + ) + cancel_failed_session_ids: set[int] = set() + for direction, sessions in ( + ("send", self._send_sessions), + ("receive", self._recv_sessions), + ): + retirement_sessions = ( + context_retirement_sessions + if direction == "send" + else generation_retirement_sessions + ) + for session in list(sessions.values()): + if id(session) in retirement_sessions: + continue + try: + # Cancellation closes future publication but leaves every + # already-exposed resource owned until terminal evidence. + session.cancel() + except Exception as error: + cancel_failed_session_ids.add(id(session)) + logger.error( + f"KvCacheTransceiverV2 shutdown failed to cancel {direction} " + f"session {session.disagg_request_id}: {error}" + ) + return cancel_failed_session_ids + + def shutdown(self) -> bool: + admission_lock = self._get_session_admission_lock() + with admission_lock: + if getattr(self, "_shutdown", False): + return True + # Retain the ownership root before any fallible cancellation or + # cleanup. A failed shutdown must remain retryable even when the + # caller releases its last reference after this attempt. + _NON_DRAINED_TRANSCEIVERS.add(self) + # Close admission before taking any session snapshot. The same lock + # covers owner enrollment through publication/worker enqueue, so a + # launch cannot resume against a worker that shutdown just destroyed. + self._shutdown_started = True + cancel_failed_session_ids = self._cancel_sessions_for_shutdown() + # A status call may be blocked waiting for receive resources to + # drain. Cancellation above lets it make progress; waiting here + # would deadlock shutdown behind the very call it must unblock. + # Keep every owner/ledger intact and let a later shutdown retry do + # the local drain after the collective stream is idle. + if getattr(self, "_active_status_calls", 0) > 0: + return False + # A prior status poll may already have crossed distributed + # consensus, or may have prepared a local terminal candidate. + # With no active status call, shutdown can locally finish those + # exact owners without entering another distributed collective. + self._drain_context_retirements() + self._drain_generation_retirements() + try: + worker_drained = self._transfer_worker.shutdown() + except Exception as e: + logger.error( + "KvCacheTransceiverV2 worker shutdown did not drain; retaining " + f"sessions and request resources for retry: {e}" + ) + return False + if not worker_drained: + return False + + # Worker drain can make a previously failing session close + # retryable. Advance the post-consensus ledgers again, but retain + # any entry whose local commit is still incomplete. + self._drain_context_retirements() + self._drain_generation_retirements() + context_retirement_sessions = { + id(progress.session) for progress in self._get_context_retirements().values() + } + generation_retirement_sessions = { + id(progress.session) for progress in self._get_generation_retirements().values() + } + + close_failed = False + for direction, sessions, requests in ( + ("send", self._send_sessions, self._send_reqs), + ("receive", self._recv_sessions, self._recv_reqs), + ): + for rid, session in list(sessions.items()): + if id(session) in cancel_failed_session_ids: + continue + retirement_sessions = ( + context_retirement_sessions + if direction == "send" + else generation_retirement_sessions + ) + if id(session) in retirement_sessions: + continue + try: + session.close() + except Exception as e: + close_failed = True + logger.error( + f"KvCacheTransceiverV2 shutdown failed to close {direction} " + f"session {session.disagg_request_id}: {e}" + ) + continue + if sessions.get(rid) is session: + sessions.pop(rid, None) + requests.pop(rid, None) + + if ( + cancel_failed_session_ids + or close_failed + or self._get_context_retirements() + or self._get_generation_retirements() + ): + return False + getattr(self, "_wait_reqs", {}).clear() + self._get_cancelled_wait_reqs().clear() + self._shutdown = True + _NON_DRAINED_TRANSCEIVERS.discard(self) + return True def __enter__(self): return self def __exit__(self, _exc_type, _exc_val, _exc_tb): - self.shutdown() + try: + drained = self.shutdown() + except Exception as shutdown_error: + if _exc_type is None: + raise + logger.error( + "KvCacheTransceiverV2 shutdown also failed while propagating " + f"{_exc_type.__name__}: {shutdown_error}" + ) + return False + if not drained: + if _exc_type is None: + raise RuntimeError("KV cache transceiver shutdown did not drain") + logger.error( + "KV cache transceiver shutdown did not drain while propagating " + f"{_exc_type.__name__}" + ) + return False def _create_kv_slice(self, req: LlmRequest) -> KVSlice: adapter = self._reuse_adapter @@ -344,6 +611,12 @@ def _gen_consensus(self, local_ids: list) -> list: all_ranks = self._gen_allgather(local_ids) if self._gen_need_sync else [local_ids] return _find_consensus_request_ids(all_ranks, sync_size) + @staticmethod + def _allgather_or_passthrough(local_value: list, allgather: Callable, need_sync: bool) -> list: + if not need_sync: + return [list(local_value)] + return list(allgather(list(local_value))) + @staticmethod def _union(all_lists: List[List[int]]) -> set: merged: set = set() @@ -362,48 +635,92 @@ def _intersection(all_lists: List[List[int]], n_ranks: int) -> set: return {rid for rid, c in cnt.items() if c == n_ranks} def _consensus_outcome( - self, to_process, cancelled, failed, completed, allgather: Callable, need_sync: bool + self, + to_process, + cancelled, + failed, + completed, + cleanup_ready, + allgather: Callable, + need_sync: bool, ): - # CANCELLED/FAILED on any rank → global; COMPLETED only when ALL ranks agree. - # Batch the three id lists into one allgather to cut the per-step collective count. - if not need_sync: - all_c, all_f, all_done = [list(cancelled)], [list(failed)], [list(completed)] - else: - packed = list(allgather([list(cancelled), list(failed), list(completed)])) - all_c = [p[0] for p in packed] - all_f = [p[1] for p in packed] - all_done = [p[2] for p in packed] - n = len(all_c) + """Agree on logical outcome without outrunning local physical drain. + + Cancellation/failure is a union because one rank's logical failure must + fail the distributed request. Cleanup readiness is an intersection: + every rank must independently prove its local source/target accessors + terminal before any rank removes its request owner. + """ + payload = [ + list(cancelled), + list(failed), + list(completed), + list(cleanup_ready), + ] + gathered = self._allgather_or_passthrough(payload, allgather, need_sync) + all_c = [rank_payload[0] for rank_payload in gathered] + all_f = [rank_payload[1] for rank_payload in gathered] + all_done = [rank_payload[2] for rank_payload in gathered] + all_ready = [rank_payload[3] for rank_payload in gathered] + n = len(gathered) global_cancelled = self._union(all_c) global_failed = self._union(all_f) global_completed = self._intersection(all_done, n) - new_cancelled = [rid for rid in to_process if rid in global_cancelled] + globally_ready = self._intersection(all_ready, n) + new_cancelled = [ + rid for rid in to_process if rid in global_cancelled and rid in globally_ready + ] cancel_set = set(new_cancelled) - new_failed = [rid for rid in to_process if rid in global_failed and rid not in cancel_set] + new_failed = [ + rid + for rid in to_process + if rid in global_failed and rid in globally_ready and rid not in cancel_set + ] terminal = cancel_set | set(new_failed) new_completed = [ - rid for rid in to_process if rid in global_completed and rid not in terminal + rid + for rid in to_process + if rid in global_completed and rid in globally_ready and rid not in terminal ] return new_cancelled, new_failed, new_completed - def _gen_consensus_outcome(self, to_process, cancelled, failed, completed): + def _gen_consensus_outcome(self, to_process, cancelled, failed, completed, cleanup_ready): return self._consensus_outcome( - to_process, cancelled, failed, completed, self._gen_allgather, self._gen_need_sync + to_process, + cancelled, + failed, + completed, + cleanup_ready, + self._gen_allgather, + self._gen_need_sync, ) - def _ctx_consensus_outcome(self, to_process, cancelled, failed, completed, timed_out): + def _ctx_consensus_outcome( + self, to_process, cancelled, failed, completed, timed_out, cleanup_ready + ): # TP first, then PP. timed_out is local-only (back-off signal). c, f, d = self._consensus_outcome( to_process, cancelled, failed, completed, + cleanup_ready, self._dist.tp_allgather, self._ctx_need_tp_sync, ) if self._ctx_need_pp_sync: pp_allgather: Callable = getattr(self._dist, "pp_allgather") - c, f, d = self._consensus_outcome(to_process, c, f, d, pp_allgather, True) + # A TP group appears ready to the PP stage only after the first + # consensus admitted one terminal outcome for that request. + c, f, d = self._consensus_outcome( + to_process, + c, + f, + d, + c + f + d, + pp_allgather, + True, + ) return c, f, d, timed_out def _collect_done(self, sessions: dict, reqs: dict): @@ -420,7 +737,7 @@ def _build_to_process( self, sessions: dict, consensus: list, wait_num: int, block_all: bool ) -> list: if block_all: - return list(sessions.keys()) + return list(dict.fromkeys([*consensus, *sessions.keys()])) to_process = list(consensus) for rid in sessions: if len(to_process) >= wait_num: @@ -429,12 +746,485 @@ def _build_to_process( to_process.append(rid) return to_process - def _close_failed_sessions(self, sessions: dict, reqs: dict, failed: list): - for rid in failed: - reqs[rid].state = LlmRequestState.DISAGG_TRANS_ERROR - sessions[rid].close() - del reqs[rid] - del sessions[rid] + def _get_context_retirements(self) -> Dict[int, _ContextRetirementProgress]: + """Return the context retirement ledger, including object.__new__ fixtures.""" + + retirements = getattr(self, "_context_retirements", None) + if retirements is None: + retirements = {} + self._context_retirements = retirements + return retirements + + def _get_generation_retirements(self) -> Dict[int, _GenerationRetirementProgress]: + """Return the generation retirement ledger, including object.__new__ fixtures.""" + + retirements = getattr(self, "_generation_retirements", None) + if retirements is None: + retirements = {} + self._generation_retirements = retirements + return retirements + + def _get_cancelled_wait_reqs(self) -> Dict[int, LlmRequest]: + """Return wait-owner cancellation markers, including test fixtures.""" + cancelled = getattr(self, "_cancelled_wait_reqs", None) + if cancelled is None: + cancelled = {} + self._cancelled_wait_reqs = cancelled + return cancelled + + @staticmethod + def _validate_exact_map_entry(mapping: dict, rid: int, expected, label: str) -> None: + """Validate the commit owner before advertising cleanup readiness.""" + if mapping.get(rid) is not expected: + raise RuntimeError(f"{label} {rid} changed identity before retirement commit") + + @staticmethod + def _retire_exact_map_entry(mapping: dict, rid: int, expected, label: str) -> None: + """Idempotently remove one exact owner during local shutdown retry.""" + + current = mapping.get(rid) + if current is not None and current is not expected: + raise RuntimeError(f"{label} {rid} changed identity during retirement") + if current is expected: + mapping.pop(rid, None) + + def _enroll_context_retirements( + self, + outcomes: tuple[tuple[str, list], ...], + sessions: dict, + reqs: dict, + *, + mark_complete: bool, + ) -> None: + """Persist exact owners before fallible context-side preparation.""" + + retirements = self._get_context_retirements() + with self._get_session_admission_lock(): + for outcome, rids in outcomes: + for rid in rids: + session = sessions.get(rid) + req = reqs.get(rid) + if session is None or req is None: + continue + existing = retirements.get(rid) + if existing is not None: + if ( + existing.session is not session + or existing.request is not req + or existing.local_outcome != outcome + ): + raise RuntimeError( + f"Conflicting local context terminal candidate for request {rid}" + ) + existing.mark_complete = existing.mark_complete or mark_complete + continue + progress = _ContextRetirementProgress( + rid=rid, + session=session, + request=req, + local_outcome=outcome, + mark_complete=mark_complete, + ) + # Cancellation may retire this snapshot while the status + # call is outside the admission lock. Keep a no-op + # provisional record for collective symmetry, but never + # mutate a replacement owner that reused the integer RID. + exact_owner = ( + self._send_reqs.get(rid) is req and self._send_sessions.get(rid) is session + ) + if not exact_owner: + progress.report_result = False + progress.session_closed = True + progress.maps_validated = True + progress.state_updated = True + progress.request_map_retired = True + progress.session_map_retired = True + retirements[rid] = progress + + def _prepare_context_retirement(self, progress: _ContextRetirementProgress) -> None: + """Complete fallible local cleanup before the distributed commit.""" + + if not progress.session_closed: + progress.session.close() + progress.session_closed = True + if not progress.maps_validated: + self._validate_exact_map_entry( + self._send_reqs, progress.rid, progress.request, "send request" + ) + self._validate_exact_map_entry( + self._send_sessions, progress.rid, progress.session, "send session" + ) + progress.maps_validated = True + + @staticmethod + def _context_retirement_prepared(progress: _ContextRetirementProgress) -> bool: + return progress.session_closed and progress.maps_validated + + def _prepare_context_retirements(self) -> None: + """Retry every provisional context cleanup without adding collectives.""" + with self._get_session_admission_lock(): + for rid, progress in list(self._get_context_retirements().items()): + try: + self._prepare_context_retirement(progress) + except Exception as error: + logger.error( + "Context transfer cleanup failed locally for request " + f"{rid}; retaining exact phase progress for retry: {error}" + ) + + def _context_candidate_payload( + self, to_process: list[int] + ) -> tuple[list[int], list[int], list[int], list[int]]: + cancelled: list[int] = [] + failed: list[int] = [] + completed: list[int] = [] + cleanup_ready: list[int] = [] + with self._get_session_admission_lock(): + retirements = self._get_context_retirements() + for rid in to_process: + progress = retirements.get(rid) + if progress is None: + continue + if progress.local_outcome == "cancelled": + cancelled.append(rid) + elif progress.local_outcome == "failed": + failed.append(rid) + else: + completed.append(rid) + if self._context_retirement_prepared(progress): + cleanup_ready.append(rid) + return cancelled, failed, completed, cleanup_ready + + def _commit_context_retirement(self, progress: _ContextRetirementProgress) -> None: + """Publish one consensus decision using asserted-infallible operations.""" + + if not progress.state_updated: + outcome = progress.outcome or progress.local_outcome + target_state = None + if outcome == "completed": + if progress.mark_complete: + target_state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + else: + target_state = LlmRequestState.DISAGG_TRANS_ERROR + if target_state is not None and progress.request.state != target_state: + progress.request.state = target_state + progress.state_updated = True + + if not progress.request_map_retired: + self._retire_exact_map_entry( + self._send_reqs, progress.rid, progress.request, "send request" + ) + progress.request_map_retired = True + if not progress.session_map_retired: + self._retire_exact_map_entry( + self._send_sessions, progress.rid, progress.session, "send session" + ) + progress.session_map_retired = True + + @staticmethod + def _context_retirement_committed(progress: _ContextRetirementProgress) -> bool: + return ( + progress.state_updated and progress.request_map_retired and progress.session_map_retired + ) + + def _commit_context_decisions( + self, + outcomes: tuple[tuple[str, list[int]], ...], + ) -> tuple[list[int], list[int]]: + """Commit the outcome-consensus result without another sync round.""" + completed: list[int] = [] + failed: list[int] = [] + try: + with self._get_session_admission_lock(): + retirements = self._get_context_retirements() + for outcome, rids in outcomes: + for rid in rids: + progress = retirements.get(rid) + if progress is None or not self._context_retirement_prepared(progress): + raise RuntimeError( + f"Context outcome consensus admitted an unprepared request {rid}" + ) + if progress.outcome is not None and progress.outcome != outcome: + raise RuntimeError(f"Context request {rid} changed global outcome") + progress.outcome = outcome + self._commit_context_retirement(progress) + if not self._context_retirement_committed(progress): + raise RuntimeError(f"Context request {rid} commit was incomplete") + retirements.pop(rid) + if not progress.report_result: + continue + if outcome == "completed": + completed.append(rid) + else: + failed.append(rid) + except Exception as error: + self._latch_retirement_fault(error) + raise + return completed, failed + + def _drain_context_retirements(self) -> tuple[list[int], list[int]]: + """Locally force decided context retirement during shutdown.""" + + completed: list[int] = [] + failed: list[int] = [] + retirements = self._get_context_retirements() + with self._get_session_admission_lock(): + for rid, progress in list(retirements.items()): + if retirements.get(rid) is not progress: + continue + try: + self._prepare_context_retirement(progress) + if progress.outcome is None: + progress.outcome = progress.local_outcome + self._commit_context_retirement(progress) + except Exception as error: + logger.error( + "Context transfer retirement failed locally for request " + f"{rid}; retaining exact phase progress for retry: {error}" + ) + continue + if retirements.get(rid) is progress: + retirements.pop(rid, None) + if not progress.report_result: + continue + if (progress.outcome or progress.local_outcome) == "completed": + completed.append(rid) + else: + failed.append(rid) + return completed, failed + + def _enroll_generation_retirements( + self, + outcomes: tuple[tuple[str, list], ...], + sessions: dict, + reqs: dict, + ) -> None: + """Persist exact owners before fallible generation-side preparation.""" + + retirements = self._get_generation_retirements() + with self._get_session_admission_lock(): + for outcome, rids in outcomes: + for rid in rids: + session = sessions.get(rid) + req = reqs.get(rid) + if session is None or req is None: + continue + existing = retirements.get(rid) + if existing is not None: + if ( + existing.session is not session + or existing.request is not req + or existing.local_outcome != outcome + ): + raise RuntimeError( + f"Conflicting local generation terminal candidate for request {rid}" + ) + continue + progress = _GenerationRetirementProgress( + rid=rid, + session=session, + request=req, + local_outcome=outcome, + ) + exact_owner = ( + self._recv_reqs.get(rid) is req and self._recv_sessions.get(rid) is session + ) + if not exact_owner: + progress.report_result = False + progress.kv_size_set = True + progress.aux_applied = True + progress.history_checked = True + progress.session_closed = True + progress.maps_validated = True + progress.state_updated = True + progress.request_map_retired = True + progress.session_map_retired = True + retirements[rid] = progress + + def _prepare_generation_retirement(self, progress: _GenerationRetirementProgress) -> None: + """Complete fallible local generation cleanup before commit.""" + + if progress.local_outcome == "completed": + if not progress.kv_size_set: + progress.request.set_kv_cache_size( + getattr(progress.request, "py_kv_cache_xfer_bytes", 0) + ) + progress.kv_size_set = True + if not progress.aux_applied: + if self._need_aux_transfer(progress.request): + self._apply_aux(progress.session, progress.request) + progress.aux_applied = True + if not progress.history_checked: + self._assert_disagg_history_declared(progress.request) + progress.history_checked = True + else: + # These phases are not part of cancelled/failed retirement. + progress.kv_size_set = True + progress.aux_applied = True + progress.history_checked = True + + if not progress.session_closed: + progress.session.close() + progress.session_closed = True + if not progress.maps_validated: + self._validate_exact_map_entry( + self._recv_reqs, progress.rid, progress.request, "receive request" + ) + self._validate_exact_map_entry( + self._recv_sessions, progress.rid, progress.session, "receive session" + ) + progress.maps_validated = True + + def _commit_generation_retirement(self, progress: _GenerationRetirementProgress) -> None: + """Publish one consensus decision using asserted-infallible operations.""" + + if not progress.state_updated: + outcome = progress.outcome or progress.local_outcome + target_state = None + if outcome == "completed": + target_state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + elif outcome == "failed": + target_state = LlmRequestState.DISAGG_TRANS_ERROR + if target_state is not None and progress.request.state != target_state: + progress.request.state = target_state + progress.state_updated = True + + if not progress.request_map_retired: + self._retire_exact_map_entry( + self._recv_reqs, progress.rid, progress.request, "receive request" + ) + progress.request_map_retired = True + if not progress.session_map_retired: + self._retire_exact_map_entry( + self._recv_sessions, progress.rid, progress.session, "receive session" + ) + progress.session_map_retired = True + + @staticmethod + def _generation_retirement_prepared(progress: _GenerationRetirementProgress) -> bool: + return ( + progress.kv_size_set + and progress.aux_applied + and progress.history_checked + and progress.session_closed + and progress.maps_validated + ) + + @staticmethod + def _generation_retirement_committed(progress: _GenerationRetirementProgress) -> bool: + return ( + progress.state_updated and progress.request_map_retired and progress.session_map_retired + ) + + def _prepare_generation_retirements(self) -> None: + """Retry every provisional generation cleanup without extra collectives.""" + with self._get_session_admission_lock(): + for rid, progress in list(self._get_generation_retirements().items()): + try: + self._prepare_generation_retirement(progress) + except Exception as error: + logger.error( + "Generation transfer cleanup failed locally for request " + f"{rid}; retaining exact phase progress for retry: {error}" + ) + + def _generation_candidate_payload( + self, to_process: list[int] + ) -> tuple[list[int], list[int], list[int], list[int]]: + cancelled: list[int] = [] + failed: list[int] = [] + completed: list[int] = [] + cleanup_ready: list[int] = [] + with self._get_session_admission_lock(): + retirements = self._get_generation_retirements() + for rid in to_process: + progress = retirements.get(rid) + if progress is None: + continue + if progress.local_outcome == "cancelled": + cancelled.append(rid) + elif progress.local_outcome == "failed": + failed.append(rid) + else: + completed.append(rid) + if self._generation_retirement_prepared(progress): + cleanup_ready.append(rid) + return cancelled, failed, completed, cleanup_ready + + def _commit_generation_decisions( + self, + outcomes: tuple[tuple[str, list[int]], ...], + ) -> tuple[list[int], list[int], list[LlmRequest]]: + """Commit the outcome-consensus result without another sync round.""" + completed: list[int] = [] + failed: list[int] = [] + cancelled: list[LlmRequest] = [] + try: + with self._get_session_admission_lock(): + retirements = self._get_generation_retirements() + for outcome, rids in outcomes: + for rid in rids: + progress = retirements.get(rid) + if progress is None or not self._generation_retirement_prepared(progress): + raise RuntimeError( + f"Generation outcome consensus admitted an unprepared request {rid}" + ) + if progress.outcome is not None and progress.outcome != outcome: + raise RuntimeError(f"Generation request {rid} changed global outcome") + progress.outcome = outcome + self._commit_generation_retirement(progress) + if not self._generation_retirement_committed(progress): + raise RuntimeError(f"Generation request {rid} commit was incomplete") + retirements.pop(rid) + if not progress.report_result: + continue + if outcome == "completed": + completed.append(rid) + elif outcome == "failed": + failed.append(rid) + else: + cancelled.append(progress.request) + except Exception as error: + self._latch_retirement_fault(error) + raise + return completed, failed, cancelled + + def _drain_generation_retirements( + self, + ) -> tuple[list[int], list[int], list[LlmRequest]]: + """Locally force decided generation retirement during shutdown.""" + + completed: list[int] = [] + failed: list[int] = [] + cancelled: list[LlmRequest] = [] + retirements = self._get_generation_retirements() + with self._get_session_admission_lock(): + for rid, progress in list(retirements.items()): + if retirements.get(rid) is not progress: + continue + try: + self._prepare_generation_retirement(progress) + if progress.outcome is None: + progress.outcome = progress.local_outcome + self._commit_generation_retirement(progress) + except Exception as error: + logger.error( + "Generation transfer retirement failed locally for request " + f"{rid}; retaining exact phase progress for retry: {error}" + ) + continue + if retirements.get(rid) is progress: + retirements.pop(rid, None) + if not progress.report_result: + continue + outcome = progress.outcome or progress.local_outcome + if outcome == "completed": + completed.append(rid) + elif outcome == "failed": + failed.append(rid) + else: + cancelled.append(progress.request) + return completed, failed, cancelled def _apply_aux(self, session, req: LlmRequest): """Unpack aux tokens from session into request's context_phase_params.""" @@ -481,97 +1271,283 @@ def _finalize_send(self, req: LlmRequest, session: TxSessionBase): @nvtx_range("KvCacheTransceiverV2.respond_and_send_async") def respond_and_send_async(self, req: LlmRequest): - self._ever_had_send_session = True - session = self._get_or_create_send_session(req) - req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS - session.send(self._create_kv_slice(req)) - self._finalize_send(req, session) + with self._get_session_admission_lock(): + if getattr(self, "_shutdown_started", False): + raise RuntimeError("KV cache transceiver is shutting down") + rid = get_unique_rid(req) + assert rid is not None + wait_owner = getattr(self, "_wait_reqs", {}).get(rid) + if wait_owner is not None: + if wait_owner is not req: + raise RuntimeError( + f"waiting request {rid} already has a different live source owner" + ) + raise RuntimeError(f"send request {rid} is still waiting for scheduler promotion") + cancelled_wait_owner = self._get_cancelled_wait_reqs().get(rid) + if cancelled_wait_owner is not None: + raise RuntimeError(f"send request {rid} has a pending wait-owner cancellation") + if rid in self._get_context_retirements(): + raise RuntimeError(f"send request {rid} has a pending terminal retirement") + existing_owner = self._send_reqs.get(rid) + if existing_owner is not None and existing_owner is not req: + raise RuntimeError(f"send request {rid} already has a different live source owner") + owner_newly_enrolled = existing_owner is None + self._ever_had_send_session = True + # Enroll the source owner before a worker can launch gather/NIXL. + self._send_reqs[rid] = req + session = None + try: + session = self._get_or_create_send_session(req) + req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + session.send(self._create_kv_slice(req)) + self._finalize_send(req, session) + except Exception: + if session is None: + if owner_newly_enrolled: + self._send_reqs.pop(rid, None) + else: + try: + session.cancel() + if not session.has_transferring_tasks(): + session.close() + self._send_sessions.pop(rid, None) + self._send_reqs.pop(rid, None) + except Exception as cleanup_error: + logger.warning( + f"respond_and_send_async: retaining rid={rid} after " + f"exception-path cleanup could not prove drain: {cleanup_error}" + ) + raise @nvtx_range("KvCacheTransceiverV2.request_and_receive_sync") def request_and_receive_sync(self, req: LlmRequest): + admission_lock = self._get_session_admission_lock() + with admission_lock: + if getattr(self, "_shutdown_started", False): + raise RuntimeError("KV cache transceiver is shutting down") + admitted = self._request_and_receive_sync_admitted(req) + + if admitted is None: + return + session, kv_slice = admitted + rid = get_unique_rid(req) + assert rid is not None + + try: + result = session.wait_complete(blocking=True) + except Exception: + with admission_lock: + if self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + session.cancel() + if not session.has_transferring_tasks(): + session.close() + self._recv_sessions.pop(rid, None) + self._recv_reqs.pop(rid, None) + raise + + # Cancellation and shutdown use the same gate. They may have retired + # this session while wait_complete() was blocked, so only the exact + # enrolled owner may consume AUX data or remove the map entries. + with admission_lock: + owns_session = ( + self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req + ) + if not owns_session: + if result != WaitResult.COMPLETED or session.status == SessionStatus.CANCELLED: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + return + try: + if result == WaitResult.COMPLETED: + # KV-transfer timing setters deferred to #15871 (clock-source consistency); + # size only. + req.set_kv_cache_size( + self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor + ) + if self._need_aux_transfer(req): + self._apply_aux(session, req) + self._assert_disagg_history_declared(req) + req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + else: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + except Exception: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + raise + finally: + session.close() + if self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req: + self._recv_sessions.pop(rid, None) + self._recv_reqs.pop(rid, None) + + def _request_and_receive_sync_admitted(self, req: LlmRequest): rid = get_unique_rid(req) + assert rid is not None + if rid in self._get_generation_retirements(): + raise RuntimeError(f"receive request {rid} has a pending terminal retirement") if rid in self._recv_sessions: + if self._recv_reqs.get(rid) is not req: + raise RuntimeError( + f"receive request {rid} already has a different live destination owner" + ) logger.warning( f"request_and_receive_sync: rid={rid} already has a recv session, skipping" ) return req.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + self._recv_reqs[rid] = req session = None try: session = self._transfer_worker.create_rx_session(req) self._recv_sessions[rid] = session - self._recv_reqs[rid] = req kv_slice = self._create_kv_slice(req) session.receive(kv_slice) - result = session.wait_complete(blocking=True) - - if result == WaitResult.COMPLETED: - # KV-transfer timing setters deferred to #15871 (clock-source consistency); size only. - req.set_kv_cache_size(self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor) - if self._need_aux_transfer(req): - self._apply_aux(session, req) - self._assert_disagg_history_declared(req) - req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - else: - req.state = LlmRequestState.DISAGG_TRANS_ERROR + return session, kv_slice except Exception: req.state = LlmRequestState.DISAGG_TRANS_ERROR + if session is None: + if self._recv_reqs.get(rid) is req: + self._recv_reqs.pop(rid, None) + else: + session.cancel() + if not session.has_transferring_tasks(): + session.close() + if self._recv_sessions.get(rid) is session: + self._recv_sessions.pop(rid, None) + if self._recv_reqs.get(rid) is req: + self._recv_reqs.pop(rid, None) raise - finally: - if session is not None: - session.close() - self._recv_sessions.pop(rid, None) - self._recv_reqs.pop(rid, None) @nvtx_range("KvCacheTransceiverV2.request_and_receive_async") def request_and_receive_async(self, req: LlmRequest): - self._ever_had_recv_session = True - rid = get_unique_rid(req) - if rid in self._recv_sessions: - logger.warning( - f"request_and_receive_async: rid={rid} already has a recv session, skipping" - ) - return - req.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS - session = self._transfer_worker.create_rx_session(req) - self._recv_sessions[rid] = session - kv_slice = self._create_kv_slice(req) - req.py_kv_cache_xfer_bytes = self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor - session.receive(kv_slice) - self._recv_reqs[rid] = req + with self._get_session_admission_lock(): + if getattr(self, "_shutdown_started", False): + raise RuntimeError("KV cache transceiver is shutting down") + self._ever_had_recv_session = True + rid = get_unique_rid(req) + assert rid is not None + if rid in self._get_generation_retirements(): + raise RuntimeError(f"receive request {rid} has a pending terminal retirement") + if rid in self._recv_sessions: + if self._recv_reqs.get(rid) is not req: + raise RuntimeError( + f"receive request {rid} already has a different live destination owner" + ) + logger.warning( + f"request_and_receive_async: rid={rid} already has a recv session, skipping" + ) + return + req.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + # Enroll the request owner before session construction can allocate an + # AUX slot and before receive() can publish KV/AUX target addresses. + # If either step raises, keep every owner enrolled so cleanup continues + # through the normal physical-drain gate rather than dropping the only + # Python-level owner during exception unwinding. + self._recv_reqs[rid] = req + session = None + try: + session = self._transfer_worker.create_rx_session(req) + self._recv_sessions[rid] = session + kv_slice = self._create_kv_slice(req) + req.py_kv_cache_xfer_bytes = ( + self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor + ) + session.receive(kv_slice) + except Exception: + if session is None: + # Constructor rollback owns any pre-publication allocation. + self._recv_reqs.pop(rid, None) + else: + try: + session.cancel() + drained = not session.has_transferring_tasks() + if drained: + session.close() + self._recv_sessions.pop(rid, None) + self._recv_reqs.pop(rid, None) + except Exception as cleanup_error: + logger.warning( + f"request_and_receive_async: retaining rid={rid} after " + f"exception-path cleanup could not prove drain: {cleanup_error}" + ) + raise def check_context_transfer_status( - self, at_least_request_num: Optional[int], mark_complete: bool = False + self, + at_least_request_num: Optional[int], + mark_complete: bool = False, + ): + with self._status_call() as admitted: + if not admitted: + return [], [] + return self._check_context_transfer_status_impl(at_least_request_num, mark_complete) + + def _check_context_transfer_status_impl( + self, + at_least_request_num: Optional[int], + mark_complete: bool, ): - # A worker that never sends KV has nothing to reconcile here, so skip the consensus. Safe - # because the flag flips together on every rank and never resets, so they all skip in step; - # gating on the live session dict instead would not be, since a cancel clears it per-rank. - # Keep the original sweep (only when tp/pp sync is on) so nothing is leaked. - if not self._ever_had_send_session: - if self._ctx_need_tp_sync or self._ctx_need_pp_sync: - self._transfer_worker.sweep_stale_req_infos() - return [], [] + # This is also the native sender's periodic control-result progress + # hook. Run it before the idle fast path so cancel-before-session + # failures are retried even when no TxSession is ever created. + self._prepare_context_retirements() + with self._get_session_admission_lock(): + self._transfer_worker.sweep_stale_req_infos() + pending_retirements = self._get_context_retirements() + if not self._ever_had_send_session and not pending_retirements: + return [], [] + if mark_complete: + for progress in pending_retirements.values(): + progress.mark_complete = True + pending_candidate_rids = list(pending_retirements) + pending_rids = set(pending_retirements) + send_sessions = { + rid: session + for rid, session in self._send_sessions.items() + if rid not in pending_rids + } + send_reqs = {rid: self._send_reqs[rid] for rid in send_sessions} block_all = at_least_request_num is None wait_num = at_least_request_num if not block_all else 0 - - local_completed, local_failed = self._collect_done(self._send_sessions, self._send_reqs) + # A durable local terminal candidate already counts as progress. Do not + # block on an unrelated request while retrying its preparation. + if pending_rids and not block_all: + wait_num = 0 + + local_completed, local_failed = self._collect_done(send_sessions, send_reqs) + local_terminal_candidates = list( + dict.fromkeys( + [ + *pending_candidate_rids, + *local_completed, + *local_failed, + ] + ) + ) to_process = self._build_to_process( - self._send_sessions, - self._ctx_consensus(local_completed + local_failed), + send_sessions, + self._ctx_consensus(local_terminal_candidates), wait_num, block_all, ) - completed, timed_out, failed, cancelled = [], [], [], [] + observed_completed: list[int] = [] + observed_failed: list[int] = [] + observed_cancelled: list[int] = [] + timed_out: list[int] = [] for rid in to_process: - session = self._send_sessions[rid] + if rid in pending_rids: + continue + session = send_sessions[rid] result = session.wait_complete(blocking=block_all) + if result is None: + # Logical cancellation can precede source gather/NIXL drain. + # Keep the source owner until wait_complete reports a terminal + # physical state. + continue if session.status == SessionStatus.CANCELLED: - cancelled.append(rid) + observed_cancelled.append(rid) elif result == WaitResult.COMPLETED: - completed.append(rid) - elif result is None: - continue + observed_completed.append(rid) elif result == WaitResult.TIMEOUT: logger.warning( f"TxSession rid={session.disagg_request_id} timed out after {self._sender_future_timeout_ms}ms" @@ -579,114 +1555,169 @@ def check_context_transfer_status( timed_out.append(rid) else: logger.warning(f"TxSession rid={session.disagg_request_id} failed") - failed.append(rid) + observed_failed.append(rid) + + # Persist the exact owner before close/map validation can fail. A + # prepared candidate remains in the ready advertisement on later polls + # until the existing outcome consensus makes the distributed decision. + self._enroll_context_retirements( + ( + ("completed", observed_completed), + ("failed", observed_failed), + ("cancelled", observed_cancelled), + ), + send_sessions, + send_reqs, + mark_complete=mark_complete, + ) + self._prepare_context_retirements() + cancelled, failed, completed, cleanup_ready = self._context_candidate_payload(to_process) # All ranks must agree on per-rid outcome to avoid req.state divergence. cancelled, failed, completed, timed_out = self._ctx_consensus_outcome( - to_process, cancelled, failed, completed, timed_out + to_process, + cancelled, + failed, + completed, + timed_out, + cleanup_ready, ) - for rid in cancelled: - self._send_sessions[rid].close() - del self._send_reqs[rid] - del self._send_sessions[rid] - - for rid in completed: - if mark_complete: - self._send_reqs[rid].state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - self._send_sessions[rid].close() - del self._send_reqs[rid] - del self._send_sessions[rid] - self._close_failed_sessions(self._send_sessions, self._send_reqs, failed) - - # Sweep orphaned RecvReqInfo entries from ADP broadcast on non-assigned - # DP ranks (entries that will never have a TxSession created for them). - self._transfer_worker.sweep_stale_req_infos() - - return completed, failed + # The outcome consensus is also the prepare barrier: cleanup_ready is + # intersected across every participating rank. Commit contains only the + # real C++ state assignment and exact built-in dict removals. + return self._commit_context_decisions( + ( + ("completed", completed), + ("failed", failed), + ("cancelled", cancelled), + ) + ) def check_gen_transfer_status(self, at_least_request_num: Optional[int]): - if not self._ever_had_recv_session and not self._gen_need_sync: - return [], [], [] + with self._status_call() as admitted: + if not admitted: + return [], [], [] + return self._check_gen_transfer_status_impl(at_least_request_num) + + def _check_gen_transfer_status_impl(self, at_least_request_num: Optional[int]): + self._prepare_generation_retirements() + with self._get_session_admission_lock(): + pending_retirements = self._get_generation_retirements() + if ( + not self._ever_had_recv_session + and not self._gen_need_sync + and not pending_retirements + ): + return [], [], [] + pending_candidate_rids = list(pending_retirements) + pending_rids = set(pending_candidate_rids) + recv_sessions = { + rid: session + for rid, session in self._recv_sessions.items() + if rid not in pending_rids + } + recv_reqs = {rid: self._recv_reqs[rid] for rid in recv_sessions} block_all = at_least_request_num is None wait_num = at_least_request_num if not block_all else 0 + if pending_rids and not block_all: + wait_num = 0 need_progress = wait_num > 0 if need_progress: - self._poll_gen_sessions_for_poll_interval(wait_num) - - local_completed, local_failed = self._collect_done(self._recv_sessions, self._recv_reqs) + self._poll_gen_sessions_for_poll_interval(wait_num, recv_sessions, recv_reqs) + + local_completed, local_failed = self._collect_done(recv_sessions, recv_reqs) + local_terminal_candidates = list( + dict.fromkeys( + [ + *pending_candidate_rids, + *local_completed, + *local_failed, + ] + ) + ) to_process = self._build_to_process( - self._recv_sessions, - self._gen_consensus(local_completed + local_failed), + recv_sessions, + self._gen_consensus(local_terminal_candidates), 0 if need_progress else wait_num, block_all, ) - completed, failed, cancelled = [], [], [] + observed_completed: list[int] = [] + observed_failed: list[int] = [] + observed_cancelled: list[int] = [] for rid in to_process: - session = self._recv_sessions[rid] + if rid in pending_rids: + continue + session = recv_sessions[rid] result = session.wait_complete(blocking=block_all) + if result not in (WaitResult.COMPLETED, WaitResult.FAILED): + # Logical cancellation/failure may precede physical drain. + # Neither a non-blocking miss nor a timeout proves that request + # cleanup can no longer race a transfer accessor. + continue if session.status == SessionStatus.CANCELLED: # Session cancelled — either by local cancel_request() (user # cancel) or by a remote CANCEL_SESSION message (e.g. CTX # server timeout). Return the req objects so the caller can # distinguish the two cases and set the appropriate state. - cancelled.append(rid) + observed_cancelled.append(rid) elif result == WaitResult.COMPLETED: - completed.append(rid) + observed_completed.append(rid) elif result == WaitResult.FAILED: - failed.append(rid) - # else: None — KV done but aux still in flight; re-poll next cycle + observed_failed.append(rid) + + self._enroll_generation_retirements( + ( + ("completed", observed_completed), + ("failed", observed_failed), + ("cancelled", observed_cancelled), + ), + recv_sessions, + recv_reqs, + ) + self._prepare_generation_retirements() + cancelled, failed, completed, cleanup_ready = self._generation_candidate_payload(to_process) # All ranks must agree on per-rid outcome to avoid req.state divergence. cancelled, failed, completed = self._gen_consensus_outcome( - to_process, cancelled, failed, completed + to_process, cancelled, failed, completed, cleanup_ready ) - cancelled_reqs = [] - for rid in cancelled: - cancelled_reqs.append(self._recv_reqs[rid]) - self._recv_sessions[rid].close() - del self._recv_reqs[rid] - del self._recv_sessions[rid] - - for rid in completed: - session = self._recv_sessions[rid] - req = self._recv_reqs[rid] - # transfer_end already stamped at completion detection above. - req.set_kv_cache_size(getattr(req, "py_kv_cache_xfer_bytes", 0)) - if self._need_aux_transfer(req): - self._apply_aux(session, req) - self._assert_disagg_history_declared(req) - req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - session.close() - del self._recv_reqs[rid] - del self._recv_sessions[rid] - if failed: + retired_completed, retired_failed, cancelled_reqs = self._commit_generation_decisions( + ( + ("completed", completed), + ("failed", failed), + ("cancelled", cancelled), + ) + ) + if retired_failed: logger.warning( f"Disagg gen transfer FAILED rank={self._dist.rank} " - f"rids={failed} gen_need_sync={self._gen_need_sync}" + f"rids={retired_failed} gen_need_sync={self._gen_need_sync}" ) - self._close_failed_sessions(self._recv_sessions, self._recv_reqs, failed) - return completed, failed, cancelled_reqs + return retired_completed, retired_failed, cancelled_reqs - def _poll_gen_sessions_for_poll_interval(self, wait_num: int) -> None: + def _poll_gen_sessions_for_poll_interval( + self, wait_num: int, sessions: dict, reqs: dict + ) -> None: poll_interval_s = (self.kv_transfer_poll_interval_ms or 0) / 1000.0 deadline = time.monotonic() + poll_interval_s while True: - completed, failed = self._collect_done(self._recv_sessions, self._recv_reqs) + completed, failed = self._collect_done(sessions, reqs) if len(completed) + len(failed) >= wait_num: return remaining_s = deadline - time.monotonic() if remaining_s <= 0: return - for session in self._recv_sessions.values(): + for session in sessions.values(): session.wait_complete(blocking=False) time.sleep(min(0.001, remaining_s)) def check_gen_transfer_complete(self): - return len(self._recv_sessions) == 0 + with self._get_session_admission_lock(): + return not self._recv_sessions and not self._get_generation_retirements() def _assert_disagg_history_declared(self, req: LlmRequest) -> None: """Verify the V2 scheduler pre-declared prompt_len as history. @@ -729,33 +1760,90 @@ def cancel_request(self, req: LlmRequest) -> bool: retry next iteration. Returns True when safe to free KV memory. """ rid = get_unique_rid(req) - - # Not yet started (generation-first wait queue). - self._wait_reqs.pop(rid, None) - - has_transferring = False - - if rid in self._send_sessions: - self._send_sessions[rid].cancel() - if self._send_sessions[rid].has_transferring_tasks(): - has_transferring = True - else: - self._send_sessions[rid].close() - del self._send_reqs[rid] - del self._send_sessions[rid] - - if rid in self._recv_sessions: - self._recv_sessions[rid].cancel() - if self._recv_sessions[rid].has_transferring_tasks(): - has_transferring = True - else: - self._recv_sessions[rid].close() - del self._recv_reqs[rid] - del self._recv_sessions[rid] - - if has_transferring: - return False - return True + assert rid is not None + with self._get_session_admission_lock(): + # Status/preparation calls take request/session snapshots outside + # this lock while they enter collectives or poll workers. Removing + # an owner in that interval would let this rank enroll a no-op + # terminal candidate while peers retire the real request. Signal + # cancellation, but retain every exact map entry until the active + # collective-bearing call exits and cancellation is retried. + status_call_active = getattr(self, "_active_status_calls", 0) > 0 + cleanup_pending = status_call_active + cancelled_wait_reqs = self._get_cancelled_wait_reqs() + context_retirement = self._get_context_retirements().get(rid) + generation_retirement = self._get_generation_retirements().get(rid) + # A stale cancellation must not remove a replacement request that + # happens to reuse the same integer request ID. + if self._wait_reqs.get(rid) is req: + if status_call_active: + cancelled_wait_reqs[rid] = req + else: + self._wait_reqs.pop(rid, None) + if cancelled_wait_reqs.get(rid) is req: + cancelled_wait_reqs.pop(rid, None) + elif ( + status_call_active + and self._wait_reqs.get(rid) is None + and self._send_reqs.get(rid) is None + and self._recv_reqs.get(rid) is None + and context_retirement is None + and generation_retirement is None + and (cancelled_wait_reqs.get(rid) is None or cancelled_wait_reqs.get(rid) is req) + ): + # prepare_context_requests() marks its collective-bearing call + # active before it can enroll the incoming owner. Preserve an + # exact cancellation marker for that otherwise-unowned window + # so the preparation cannot subsequently admit this request. + cancelled_wait_reqs[rid] = req + elif not status_call_active and cancelled_wait_reqs.get(rid) is req: + cancelled_wait_reqs.pop(rid, None) + + cleanup_pending = ( + cleanup_pending + or (context_retirement is not None and context_retirement.request is req) + or (generation_retirement is not None and generation_retirement.request is req) + ) + for direction, sessions, reqs in ( + ("send", self._send_sessions, self._send_reqs), + ("receive", self._recv_sessions, self._recv_reqs), + ): + session = sessions.get(rid) + if session is None or reqs.get(rid) is not req: + continue + retirement = context_retirement if direction == "send" else generation_retirement + if ( + retirement is not None + and retirement.request is req + and retirement.session is session + ): + # Local terminal arbitration already enrolled this exact + # owner. Its provisional ledger is now the sole path that + # may prepare teardown and publish the consensus decision. + continue + try: + session.cancel() + if status_call_active: + # The active status call may still enroll this exact + # snapshot. It owns close/map retirement for this pass. + cleanup_pending = True + continue + if session.has_transferring_tasks(): + cleanup_pending = True + continue + session.close() + except Exception as error: + cleanup_pending = True + logger.error( + f"Failed to cancel {direction} KV transfer for request {rid}; " + f"retaining its owner for retry: {error}" + ) + continue + if sessions.get(rid) is session and reqs.get(rid) is req: + reqs.pop(rid, None) + sessions.pop(rid, None) + + return not cleanup_pending def get_disaggregated_params(self) -> Dict[str, Any]: # Keep this aligned with fields populated in respond_and_send_async(). @@ -777,31 +1865,89 @@ def get_disaggregated_params(self) -> Dict[str, Any]: } def prepare_context_requests(self, requests: List[LlmRequest]): + with self._status_call() as admitted: + if not admitted: + raise RuntimeError("KV cache transceiver is shutting down") + self._prepare_context_requests_impl(requests) + + def _prepare_context_requests_impl(self, requests: List[LlmRequest]) -> None: # Place new generation-first context requests into wait state, then # use allgather consensus to promote ready requests to CONTEXT_INIT. - for req in requests: - rid = get_unique_rid(req) - if rid not in self._send_sessions: + with self._get_session_admission_lock(): + pending_admissions = {} + planned_wait_owners = dict(self._wait_reqs) + cancelled_wait_reqs = self._get_cancelled_wait_reqs() + for req in requests: + rid = get_unique_rid(req) + assert rid is not None + cancelled_wait_owner = cancelled_wait_reqs.get(rid) + if cancelled_wait_owner is not None: + if cancelled_wait_owner is not req: + raise RuntimeError( + f"waiting request {rid} has a pending cancellation for a different request" + ) + # Cancellation won after status-call admission but before + # wait enrollment. Still publish this exact owner to the + # local wait ledger so every rank enters readiness + # consensus; the marker below prevents local promotion + # until cancel_request() retries and retires both entries. + if rid in self._get_context_retirements(): + raise RuntimeError(f"send request {rid} has a pending terminal retirement") + send_owner = self._send_reqs.get(rid) + if rid in self._send_sessions: + if send_owner is not req: + raise RuntimeError( + f"send request {rid} already has a different live source owner" + ) + continue + if send_owner is not None and send_owner is not req: + raise RuntimeError( + f"send request {rid} already has a different live source owner" + ) + wait_owner = planned_wait_owners.get(rid) + if wait_owner is not None and wait_owner is not req: + raise RuntimeError( + f"waiting request {rid} already has a different live source owner" + ) + planned_wait_owners[rid] = req + pending_admissions[rid] = req + + # Validate the full batch before publishing any owner. A collision + # must not leave an earlier request stranded in the wait ledger. + for rid, req in pending_admissions.items(): self._wait_reqs[rid] = req req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER - # Nothing waiting on any rank, so skip the consensus. The waiting set is the same on every - # rank, so they all skip together. - if not self._wait_reqs: - return + if not self._wait_reqs: + return + + # Check which waiting requests have peer info while the worker is + # protected from concurrent shutdown. Promotion happens after the + # collective and therefore revalidates each exact owner below. + ready_owners = { + rid: owner + for rid, owner in self._wait_reqs.items() + if self._get_cancelled_wait_reqs().get(rid) is not owner + if self._transfer_worker.has_all_peer_req_infos_for_send(rid) + } - # Check which waiting requests have peer info locally, then allgather - # consensus so all TP/PP ranks agree before promoting. # Without consensus, background peer info arriving at different times on - # different ranks causes scheduling mismatches → hang. - local_ready = [ - rid - for rid in self._wait_reqs - if self._transfer_worker.has_all_peer_req_infos_for_send(rid) - ] - for rid in self._ctx_consensus(local_ready): - self._wait_reqs[rid].state = LlmRequestState.CONTEXT_INIT - del self._wait_reqs[rid] + # different ranks causes scheduling mismatches → hang. Do not hold the + # local admission lock across the distributed collective. + ready_ids = self._ctx_consensus(list(ready_owners)) + with self._get_session_admission_lock(): + if getattr(self, "_shutdown_started", False): + return + for rid in ready_ids: + owner = ready_owners.get(rid) + if owner is not None and self._wait_reqs.get(rid) is owner: + # Every rank applies the decision from the same ready + # snapshot. A cancellation that arrived after that + # snapshot remains marked to block launch/reuse, but must + # not make this rank reverse a distributed promotion that + # its peers have already committed. + owner.state = LlmRequestState.CONTEXT_INIT + self._wait_reqs.pop(rid, None) def _check_compatible(self): if self._mapping.cp_size != 1: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 348d47a9659c..90137026bf16 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -182,6 +182,11 @@ def create_kv_cache_transceiver( class KvCacheTransceiver(ABC): + @property + def requires_physical_drain_before_request_release(self) -> bool: + """Whether request resources must remain owned until cancel drains.""" + return False + @abstractmethod def respond_and_send_async(self, req: LlmRequest): raise NotImplementedError @@ -237,8 +242,15 @@ def get_disaggregated_params(self) -> Dict[str, Any]: def commit_blocks_for_reuse(self, req: LlmRequest) -> None: """Commit received KV blocks to the radix tree for prefix reuse. No-op by default.""" - def shutdown(self): - """Shut down the transceiver and release registered resources.""" + def shutdown(self) -> Optional[bool]: + """Shut down the transceiver and release registered resources. + + Returns: + Lifecycle-capable implementations return ``True`` only after a + complete drain and ``False`` while resources remain owned. Legacy + implementations may return ``None``; callers must not interpret it + as transport-quiescence evidence. + """ class BindKvCacheTransceiver(KvCacheTransceiver): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4f4fc081a1d1..529d393be981 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -379,6 +379,100 @@ class BatchStatePP(BatchState): microbatch_id: int = -1 +_PYTHON_NATIVE_TRANSCEIVER_OWNER = "python_native_transceiver" +# At shutdown, allow an asynchronous connector roughly one second to publish +# its one-shot completion while keeping the number of collective polls fixed. +_SHUTDOWN_CONNECTOR_COMPLETION_POLLS = 20 +_SHUTDOWN_CONNECTOR_COMPLETION_POLL_INTERVAL_S = 0.05 + + +class _TransferTerminalOutcome(StrEnum): + """Request-level logical result retained across provider completion.""" + + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + + +@dataclasses.dataclass +class _TransferTerminalProgress: + """First terminal result and exact response phases for one request.""" + + request: LlmRequest + outcome: _TransferTerminalOutcome + claimant: str + response_creation_in_doubt: bool = False + response_created: bool = False + response: Optional[LlmResponse] = None + response_enriched: bool = False + response_buffered: bool = False + response_publication_in_doubt: bool = False + response_published: bool = False + response_discarded: bool = False + teardown_complete: bool = False + + +@dataclasses.dataclass +class _RequestCancellationProgress: + """Exact-request progress through user-cancellation ownership phases.""" + + request: LlmRequest + cancel_id: int + transport_drained: bool = False + transceiver_release_resolved: bool = False + all_transfers_complete: Optional[bool] = None + logical_cancelled: bool = False + decoding_iter_updated: bool = False + termination_handoff_complete: bool = False + + +@dataclasses.dataclass +class _ConnectorCompletionProgress: + """Retry state after a connector consumes its sole completion notice.""" + + request: LlmRequest + was_active: bool + owner_released: bool = False + all_transfers_complete: Optional[bool] = None + active_request_removed: bool = False + request_terminated: bool = False + + +@dataclasses.dataclass +class _NativeContextCompletionProgress: + """Retry state after the native transceiver retires a send session.""" + + request: LlmRequest + was_active: bool + was_deferred: bool + transfer_failed: bool + was_already_terminated: bool + owner_released: bool = False + all_transfers_complete: Optional[bool] = None + active_request_removed: bool = False + request_terminated: bool = False + terminated_marker_forgotten: bool = False + + +@dataclasses.dataclass +class _ResourceManagerShutdownProgress: + """One local manager's non-replayable shutdown phase.""" + + manager: object + completed: bool = False + in_doubt: bool = False + failure: Optional[str] = None + + +@dataclasses.dataclass +class _ExecutorShutdownPhaseProgress: + """One non-replayable executor-finalization phase.""" + + completed: bool = False + in_doubt: bool = False + failure: Optional[str] = None + + class AsyncTransferManager: """ Handle asynchronous transfer of KV cache after a request has completed. @@ -393,21 +487,89 @@ class AsyncTransferManager: class RequestTransferMetadata: - def __init__(self, block_id: Optional[int]): - self.block_id = block_id + def __init__(self, pinned_block_ids: Optional[List[int]]): + self.pinned_block_ids = pinned_block_ids self.counter = 0 - - def start_transfer(self): + # Named owners make one provider's completion independently + # retireable without guessing which anonymous counter belongs to + # another provider. Existing callers remain anonymous. + self.owners: set[str] = set() + + def start_transfer(self, owner: Optional[str] = None): + if owner is not None: + if owner in self.owners: + raise RuntimeError( + f"transfer owner {owner!r} is already active") + self.owners.add(owner) self.counter += 1 - def end_transfer(self) -> bool: + def will_end_all_transfers(self, owner: Optional[str] = None) -> bool: + """Validate one owner retirement without mutating the ledger.""" + if owner is not None: + if owner not in self.owners: + raise RuntimeError( + f"transfer owner {owner!r} is not active") + elif self.counter <= len(self.owners): + raise RuntimeError("no anonymous transfer owner is active") + return self.counter == 1 + + def end_transfer(self, owner: Optional[str] = None) -> bool: """ Returns: bool: True if there are no more transfers for this request """ + self.will_end_all_transfers(owner) + if owner is not None: + self.owners.remove(owner) self.counter -= 1 return self.counter == 0 + def has_owner(self, owner: str) -> bool: + return owner in self.owners + + @dataclasses.dataclass + class TransferAdmissionProgress: + """Durable first-owner admission before any external mutation.""" + + request: LlmRequest + owner: Optional[str] + manager_plan: Tuple[ResourceManagerType, ...] + completed_managers: set[ResourceManagerType] = dataclasses.field( + default_factory=set) + manager_in_doubt: Optional[ResourceManagerType] = None + state_updated: bool = False + state_update_in_doubt: bool = False + pin_started: bool = False + pin_complete: bool = False + pin_in_doubt: bool = False + pinned_block_ids: Optional[List[int]] = None + metadata: Optional[ + "AsyncTransferManager.RequestTransferMetadata"] = None + owner_enrolled: bool = False + metadata_published: bool = False + request_published: bool = False + + def in_doubt_reason(self) -> Optional[str]: + if self.manager_in_doubt is not None: + return f"resource manager {self.manager_in_doubt.value}" + if self.state_update_in_doubt: + return "request state update" + if self.pin_in_doubt: + return "KV block pin" + return None + + @dataclasses.dataclass + class TransferCompletionProgress: + """Durable final-owner release across fallible completion phases.""" + + request: LlmRequest + metadata: "AsyncTransferManager.RequestTransferMetadata" + owner: Optional[str] + unpin_complete: bool = False + unpin_in_doubt: bool = False + owner_released: bool = False + state_updated: bool = False + def __init__(self, resource_manager: "ResourceManager", should_store_blocks: bool = True): @@ -424,10 +586,92 @@ def __init__(self, self._request_transfer_metadata: Dict[ int, self.RequestTransferMetadata] = dict() + # Published before first-owner setup mutates any external manager, + # request state, or KV pin. An exception therefore retains the exact + # request and the last potentially ambiguous phase for fail-closed + # response handling and shutdown veto. + self._transfer_admission_progress: Dict[ + int, self.TransferAdmissionProgress] = dict() + self._transfer_completion_progress: Dict[ + int, self.TransferCompletionProgress] = dict() + self._admission_closed = False + def requests_in_transfer(self) -> Dict[int, LlmRequest]: return self._requests_in_transfer - def start_transfer(self, request: LlmRequest): + def requests_in_admission(self) -> Dict[int, LlmRequest]: + return { + req_id: progress.request + for req_id, progress in self._transfer_admission_progress.items() + } + + def has_pending_admission(self, request: LlmRequest) -> bool: + progress = self._transfer_admission_progress.get(request.py_request_id) + return progress is not None and progress.request is request + + def has_in_doubt_admissions(self) -> bool: + return any(progress.in_doubt_reason() is not None + for progress in self._transfer_admission_progress.values()) + + def begin_shutdown(self) -> None: + """Close first/subsequent-owner admission before teardown checks.""" + self._admission_closed = True + + def _resume_first_transfer_admission( + self, progress: TransferAdmissionProgress) -> None: + request = progress.request + req_id = request.py_request_id + in_doubt_reason = progress.in_doubt_reason() + if in_doubt_reason is not None: + raise RuntimeError( + f"transfer admission for request {req_id} is in doubt after " + f"{in_doubt_reason}; refusing to replay the phase") + + for resource_mgr_type in progress.manager_plan: + if resource_mgr_type in progress.completed_managers: + continue + progress.manager_in_doubt = resource_mgr_type + self.resource_manager.release_resource_for_transfer( + request, resource_mgr_type) + progress.completed_managers.add(resource_mgr_type) + progress.manager_in_doubt = None + + if not progress.state_updated: + try: + request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + except Exception: + progress.state_update_in_doubt = True + raise + progress.state_updated = True + + if not progress.pin_complete: + if self.should_store_blocks: + progress.pin_started = True + try: + progress.pinned_block_ids = ( + self.kv_cache_manager.store_blocks_for_reuse( + request, True)) + except Exception: + progress.pin_in_doubt = True + raise + progress.pin_complete = True + + if progress.metadata is None: + progress.metadata = self.RequestTransferMetadata( + progress.pinned_block_ids) + if not progress.owner_enrolled: + progress.metadata.start_transfer(progress.owner) + progress.owner_enrolled = True + if not progress.metadata_published: + self._request_transfer_metadata[req_id] = progress.metadata + progress.metadata_published = True + if not progress.request_published: + self._requests_in_transfer[req_id] = request + progress.request_published = True + if self._transfer_admission_progress.get(req_id) is progress: + self._transfer_admission_progress.pop(req_id) + + def start_transfer(self, request: LlmRequest, owner: Optional[str] = None): """ Called when a Cache transceiver or connector transfer is started. 1. Increment the counter for the request. @@ -437,30 +681,62 @@ def start_transfer(self, request: LlmRequest): req_id = request.py_request_id - if req_id not in self._requests_in_transfer: - for resource_mgr_type in ( - ResourceManagerType.SEQ_SLOT_MANAGER, - ResourceManagerType.SPEC_RESOURCE_MANAGER): - if resource_mgr_type in self.resource_manager.resource_managers and self.resource_manager.resource_managers[ - resource_mgr_type] is not None: - self.resource_manager.resource_managers[ - resource_mgr_type].free_resources(request) + if self._admission_closed: + raise RuntimeError( + f"transfer admission is closed for request {req_id}") + + # Prefer the durable admission record over partially published commit + # maps so retry can finish metadata/request publication without + # replaying manager release, state mutation, pinning, or owner enroll. + progress = self._transfer_admission_progress.get(req_id) + if progress is not None: + if progress.request is not request: + raise RuntimeError( + f"request {req_id} already has a different transfer " + "admission") + if progress.owner != owner: + raise RuntimeError( + f"request {req_id} transfer admission belongs to owner " + f"{progress.owner!r}, not {owner!r}") + self._resume_first_transfer_admission(progress) + return - request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + completion = self._transfer_completion_progress.get(req_id) + if completion is not None: + if completion.request is not request: + raise RuntimeError( + f"request {req_id} already has a different transfer " + "completion") + raise RuntimeError( + f"request {req_id} final transfer completion is incomplete") - if self.should_store_blocks: - block_id = self.kv_cache_manager.store_blocks_for_reuse( - request, True) - else: - block_id = None + existing_request = self._requests_in_transfer.get(req_id) + if existing_request is not None and existing_request is not request: + raise RuntimeError( + f"request {req_id} already has a different transfer owner") - self._requests_in_transfer[req_id] = request - self._request_transfer_metadata[ - req_id] = self.RequestTransferMetadata(block_id) + if existing_request is not None: + self._request_transfer_metadata[req_id].start_transfer(owner) + return - self._request_transfer_metadata[req_id].start_transfer() + manager_plan = tuple( + resource_mgr_type + for resource_mgr_type in (ResourceManagerType.SEQ_SLOT_MANAGER, + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if self.resource_manager.resource_managers.get( + resource_mgr_type) is not None) + progress = self.TransferAdmissionProgress( + request=request, + owner=owner, + manager_plan=manager_plan, + ) + self._transfer_admission_progress[req_id] = progress + + self._resume_first_transfer_admission(progress) - def end_transfer(self, request: LlmRequest) -> bool: + def end_transfer(self, + request: LlmRequest, + owner: Optional[str] = None) -> bool: """ Called after a send of KV cache is complete. 1. Decrements counter for request. @@ -469,33 +745,182 @@ def end_transfer(self, request: LlmRequest) -> bool: Returns: bool: True if the request should be terminated after call to end_transfer """ - try: - transfer_metadata = self._request_transfer_metadata[ - request.py_request_id] - except KeyError: - logger.warning( - f"Request {request.py_request_id} not found in transfer manager" + req_id = request.py_request_id + admission = self._transfer_admission_progress.get(req_id) + if admission is not None: + if admission.request is not request: + raise RuntimeError( + f"request {req_id} has a different active transfer " + "admission") + reason = admission.in_doubt_reason() + detail = f"; {reason} outcome is in doubt" if reason else "" + raise RuntimeError( + f"request {req_id} transfer admission is incomplete{detail}") + completion = self._transfer_completion_progress.get(req_id) + if completion is not None: + if completion.request is not request: + raise RuntimeError( + f"request {req_id} has a different active transfer " + "completion") + if completion.owner != owner: + raise RuntimeError( + f"request {req_id} transfer completion belongs to owner " + f"{completion.owner!r}, not {owner!r}") + transfer_metadata = completion.metadata + else: + transfer_metadata = None + existing_request = self._requests_in_transfer.get(req_id) + if existing_request is None: + if completion is None: + logger.warning( + f"Request {req_id} not found in transfer manager") + return False + if existing_request is not request: + if existing_request is not None: + raise RuntimeError( + f"request {req_id} has a different active transfer owner") + + if transfer_metadata is None: + transfer_metadata = self._request_transfer_metadata[req_id] + final_transfer = transfer_metadata.will_end_all_transfers(owner) + if not final_transfer: + return transfer_metadata.end_transfer(owner) + completion = self.TransferCompletionProgress( + request=request, + metadata=transfer_metadata, + owner=owner, ) - return False + self._transfer_completion_progress[req_id] = completion - if transfer_metadata.end_transfer(): - self._requests_in_transfer.pop(request.py_request_id) - self._request_transfer_metadata.pop(request.py_request_id) - - if self.should_store_blocks: + # The completion record is published before the first external + # mutation. Each successful phase is then skipped on retry, while the + # request maps keep the exact allocator owner rooted until commit. + if completion.unpin_in_doubt: + raise RuntimeError( + f"request {req_id} final transfer unpin outcome is in doubt; " + "refusing to replay the phase") + if not completion.unpin_complete and self.should_store_blocks: + try: self.kv_cache_manager.unpin_blocks_by_id( - transfer_metadata.block_id) + transfer_metadata.pinned_block_ids) + except Exception: + completion.unpin_in_doubt = True + raise + completion.unpin_complete = True + if not completion.state_updated: # We don't want to overwrite any error state. - if request.state != LlmRequestState.DISAGG_TRANS_ERROR: + if request.state not in (LlmRequestState.DISAGG_TRANS_ERROR, + LlmRequestState.DISAGG_CONTEXT_COMPLETE): request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + completion.state_updated = True + + if not completion.owner_released: + if not transfer_metadata.end_transfer(owner): + raise RuntimeError( + f"request {req_id} final transfer owner did not retire") + completion.owner_released = True + + published_request = self._requests_in_transfer.get(req_id) + if published_request is not None and published_request is not request: + raise RuntimeError( + f"request {req_id} final transfer commit found a different request" + ) + if published_request is request: + self._requests_in_transfer.pop(req_id) + + published_metadata = self._request_transfer_metadata.get(req_id) + if (published_metadata is not None + and published_metadata is not transfer_metadata): + raise RuntimeError( + f"request {req_id} final transfer commit found different metadata" + ) + if published_metadata is transfer_metadata: + self._request_transfer_metadata.pop(req_id) + + if self._transfer_completion_progress.get(req_id) is completion: + self._transfer_completion_progress.pop(req_id) + return True + + def has_transfer_owner(self, request: LlmRequest, owner: str) -> bool: + """Whether the exact request still has the named provider owner.""" + req_id = request.py_request_id + published_request = self._requests_in_transfer.get(req_id) + transfer_metadata = self._request_transfer_metadata.get(req_id) + completion = self._transfer_completion_progress.get(req_id) + return ((published_request is request and transfer_metadata is not None + and transfer_metadata.has_owner(owner)) + or (completion is not None and completion.request is request + and completion.owner == owner)) + + def has_any_transfer_owner(self, request: LlmRequest) -> bool: + """Whether the exact request still has any enrolled provider owner.""" + req_id = request.py_request_id + completion = self._transfer_completion_progress.get(req_id) + return (self._requests_in_transfer.get(req_id) is request + or (completion is not None and completion.request is request)) + + def is_final_transfer_owner(self, + request: LlmRequest, + owner: Optional[str] = None) -> bool: + """Whether retiring ``owner`` would retire the request's last owner. + + This is a read-only arbitration query. In particular, it lets the + anonymous connector distinguish connector-only completion from a mixed + connector plus legacy-C++ transfer without guessing which provider an + anonymous count belongs to. + """ + req_id = request.py_request_id + admission = self._transfer_admission_progress.get(req_id) + if admission is not None: + if admission.request is not request: + raise RuntimeError( + f"request {req_id} has a different active transfer " + "admission") + raise RuntimeError( + f"request {req_id} transfer admission is incomplete") + completion = self._transfer_completion_progress.get(req_id) + if completion is not None: + if completion.request is not request: + raise RuntimeError( + f"request {req_id} has a different active transfer " + "completion") + if completion.owner != owner: + raise RuntimeError( + f"request {req_id} transfer completion belongs to owner " + f"{completion.owner!r}, not {owner!r}") return True - return False + existing_request = self._requests_in_transfer.get(req_id) + if existing_request is None: + return False + if existing_request is not request: + raise RuntimeError( + f"request {req_id} has a different active transfer owner") + return self._request_transfer_metadata[req_id].will_end_all_transfers( + owner) + + def requests_with_owner(self, owner: str) -> Dict[int, LlmRequest]: + """Snapshot requests whose transfer count includes ``owner``.""" + requests = { + req_id: request + for req_id, request in self._requests_in_transfer.items() + if self._request_transfer_metadata.get(req_id) is not None + and self._request_transfer_metadata[req_id].has_owner(owner) + } + requests.update({ + req_id: completion.request + for req_id, completion in + self._transfer_completion_progress.items() + if completion.owner == owner + }) + return requests def has_any_inflight_requests(self) -> bool: - return len(self._requests_in_transfer) > 0 + return bool(self._requests_in_transfer + or self._transfer_admission_progress + or self._transfer_completion_progress) class PyExecutor: @@ -703,11 +1128,56 @@ def __init__( # responses and flushing them at a synchronised point in the executor # loop avoids the mismatch. self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + # Exact O(1) association between a buffered response and the terminal + # ledger that owns it. The terminal roots the response, so ``id`` cannot + # be reused while this entry exists. + self._pending_transfer_response_terminals: Dict[ + int, _TransferTerminalProgress] = {} # Same buffer-then-synced-flush pattern as _pending_transfer_responses # above: _handle_responses and _append_iter_stats are reached from # per-rank-divergent gates, so their tp_allgather collectives are # lifted to _handle_kv_transfer_timeouts_synced / _flush_iter_stats_synced. self._pending_timed_out_requests: List[LlmRequest] = [] + # Requests may become logically terminal while a transceiver still + # owns their memory. Keep a strong root until runtime cancellation or + # post-loop shutdown satisfies that implementation's drain contract. + self._deferred_transfer_terminations: Dict[int, LlmRequest] = {} + # PP=1 partial-reuse requests can be freed before their asynchronous + # transfer owner retires. Shutdown still retains those request objects + # for owner retirement, but must not free their resources a second time. + self._deferred_transfer_terminations_already_terminated: set[int] = set( + ) + # Requests whose normal resource termination completed while any + # AsyncTransferManager provider still held a lease. This exact marker + # prevents the last provider from terminating the request twice. + self._terminated_transfer_requests: Dict[int, LlmRequest] = {} + # A request remains here after its ResourceManager release succeeds + # until all remaining PyExecutor bookkeeping also completes. This + # makes a retry resume after resource release instead of freeing the + # same exact request twice. + self._request_resource_termination_progress: Dict[int, LlmRequest] = {} + # A user-cancel attempt can drain its transport and then fail while + # releasing the transceiver owner. Retain the exact + # request and resume at the failed phase instead of relying on the + # request to remain in active_requests. + self._pending_request_cancellations: Dict[ + int, _RequestCancellationProgress] = {} + # Connector completion notifications are one-shot. Keep phase-aware + # progress so unpin or request-termination failures can be retried + # without recreating a response or releasing the owner twice. + self._pending_connector_completions: Dict[ + int, _ConnectorCompletionProgress] = {} + # Python-native context completion notifications are also one-shot: + # check_context_transfer_status retires the transport session before + # returning its request ID. Retain the exact request and phase progress + # until response preparation, owner release, and final teardown finish. + self._pending_native_context_completions: Dict[ + int, _NativeContextCompletionProgress] = {} + # Native and connector notifications can retire in either order. The + # first logical result and its exact response object therefore outlive + # either provider-local retry record until all sibling owners retire. + self._pending_transfer_terminals: Dict[int, + _TransferTerminalProgress] = {} self._pending_iter_stats_dict: Optional[Dict] = None # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen; # updated from observed request types. @@ -1075,43 +1545,277 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() - def _end_transfer_and_maybe_terminate(self, request: LlmRequest): + def _get_pending_transfer_terminals( + self) -> Dict[int, _TransferTerminalProgress]: + terminals = getattr(self, "_pending_transfer_terminals", None) + if terminals is None: + terminals = {} + self._pending_transfer_terminals = terminals + return terminals + + def _get_transfer_terminal( + self, request: LlmRequest) -> Optional[_TransferTerminalProgress]: + """Return the exact request's first logical terminal result.""" + terminal = self._get_pending_transfer_terminals().get(id(request)) + if terminal is not None and terminal.request is not request: + raise RuntimeError( + "transfer terminal identity collision for request " + f"{request.py_request_id}") + return terminal + + def _claim_transfer_terminal(self, request: LlmRequest, + outcome: _TransferTerminalOutcome, + claimant: str) -> _TransferTerminalProgress: + """Publish the first request-level terminal result exactly once.""" + terminals = self._get_pending_transfer_terminals() + request_key = id(request) + terminal = terminals.get(request_key) + if terminal is not None: + if terminal.request is not request: + raise RuntimeError( + "transfer terminal identity collision for request " + f"{request.py_request_id}") + return terminal + terminal = _TransferTerminalProgress(request=request, + outcome=outcome, + claimant=claimant) + terminals[request_key] = terminal + return terminal + + def _get_pending_transfer_response_terminals( + self) -> Dict[int, _TransferTerminalProgress]: + response_terminals = getattr(self, + "_pending_transfer_response_terminals", + None) + if response_terminals is None: + response_terminals = {} + self._pending_transfer_response_terminals = response_terminals + return response_terminals + + def _create_transfer_terminal_response( + self, terminal: _TransferTerminalProgress) -> None: + """Create one terminal response, refusing an ambiguous replay.""" + if terminal.response_created: + return + if terminal.response_creation_in_doubt: + raise RuntimeError( + "terminal response creation outcome is in doubt for request " + f"{terminal.request.py_request_id}; refusing to replay") + + terminal.response_creation_in_doubt = True + terminal.response = terminal.request.create_response( + False, self.dist.rank) + terminal.response_created = True + terminal.response_creation_in_doubt = False + + def _buffer_transfer_terminal_response( + self, terminal: _TransferTerminalProgress) -> None: + """Root one response and its exact terminal until publish or discard.""" + response = terminal.response + if response is None or terminal.response_buffered: + return + response_terminals = self._get_pending_transfer_response_terminals() + response_key = id(response) + existing = response_terminals.get(response_key) + if existing is not None and (existing is not terminal + or existing.response is not response): + raise RuntimeError( + "transfer response identity collision for request " + f"{terminal.request.py_request_id}") + pending_responses = getattr(self, "_pending_transfer_responses", None) + if pending_responses is None: + pending_responses = [] + self._pending_transfer_responses = pending_responses + pending_responses.append((terminal.request.py_request_id, response)) + response_terminals[response_key] = terminal + terminal.response_buffered = True + + def _begin_transfer_response_publication( + self, terminals: Iterable[_TransferTerminalProgress]) -> None: + """Mark publication ambiguous before entering the side-effecting sink.""" + terminals = list(terminals) + for terminal in terminals: + if terminal.response_publication_in_doubt: + raise RuntimeError( + "terminal response publication outcome is in doubt for " + f"request {terminal.request.py_request_id}; refusing to " + "replay") + for terminal in terminals: + terminal.response_publication_in_doubt = True + + def _commit_transfer_response_publication( + self, terminals: Iterable[_TransferTerminalProgress]) -> None: + """Commit successful publication and remove its local response roots.""" + terminals = list(terminals) + response_keys = { + id(terminal.response) + for terminal in terminals if terminal.response is not None + } + if response_keys: + pending_responses = getattr(self, "_pending_transfer_responses", []) + self._pending_transfer_responses = [ + entry for entry in pending_responses + if id(entry[1]) not in response_keys + ] + response_terminals = self._get_pending_transfer_response_terminals() + for terminal in terminals: + response = terminal.response + if response is not None: + response_key = id(response) + if response_terminals.get(response_key) is terminal: + response_terminals.pop(response_key) + terminal.response_publication_in_doubt = False + terminal.response_published = True + self._retire_transfer_terminal_if_complete(terminal.request) + + def _prepare_transfer_terminal_response( + self, request: LlmRequest, outcome: _TransferTerminalOutcome, + claimant: str) -> _TransferTerminalProgress: + """Create, enrich, and buffer the winning response fail-closed.""" + terminal = self._claim_transfer_terminal(request, outcome, claimant) + if terminal.outcome != outcome: + return terminal + if terminal.response_published or terminal.response_discarded: + return terminal + + self._create_transfer_terminal_response(terminal) + if terminal.response is None: + terminal.response_enriched = True + terminal.response_buffered = True + terminal.response_published = True + return terminal + if not terminal.response_enriched: + terminal.response.result.cached_tokens = request.cached_tokens + self._maybe_attach_ctx_usage(request, terminal.response) + terminal.response_enriched = True + if not terminal.response_buffered: + # With ADP, _enqueue_responses does a tp_gather collective. Calling + # it from a provider callback would deadlock because only the + # owning DP rank reaches that callback. + self._buffer_transfer_terminal_response(terminal) + return terminal + + def _retire_transfer_terminal_if_complete(self, + request: LlmRequest) -> None: + """Forget a terminal result after teardown and the final lease.""" + terminal = self._get_transfer_terminal(request) + response_unsettled = (terminal is not None + and not terminal.response_published + and not terminal.response_discarded) + if (terminal is None or not terminal.teardown_complete + or response_unsettled + or self._has_async_transfer_owner(request)): + return + terminals = self._get_pending_transfer_terminals() + if terminals.get(id(request)) is terminal: + terminals.pop(id(request)) + + def _mark_transfer_terminal_teardown_complete(self, + request: LlmRequest) -> None: + terminal = self._get_transfer_terminal(request) + if terminal is None: + return + terminal.teardown_complete = True + self._retire_transfer_terminal_if_complete(request) + + def _release_transfer_owner(self, + request: LlmRequest, + transfer_owner: Optional[str] = None) -> bool: + """Release one transfer owner without changing a cancelled outcome.""" + # AsyncTransferManager maps a final successful owner retirement to + # CONTEXT_COMPLETE. A connector may be the last sibling owner of a + # request that the user already cancelled; keep that terminal state + # stable across both successful and retryable owner release. + preserve_cancelled_state = (getattr(request, + "is_finished_due_to_cancellation", + False) is True) + original_state = request.state + if preserve_cancelled_state: + request.state = LlmRequestState.DISAGG_TRANS_ERROR + try: + if transfer_owner is None: + all_transfers_complete = ( + self.async_transfer_manager.end_transfer(request)) + else: + all_transfers_complete = self.async_transfer_manager.end_transfer( + request, owner=transfer_owner) + finally: + if preserve_cancelled_state: + request.state = original_state + return all_transfers_complete + + def _prepare_transfer_completion_response(self, + request: LlmRequest) -> None: + """Create and buffer the response needed before owner retirement.""" + self._prepare_transfer_terminal_response( + request, _TransferTerminalOutcome.SUCCEEDED, + _PYTHON_NATIVE_TRANSCEIVER_OWNER) + + def _end_transfer_and_maybe_terminate(self, + request: LlmRequest, + transfer_owner: Optional[str] = None): transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR - if self.kv_cache_transceiver and request in self.active_requests: + request_was_active = (self.kv_cache_transceiver + and request in self.active_requests) + request_was_already_terminated = ( + PyExecutor._was_transfer_request_terminated(self, request)) + if transfer_failed: + PyExecutor._claim_transfer_terminal( + self, request, _TransferTerminalOutcome.FAILED, transfer_owner + or "legacy_transceiver") + if request_was_active: if transfer_failed: - # End only the transfer that just became terminal. Keep the - # request active so the synchronized error path can emit an - # error response after every async transfer releases ownership. - self.async_transfer_manager.end_transfer(request) + # Keep the request active so the synchronized error path can + # emit its response after every sibling owner retires. + PyExecutor._release_transfer_owner(self, request, + transfer_owner) return # Fast-transfer: KV transfer completed in the same iteration # before _handle_responses could run. Create the response now # while state is still TRANS_IN_PROGRESS (required by C++ # createResult). Then proceed with end_transfer + termination. - response = request.create_response(False, self.dist.rank) - if response: - response.result.cached_tokens = request.cached_tokens - self._maybe_attach_ctx_usage(request, response) - # Buffer the response instead of enqueueing immediately. - # With ADP, _enqueue_responses does a tp_gather collective. - # Calling it here would deadlock because only the owning DP - # rank reaches this point; the other DP rank never enters - # the matching collective. The buffer is flushed later at - # _flush_pending_transfer_responses where all ranks - # participate. - self._pending_transfer_responses.append( - (request.py_request_id, response)) - if self.async_transfer_manager.end_transfer(request): + PyExecutor._prepare_transfer_completion_response(self, request) + if PyExecutor._release_transfer_owner(self, request, + transfer_owner): self.active_requests.remove(request) - self._terminate_request(request) + if not request_was_already_terminated: + self._terminate_request(request) + else: + PyExecutor._mark_transfer_terminal_teardown_complete( + self, request) + if hasattr(self, "_terminated_transfer_requests"): + PyExecutor._forget_terminated_transfer_request_if_unowned( + self, request) + PyExecutor._retire_transfer_terminal_if_complete(self, request) return - if self.async_transfer_manager.end_transfer(request): + if PyExecutor._release_transfer_owner(self, request, transfer_owner): if transfer_failed: + if (request_was_already_terminated + and hasattr(self, "_terminated_transfer_requests")): + PyExecutor._mark_transfer_terminal_teardown_complete( + self, request) + PyExecutor._forget_terminated_transfer_request_if_unowned( + self, request) + PyExecutor._retire_transfer_terminal_if_complete( + self, request) return - # Skip if the PP=1 early path already terminated this request; - # under PP>1 that path is off, so terminate here on transfer-complete. - if not self.force_terminate_ctx_for_partial_reuse: + # The PP=1 partial-reuse path may have released request resources + # while the transfer lease remained live. The exact-request marker + # distinguishes that case from an inactive request which still + # needs its one final teardown. + legacy_pp1_already_terminated = ( + not hasattr(self, "_terminated_transfer_requests") and getattr( + self, "force_terminate_ctx_for_partial_reuse", False)) + if (not request_was_already_terminated + and not legacy_pp1_already_terminated): self._terminate_request(request) + else: + PyExecutor._mark_transfer_terminal_teardown_complete( + self, request) + if hasattr(self, "_terminated_transfer_requests"): + PyExecutor._forget_terminated_transfer_request_if_unowned( + self, request) + PyExecutor._retire_transfer_terminal_if_complete(self, request) def _flush_pending_transfer_responses(self): """Enqueue buffered transfer-completion responses. @@ -1119,13 +1823,106 @@ def _flush_pending_transfer_responses(self): Must be called at a point where ALL DP ranks execute in lockstep so that the tp_gather inside _enqueue_responses does not deadlock. """ - responses = self._pending_transfer_responses - self._pending_transfer_responses = [] + responses = list(self._pending_transfer_responses) + response_terminals = self._get_pending_transfer_response_terminals() + terminals = [] + for _request_id, response in responses: + terminal = response_terminals.get(id(response)) + if terminal is None or terminal.response is not response: + raise RuntimeError( + "buffered transfer response has no exact terminal owner") + terminals.append(terminal) + self._begin_transfer_response_publication(terminals) if responses or self.enable_attention_dp: # Even when this rank has no responses we must participate in the # collective when ADP is enabled so that the other rank's gather # can complete. self._enqueue_responses(responses) + self._commit_transfer_response_publication(terminals) + + def _discard_pending_transfer_responses_after_shutdown(self) -> None: + """Discard responses remaining after the rank-synchronous loop stopped. + + ``_enqueue_responses`` may enter TP collectives. Once the executor loop + has exited, a rank-local teardown failure can make any new collective + asymmetric, so shutdown is an ownership drain rather than a response + delivery boundary. Normal execution flushes these responses from a + rank-synchronous point in the loop; post-loop shutdown only drops the + local payload after recording the loss. + """ + responses = getattr(self, "_pending_transfer_responses", None) + if not responses: + return + logger.warning( + f"Discarding {len(responses)} KV-transfer completion response(s) " + "remaining after the executor loop stopped; shutdown cannot safely " + "enter response collectives") + self._pending_transfer_responses = [] + response_terminals = self._get_pending_transfer_response_terminals() + for _request_id, response in responses: + terminal = response_terminals.pop(id(response), None) + if terminal is None or terminal.response is not response: + logger.error( + "Discarded transfer response had no exact terminal owner") + continue + terminal.response_publication_in_doubt = False + terminal.response_discarded = True + self._retire_transfer_terminal_if_complete(terminal.request) + + def _discard_unpublishable_transfer_terminals_after_shutdown(self) -> None: + """Settle terminal responses that never reached the post-loop buffer. + + Once the rank-synchronous executor loop has stopped, no terminal + response can safely enter TP/attention-DP publication collectives. A + provider may nevertheless retire its final owner during shutdown and + leave a failed terminal with no response object or buffer entry. Mark + every such outcome as explicitly discarded so request teardown cannot + be skipped merely because response construction never started. + """ + terminals = list( + getattr(self, "_pending_transfer_terminals", {}).values()) + newly_discarded = 0 + response_terminals = self._get_pending_transfer_response_terminals() + for terminal in terminals: + if (terminal.response_published or terminal.response_discarded + or terminal.response_creation_in_doubt + or self._has_async_transfer_owner(terminal.request) or + self._has_pending_ownership_transition(terminal.request)): + continue + response = terminal.response + if response is not None: + response_key = id(response) + if response_terminals.get(response_key) is terminal: + response_terminals.pop(response_key, None) + terminal.response_publication_in_doubt = False + terminal.response_discarded = True + newly_discarded += 1 + self._retire_transfer_terminal_if_complete(terminal.request) + if newly_discarded: + logger.warning( + f"Discarding {newly_discarded} ownerless KV-transfer terminal " + "response(s) that became final after the executor loop stopped") + + def _drain_discarded_transfer_terminals_after_shutdown(self) -> None: + """Finish local teardown after post-loop response disposition.""" + for terminal in list( + getattr(self, "_pending_transfer_terminals", {}).values()): + if ((not terminal.response_discarded + and not terminal.response_published) + or terminal.teardown_complete + or self._has_async_transfer_owner(terminal.request) or + self._has_pending_ownership_transition(terminal.request)): + continue + request = terminal.request + if request in self.active_requests: + self.active_requests.remove(request) + if self._was_transfer_request_terminated(request): + self._mark_transfer_terminal_teardown_complete(request) + else: + self._terminate_request_after_worker_shutdown(request) + self._mark_transfer_terminal_teardown_complete(request) + self._forget_terminated_transfer_request_if_unowned(request) + self._retire_transfer_terminal_if_complete(request) def _handle_kv_transfer_timeouts_synced(self): """ADP-safe drain of the KV-transfer-timeout consensus collective. @@ -1428,47 +2225,142 @@ def cancel_request(self, id: int): """ self.executor_request_queue.enqueue_cancel_request(id) - def shutdown(self): - """ - Signals the server to shutdown. - """ - self.executor_request_queue.enqueue_shutdown_request() - self.shutdown_event.wait() - if self.hang_detector.detected(): - # Early return here to avoid waiting for hanging threads. - # Since `on_detected` has sent the error message as response, - # this worker will be asked to shutdown immediately. - # Since the whole process will shutdown after this `shutdown` call, - # All threads and memory pools will be freed properly. - logger.error("Hang detected, shutting down immediately.") + def _get_resource_manager_shutdown_progress( + self) -> tuple[_ResourceManagerShutdownProgress, ...]: + """Return a stable exact-manager teardown plan across shutdown retry.""" + managers = tuple( + manager + for manager in self.resource_manager.resource_managers.values() + if manager is not None) + progress = getattr(self, "_resource_manager_shutdown_progress", None) + if progress is None: + progress = tuple( + _ResourceManagerShutdownProgress(manager) + for manager in managers) + self._resource_manager_shutdown_progress = progress + return progress + if (len(progress) != len(managers) or any( + entry.manager is not manager + for entry, manager in zip(progress, managers, strict=True))): + raise RuntimeError( + "Resource-manager shutdown plan changed during teardown") + return progress + + def _shutdown_resource_managers_rank_uniform(self) -> None: + """Shut local managers once until failure, then publish the rank outcome.""" + first_error: Optional[Exception] = None + try: + progress = self._get_resource_manager_shutdown_progress() + except Exception as error: + progress = () + first_error = error + + for manager_index, entry in enumerate(progress): + if entry.completed: + continue + if entry.in_doubt: + error = RuntimeError( + "Resource-manager shutdown outcome is in doubt at manager " + f"index {manager_index}: {entry.failure}") + if first_error is None: + first_error = error + # Manager order is a dependency order (KV is intentionally + # last). Never destroy a later dependency while an earlier + # manager may still be live after an ambiguous mutation. + break + + # Mark ambiguity before entering arbitrary manager code. A retry + # skips both a successful manager and one whose mutation boundary + # is unknown after an exception. + entry.in_doubt = True + try: + entry.manager.shutdown() + except Exception as error: + entry.failure = f"{type(error).__name__}: {error}" + if first_error is None: + first_error = error + break + else: + entry.completed = True + entry.in_doubt = False + entry.failure = None + + local_failed = first_error is not None + try: + if self._dist_size(self.dist, "world_size") > 1: + any_rank_failed = bool( + self.dist.allreduce(int(local_failed), op=ReduceOp.MAX)) + else: + any_rank_failed = local_failed + except Exception as consensus_error: + if first_error is not None: + raise first_error from consensus_error + raise + if any_rank_failed: + if first_error is not None: + raise first_error + raise RuntimeError( + "Resource-manager shutdown failed on another rank; refusing " + "to continue executor teardown") + + def _run_executor_shutdown_phase(self, name: str, + operation: Callable[[], None]) -> None: + """Run one fallible finalization phase at most once.""" + progress_by_name = getattr(self, "_executor_shutdown_phase_progress", + None) + if progress_by_name is None: + progress_by_name = {} + self._executor_shutdown_phase_progress = progress_by_name + progress = progress_by_name.setdefault(name, + _ExecutorShutdownPhaseProgress()) + if progress.completed: return - self.worker_thread.join() - if self.dist.pp_size > 1: - self.executed_batch_queue.put(None) - self.broadcast_sample_state_handler.join() - # Signal non-rank-0 sleep/wakeup listener threads to exit. This runs - # after the worker thread has joined, which guarantees that the non-rank-0 - # executor loops have already processed the shutdown broadcast and are - # no longer driving NCCL, so the send cannot deadlock. - self._shutdown_sleep_wakeup_listeners() - self.worker_started = False - # Release CUDA graphs before resource managers free their GPU memory. - # Resource managers (e.g. SuffixAutomatonManager) allocate GPU workspace - # that is referenced by raw pointers inside captured CUDA graphs. If - # the workspace is freed first (and returned to the driver via - # empty_cache), the subsequent CUDA graph teardown can trigger a - # device-wide cudaErrorIllegalAddress when the driver touches metadata - # for the now-freed memory regions. - for engine in (self.model_engine, self.draft_model_engine): - if engine is not None and hasattr(engine, '_release_cuda_graphs'): - engine._release_cuda_graphs() - # Ensure graph destruction has fully completed on device before - # resource managers start freeing GPU-backed workspaces. - if torch.cuda.is_available(): - torch.cuda.synchronize() - for manager in self.resource_manager.resource_managers.values(): - if manager: - manager.shutdown() + if progress.in_doubt: + raise RuntimeError( + f"Executor shutdown phase {name!r} outcome is in doubt: " + f"{progress.failure}") + + progress.in_doubt = True + try: + operation() + except Exception as error: + progress.failure = f"{type(error).__name__}: {error}" + raise + else: + progress.completed = True + progress.in_doubt = False + progress.failure = None + + def _prepare_final_executor_shutdown(self) -> None: + """Complete fallible finalization before releasing engine objects.""" + + def stop_sampler() -> None: + if (isinstance(self.sampler, AsyncWorkerMixin) + and self.sampler.async_worker_enabled()): + self.sampler.async_worker_stop() + + def stop_dwdp_manager() -> None: + if self.dwdp_manager is not None: + self.dwdp_manager.__exit__(None, None, None) + self.dwdp_manager = None + + first_error: Optional[Exception] = None + # Keep the fixed phase sequence rank-uniform even after a local error. + # DWDP communicator cleanup may itself enter collective lifecycle work. + for name, operation in ( + ("sampler", stop_sampler), + ("dwdp_manager", stop_dwdp_manager), + ): + try: + self._run_executor_shutdown_phase(name, operation) + except Exception as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error + + def _release_final_executor_objects(self) -> None: + """Release non-fallible executor roots after every shutdown vote.""" # Note: do NOT call engine.cleanup() here. PyExecutor.shutdown() is # also invoked mid-init by configure_kv_cache_capacity() in # tensorrt_llm/_torch/pyexecutor/_util.py — the warmup pass calls @@ -1486,13 +2378,354 @@ def shutdown(self): keys = list(self.virtual_memory_pools.keys()) for key in keys: del self.virtual_memory_pools[key] - # Stop the sampler's async worker, if it was used - if (isinstance(self.sampler, AsyncWorkerMixin) - and self.sampler.async_worker_enabled()): - self.sampler.async_worker_stop() - if self.dwdp_manager is not None: - self.dwdp_manager.__exit__(None, None, None) - self.dwdp_manager = None + + def shutdown(self): + """ + Signals the server to shutdown. + """ + if self.worker_started: + self.executor_request_queue.enqueue_shutdown_request() + self.shutdown_event.wait() + if self.hang_detector.detected(): + # Early return here to avoid waiting for hanging threads. + # Since `on_detected` has sent the error message as response, + # this worker will be asked to shutdown immediately. + # Since the whole process will shutdown after this `shutdown` call, + # All threads and memory pools will be freed properly. + logger.error("Hang detected, shutting down immediately.") + return + self.worker_thread.join() + if self.dist.pp_size > 1: + self.executed_batch_queue.put(None) + self.broadcast_sample_state_handler.join() + # Signal non-rank-0 sleep/wakeup listener threads to exit. This + # runs after the worker thread has joined, which guarantees that + # non-rank-0 executor loops no longer drive NCCL. + self._shutdown_sleep_wakeup_listeners() + self.worker_started = False + pre_connector_error: Optional[Exception] = None + + def retain_pre_connector_error(error: Exception) -> None: + nonlocal pre_connector_error + if pre_connector_error is None: + pre_connector_error = error + logger.error( + "KV ownership shutdown failed locally before connector " + f"consensus; continuing required local drains: {error}") + + # Release CUDA graphs before resource managers free their GPU memory. + # Resource managers (e.g. SuffixAutomatonManager) allocate GPU workspace + # that is referenced by raw pointers inside captured CUDA graphs. If + # the workspace is freed first (and returned to the driver via + # empty_cache), the subsequent CUDA graph teardown can trigger a + # device-wide cudaErrorIllegalAddress when the driver touches metadata + # for the now-freed memory regions. + def release_cuda_graphs() -> None: + for engine in (self.model_engine, self.draft_model_engine): + if engine is not None and hasattr(engine, + '_release_cuda_graphs'): + engine._release_cuda_graphs() + + try: + self._run_executor_shutdown_phase("cuda_graphs", + release_cuda_graphs) + except Exception as error: + retain_pre_connector_error(error) + # Ensure graph destruction has fully completed on device before + # resource managers start freeing GPU-backed workspaces. + + def synchronize_cuda() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + try: + self._run_executor_shutdown_phase("cuda_synchronize", + synchronize_cuda) + except Exception as error: + retain_pre_connector_error(error) + try: + is_python_native_transceiver = ( + self.kv_cache_transceiver is not None + and self._requires_physical_transfer_drain()) + except Exception as error: + retain_pre_connector_error(error) + # Capability lookup itself is ambiguous. Conservatively require an + # explicit drain proof from any configured transceiver. + is_python_native_transceiver = self.kv_cache_transceiver is not None + + if not is_python_native_transceiver: + # The ownership drain, veto ledger, and rank-uniform teardown below + # are a Python-native protocol. Preserve the default C++ and + # non-disaggregated executor's existing local shutdown contract. + if pre_connector_error is not None: + raise pre_connector_error + for manager in self.resource_manager.resource_managers.values(): + if manager: + manager.shutdown() + self._prepare_final_executor_shutdown() + self._release_final_executor_objects() + return + + transfer_manager = getattr(self, "async_transfer_manager", None) + if transfer_manager is not None: + try: + transfer_manager.begin_shutdown() + except Exception as error: + retain_pre_connector_error(error) + native_shutdown_proven = False + if any(terminal.response_creation_in_doubt for terminal in getattr( + self, "_pending_transfer_terminals", {}).values()): + retain_pre_connector_error( + RuntimeError( + "Terminal response creation outcome is in doubt; refusing " + "to shut down resource managers")) + + if is_python_native_transceiver: + # The executor loop can stop before its final status poll. Capture + # every still-enrolled Python-native owner before transport + # shutdown so a proven drain can retire that exact contribution. + # The default C++ transceiver keeps its legacy shutdown behavior + # and never enters this lifecycle-capable path. + defer_transceiver_owners = getattr( + self, "_defer_transceiver_owners_for_shutdown", None) + if defer_transceiver_owners is not None: + try: + defer_transceiver_owners() + except Exception as error: + retain_pre_connector_error(error) + try: + transceiver_shutdown_result = ( + self.kv_cache_transceiver.shutdown()) + except Exception as error: + retain_pre_connector_error(error) + else: + if transceiver_shutdown_result is False: + retain_pre_connector_error( + RuntimeError( + "KV cache transceiver still owns active transfer " + "targets; refusing to shut down resource managers")) + elif transceiver_shutdown_result is not True: + retain_pre_connector_error( + RuntimeError( + "Python-native KV cache transceiver did not prove " + "physical drain; refusing to shut down resource " + "managers")) + else: + native_shutdown_proven = True + drain_native_completions = getattr( + self, "_drain_pending_native_context_completions", None) + if (drain_native_completions is not None + and native_shutdown_proven): + try: + drain_native_completions( + terminate_after_worker_shutdown=True) + except Exception as error: + retain_pre_connector_error(error) + if native_shutdown_proven: + try: + self._discard_pending_transfer_responses_after_shutdown() + except Exception as error: + retain_pre_connector_error(error) + if getattr(self, "_pending_native_context_completions", None): + retain_pre_connector_error( + RuntimeError( + "Native context completion ownership is still " + "incomplete after transceiver shutdown; refusing to " + "shut down resource managers")) + # Native shutdown closes admission and proves physical drain. Re-run + # cancellation through the same safe-to-free gate before retiring + # request ownership. + drain_deferred_terminations = getattr( + self, "_drain_deferred_transfer_terminations", None) + if (drain_deferred_terminations is not None + and native_shutdown_proven): + try: + drain_deferred_terminations( + terminate_after_worker_shutdown=True) + except Exception as error: + retain_pre_connector_error(error) + drain_pending_cancellations = getattr( + self, "_drain_pending_request_cancellations_after_shutdown", None) + if (drain_pending_cancellations is not None and native_shutdown_proven): + try: + drain_pending_cancellations() + except Exception as error: + retain_pre_connector_error(error) + if getattr(self, "_pending_request_cancellations", None): + retain_pre_connector_error( + RuntimeError( + "User cancellation ownership is still incomplete after " + "transceiver shutdown; refusing to shut down resource " + "managers")) + if getattr(self, "_deferred_transfer_terminations", None): + retain_pre_connector_error( + RuntimeError( + "KV transfer request ownership is still active after " + "transceiver shutdown; refusing to shut down resource " + "managers")) + if getattr(self, "kv_connector_manager", None) is not None: + self._drain_connector_completions_after_shutdown( + preexisting_error=pre_connector_error) + else: + drain_connector_completions = getattr( + self, "_drain_pending_connector_completions", None) + if drain_connector_completions is not None: + try: + drain_connector_completions( + terminate_after_worker_shutdown=True) + except Exception as error: + retain_pre_connector_error(error) + + # Every check after connector drain must reach one final world vote. + # Otherwise a rank-local retryable failure can return early while peers + # destroy their managers; that rank's next shutdown attempt would then + # re-enter connector collectives without them. + post_connector_error = pre_connector_error + + def retain_post_connector_error(error: Exception) -> None: + nonlocal post_connector_error + if post_connector_error is None: + post_connector_error = error + logger.error( + "KV ownership shutdown is not ready for manager teardown; " + f"continuing to the final rank-uniform vote: {error}") + + for operation in ( + self._discard_pending_transfer_responses_after_shutdown, + self._discard_unpublishable_transfer_terminals_after_shutdown, + self._drain_discarded_transfer_terminals_after_shutdown, + ): + try: + operation() + except Exception as error: + retain_post_connector_error(error) + + if getattr(self, "_pending_connector_completions", None): + retain_post_connector_error( + RuntimeError( + "Connector completion ownership is still incomplete after " + "executor shutdown; refusing to shut down resource managers" + )) + try: + transceiver_transfer_owners = (transfer_manager.requests_with_owner( + _PYTHON_NATIVE_TRANSCEIVER_OWNER) if transfer_manager + is not None else {}) + if transceiver_transfer_owners: + raise RuntimeError( + "KV cache transceiver still owns active request resources; " + "refusing to shut down resource managers") + except Exception as error: + retain_post_connector_error(error) + try: + if (transfer_manager is not None + and transfer_manager.has_any_inflight_requests()): + raise RuntimeError( + "Asynchronous transfer ownership is still active after " + "executor shutdown; refusing to shut down resource managers" + ) + except Exception as error: + retain_post_connector_error(error) + + termination_handler = getattr(self, "_disagg_pp_termination_handler", + None) + if termination_handler is not None: + try: + termination_handler.terminate_all_after_shutdown() + if termination_handler.has_pending_terminations(): + raise RuntimeError( + "PP request termination is still pending after executor " + "shutdown; refusing to shut down resource managers") + except Exception as error: + retain_post_connector_error(error) + + termination_progress = getattr( + self, "_request_resource_termination_progress", {}) + for request_key, request in list(termination_progress.items()): + if termination_progress.get(request_key) is request: + try: + self._do_terminate_request(request) + except Exception as error: + retain_post_connector_error(error) + if termination_progress: + retain_post_connector_error( + RuntimeError( + "Request resource termination is still incomplete; refusing " + "to shut down resource managers")) + + try: + unresolved_ownerless_terminals = [ + terminal for terminal in getattr( + self, "_pending_transfer_terminals", {}).values() + if not terminal.teardown_complete + and not self._has_async_transfer_owner(terminal.request) + and not self._has_pending_ownership_transition(terminal.request) + ] + if unresolved_ownerless_terminals: + raise RuntimeError( + "Ownerless KV-transfer terminal teardown is incomplete for " + f"{len(unresolved_ownerless_terminals)} request(s); refusing " + "to shut down resource managers") + except Exception as error: + retain_post_connector_error(error) + + if any(terminal.response_creation_in_doubt for terminal in getattr( + self, "_pending_transfer_terminals", {}).values()): + retain_post_connector_error( + RuntimeError( + "Terminal response creation outcome is in doubt; refusing " + "to shut down resource managers")) + + has_in_doubt_releases = getattr(self.resource_manager, + "has_in_doubt_resource_releases", None) + try: + if (has_in_doubt_releases is not None and has_in_doubt_releases()): + raise RuntimeError( + "Request resource release outcome is in doubt; refusing to " + "shut down resource managers") + except Exception as error: + retain_post_connector_error(error) + has_pending_releases = getattr(self.resource_manager, + "has_pending_resource_releases", None) + try: + if (has_pending_releases is not None and has_pending_releases()): + raise RuntimeError( + "Request resource release is still pending; refusing to shut " + "down resource managers") + except Exception as error: + retain_post_connector_error(error) + + # Sampler and DWDP teardown are fallible and potentially + # non-replayable. Complete them before the final readiness vote so one + # rank cannot release engines after a peer enters an ambiguous phase. + try: + self._prepare_final_executor_shutdown() + except Exception as error: + retain_post_connector_error(error) + + local_not_ready = post_connector_error is not None + try: + if self._dist_size(self.dist, "world_size") > 1: + any_rank_not_ready = bool( + self.dist.allreduce(int(local_not_ready), op=ReduceOp.MAX)) + else: + any_rank_not_ready = local_not_ready + except Exception as consensus_error: + if post_connector_error is not None: + raise post_connector_error from consensus_error + raise + if any_rank_not_ready: + if post_connector_error is not None: + raise post_connector_error + raise RuntimeError( + "KV ownership shutdown is not ready on another rank; refusing " + "to shut down resource managers") + + # Manager teardown is itself fallible and destructive. Every returned + # manager is durably skipped on retry, an exception remains in doubt, + # and a second world vote keeps all ranks out of engine deletion when + # any local manager failed. + self._shutdown_resource_managers_rank_uniform() + self._release_final_executor_objects() def can_enqueue_requests(self) -> bool: """ @@ -2837,7 +4070,6 @@ def _broadcast_sample_state_loop(self): # Do not wait for PP send handles here. The next # _ring_broadcast_sample_state call drains the previous isend # for the same microbatch_id before reusing that slot. - # # Waiting here can hang during shutdown. Peer ranks may have # left this loop and moved to the next executor setup, so no # rank is polling the broadcast communicator. In that state, @@ -3384,6 +4616,131 @@ def _uses_async_disagg_gen_transfer(self) -> bool: return (not self._is_disagg_gen_only_no_context_benchmark() and os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") != "1") + def _requires_physical_transfer_drain(self) -> bool: + transceiver = getattr(self, "kv_cache_transceiver", None) + return (getattr(transceiver, + "requires_physical_drain_before_request_release", False) + is True) + + def _has_async_transfer_owner(self, request: LlmRequest) -> bool: + """Whether any provider owns the exact request's transfer lease.""" + transfer_manager = getattr(self, "async_transfer_manager", None) + if transfer_manager is None: + return False + has_any_transfer_owner = getattr(transfer_manager, + "has_any_transfer_owner", None) + if has_any_transfer_owner is not None: + result = has_any_transfer_owner(request) + # Some unit-test executors use an incompletely configured Mock. + # Accept only the real manager's concrete bool, then fall back to + # the exact-request map instead of treating a Mock as truthy. + if result is True or result is False: + return result + requests_in_transfer = getattr(transfer_manager, "requests_in_transfer", + None) + if requests_in_transfer is None: + return False + return (requests_in_transfer().get(request.py_request_id) is request) + + def _has_transceiver_transfer_owner(self, request: LlmRequest) -> bool: + """Whether the Python-native transceiver owns the exact request.""" + transfer_manager = getattr(self, "async_transfer_manager", None) + if transfer_manager is None: + return False + has_transfer_owner = getattr(transfer_manager, "has_transfer_owner", + None) + return (has_transfer_owner is not None and has_transfer_owner( + request, _PYTHON_NATIVE_TRANSCEIVER_OWNER) is True) + + def _is_final_async_transfer_owner(self, + request: LlmRequest, + owner: Optional[str] = None) -> bool: + """Whether retiring one exact provider would retire the final owner.""" + transfer_manager = getattr(self, "async_transfer_manager", None) + if transfer_manager is None: + return False + is_final_transfer_owner = getattr(transfer_manager, + "is_final_transfer_owner", None) + if is_final_transfer_owner is None: + return False + result = is_final_transfer_owner(request, owner=owner) + return result is True + + def _was_transfer_request_terminated(self, request: LlmRequest) -> bool: + """Whether normal teardown ran while a transfer lease stayed live.""" + terminated_requests = getattr(self, "_terminated_transfer_requests", {}) + return terminated_requests.get(id(request)) is request + + def _record_terminated_transfer_request(self, request: LlmRequest) -> None: + """Remember exact resources freed before the final transfer owner.""" + if not self._has_async_transfer_owner(request): + return + terminated_requests = getattr(self, "_terminated_transfer_requests", + None) + if terminated_requests is None: + terminated_requests = {} + self._terminated_transfer_requests = terminated_requests + terminated_requests[id(request)] = request + + def _forget_terminated_transfer_request_if_unowned( + self, request: LlmRequest) -> None: + """Drop the early-teardown root only after every owner is retired.""" + if self._has_async_transfer_owner(request): + return + terminated_requests = getattr(self, "_terminated_transfer_requests", + None) + request_key = id(request) + if (terminated_requests is not None + and terminated_requests.get(request_key) is request): + terminated_requests.pop(request_key) + + def _defer_transceiver_owners_for_shutdown(self) -> None: + """Keep Python-native owners rooted through post-loop shutdown.""" + transfer_manager = getattr(self, "async_transfer_manager", None) + if transfer_manager is None: + return + + deferred_terminations = getattr(self, "_deferred_transfer_terminations", + None) + if deferred_terminations is None: + deferred_terminations = {} + self._deferred_transfer_terminations = deferred_terminations + already_terminated = getattr( + self, "_deferred_transfer_terminations_already_terminated", None) + if already_terminated is None: + already_terminated = set() + self._deferred_transfer_terminations_already_terminated = already_terminated + terminated_requests = getattr(self, "_terminated_transfer_requests", {}) + pending_cancellations = getattr(self, "_pending_request_cancellations", + {}) + pending_native_completions = getattr( + self, "_pending_native_context_completions", {}) + for request in transfer_manager.requests_with_owner( + _PYTHON_NATIVE_TRANSCEIVER_OWNER).values(): + request_key = id(request) + cancellation_progress = pending_cancellations.get(request_key) + if (cancellation_progress is not None + and cancellation_progress.request is request): + # The cancellation ledger already owns the exact transport + # drain -> transceiver release -> logical termination sequence. + continue + native_completion = pending_native_completions.get(request_key) + if (native_completion is not None + and native_completion.request is request): + # The status poll already consumed this session's one-shot + # completion. Its phase ledger, rather than cancellation, + # owns the remaining release and teardown work. + continue + if deferred_terminations.get(request_key) is request: + # Preserve whether an earlier error path still needs normal + # request termination; a shutdown retry must not reclassify it. + continue + deferred_terminations[request_key] = request + if terminated_requests.get(request_key) is request: + already_terminated.add(request_key) + else: + already_terminated.discard(request_key) + def _uses_kv_manager_v2(self) -> bool: explicit_flag = getattr(self, "_is_kv_manager_v2", None) if explicit_flag is not None: @@ -3718,11 +5075,211 @@ def _kv_connector_start_batch(self, scheduled_batch): self.kv_connector_manager.worker.start_load_kv( torch.cuda.current_stream()) - def _kv_connector_terminate_requests(self): + def _get_pending_connector_completions( + self) -> Dict[int, _ConnectorCompletionProgress]: + pending = getattr(self, "_pending_connector_completions", None) + if pending is None: + pending = {} + self._pending_connector_completions = pending + return pending + + def _advance_connector_completion( + self, + progress: _ConnectorCompletionProgress, + *, + terminate_after_worker_shutdown: bool = False) -> bool: + """Advance one consumed connector notification without replay.""" + request = progress.request + request_was_already_terminated = ( + self._was_transfer_request_terminated(request)) + terminal = self._get_transfer_terminal(request) + if (terminal is None and progress.was_active + and self._is_final_async_transfer_owner(request)): + # In mixed mode, connector-first completion only retires the + # connector contribution: either native or legacy-C++ transceiver + # status is authoritative while another provider remains because it + # may still report a transfer failure. + outcome = (_TransferTerminalOutcome.CANCELLED if getattr( + request, "is_finished_due_to_cancellation", False) is True else + _TransferTerminalOutcome.SUCCEEDED) + terminal = self._claim_transfer_terminal(request, outcome, + "kv_connector") + if (progress.was_active and terminal is not None + and terminal.outcome in (_TransferTerminalOutcome.SUCCEEDED, + _TransferTerminalOutcome.CANCELLED) + and not terminal.response_published): + terminal = self._prepare_transfer_terminal_response( + request, terminal.outcome, terminal.claimant) + + if not progress.owner_released: + progress.all_transfers_complete = ( + PyExecutor._release_transfer_owner(self, request)) + progress.owner_released = True + + if not progress.all_transfers_complete: + return True + + if (terminal is not None + and terminal.outcome == _TransferTerminalOutcome.FAILED + and request in self.active_requests + and not terminate_after_worker_shutdown): + # The synchronized error path still owns logical response and + # teardown. The connector notification itself is fully consumed. + return True + + if not progress.active_request_removed: + if request in self.active_requests: + self.active_requests.remove(request) + progress.active_request_removed = True + + if not progress.request_terminated: + if not request_was_already_terminated: + if terminate_after_worker_shutdown: + self._terminate_request_after_worker_shutdown(request) + else: + self._terminate_request(request) + else: + self._mark_transfer_terminal_teardown_complete(request) + progress.request_terminated = True + self._forget_terminated_transfer_request_if_unowned(request) + self._retire_transfer_terminal_if_complete(request) + return True + + def _drain_pending_connector_completions( + self, *, terminate_after_worker_shutdown: bool = False) -> None: + pending = self._get_pending_connector_completions() + for request_key, progress in list(pending.items()): + if pending.get(request_key) is not progress: + continue + try: + completed = self._advance_connector_completion( + progress, + terminate_after_worker_shutdown= + terminate_after_worker_shutdown) + except Exception as e: + logger.error( + "Failed to finish connector ownership for request " + f"{progress.request.py_request_id}; retaining exact progress " + f"for retry: {e}") + continue + if completed and pending.get(request_key) is progress: + pending.pop(request_key) + + def _record_finished_connector_requests( + self, finished_requests: Iterable[LlmRequest]) -> None: + """Enroll one polled connector batch without advancing cleanup twice.""" + pending = self._get_pending_connector_completions() + for request in finished_requests: + request_key = id(request) + existing = pending.get(request_key) + if existing is not None and existing.request is not request: + raise RuntimeError( + "connector completion identity collision for request " + f"{request.py_request_id}") + if existing is None: + was_active = request in self.active_requests + if (was_active and self._get_transfer_terminal(request) is None + and self._is_final_async_transfer_owner(request)): + outcome = (_TransferTerminalOutcome.CANCELLED if getattr( + request, "is_finished_due_to_cancellation", False) + is True else _TransferTerminalOutcome.SUCCEEDED) + self._claim_transfer_terminal(request, outcome, + "kv_connector") + pending[request_key] = _ConnectorCompletionProgress( + request=request, + was_active=was_active, + ) + + def _kv_connector_terminate_requests( + self, *, terminate_after_worker_shutdown: bool = False): if self.kv_connector_manager: - reqs_to_terminate = self.kv_connector_manager.get_finished() - for req in reqs_to_terminate: - self._end_transfer_and_maybe_terminate(req) + self._record_finished_connector_requests( + self.kv_connector_manager.get_finished()) + self._drain_pending_connector_completions( + terminate_after_worker_shutdown=terminate_after_worker_shutdown) + + def _drain_connector_completions_after_shutdown( + self, preexisting_error: Optional[Exception] = None) -> None: + """Give one-shot connector completions a bounded shutdown drain.""" + # get_finished contains connector-wide collectives. Keep the poll count + # rank-uniform instead of breaking on local owner state, which may + # diverge when one rank encounters a retryable cleanup failure. + first_error = preexisting_error + + def retain_error(error: Exception) -> None: + nonlocal first_error + if first_error is None: + first_error = error + logger.error("Connector shutdown completion drain failed locally; " + f"continuing rank-uniform polling: {error}") + + transfer_manager = getattr(self, "async_transfer_manager", None) + try: + needs_grace = bool(self._get_pending_connector_completions()) + if transfer_manager is not None: + needs_grace = (needs_grace + or transfer_manager.has_any_inflight_requests()) + except Exception as error: + retain_error(error) + needs_grace = True + + for poll_index in range(_SHUTDOWN_CONNECTOR_COMPLETION_POLLS): + if poll_index > 0 and needs_grace: + try: + time.sleep(_SHUTDOWN_CONNECTOR_COMPLETION_POLL_INTERVAL_S) + except Exception as error: + retain_error(error) + + try: + finished_requests = self.kv_connector_manager.get_finished() + except Exception as error: + retain_error(error) + else: + try: + self._record_finished_connector_requests(finished_requests) + except Exception as error: + retain_error(error) + + # Completion progress records every successful side effect, so a + # later poll can resume without replaying owner release or request + # teardown even when enrollment or polling failed locally. + try: + self._drain_pending_connector_completions( + terminate_after_worker_shutdown=True) + except Exception as error: + retain_error(error) + + try: + needs_grace = bool(self._get_pending_connector_completions()) + if transfer_manager is not None: + needs_grace = (needs_grace or + transfer_manager.has_any_inflight_requests()) + except Exception as error: + retain_error(error) + needs_grace = True + + try: + local_incomplete = bool(self._get_pending_connector_completions()) + if transfer_manager is not None: + local_incomplete = ( + local_incomplete + or transfer_manager.has_any_inflight_requests()) + except Exception as error: + retain_error(error) + local_incomplete = True + + local_failed = first_error is not None or local_incomplete + if self._dist_size(self.dist, "world_size") > 1: + any_failed = bool( + self.dist.allreduce(int(local_failed), op=ReduceOp.MAX)) + else: + any_failed = local_failed + if any_failed: + if first_error is not None: + raise first_error + raise RuntimeError( + "Connector shutdown completion drain failed on another rank or " + "left ownership incomplete") def _kv_connector_wait_for_save(self): if self.kv_connector_manager is not None: @@ -3872,6 +5429,11 @@ def _handle_disagg_cache_errors_synced(self): otherwise the downstream ``tp_gather`` in ``_enqueue_responses`` deadlocks or leaves peer replicas running. """ + # Error responses remove requests from active_requests immediately, + # but physical transfer cancellation may need more than one bounded + # poll. Retry those ownership releases once per executor iteration. + PyExecutor._drain_deferred_transfer_terminations(self) + if not self.kv_cache_transceiver: return @@ -6042,8 +7604,21 @@ def kv_connector_request_finished(req: LlmRequest): req.py_request_id) # Order is important here: we need to start the transfer before responding # to make sure the blocks are stored for reuse before they are sent. - self.async_transfer_manager.start_transfer(req) - self.kv_cache_transceiver.respond_and_send_async(req) + # Python native gets an exact named contribution so it can + # retire independently from an anonymous connector save. + # The default C++ transceiver keeps its legacy anonymous + # counter contract. + if self._requires_physical_transfer_drain(): + self.async_transfer_manager.start_transfer( + req, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER) + else: + self.async_transfer_manager.start_transfer(req) + try: + self.kv_cache_transceiver.respond_and_send_async(req) + except Exception: + if self._requires_physical_transfer_drain(): + self._reconcile_failed_native_transfer_launch(req) + raise if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: req.py_kv_transfer_start_time = time.monotonic() @@ -6078,6 +7653,41 @@ def _is_disagg_error_cleanup_blocked(self, request: LlmRequest) -> bool: return True return False + def _reconcile_failed_native_transfer_launch(self, + request: LlmRequest) -> None: + """Fail closed after native owner enrollment but failed launch.""" + request.state = LlmRequestState.DISAGG_TRANS_ERROR + deferred = getattr(self, "_deferred_transfer_terminations", None) + if deferred is None: + deferred = {} + self._deferred_transfer_terminations = deferred + request_key = id(request) + existing = deferred.get(request_key) + if existing is not None and existing is not request: + raise RuntimeError( + "failed native launch identity collision for request " + f"{request.py_request_id}") + deferred[request_key] = request + + try: + is_drained = self.kv_cache_transceiver.cancel_request(request) + except Exception as e: + logger.error( + "Failed to cancel native transfer after launch error for " + f"request {request.py_request_id}; retaining exact request for " + f"shutdown retry: {e}") + return + if not is_drained: + return + + try: + self._finalize_deferred_transfer_termination(request) + except Exception as e: + logger.error( + "Failed to retire native transfer owner after launch error " + f"for request {request.py_request_id}; retaining exact request " + f"for retry: {e}") + def _get_disagg_reqs_in_error_state(self): return [ req for req in self.active_requests @@ -6085,6 +7695,204 @@ def _get_disagg_reqs_in_error_state(self): and not self._is_disagg_error_cleanup_blocked(req) ] + def _get_pending_native_context_completions( + self) -> Dict[int, _NativeContextCompletionProgress]: + pending = getattr(self, "_pending_native_context_completions", None) + if pending is None: + pending = {} + self._pending_native_context_completions = pending + return pending + + def _get_native_context_completion( + self, + request: LlmRequest) -> Optional[_NativeContextCompletionProgress]: + """Return the exact one-shot native completion owner, if present.""" + progress = getattr(self, "_pending_native_context_completions", + {}).get(id(request)) + if progress is not None and progress.request is request: + return progress + return None + + def _advance_native_context_completion( + self, + progress: _NativeContextCompletionProgress, + *, + terminate_after_worker_shutdown: bool = False) -> bool: + """Advance one consumed native completion notification without replay.""" + request = progress.request + if progress.was_deferred: + if not self._finalize_deferred_transfer_termination( + request, + terminate_after_worker_shutdown= + terminate_after_worker_shutdown): + raise RuntimeError( + "deferred native completion lost its exact request " + f"ledger for request {request.py_request_id}") + return True + + if progress.transfer_failed: + terminal = self._claim_transfer_terminal( + request, _TransferTerminalOutcome.FAILED, + _PYTHON_NATIVE_TRANSCEIVER_OWNER) + elif progress.was_active: + terminal = self._prepare_transfer_terminal_response( + request, _TransferTerminalOutcome.SUCCEEDED, + _PYTHON_NATIVE_TRANSCEIVER_OWNER) + else: + terminal = self._get_transfer_terminal(request) + + if not progress.owner_released: + progress.all_transfers_complete = PyExecutor._release_transfer_owner( + self, request, _PYTHON_NATIVE_TRANSCEIVER_OWNER) + progress.owner_released = True + + # Another provider owns the final teardown. This native contribution is + # nevertheless complete and must not be released again on a later poll. + if not progress.all_transfers_complete: + return True + + terminal_failed = (terminal is not None and terminal.outcome + == _TransferTerminalOutcome.FAILED) + terminal_cancelled = (terminal is not None and terminal.outcome + == _TransferTerminalOutcome.CANCELLED) + if progress.was_active: + if terminal_failed or terminal_cancelled: + # If the logical error/cancel path still owns the active root, + # leave response delivery to it. Once that path removes the + # request (or after loop shutdown), this final provider owns + # the remaining physical teardown. + if (terminate_after_worker_shutdown + or request not in self.active_requests): + if not progress.active_request_removed: + if request in self.active_requests: + self.active_requests.remove(request) + progress.active_request_removed = True + if not progress.request_terminated: + if not progress.was_already_terminated: + if terminate_after_worker_shutdown: + self._terminate_request_after_worker_shutdown( + request) + else: + self._terminate_request(request) + else: + self._mark_transfer_terminal_teardown_complete( + request) + progress.request_terminated = True + if (hasattr(self, "_terminated_transfer_requests") + and not progress.terminated_marker_forgotten): + self._forget_terminated_transfer_request_if_unowned( + request) + progress.terminated_marker_forgotten = True + return True + if not progress.active_request_removed: + if request in self.active_requests: + self.active_requests.remove(request) + progress.active_request_removed = True + elif terminal_failed or terminal_cancelled: + if (progress.was_already_terminated + and not progress.terminated_marker_forgotten): + self._mark_transfer_terminal_teardown_complete(request) + self._forget_terminated_transfer_request_if_unowned(request) + progress.terminated_marker_forgotten = True + self._retire_transfer_terminal_if_complete(request) + return True + + legacy_pp1_already_terminated = ( + not hasattr(self, "_terminated_transfer_requests") + and getattr(self, "force_terminate_ctx_for_partial_reuse", False)) + if not progress.request_terminated: + if (not progress.was_already_terminated + and not legacy_pp1_already_terminated): + if terminate_after_worker_shutdown: + self._terminate_request_after_worker_shutdown(request) + else: + self._terminate_request(request) + else: + self._mark_transfer_terminal_teardown_complete(request) + progress.request_terminated = True + + if (hasattr(self, "_terminated_transfer_requests") + and not progress.terminated_marker_forgotten): + self._forget_terminated_transfer_request_if_unowned(request) + progress.terminated_marker_forgotten = True + self._retire_transfer_terminal_if_complete(request) + return True + + def _drain_pending_native_context_completions( + self, + *, + terminate_after_worker_shutdown: bool = False, + request_keys: Optional[set[int]] = None) -> None: + pending = self._get_pending_native_context_completions() + for request_key, progress in list(pending.items()): + if request_keys is not None and request_key not in request_keys: + continue + if pending.get(request_key) is not progress: + continue + try: + completed = self._advance_native_context_completion( + progress, + terminate_after_worker_shutdown= + terminate_after_worker_shutdown) + except Exception as e: + logger.error( + "Failed to finish native context ownership for request " + f"{progress.request.py_request_id}; retaining exact " + f"progress for retry: {e}") + continue + if completed and pending.get(request_key) is progress: + pending.pop(request_key) + + def _ingest_native_context_completions( + self, finished_request_ids: Iterable[int], + error_request_ids: Iterable[int], + requests_in_transfer: Dict[int, LlmRequest]) -> None: + pending = self._get_pending_native_context_completions() + error_request_ids = set(error_request_ids) + for request_id in list(finished_request_ids) + list(error_request_ids): + request = requests_in_transfer.get(request_id) + if request is None: + logger.warning( + f"Request {request_id} not found in transfer manager") + continue + request_key = id(request) + existing = pending.get(request_key) + if existing is not None and existing.request is not request: + raise RuntimeError( + "native completion identity collision for request " + f"{request.py_request_id}") + if existing is not None: + continue + pending_cancellations = getattr(self, + "_pending_request_cancellations", + {}) + cancellation = pending_cancellations.get(request_key) + if cancellation is not None and cancellation.request is request: + # Cancellation was admitted before this status notification. + # The one-shot status is positive transport-drain evidence, but + # cancellation remains the sole named-owner/terminal ledger. + cancellation.transport_drained = True + continue + deferred = getattr(self, "_deferred_transfer_terminations", {}) + transfer_failed = (request_id in error_request_ids or request.state + == LlmRequestState.DISAGG_TRANS_ERROR) + if transfer_failed: + self._claim_transfer_terminal(request, + _TransferTerminalOutcome.FAILED, + _PYTHON_NATIVE_TRANSCEIVER_OWNER) + elif request in self.active_requests: + self._claim_transfer_terminal( + request, _TransferTerminalOutcome.SUCCEEDED, + _PYTHON_NATIVE_TRANSCEIVER_OWNER) + pending[request_key] = _NativeContextCompletionProgress( + request=request, + was_active=request in self.active_requests, + was_deferred=deferred.get(request_key) is request, + transfer_failed=transfer_failed, + was_already_terminated=self._was_transfer_request_terminated( + request), + ) + def _check_cache_transfer_errors(self, error_msg_prefix: str): """Check and handle cache transfer errors. @@ -6102,24 +7910,47 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): @nvtx_range("_check_disagg_ctx_cache_transfer_status") def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): + native_owner_mode = self._requires_physical_transfer_drain() + pending_native = self._get_pending_native_context_completions() + pending_request_keys = set(pending_native) + pending_request_ids = { + progress.request.py_request_id + for progress in pending_native.values() + } + if native_owner_mode and pending_native: + # Evidence already consumed by a prior poll counts as progress for + # a blocking caller. Retry it before asking the transport to wait + # for an unrelated new completion. + self._drain_pending_native_context_completions() + atLeastNum = 0 + finished_requests, error_requests = self.kv_cache_transceiver.check_context_transfer_status( atLeastNum) + transfer_owner = (_PYTHON_NATIVE_TRANSCEIVER_OWNER + if native_owner_mode else None) - completed_req_ids = set(finished_requests + error_requests) + completed_req_ids = (set(finished_requests + error_requests) + | pending_request_ids) requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) - for request_id in completed_req_ids: - - if request_id not in requests_in_transfer: - logger.warning( - f"Request {request_id} not found in transfer manager") - continue - - request = requests_in_transfer[request_id] - - self._end_transfer_and_maybe_terminate(request) + if native_owner_mode: + self._ingest_native_context_completions(finished_requests, + error_requests, + requests_in_transfer) + self._drain_pending_native_context_completions( + request_keys=set(pending_native) - pending_request_keys) + else: + for request_id in set(finished_requests + error_requests): + if request_id not in requests_in_transfer: + logger.warning( + f"Request {request_id} not found in transfer manager") + continue + request = requests_in_transfer[request_id] + if not self._finalize_deferred_transfer_termination(request): + self._end_transfer_and_maybe_terminate( + request, transfer_owner=transfer_owner) # The set of requests in transfer may have changed since we terminated some requests. requests_in_transfer = self.async_transfer_manager.requests_in_transfer( @@ -6132,6 +7963,17 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): or request_id in self._disagg_timed_out_ctx_cancelled_ids): continue + if self._requires_physical_transfer_drain(): + is_cancelled = self._request_kv_transfer_cancellation(request) + if not is_cancelled: + continue + request.py_kv_transfer_start_time = None + request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + if not self._finalize_deferred_transfer_termination(request): + self._end_transfer_and_maybe_terminate( + request, transfer_owner=transfer_owner) + continue + is_cancelled = self._request_kv_transfer_cancellation(request) if not is_cancelled: continue @@ -6444,6 +8286,7 @@ def _handle_errors(self, server health. """ error_responses: Dict[int, LlmResponse] = {} + new_error_response_request_keys = set() error_msg = error_msg or "error" budget_fatal = (self._error_budget.consume(error_msg) @@ -6510,23 +8353,132 @@ def _handle_errors(self, failed_requests = (list(self.active_requests) if requests is None else requests) + pending_ownership_request_keys = set() + competing_terminal_request_keys = set() + deferred_request_keys = set() + drained_deferred_request_keys = set() + provider_delegated_request_keys = set() + deferred_terminations = getattr(self, "_deferred_transfer_terminations", + None) + if deferred_terminations is None: + deferred_terminations = {} + self._deferred_transfer_terminations = deferred_terminations + already_terminated = getattr( + self, "_deferred_transfer_terminations_already_terminated", None) + if already_terminated is None: + already_terminated = set() + self._deferred_transfer_terminations_already_terminated = already_terminated + context_transfer_state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + for request in failed_requests: + request_key = id(request) + terminal = self._claim_transfer_terminal( + request, _TransferTerminalOutcome.FAILED, "request_error") + if terminal.outcome != _TransferTerminalOutcome.FAILED: + # A provider completion or cancellation already owns the first + # logical result through every remaining sibling lease. + competing_terminal_request_keys.add(request_key) + pending_ownership_request_keys.add(request_key) + continue + if self._has_pending_ownership_transition(request): + pending_ownership_request_keys.add(request_key) + continue + already_deferred = deferred_terminations.get(request_key) is request + is_context_transfer = request.state == context_transfer_state + has_transceiver_owner = self._has_transceiver_transfer_owner( + request) + if not self._requires_physical_transfer_drain(): + continue + try: + is_drained = self.kv_cache_transceiver.cancel_request(request) + except Exception as e: + logger.error( + f"Failed to cancel KV transfer for request " + f"{request.py_request_id}; retaining request resources until " + f"drain can be retried: {e}") + is_drained = False + if not is_drained: + deferred_terminations[request_key] = request + already_terminated.discard(request_key) + deferred_request_keys.add(request_key) + elif (already_deferred or is_context_transfer + or has_transceiver_owner): + # Context transfers also hold an AsyncTransferManager entry. + # Keep the request in the map through logical error handling + # so finalization releases that entry exactly once. + deferred_terminations[request_key] = request + already_terminated.discard(request_key) + drained_deferred_request_keys.add(request_key) + + force_terminate = getattr(self, "force_terminate_ctx_for_partial_reuse", + False) + for request in failed_requests: + request_key = id(request) + if (request_key in pending_ownership_request_keys + or request_key in deferred_request_keys + or request_key in drained_deferred_request_keys): + continue + if self._has_async_transfer_owner(request) and not force_terminate: + # Provider completion is the only path that can prove the + # shared transfer lease is final. Keep resources until then. + provider_delegated_request_keys.add(request_key) + for request in failed_requests: + request_key = id(request) + if request_key in competing_terminal_request_keys: + continue + terminal = self._get_transfer_terminal(request) + if terminal is None or terminal.outcome != _TransferTerminalOutcome.FAILED: + continue req_id = request.py_request_id request.state = LlmRequestState.GENERATION_COMPLETE - error_responses[req_id] = LlmResponse( - request_id=req_id, - error_msg=error_msg, - client_id=request.py_client_id) - if requests is None: - self.active_requests.clear() - else: - self.active_requests = [ - request for request in self.active_requests - if request not in requests - ] + if not terminal.response_published: + if not terminal.response_created: + terminal.response = LlmResponse( + request_id=req_id, + error_msg=error_msg, + client_id=request.py_client_id) + terminal.response_created = True + self._buffer_transfer_terminal_response(terminal) + error_responses[req_id] = terminal.response + new_error_response_request_keys.add(request_key) + if self._has_async_transfer_owner(request): + # Prevent a later final owner retirement from translating the + # failed request into CONTEXT_COMPLETE. + request.state = LlmRequestState.DISAGG_TRANS_ERROR + publishing_error_terminals = [ + terminal for request in failed_requests + if id(request) in new_error_response_request_keys + if (terminal := self._get_transfer_terminal(request)) is not None + ] + self._begin_transfer_response_publication(publishing_error_terminals) self._enqueue_responses(list(error_responses.items())) + self._commit_transfer_response_publication(publishing_error_terminals) + published_error_request_keys = set() for request in failed_requests: - self._terminate_request(request) + terminal = self._get_transfer_terminal(request) + if (terminal is not None + and terminal.outcome == _TransferTerminalOutcome.FAILED + and terminal.response_published): + published_error_request_keys.add(id(request)) + self.active_requests = [ + request for request in self.active_requests + if id(request) not in published_error_request_keys + ] + for request in failed_requests: + if id(request) in published_error_request_keys: + self._retire_transfer_terminal_if_complete(request) + for request in failed_requests: + request_key = id(request) + if request_key in competing_terminal_request_keys: + continue + if (request_key in pending_ownership_request_keys + or request_key in deferred_request_keys + or request_key in provider_delegated_request_keys): + continue + if request_key in drained_deferred_request_keys: + self._finalize_deferred_transfer_termination(request) + else: + self._terminate_request(request) if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() @@ -6543,13 +8495,38 @@ def _terminate_request(self, request: LlmRequest): self._do_terminate_request(request) def _do_terminate_request(self, request: LlmRequest): - self.resource_manager.free_resources(request) + transfer_manager = getattr(self, "async_transfer_manager", None) + if (transfer_manager is not None + and transfer_manager.has_pending_admission(request) is True): + raise RuntimeError( + "transfer admission is incomplete for request " + f"{request.py_request_id}; refusing to release resources") + + request_key = id(request) + termination_progress = getattr( + self, "_request_resource_termination_progress", None) + if termination_progress is None: + termination_progress = {} + self._request_resource_termination_progress = termination_progress + released_request = termination_progress.get(request_key) + if released_request is not None and released_request is not request: + raise RuntimeError( + "request resource termination identity collision for request " + f"{request.py_request_id}") + if released_request is None: + self.resource_manager.free_resources(request) + termination_progress[request_key] = request + self._prefetched_request_ids.discard(request.py_request_id) self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) self._disagg_timed_out_gen_cancelled_ids.discard(request.py_request_id) if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) + self._record_terminated_transfer_request(request) + self._mark_transfer_terminal_teardown_complete(request) + if termination_progress.get(request_key) is request: + termination_progress.pop(request_key) def _is_request_in_transmission(self, request) -> bool: """Check if a request is currently in transmission state.""" @@ -6567,6 +8544,12 @@ def _try_cancel_request(self, request) -> bool: if self.kv_cache_transceiver is None: return True + # The Python-native transceiver can still own a live send/receive + # session after another path has moved the request out of a TRANS + # state. Its cancellation result is the physical-drain proof. + if self._requires_physical_transfer_drain(): + return self._request_kv_transfer_cancellation(request) + async_transfer_manager = getattr(self, "async_transfer_manager", None) if (getattr(request, "is_context_only_request", False) is True and async_transfer_manager is not None and request.py_request_id @@ -6584,9 +8567,233 @@ def _try_cancel_request(self, request) -> bool: return self._request_kv_transfer_cancellation(request) + def _resolve_cancelled_context_transceiver_owner( + self, request: LlmRequest) -> bool: + """Retire a drained Python-native owner or await another provider.""" + transfer_manager = getattr(self, "async_transfer_manager", None) + if (request.is_generation_only_request() or transfer_manager is None + or not self._has_transceiver_transfer_owner(request)): + return not self._has_async_transfer_owner(request) + + # AsyncTransferManager normally maps its final successful owner to + # CONTEXT_COMPLETE. User cancellation must remain cancellation/error, + # including when owner release raises and is retried next iteration. + original_state = request.state + request.state = LlmRequestState.DISAGG_TRANS_ERROR + try: + return transfer_manager.end_transfer( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER) + except Exception: + request.state = original_state + raise + + def _get_pending_request_cancellations( + self) -> Dict[int, _RequestCancellationProgress]: + pending = getattr(self, "_pending_request_cancellations", None) + if pending is None: + pending = {} + self._pending_request_cancellations = pending + return pending + + def _advance_request_cancellation( + self, progress: _RequestCancellationProgress) -> bool: + """Advance one exact request through cancellation ownership phases.""" + request = progress.request + terminal = self._claim_transfer_terminal( + request, _TransferTerminalOutcome.CANCELLED, "user_cancellation") + if terminal.outcome != _TransferTerminalOutcome.CANCELLED: + return True + if not progress.transport_drained: + if not self._try_cancel_request(request): + return False + progress.transport_drained = True + + if not progress.transceiver_release_resolved: + progress.all_transfers_complete = ( + self._resolve_cancelled_context_transceiver_owner(request)) + progress.transceiver_release_resolved = True + + if not progress.logical_cancelled: + request.py_kv_transfer_timed_out = False + request.finish_by_reason(FinishReason.CANCELLED) + progress.logical_cancelled = True + + if not progress.decoding_iter_updated: + request.decoding_iter = request.py_decoding_iter + progress.decoding_iter_updated = True + return True + + def _complete_request_cancellation_termination( + self, + progress: _RequestCancellationProgress, + *, + terminate_after_worker_shutdown: bool = False) -> None: + """Terminate, or hand teardown to the last remaining provider.""" + if progress.termination_handoff_complete: + return + + request = progress.request + if progress.all_transfers_complete: + if not self._was_transfer_request_terminated(request): + if terminate_after_worker_shutdown: + self._terminate_request_after_worker_shutdown(request) + else: + self._terminate_request(request) + else: + self._mark_transfer_terminal_teardown_complete(request) + self._forget_terminated_transfer_request_if_unowned(request) + self._retire_transfer_terminal_if_complete(request) + + # When another owner remains, its completion path owns final teardown. + progress.termination_handoff_complete = True + + def _drain_pending_request_cancellations_after_shutdown(self) -> None: + """Finish exact pending cancellations after executor loops stop.""" + pending = self._get_pending_request_cancellations() + for request_key, progress in list(pending.items()): + if pending.get(request_key) is not progress: + continue + try: + completed = self._advance_request_cancellation(progress) + if completed and not progress.termination_handoff_complete: + self._complete_request_cancellation_termination( + progress, terminate_after_worker_shutdown=True) + except Exception as e: + logger.error( + "Failed to finish cancellation ownership for request " + f"{progress.cancel_id} after shutdown; retaining exact " + f"progress for retry: {e}") + continue + if (progress.termination_handoff_complete + and pending.get(request_key) is progress): + pending.pop(request_key) + + def _has_pending_ownership_transition(self, request: LlmRequest) -> bool: + """Whether response/resource cleanup must wait for a retry phase.""" + transfer_manager = getattr(self, "async_transfer_manager", None) + if (transfer_manager is not None + and transfer_manager.has_pending_admission(request) is True): + return True + request_key = id(request) + pending_cancellations = getattr(self, "_pending_request_cancellations", + {}) + if (pending_cancellations.get(request_key) is not None + and pending_cancellations[request_key].request is request): + return True + pending_native = getattr(self, "_pending_native_context_completions", + {}) + if (pending_native.get(request_key) is not None + and pending_native[request_key].request is request): + return True + pending_connector = getattr(self, "_pending_connector_completions", {}) + return (pending_connector.get(request_key) is not None + and pending_connector[request_key].request is request) + + def _has_pending_cancelled_transfer_owner(self, + request: LlmRequest) -> bool: + """Whether a cancelled request still has a provider owner.""" + if not request.is_finished_due_to_cancellation: + return False + return self._has_async_transfer_owner(request) + + def _finalize_deferred_transfer_termination( + self, + request: LlmRequest, + *, + terminate_after_worker_shutdown: bool = False) -> bool: + """Release a drained Python-native owner and terminate if appropriate.""" + deferred_terminations = getattr(self, "_deferred_transfer_terminations", + None) + request_key = id(request) + if (not deferred_terminations + or deferred_terminations.get(request_key) is not request): + return False + already_terminated = getattr( + self, "_deferred_transfer_terminations_already_terminated", set()) + request_was_already_terminated = ( + request_key in already_terminated + or self._was_transfer_request_terminated(request)) + + should_terminate = True + if self._has_async_transfer_owner(request): + # AsyncTransferManager normally transitions a successful context + # send to CONTEXT_COMPLETE. Latch the error before releasing the + # transceiver contribution so a sibling provider cannot resurrect + # it. + request.state = LlmRequestState.DISAGG_TRANS_ERROR + force_terminate = getattr(self, + "force_terminate_ctx_for_partial_reuse", + False) + if self._has_transceiver_transfer_owner(request): + transfer_manager = self.async_transfer_manager + all_transfers_complete = transfer_manager.end_transfer( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER) + should_terminate = all_transfers_complete or force_terminate + else: + # Transceiver ownership may have been retired by an earlier + # retry; a remaining provider owns final teardown. + should_terminate = force_terminate + + if should_terminate: + if not request_was_already_terminated: + if terminate_after_worker_shutdown: + self._terminate_request_after_worker_shutdown(request) + else: + self._terminate_request(request) + else: + self._mark_transfer_terminal_teardown_complete(request) + deferred_terminations.pop(request_key) + already_terminated.discard(request_key) + self._forget_terminated_transfer_request_if_unowned(request) + self._retire_transfer_terminal_if_complete(request) + return True + + def _terminate_request_after_worker_shutdown(self, + request: LlmRequest) -> None: + """Terminate directly when no executor iteration can run PP consensus.""" + termination_handler = getattr(self, "_disagg_pp_termination_handler", + None) + if termination_handler is not None and not request.is_dummy_request: + termination_handler.terminate_after_shutdown(request) + else: + self._terminate_request(request) + + def _drain_deferred_transfer_terminations( + self, *, terminate_after_worker_shutdown: bool = False) -> None: + """Drain transport, then retire exact deferred native owners.""" + deferred_terminations = getattr(self, "_deferred_transfer_terminations", + None) + if not deferred_terminations or self.kv_cache_transceiver is None: + return + + for request_key, request in list(deferred_terminations.items()): + try: + is_drained = self.kv_cache_transceiver.cancel_request(request) + except Exception as e: + logger.error( + "Failed to retry KV transfer cancellation for request " + f"{request.py_request_id}; retaining request resources: {e}" + ) + continue + if not is_drained: + continue + if deferred_terminations.get(request_key) is not request: + continue + try: + self._finalize_deferred_transfer_termination( + request, + terminate_after_worker_shutdown= + terminate_after_worker_shutdown) + except Exception as e: + logger.error( + "Failed to retire deferred KV transceiver owner for request " + f"{request.py_request_id}; retaining exact request for " + f"retry: {e}") + @nvtx_range("_handle_canceled_requests") def _handle_canceled_requests(self): - if len(self.canceled_req_ids) == 0: + pending = self._get_pending_request_cancellations() + if len(self.canceled_req_ids) == 0 and not pending: return # Create set from list of canceled request ids to speed up canceled test @@ -6595,23 +8802,54 @@ def _handle_canceled_requests(self): # Remove canceled requests from the waiting queue self.waiting_queue.remove_by_ids(canceled_req_ids_set) - still_pending_canceled_ids = [] for request in self.active_requests: req_id = request.py_request_id if not request.is_child else request.parent_request_id if req_id not in canceled_req_ids_set: continue + request_key = id(request) + existing = pending.get(request_key) + if existing is not None and existing.request is not request: + raise RuntimeError( + "request cancellation identity collision for request " + f"{request.py_request_id}") + if existing is None: + terminal = self._get_transfer_terminal(request) + if (terminal is not None and terminal.outcome + != _TransferTerminalOutcome.CANCELLED): + # A consumed provider status or an earlier request error + # already owns the logical result through sibling cleanup. + continue + terminal = self._claim_transfer_terminal( + request, _TransferTerminalOutcome.CANCELLED, + "user_cancellation") + if terminal.outcome != _TransferTerminalOutcome.CANCELLED: + continue + pending[request_key] = _RequestCancellationProgress( + request=request, cancel_id=req_id) - is_cancelled = self._try_cancel_request(request) - if is_cancelled: - # Mark requests as finished, then, we reuse all existing code - # to clean up the KV cache resources. - request.py_kv_transfer_timed_out = False - request.finish_by_reason(FinishReason.CANCELLED) - request.decoding_iter = request.py_decoding_iter - else: - still_pending_canceled_ids.append(req_id) + for request_key, progress in list(pending.items()): + if pending.get(request_key) is not progress: + continue + try: + completed = self._advance_request_cancellation(progress) + except Exception as e: + logger.error( + f"Failed to cancel request {progress.cancel_id}; retaining " + f"cancellation for retry: {e}") + continue + if completed and pending.get(request_key) is progress: + pending.pop(request_key) - # Clear list of requests marked for cancellation and add back those that failed to cancel. + # Keep one ID while any exact child/request cancellation still has an + # incomplete phase. IDs with no active request retain legacy behavior + # and are cleared after waiting-queue removal. + still_pending_canceled_ids = [] + seen_cancel_ids = set() + for progress in pending.values(): + if progress.cancel_id in seen_cancel_ids: + continue + seen_cancel_ids.add(progress.cancel_id) + still_pending_canceled_ids.append(progress.cancel_id) self.canceled_req_ids.clear() self.canceled_req_ids.extend(still_pending_canceled_ids) @@ -6734,6 +8972,7 @@ def _emit_first_token_responses(self, prev_scheduled_requests): @nvtx_range("_handle_responses") def _handle_responses(self, emit_first_iter: bool = True): new_responses = [] + terminal_response_request_keys = set() requests_to_terminate = [] # Requests terminated by _check_disagg_ctx_cache_transfer_status (DISAGG_CONTEXT_COMPLETE); # included in the return value for stats but not re-terminated here. @@ -6753,7 +8992,28 @@ def _handle_responses(self, emit_first_iter: bool = True): requests_to_terminate.append(request) continue - # Check if a generation request needs cleanup due to KV cache transfer timeout. + if self._has_pending_ownership_transition(request): + # A retry ledger is the exact owner while a one-way transport, + # owner-release, or termination phase is incomplete. Do not + # emit a competing terminal response or drop the active root. + new_active_requests.append(request) + continue + + terminal = self._get_transfer_terminal(request) + native_cancel_can_publish = ( + terminal is not None + and terminal.outcome == _TransferTerminalOutcome.CANCELLED + and self._requires_physical_transfer_drain()) + if (terminal is not None and self._has_async_transfer_owner(request) + and not native_cancel_can_publish): + # Legacy providers retain response creation until their final + # completion callback. For native cancellation, the logical + # result is already final, and the exact terminal/manager root + # can keep physical teardown gated after response publication. + new_active_requests.append(request) + continue + + # Check if generation request needs cleanup due to KV cache transfer timeout if request.py_kv_transfer_timed_out: if self._is_disagg_inflight_cancel_active(): if (request.is_disagg_generation_transmission_in_progress @@ -6773,9 +9033,9 @@ def _handle_responses(self, emit_first_iter: bool = True): # Defer it until the rank-uniform vote below. timed_out_requests.append(request) else: - # Legacy transceivers cannot cancel every in-flight - # transfer. Keep polling instead of dropping ownership of - # the request and its KV resources. + # Ownership-safe cancellation is retryable: keep the + # request visible to the next iteration until every + # receive target is physically drained. new_active_requests.append(request) continue @@ -6828,14 +9088,36 @@ def _handle_responses(self, emit_first_iter: bool = True): if request.return_perf_metrics: # Response creation may finalize and copy scalar ctx GPU totals. self.perf_manager.compute_batch_gpu_times([request]) - response = request.create_response(False, self.dist.rank) + terminal_owns_response = ( + terminal is not None and terminal.outcome + in (_TransferTerminalOutcome.SUCCEEDED, + _TransferTerminalOutcome.CANCELLED)) + if terminal_owns_response and terminal.response_created: + response = (None if terminal.response_published else + terminal.response) + else: + if terminal_owns_response: + self._create_transfer_terminal_response(terminal) + response = terminal.response + else: + response = request.create_response( + False, self.dist.rank) if response: request_done = request.is_finished - response.result.cached_tokens = request.cached_tokens - self._maybe_attach_ctx_usage(request, response) + if (not terminal_owns_response + or not terminal.response_enriched): + response.result.cached_tokens = request.cached_tokens + self._maybe_attach_ctx_usage(request, response) + if terminal_owns_response: + terminal.response_enriched = True response.result.per_pos_drafted = request.py_per_pos_drafted response.result.per_pos_accepted = request.py_per_pos_accepted new_responses.append((req_id, response)) + if terminal_owns_response: + self._buffer_transfer_terminal_response(terminal) + terminal_response_request_keys.add(id(request)) + elif terminal_owns_response and terminal.response_published: + request_done = request.is_finished if request_done: # PP=1-only early termination; _end_transfer_and_maybe_terminate @@ -6848,6 +9130,12 @@ def _handle_responses(self, emit_first_iter: bool = True): requests_finished_by_transfer.append(request) elif force_terminate_for_partial_reuse: requests_to_terminate.append(request) + elif self._has_pending_cancelled_transfer_owner(request): + # A provider still holds the shared transfer lease. Its + # completion path performs the one remaining termination. + # The request is already logically finished, so report it + # once for iteration completion and request stats now. + requests_finished_by_transfer.append(request) elif not request.is_disagg_context_transmission_state: requests_to_terminate.append(request) else: @@ -6856,7 +9144,14 @@ def _handle_responses(self, emit_first_iter: bool = True): self.active_requests.clear() self.active_requests.extend(new_active_requests) # Request should be terminated after enqueueing response to ensure we can enqueue response successfully. + publishing_terminals = [ + terminal for request_key in terminal_response_request_keys + if (terminal := self._get_pending_transfer_terminals().get( + request_key)) is not None + ] + self._begin_transfer_response_publication(publishing_terminals) self._enqueue_responses(new_responses) + self._commit_transfer_response_publication(publishing_terminals) for request in requests_to_terminate: self._terminate_request(request) if (self.kv_cache_transceiver is not None and self.enable_attention_dp @@ -7032,6 +9327,39 @@ def __init__(self, dist, terminator_func: Callable[[LlmRequest], None]): def terminate(self, request: LlmRequest): self._pending_termination[request.py_request_id] = request + def terminate_after_shutdown(self, request: LlmRequest): + """Free a request after the executor loop can no longer run consensus.""" + self._terminator_func(request) + if self._pending_termination.get(request.py_request_id) is request: + self._pending_termination.pop(request.py_request_id) + + def terminate_all_after_shutdown(self) -> None: + """Finish locally pending requests after all executor loops stop.""" + errors = [] + # Do not wait for the ring's final control send here. Its matching + # next-iteration receive may never be posted after executor loops stop. + # The handle carries request IDs only and owns no request/KV resource; + # retaining it on this handler is safer than blocking teardown. + + for request_id, request in list(self._pending_termination.items()): + if self._pending_termination.get(request_id) is not request: + continue + try: + self._terminator_func(request) + except Exception as error: + errors.append(f"request {request_id}: {error}") + continue + if self._pending_termination.get(request_id) is request: + self._pending_termination.pop(request_id) + + if errors: + raise RuntimeError( + "PP request termination did not drain after shutdown: " + + "; ".join(errors)) + + def has_pending_terminations(self) -> bool: + return bool(self._pending_termination) + @nvtx_range("_disagg_pp_termination_handler_sync") def terminate_pending_requests(self): """ diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f82207881c8f..316a4317c742 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -333,6 +333,38 @@ def log_memory_usage(stage: str): ) +def _teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, resources) -> None: + """Preserve manager-owned memory through transport and owner teardown.""" + transceiver = getattr(py_executor, "kv_cache_transceiver", None) + if transceiver is not None: + shutdown_result = transceiver.shutdown() + requires_physical_drain = (getattr( + transceiver, "requires_physical_drain_before_request_release", + False) is True) + if requires_physical_drain and shutdown_result is not True: + raise RuntimeError( + "KV cache transceiver did not prove physical drain; refusing " + "to tear down KV cache managers") + if shutdown_result is False: + raise RuntimeError( + "KV cache transceiver still owns active transfer targets; " + "refusing to tear down KV cache managers") + + transfer_manager = getattr(py_executor, "async_transfer_manager", None) + has_inflight_requests = getattr(transfer_manager, + "has_any_inflight_requests", None) + if (has_inflight_requests is not None and has_inflight_requests() is True): + # This helper has no request-level owner-retirement phase. Even an + # explicit transceiver drain cannot retire a connector sibling or clear + # the AsyncTransferManager ledger, and a legacy ``None`` result proves + # nothing about either. Capacity warmup is expected to have no owners. + raise RuntimeError( + "Asynchronous transfer ownership is still active; refusing to " + "tear down KV cache managers") + kv_cache_creator.teardown_managers(resources) + + def create_py_executor( llm_args: TorchLlmArgs, checkpoint_dir: Optional[str] = None, @@ -1026,12 +1058,8 @@ def drafting_loop_wrapper(model): # that NIXL-registered (pinned) GPU memory is deregistered first; # otherwise the old KV cache memory stays pinned and the subsequent # KV cache allocation will OOM. - try: - if hasattr(py_executor, 'kv_cache_transceiver' - ) and py_executor.kv_cache_transceiver is not None: - py_executor.kv_cache_transceiver.shutdown() - finally: - kv_cache_creator.teardown_managers(resources) + _teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, resources) # Release Phase-1 CUDA graph pools before final KV allocation to avoid overshoot. for eng in [model_engine, draft_model_engine]: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 157e2d7d720a..61fd9a7e4223 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -16,6 +16,7 @@ import enum import math import os +import threading from abc import ABC, abstractmethod from collections import OrderedDict, defaultdict, deque from dataclasses import dataclass @@ -1069,7 +1070,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): def store_blocks_for_reuse(self, request: LlmRequest, - pin_blocks: bool = False): + pin_blocks: bool = False) -> List[int]: return self.impl.store_blocks_for_reuse(request.py_request_id, request, pin_blocks) @@ -1346,8 +1347,8 @@ def commit_and_get_block_hashes( self.impl.commit_and_get_block_hashes_for_request( request, window_size)) - def unpin_blocks_by_id(self, kv_cache_block_id: int): - self.impl.unpin_blocks_by_id(kv_cache_block_id) + def unpin_blocks_by_id(self, block_ids: List[int]) -> None: + self.impl.unpin_blocks_by_id(block_ids) def get_last_block_id(self, request_id: int) -> int: return self.impl.get_last_block_id(request_id) @@ -2564,11 +2565,43 @@ def free_resources(self, request: "LlmRequest") -> None: self.on_request_finish(request) +@dataclass +class _ResourceReleaseProgress: + request: LlmRequest + release_plan: Tuple[BaseResourceManager, ...] + next_manager: int = 0 + in_doubt_manager: Optional[int] = None + manager_call_active: bool = False + failure: Optional[str] = None + + +@dataclass +class _TransferResourceReleaseProgress: + request: LlmRequest + completed_managers: Set[ResourceManagerType] + in_doubt_manager: Optional[ResourceManagerType] = None + manager_call_active: bool = False + failure: Optional[str] = None + + class ResourceManager: def __init__(self, resource_managers: dict[ResourceManagerType, BaseResourceManager]): self.resource_managers = OrderedDict(resource_managers) + # Keep the exact request and a stable teardown plan rooted while a + # multi-manager release is incomplete. A retry never invokes managers + # that already returned. A manager that raised remains in doubt and + # blocks teardown because its internal mutation boundary is unknown. + self._resource_release_progress: Dict[int, + _ResourceReleaseProgress] = {} + # Async transfer admission releases seq/spec resources early. Record + # each successful manager here so final request teardown never invokes + # it again, while an exception remains fail-closed and exact-request + # rooted. + self._transfer_resource_release_progress: Dict[ + int, _TransferResourceReleaseProgress] = {} + self._resource_release_lock = threading.RLock() def __call__(self, type: ResourceManagerType): return self.resource_managers[type] @@ -2604,10 +2637,192 @@ def update_resources( resource_manager.update_resources(scheduled_batch) def free_resources(self, request: LlmRequest): - for resource_type, resource_manager in reversed( - self.resource_managers.items()): - if hasattr(resource_manager, "free_resources"): - resource_manager.free_resources(request) + request_key = id(request) + while True: + with self._resource_release_lock: + transfer_progress = self._transfer_resource_release_progress.get( + request_key) + if (transfer_progress is not None + and transfer_progress.request is not request): + raise RuntimeError( + "transfer resource release identity collision for request " + f"{request.py_request_id}") + if (transfer_progress is not None + and transfer_progress.in_doubt_manager is not None): + if transfer_progress.manager_call_active: + raise RuntimeError( + "transfer resource release is already in progress for " + f"request {request.py_request_id} at manager " + f"{transfer_progress.in_doubt_manager.value}") + raise RuntimeError( + "transfer resource release outcome is in doubt for request " + f"{request.py_request_id} at manager " + f"{transfer_progress.in_doubt_manager.value}: " + f"{transfer_progress.failure}") + + progress = self._resource_release_progress.get(request_key) + if progress is not None and progress.request is not request: + raise RuntimeError( + "resource release identity collision for request " + f"{request.py_request_id}") + + if progress is None: + release_plan = tuple( + resource_manager + for resource_mgr_type, resource_manager in reversed( + self.resource_managers.items()) + if (transfer_progress is None or resource_mgr_type + not in transfer_progress.completed_managers) + if hasattr(resource_manager, "free_resources")) + progress = _ResourceReleaseProgress(request, release_plan) + self._resource_release_progress[request_key] = progress + + if progress.in_doubt_manager is not None: + if progress.manager_call_active: + raise RuntimeError( + "resource release is already in progress for request " + f"{request.py_request_id} at manager index " + f"{progress.in_doubt_manager}") + raise RuntimeError( + "resource release outcome is in doubt for request " + f"{request.py_request_id} at manager index " + f"{progress.in_doubt_manager}: {progress.failure}") + + if progress.next_manager == len(progress.release_plan): + if self._resource_release_progress.get( + request_key) is progress: + self._resource_release_progress.pop(request_key) + if (transfer_progress is not None + and self._transfer_resource_release_progress.get( + request_key) is transfer_progress): + self._transfer_resource_release_progress.pop( + request_key) + return + + # Reserve one exact manager attempt under synchronization, then + # release the global metadata lock before arbitrary allocator or + # callback code runs. A concurrent call for this same request + # fails closed, while unrelated requests can continue. + manager_index = progress.next_manager + manager = progress.release_plan[manager_index] + progress.in_doubt_manager = manager_index + progress.manager_call_active = True + + try: + manager.free_resources(request) + except Exception as error: + with self._resource_release_lock: + if (self._resource_release_progress.get(request_key) + is progress): + progress.manager_call_active = False + progress.failure = f"{type(error).__name__}: {error}" + raise + + with self._resource_release_lock: + if self._resource_release_progress.get( + request_key) is not progress: + raise RuntimeError( + "resource release progress disappeared for request " + f"{request.py_request_id}") + if (progress.in_doubt_manager != manager_index + or not progress.manager_call_active): + raise RuntimeError( + "resource release progress changed during callback for " + f"request {request.py_request_id}") + progress.next_manager += 1 + progress.in_doubt_manager = None + progress.manager_call_active = False + progress.failure = None + + def release_resource_for_transfer( + self, request: LlmRequest, + resource_mgr_type: ResourceManagerType) -> None: + """Release one manager early and remember it through final teardown.""" + request_key = id(request) + with self._resource_release_lock: + final_progress = self._resource_release_progress.get(request_key) + if final_progress is not None: + if final_progress.request is not request: + raise RuntimeError( + "resource release identity collision for request " + f"{request.py_request_id}") + raise RuntimeError( + "final resource release already started for request " + f"{request.py_request_id}; refusing an early transfer " + "release") + + progress = self._transfer_resource_release_progress.get(request_key) + if progress is not None and progress.request is not request: + raise RuntimeError( + "transfer resource release identity collision for request " + f"{request.py_request_id}") + if progress is None: + progress = _TransferResourceReleaseProgress(request, set()) + self._transfer_resource_release_progress[request_key] = progress + + if progress.in_doubt_manager is not None: + if progress.manager_call_active: + raise RuntimeError( + "transfer resource release is already in progress for " + f"request {request.py_request_id} at manager " + f"{progress.in_doubt_manager.value}") + raise RuntimeError( + "transfer resource release outcome is in doubt for request " + f"{request.py_request_id} at manager " + f"{progress.in_doubt_manager.value}: {progress.failure}") + if resource_mgr_type in progress.completed_managers: + return + + manager = self.resource_managers.get(resource_mgr_type) + if manager is None or not hasattr(manager, "free_resources"): + progress.completed_managers.add(resource_mgr_type) + return + + progress.in_doubt_manager = resource_mgr_type + progress.manager_call_active = True + + try: + manager.free_resources(request) + except Exception as error: + with self._resource_release_lock: + if self._transfer_resource_release_progress.get( + request_key) is progress: + progress.manager_call_active = False + progress.failure = f"{type(error).__name__}: {error}" + raise + + with self._resource_release_lock: + if self._transfer_resource_release_progress.get( + request_key) is not progress: + raise RuntimeError( + "transfer resource release progress disappeared for request " + f"{request.py_request_id}") + if (progress.in_doubt_manager != resource_mgr_type + or not progress.manager_call_active): + raise RuntimeError( + "transfer resource release progress changed during callback " + f"for request {request.py_request_id}") + progress.completed_managers.add(resource_mgr_type) + progress.in_doubt_manager = None + progress.manager_call_active = False + progress.failure = None + + def has_in_doubt_resource_releases(self) -> bool: + """Whether a resource manager raised after release was attempted.""" + with self._resource_release_lock: + return (any( + progress.in_doubt_manager is not None + and not progress.manager_call_active + for progress in self._resource_release_progress.values()) + or any(progress.in_doubt_manager is not None + and not progress.manager_call_active for progress in + self._transfer_resource_release_progress.values())) + + def has_pending_resource_releases(self) -> bool: + """Whether an exact request is rooted in either release ledger.""" + with self._resource_release_lock: + return bool(self._resource_release_progress + or self._transfer_resource_release_progress) def reorder_pipeline(self, resource_manager_list: list[ResourceManagerType]): diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index bf8c6d131d50..63160dc9127a 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -39,6 +39,8 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_compression_manager.py - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py + - unittest/_torch/executor/test_py_executor_teardown.py + - unittest/_torch/executor/test_py_executor_transfer_termination.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py - unittest/_torch/modules/dwdp/test_dwdp_mapping.py @@ -73,6 +75,9 @@ l0_a10: - unittest/disaggregated/test_extractor.py - unittest/disaggregated/test_peer.py - unittest/disaggregated/test_bounce.py + - unittest/disaggregated/test_receive_lifecycle.py + - unittest/disaggregated/test_receive_ownership_wiring.py + - unittest/disaggregated/test_transceiver_bounded_polling.py - unittest/disaggregated/region/test_block.py - unittest/disaggregated/test_mamba_transfer.py - unittest/tools diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 2874ce258d48..633c015cff35 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -76,6 +76,9 @@ l0_h100: - unittest/disaggregated/test_extractor.py - unittest/disaggregated/test_peer.py - unittest/disaggregated/test_bounce.py + - unittest/disaggregated/test_receive_lifecycle.py + - unittest/disaggregated/test_receive_ownership_wiring.py + - unittest/disaggregated/test_transceiver_bounded_polling.py - unittest/disaggregated/region/test_block.py - unittest/disaggregated/region/test_aux.py - unittest/disaggregated/region/test_page.py diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/executor/test_async_transfer_manager.py index 1f2f9013d903..bd73e88f2e8d 100644 --- a/tests/unittest/_torch/executor/test_async_transfer_manager.py +++ b/tests/unittest/_torch/executor/test_async_transfer_manager.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,10 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import threading from unittest.mock import MagicMock +import pytest + from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager -from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm.bindings import LlmRequestState @@ -28,33 +31,56 @@ def create_mock_request(request_id: int): return request +class _CompletionStateFailingRequest: + """Request whose first completion-state update fails before or after mutation.""" + + def __init__(self, request_id: int, *, mutate_before_raise: bool): + self.py_request_id = request_id + self._state = LlmRequestState.GENERATION_IN_PROGRESS + self._mutate_before_raise = mutate_before_raise + self._completion_failure_pending = True + self.completion_state_set_calls = 0 + + @property + def state(self): + return self._state + + @state.setter + def state(self, value): + if value == LlmRequestState.DISAGG_CONTEXT_COMPLETE: + self.completion_state_set_calls += 1 + if self._completion_failure_pending: + self._completion_failure_pending = False + if self._mutate_before_raise: + self._state = value + raise RuntimeError("completion state update failed") + self._state = value + + def create_mock_resource_manager( kv_cache_manager=None, seq_slot_manager=None, spec_resource_manager=None, ): """Create a mock ResourceManager with the specified resource managers.""" - resource_manager = MagicMock() - resource_manager.resource_managers = {} + resource_managers = {} if kv_cache_manager is not None: - resource_manager.resource_managers[ResourceManagerType.KV_CACHE_MANAGER] = kv_cache_manager + resource_managers[ResourceManagerType.KV_CACHE_MANAGER] = kv_cache_manager if seq_slot_manager is not None: - resource_manager.resource_managers[ResourceManagerType.SEQ_SLOT_MANAGER] = seq_slot_manager + resource_managers[ResourceManagerType.SEQ_SLOT_MANAGER] = seq_slot_manager if spec_resource_manager is not None: - resource_manager.resource_managers[ResourceManagerType.SPEC_RESOURCE_MANAGER] = ( - spec_resource_manager - ) + resource_managers[ResourceManagerType.SPEC_RESOURCE_MANAGER] = spec_resource_manager - return resource_manager + return ResourceManager(resource_managers) def test_start_transfer_single_request(): """Test starting a single transfer.""" kv_cache_manager = MagicMock() - kv_cache_manager.store_blocks_for_reuse.return_value = 100 + kv_cache_manager.store_blocks_for_reuse.return_value = [100] seq_slot_manager = MagicMock() resource_manager = create_mock_resource_manager( kv_cache_manager=kv_cache_manager, seq_slot_manager=seq_slot_manager @@ -69,7 +95,7 @@ def test_start_transfer_single_request(): transfer_metadata = manager._request_transfer_metadata[42] - assert transfer_metadata.block_id == 100 + assert transfer_metadata.pinned_block_ids == [100] assert transfer_metadata.counter == 1 # Check state was updated @@ -82,13 +108,13 @@ def test_start_transfer_single_request(): seq_slot_manager.free_resources.assert_called_once_with(request) manager.end_transfer(request) - kv_cache_manager.unpin_blocks_by_id.assert_called_once() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) def test_start_transfer_multiple_transfers_same_request(): """Test starting multiple transfers for the same request.""" kv_cache_manager = MagicMock() - kv_cache_manager.store_blocks_for_reuse.return_value = 100 + kv_cache_manager.store_blocks_for_reuse.return_value = [100] resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) manager = AsyncTransferManager(resource_manager) @@ -109,13 +135,328 @@ def test_start_transfer_multiple_transfers_same_request(): kv_cache_manager.unpin_blocks_by_id.assert_not_called() manager.end_transfer(request) - kv_cache_manager.unpin_blocks_by_id.assert_called_once() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_final_transfer_owner_query_distinguishes_mixed_anonymous_owners(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + + manager.start_transfer(request) + manager.start_transfer(request) + + assert not manager.is_final_transfer_owner(request) + assert not manager.end_transfer(request) + assert manager.is_final_transfer_owner(request) + assert manager.end_transfer(request) + + +def test_named_transfer_owner_can_retire_independently(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + + manager.start_transfer(request, owner="python_native") + manager.start_transfer(request) + + assert manager.requests_with_owner("python_native") == {42: request} + assert manager.has_transfer_owner(request, "python_native") + assert not manager.end_transfer(request, owner="python_native") + assert manager.requests_with_owner("python_native") == {} + assert 42 in manager.requests_in_transfer() + kv_cache_manager.unpin_blocks_by_id.assert_not_called() + + assert manager.end_transfer(request) + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_start_transfer_rejects_different_request_with_same_id(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + replacement = create_mock_request(42) + + manager.start_transfer(request, owner="python_native") + + with pytest.raises(RuntimeError, match="different transfer owner"): + manager.start_transfer(replacement, owner="connector") + + metadata = manager._request_transfer_metadata[42] + assert manager.requests_in_transfer() == {42: request} + assert metadata.counter == 1 + assert metadata.owners == {"python_native"} + kv_cache_manager.store_blocks_for_reuse.assert_called_once_with(request, True) + + +def test_end_transfer_rejects_different_request_with_same_id(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + replacement = create_mock_request(42) + + manager.start_transfer(request, owner="python_native") + + with pytest.raises(RuntimeError, match="different active transfer owner"): + manager.end_transfer(replacement, owner="python_native") + + metadata = manager._request_transfer_metadata[42] + assert manager.requests_in_transfer() == {42: request} + assert metadata.counter == 1 + assert metadata.owners == {"python_native"} + kv_cache_manager.unpin_blocks_by_id.assert_not_called() + + +def test_final_transfer_unpin_failure_is_retained_without_replay(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + unpinned_blocks = [] + + def mutate_then_raise(block_ids): + unpinned_blocks.extend(block_ids) + raise RuntimeError("unpin failed after mutation") + + kv_cache_manager.unpin_blocks_by_id.side_effect = mutate_then_raise + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + + manager.start_transfer(request, owner="python_native") + + with pytest.raises(RuntimeError, match="unpin failed after mutation"): + manager.end_transfer(request, owner="python_native") + + metadata = manager._request_transfer_metadata[42] + assert manager.requests_in_transfer() == {42: request} + assert metadata.counter == 1 + assert metadata.owners == {"python_native"} + assert request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + assert unpinned_blocks == [100] + + with pytest.raises(RuntimeError, match="unpin outcome is in doubt"): + manager.end_transfer(request, owner="python_native") + + progress = manager._transfer_completion_progress[42] + assert progress.unpin_in_doubt + assert not progress.unpin_complete + assert manager.requests_in_transfer() == {42: request} + assert manager.has_any_inflight_requests() + assert kv_cache_manager.unpin_blocks_by_id.call_count == 1 + assert unpinned_blocks == [100] + + +@pytest.mark.parametrize("mutate_before_raise", [False, True]) +def test_final_transfer_state_failure_resumes_without_replaying_unpin( + mutate_before_raise, +): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = _CompletionStateFailingRequest(42, mutate_before_raise=mutate_before_raise) + + manager.start_transfer(request, owner="python_native") + + with pytest.raises(RuntimeError, match="completion state update failed"): + manager.end_transfer(request, owner="python_native") + + metadata = manager._request_transfer_metadata[42] + progress = manager._transfer_completion_progress[42] + assert manager.requests_in_transfer() == {42: request} + assert metadata.counter == 1 + assert metadata.owners == {"python_native"} + assert progress.request is request + assert progress.metadata is metadata + assert progress.unpin_complete + assert not progress.state_updated + assert not progress.owner_released + assert manager.has_any_inflight_requests() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + with pytest.raises(RuntimeError, match="completion is incomplete"): + manager.start_transfer(request, owner="connector") + + assert manager.end_transfer(request, owner="python_native") + + assert manager.requests_in_transfer() == {} + assert manager._request_transfer_metadata == {} + assert manager._transfer_completion_progress == {} + assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + expected_state_set_calls = 1 if mutate_before_raise else 2 + assert request.completion_state_set_calls == expected_state_set_calls + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_final_transfer_partial_map_commit_resumes_exact_completion(): + class _FailAfterPop(dict): + def __init__(self, values): + super().__init__(values) + self._failure_pending = True + + def pop(self, key): + value = super().pop(key) + if self._failure_pending: + self._failure_pending = False + raise RuntimeError("post-owner cleanup failed") + return value + + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + + manager.start_transfer(request, owner="python_native") + metadata = manager._request_transfer_metadata[42] + metadata.end_transfer = MagicMock(wraps=metadata.end_transfer) + manager._request_transfer_metadata = _FailAfterPop(manager._request_transfer_metadata) + + with pytest.raises(RuntimeError, match="post-owner cleanup failed"): + manager.end_transfer(request, owner="python_native") + + progress = manager._transfer_completion_progress[42] + # Request-map retirement committed before metadata-map retirement mutated + # and raised. The durable completion record is now the only exact owner. + assert manager._requests_in_transfer == {} + assert manager._request_transfer_metadata == {} + assert metadata.counter == 0 + assert metadata.owners == set() + assert progress.unpin_complete + assert progress.state_updated + assert progress.owner_released + assert manager.has_transfer_owner(request, "python_native") + assert manager.requests_with_owner("python_native") == {42: request} + assert manager.has_any_transfer_owner(request) + assert manager.has_any_inflight_requests() + assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + + assert manager.end_transfer(request, owner="python_native") + + assert manager.requests_in_transfer() == {} + assert manager._request_transfer_metadata == {} + assert manager._transfer_completion_progress == {} + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + metadata.end_transfer.assert_called_once_with("python_native") + + +def test_resource_release_failure_is_retained_as_in_doubt(): + first_manager = MagicMock() + failing_manager = MagicMock() + failing_manager.free_resources.side_effect = RuntimeError( + "release may have partially completed" + ) + manager = ResourceManager( + { + ResourceManagerType.SEQ_SLOT_MANAGER: failing_manager, + ResourceManagerType.SPEC_RESOURCE_MANAGER: first_manager, + } + ) + request = create_mock_request(42) + + with pytest.raises(RuntimeError, match="partially completed"): + manager.free_resources(request) + + progress = manager._resource_release_progress[id(request)] + assert progress.request is request + assert progress.next_manager == 1 + assert progress.in_doubt_manager == 1 + assert manager.has_in_doubt_resource_releases() + + with pytest.raises(RuntimeError, match="outcome is in doubt"): + manager.free_resources(request) + + first_manager.free_resources.assert_called_once_with(request) + failing_manager.free_resources.assert_called_once_with(request) + + +def test_successful_resource_release_clears_progress(): + resource = MagicMock() + manager = ResourceManager({ResourceManagerType.SEQ_SLOT_MANAGER: resource}) + request = create_mock_request(42) + + manager.free_resources(request) + + resource.free_resources.assert_called_once_with(request) + assert manager._resource_release_progress == {} + assert not manager.has_in_doubt_resource_releases() + + +def test_resource_release_callback_does_not_hold_global_lock(): + callback_started = threading.Event() + callback_continue = threading.Event() + errors = [] + blocking_manager = MagicMock() + + def block_first_request(request): + if request.py_request_id != 1: + return + callback_started.set() + if not callback_continue.wait(timeout=5): + raise RuntimeError("test callback timed out") + + blocking_manager.free_resources.side_effect = block_first_request + manager = ResourceManager( + { + ResourceManagerType.SEQ_SLOT_MANAGER: blocking_manager, + } + ) + first_request = create_mock_request(1) + second_request = create_mock_request(2) + + def release_first_request(): + try: + manager.free_resources(first_request) + except Exception as error: + errors.append(error) + + thread = threading.Thread(target=release_first_request) + thread.start() + assert callback_started.wait(timeout=5) + + manager.free_resources(second_request) + callback_continue.set() + thread.join(timeout=5) + + assert not thread.is_alive() + assert errors == [] + assert blocking_manager.free_resources.call_count == 2 + assert not manager.has_pending_resource_releases() + + +def test_resource_release_callback_reentry_fails_without_deadlock(): + request = create_mock_request(42) + resource = MagicMock() + manager = ResourceManager({ResourceManagerType.SEQ_SLOT_MANAGER: resource}) + reentry_errors = [] + + def reenter(_request): + try: + manager.free_resources(request) + except RuntimeError as error: + reentry_errors.append(str(error)) + + resource.free_resources.side_effect = reenter + + manager.free_resources(request) + + assert len(reentry_errors) == 1 + assert "already in progress" in reentry_errors[0] + resource.free_resources.assert_called_once_with(request) + assert not manager.has_pending_resource_releases() def test_transfer_without_storing_blocks(): """Test starting a transfer with should_store_blocks=False.""" kv_cache_manager = MagicMock() - kv_cache_manager.store_blocks_for_reuse.return_value = 0 + kv_cache_manager.store_blocks_for_reuse.return_value = [] spec_resource_manager = MagicMock() resource_manager = create_mock_resource_manager( kv_cache_manager=kv_cache_manager, spec_resource_manager=spec_resource_manager @@ -128,7 +469,7 @@ def test_transfer_without_storing_blocks(): # Check request is tracked assert 42 in manager._requests_in_transfer transfer_metadata = manager._request_transfer_metadata[42] - assert transfer_metadata.block_id is None # No block stored + assert transfer_metadata.pinned_block_ids is None # No blocks stored assert transfer_metadata.counter == 1 # Check KV cache manager was NOT called @@ -143,7 +484,7 @@ def test_transfer_without_storing_blocks(): def test_end_transfer_preserves_error_state(): """Test that end_transfer does not overwrite error state.""" kv_cache_manager = MagicMock() - kv_cache_manager.store_blocks_for_reuse.return_value = 100 + kv_cache_manager.store_blocks_for_reuse.return_value = [100] resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) manager = AsyncTransferManager(resource_manager) @@ -162,7 +503,7 @@ def test_end_transfer_preserves_error_state(): def test_requests_in_transfer(): """Test that requests_in_transfer returns correct mapping.""" kv_cache_manager = MagicMock() - kv_cache_manager.store_blocks_for_reuse.return_value = 100 + kv_cache_manager.store_blocks_for_reuse.return_value = [100] resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) manager = AsyncTransferManager(resource_manager) @@ -180,3 +521,184 @@ def test_requests_in_transfer(): assert in_transfer[1] is request1 assert in_transfer[2] is request2 assert in_transfer[3] is request3 + + +def test_first_owner_admission_is_published_before_external_mutations(): + request = create_mock_request(42) + kv_cache_manager = MagicMock() + seq_slot_manager = MagicMock() + spec_resource_manager = MagicMock() + resource_manager = create_mock_resource_manager( + kv_cache_manager=kv_cache_manager, + seq_slot_manager=seq_slot_manager, + spec_resource_manager=spec_resource_manager, + ) + manager = AsyncTransferManager(resource_manager) + + def assert_seq_release_is_rooted(released_request): + assert released_request is request + assert manager.requests_in_admission() == {42: request} + progress = manager._transfer_admission_progress[42] + assert progress.manager_in_doubt == ResourceManagerType.SEQ_SLOT_MANAGER + assert progress.completed_managers == set() + + def assert_spec_release_is_rooted(released_request): + assert released_request is request + assert manager.requests_in_admission() == {42: request} + progress = manager._transfer_admission_progress[42] + assert progress.manager_in_doubt == ResourceManagerType.SPEC_RESOURCE_MANAGER + assert progress.completed_managers == {ResourceManagerType.SEQ_SLOT_MANAGER} + + def assert_pin_is_rooted(pinned_request, pin_blocks): + assert pinned_request is request + assert pin_blocks is True + assert manager.requests_in_admission() == {42: request} + progress = manager._transfer_admission_progress[42] + assert progress.manager_in_doubt is None + assert progress.completed_managers == { + ResourceManagerType.SEQ_SLOT_MANAGER, + ResourceManagerType.SPEC_RESOURCE_MANAGER, + } + assert progress.state_updated + assert progress.pin_started + assert not progress.pin_complete + return [100, 101] + + seq_slot_manager.free_resources.side_effect = assert_seq_release_is_rooted + spec_resource_manager.free_resources.side_effect = assert_spec_release_is_rooted + kv_cache_manager.store_blocks_for_reuse.side_effect = assert_pin_is_rooted + + manager.start_transfer(request, owner="python_native") + + assert manager.requests_in_admission() == {} + assert manager.requests_in_transfer() == {42: request} + metadata = manager._request_transfer_metadata[42] + assert metadata.pinned_block_ids == [100, 101] + assert metadata.owners == {"python_native"} + assert manager.end_transfer(request, owner="python_native") + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100, 101]) + + +def test_partial_transfer_release_failure_is_fail_closed_and_not_replayed(): + request = create_mock_request(42) + kv_cache_manager = MagicMock() + seq_slot_manager = MagicMock() + spec_resource_manager = MagicMock() + spec_resource_manager.free_resources.side_effect = RuntimeError( + "spec release may have completed" + ) + resource_manager = create_mock_resource_manager( + kv_cache_manager=kv_cache_manager, + seq_slot_manager=seq_slot_manager, + spec_resource_manager=spec_resource_manager, + ) + manager = AsyncTransferManager(resource_manager) + + with pytest.raises(RuntimeError, match="may have completed"): + manager.start_transfer(request, owner="python_native") + + progress = manager._transfer_admission_progress[42] + assert progress.request is request + assert progress.completed_managers == {ResourceManagerType.SEQ_SLOT_MANAGER} + assert progress.manager_in_doubt == ResourceManagerType.SPEC_RESOURCE_MANAGER + assert manager.requests_in_admission() == {42: request} + assert manager.has_in_doubt_admissions() + assert manager.has_any_inflight_requests() + assert resource_manager.has_in_doubt_resource_releases() + + with pytest.raises(RuntimeError, match="refusing to replay"): + manager.start_transfer(request, owner="python_native") + with pytest.raises(RuntimeError, match="admission is incomplete"): + manager.end_transfer(request, owner="python_native") + with pytest.raises(RuntimeError, match="different transfer admission"): + manager.start_transfer(create_mock_request(42), owner="python_native") + with pytest.raises(RuntimeError, match="outcome is in doubt"): + resource_manager.free_resources(request) + + seq_slot_manager.free_resources.assert_called_once_with(request) + spec_resource_manager.free_resources.assert_called_once_with(request) + kv_cache_manager.store_blocks_for_reuse.assert_not_called() + kv_cache_manager.free_resources.assert_not_called() + + +def test_pin_failure_retains_exact_admission_and_refuses_replay(): + request = create_mock_request(42) + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.side_effect = RuntimeError("pin outcome unknown") + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + + with pytest.raises(RuntimeError, match="pin outcome unknown"): + manager.start_transfer(request, owner="python_native") + + progress = manager._transfer_admission_progress[42] + assert progress.request is request + assert progress.pin_started + assert not progress.pin_complete + assert progress.pin_in_doubt + assert progress.pinned_block_ids is None + assert manager.requests_in_admission() == {42: request} + assert manager.has_in_doubt_admissions() + assert manager.has_any_inflight_requests() + + with pytest.raises(RuntimeError, match="refusing to replay"): + manager.start_transfer(request, owner="python_native") + with pytest.raises(RuntimeError, match="admission is incomplete"): + manager.end_transfer(request, owner="python_native") + + kv_cache_manager.store_blocks_for_reuse.assert_called_once_with(request, True) + kv_cache_manager.unpin_blocks_by_id.assert_not_called() + + +def test_empty_pinned_block_list_is_a_valid_completed_pin(): + request = create_mock_request(42) + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + + manager.start_transfer(request) + + progress = manager._request_transfer_metadata[42] + assert progress.pinned_block_ids == [] + assert manager.end_transfer(request) + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([]) + + +def test_early_transfer_release_is_not_repeated_by_final_release(): + request = create_mock_request(42) + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + seq_slot_manager = MagicMock() + spec_resource_manager = MagicMock() + resource_manager = create_mock_resource_manager( + kv_cache_manager=kv_cache_manager, + seq_slot_manager=seq_slot_manager, + spec_resource_manager=spec_resource_manager, + ) + manager = AsyncTransferManager(resource_manager) + + manager.start_transfer(request) + assert resource_manager.has_pending_resource_releases() + assert manager.end_transfer(request) + resource_manager.free_resources(request) + + seq_slot_manager.free_resources.assert_called_once_with(request) + spec_resource_manager.free_resources.assert_called_once_with(request) + kv_cache_manager.free_resources.assert_called_once_with(request) + assert resource_manager._transfer_resource_release_progress == {} + assert resource_manager._resource_release_progress == {} + assert not resource_manager.has_pending_resource_releases() + + +def test_transfer_admission_closes_before_shutdown(): + request = create_mock_request(42) + kv_cache_manager = MagicMock() + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + + manager.begin_shutdown() + + with pytest.raises(RuntimeError, match="admission is closed"): + manager.start_transfer(request) + kv_cache_manager.store_blocks_for_reuse.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py index 9af481b2f686..c39345f04701 100644 --- a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py +++ b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -80,6 +80,9 @@ def __init__(self, kv_cache_manager, async_transfer_manager, kv_cache_transceive def _check_disagg_ctx_cache_transfer_status(self, _): return None + def _requires_physical_transfer_drain(self): + return PyExecutor._requires_physical_transfer_drain(self) + class TestSendKvAsyncReleasesIndexSlot: """Test that `_send_kv_async` releases the IndexMapper slot before the diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 3d9a3f83dcc3..0cb6249a29f5 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -288,6 +288,7 @@ def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): py_kv_transfer_timed_out=True, py_decoding_iter=3, finish_by_reason=Mock(), + is_generation_only_request=Mock(return_value=False), ) executor = object.__new__(PyExecutor) executor.active_requests = [request] @@ -301,6 +302,11 @@ def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): executor.async_transfer_manager.requests_in_transfer.return_value = { request.py_request_id: request } + executor.async_transfer_manager.has_transfer_owner.return_value = False + executor.async_transfer_manager.has_any_transfer_owner.side_effect = lambda candidate: ( + executor.async_transfer_manager.requests_in_transfer().get(candidate.py_request_id) + is candidate + ) monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) PyExecutor._handle_canceled_requests(executor) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 101edffbe155..570e9b27796f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -362,13 +362,12 @@ def test_pp_gt_1_terminates_on_transfer_complete(self): _end_transfer_and_maybe_terminate must then terminate it exactly once when the transfer completes (force_terminate_ctx_for_partial_reuse=False).""" req = Mock() - executor = types.SimpleNamespace( - kv_cache_transceiver=Mock(), - active_requests=[], # already removed by _handle_responses - async_transfer_manager=Mock(), - force_terminate_ctx_for_partial_reuse=False, - _terminate_request=Mock(), - ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.active_requests = [] # already removed by _handle_responses + executor.async_transfer_manager = Mock() + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = Mock() executor.async_transfer_manager.end_transfer.return_value = True PyExecutor._end_transfer_and_maybe_terminate(executor, req) @@ -381,13 +380,12 @@ def test_pp1_does_not_double_terminate_on_transfer_complete(self): skip re-terminating it (force_terminate_ctx_for_partial_reuse=True) to avoid a double free_resources (nvbug/5961736).""" req = Mock() - executor = types.SimpleNamespace( - kv_cache_transceiver=Mock(), - active_requests=[], # already removed + terminated by early path - async_transfer_manager=Mock(), - force_terminate_ctx_for_partial_reuse=True, - _terminate_request=Mock(), - ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.active_requests = [] # already removed + terminated by early path + executor.async_transfer_manager = Mock() + executor.force_terminate_ctx_for_partial_reuse = True + executor._terminate_request = Mock() executor.async_transfer_manager.end_transfer.return_value = True PyExecutor._end_transfer_and_maybe_terminate(executor, req) diff --git a/tests/unittest/_torch/executor/test_py_executor_teardown.py b/tests/unittest/_torch/executor/test_py_executor_teardown.py new file mode 100644 index 000000000000..eca36c23640f --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_executor_teardown.py @@ -0,0 +1,248 @@ +# 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 threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor import py_executor_creator +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + +class _ShutdownOwnershipStub: + """Minimal owner of the attributes exercised by ``PyExecutor.shutdown``.""" + + shutdown = PyExecutor.shutdown + _requires_physical_transfer_drain = PyExecutor._requires_physical_transfer_drain + _discard_pending_transfer_responses_after_shutdown = ( + PyExecutor._discard_pending_transfer_responses_after_shutdown + ) + _discard_unpublishable_transfer_terminals_after_shutdown = ( + PyExecutor._discard_unpublishable_transfer_terminals_after_shutdown + ) + _drain_discarded_transfer_terminals_after_shutdown = ( + PyExecutor._drain_discarded_transfer_terminals_after_shutdown + ) + _get_pending_transfer_response_terminals = PyExecutor._get_pending_transfer_response_terminals + _has_async_transfer_owner = PyExecutor._has_async_transfer_owner + _has_pending_ownership_transition = PyExecutor._has_pending_ownership_transition + _dist_size = staticmethod(PyExecutor._dist_size) + _get_resource_manager_shutdown_progress = PyExecutor._get_resource_manager_shutdown_progress + _shutdown_resource_managers_rank_uniform = PyExecutor._shutdown_resource_managers_rank_uniform + _run_executor_shutdown_phase = PyExecutor._run_executor_shutdown_phase + _prepare_final_executor_shutdown = PyExecutor._prepare_final_executor_shutdown + _release_final_executor_objects = PyExecutor._release_final_executor_objects + + def __init__(self, terminal_shutdown_result): + self.executor_request_queue = Mock() + self.shutdown_event = threading.Event() + self.shutdown_event.set() + self.hang_detector = Mock() + self.hang_detector.detected.return_value = False + self.worker_thread = Mock() + self.dist = SimpleNamespace(pp_size=1) + self._shutdown_sleep_wakeup_listeners = Mock() + self.worker_started = True + self._pending_transfer_responses = [] + self._pending_transfer_terminals = {} + self.model_engine = None + self.draft_model_engine = None + self.kv_cache_transceiver = Mock() + self.kv_cache_transceiver.requires_physical_drain_before_request_release = True + self.kv_cache_transceiver.shutdown.side_effect = [ + False, + terminal_shutdown_result, + ] + self.manager = Mock() + self.resource_manager = SimpleNamespace(resource_managers={"manager": self.manager}) + self.virtual_memory_pools = {} + self.sampler = object() + self.dwdp_manager = None + + +def test_native_shutdown_retries_before_releasing_resource_managers(monkeypatch): + """A non-drained Python-native transceiver vetoes manager teardown.""" + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + stub = _ShutdownOwnershipStub(True) + + with pytest.raises(RuntimeError, match="still owns active transfer targets"): + stub.shutdown() + + stub.manager.shutdown.assert_not_called() + stub.shutdown() + + stub.manager.shutdown.assert_called_once_with() + assert stub.kv_cache_transceiver.shutdown.call_count == 2 + stub.executor_request_queue.enqueue_shutdown_request.assert_called_once_with() + stub.worker_thread.join.assert_called_once_with() + stub._shutdown_sleep_wakeup_listeners.assert_called_once_with() + + +def test_native_shutdown_rejects_legacy_none_result(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + stub = _ShutdownOwnershipStub(None) + + with pytest.raises(RuntimeError, match="still owns active transfer targets"): + stub.shutdown() + with pytest.raises(RuntimeError, match="did not prove physical drain"): + stub.shutdown() + + stub.manager.shutdown.assert_not_called() + + +@pytest.mark.parametrize("failing_phase", ["sampler", "dwdp_manager"]) +def test_native_finalization_failure_is_not_replayed(monkeypatch, failing_phase): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + + class FailingSampler: + def __init__(self): + self.stop_calls = 0 + + @staticmethod + def async_worker_enabled(): + return True + + def async_worker_stop(self): + self.stop_calls += 1 + raise RuntimeError("sampler finalization failed") + + class TestDwdpManager: + def __init__(self, *, fail): + self.exit_calls = 0 + self.fail = fail + + def __exit__(self, *_): + self.exit_calls += 1 + if self.fail: + raise RuntimeError("DWDP finalization failed") + + stub = _ShutdownOwnershipStub(True) + stub.kv_cache_transceiver.shutdown.side_effect = None + stub.kv_cache_transceiver.shutdown.return_value = True + if failing_phase == "sampler": + sampler = FailingSampler() + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.AsyncWorkerMixin", + FailingSampler, + ) + stub.sampler = sampler + dwdp_manager = TestDwdpManager(fail=False) + stub.dwdp_manager = dwdp_manager + failed_object = sampler + else: + dwdp_manager = TestDwdpManager(fail=True) + stub.dwdp_manager = dwdp_manager + failed_object = dwdp_manager + + with pytest.raises(RuntimeError, match="finalization failed"): + stub.shutdown() + with pytest.raises(RuntimeError, match="outcome is in doubt"): + stub.shutdown() + + if failing_phase == "sampler": + assert failed_object.stop_calls == 1 + else: + assert failed_object.exit_calls == 1 + assert dwdp_manager.exit_calls == 1 + stub.manager.shutdown.assert_not_called() + assert hasattr(stub, "model_engine") + stub._shutdown_sleep_wakeup_listeners.assert_called_once_with() + + +def test_estimation_teardown_retries_only_after_transceiver_drains(): + """Lifecycle-capable teardown requires explicit positive drain proof.""" + transceiver = Mock() + transceiver.requires_physical_drain_before_request_release = True + transceiver.shutdown.side_effect = [False, True] + py_executor = SimpleNamespace(kv_cache_transceiver=transceiver) + kv_cache_creator = Mock() + resources = {} + + with pytest.raises(RuntimeError, match="did not prove physical drain"): + py_executor_creator._teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, resources + ) + + kv_cache_creator.teardown_managers.assert_not_called() + py_executor_creator._teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, resources + ) + + kv_cache_creator.teardown_managers.assert_called_once_with(resources) + assert transceiver.shutdown.call_count == 2 + + +def test_estimation_teardown_rejects_native_none_result(): + transceiver = Mock() + transceiver.requires_physical_drain_before_request_release = True + transceiver.shutdown.return_value = None + py_executor = SimpleNamespace(kv_cache_transceiver=transceiver) + kv_cache_creator = Mock() + + with pytest.raises(RuntimeError, match="did not prove physical drain"): + py_executor_creator._teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, {} + ) + + kv_cache_creator.teardown_managers.assert_not_called() + + +def test_estimation_teardown_preserves_legacy_none_result(): + transceiver = Mock() + transceiver.requires_physical_drain_before_request_release = False + transceiver.shutdown.return_value = None + py_executor = SimpleNamespace(kv_cache_transceiver=transceiver) + kv_cache_creator = Mock() + resources = {} + + py_executor_creator._teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, resources + ) + + kv_cache_creator.teardown_managers.assert_called_once_with(resources) + + +@pytest.mark.parametrize("shutdown_result", [True, None]) +def test_estimation_teardown_rejects_inflight_transfer_owner(shutdown_result): + transceiver = Mock() + transceiver.requires_physical_drain_before_request_release = False + transceiver.shutdown.return_value = shutdown_result + transfer_manager = Mock() + transfer_manager.has_any_inflight_requests.return_value = True + py_executor = SimpleNamespace( + kv_cache_transceiver=transceiver, + async_transfer_manager=transfer_manager, + ) + kv_cache_creator = Mock() + + with pytest.raises(RuntimeError, match="transfer ownership is still active"): + py_executor_creator._teardown_kv_cache_managers_after_transceiver_shutdown( + py_executor, kv_cache_creator, {} + ) + + transceiver.shutdown.assert_called_once_with() + transfer_manager.has_any_inflight_requests.assert_called_once_with() + kv_cache_creator.teardown_managers.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py new file mode 100644 index 000000000000..40866fbb56fe --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py @@ -0,0 +1,2444 @@ +# 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. + +from queue import Queue +from types import SimpleNamespace +from unittest.mock import Mock, call + +import pytest + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import ( + _PYTHON_NATIVE_TRANSCEIVER_OWNER, + _SHUTDOWN_CONNECTOR_COMPLETION_POLL_INTERVAL_S, + _SHUTDOWN_CONNECTOR_COMPLETION_POLLS, + DisaggPPTerminationHandler, + PyExecutor, + _TransferTerminalOutcome, +) + + +class _ExactTransferManager: + """Small exact-request owner ledger for executor lifecycle tests.""" + + def __init__(self, request, *, transceiver=False, anonymous=0, events=None): + self._request = request + self._transceiver = transceiver + self._anonymous = anonymous + self._events = events + self.begin_shutdown = Mock() + self.has_pending_admission = Mock(return_value=False) + self.has_transfer_owner = Mock(side_effect=self._has_transfer_owner) + self.has_any_transfer_owner = Mock(side_effect=self._has_any_transfer_owner) + self.is_final_transfer_owner = Mock(side_effect=self._is_final_transfer_owner) + self.requests_in_transfer = Mock(side_effect=self._requests_in_transfer) + self.requests_with_owner = Mock(side_effect=self._requests_with_owner) + self.end_transfer = Mock(side_effect=self._end_transfer) + self.has_any_inflight_requests = Mock(side_effect=lambda: self._has_any_owner()) + + def _has_any_owner(self): + return self._transceiver or self._anonymous > 0 + + def _has_exact_request(self, request): + return request is self._request and self._has_any_owner() + + def _has_transfer_owner(self, request, owner): + return ( + self._has_exact_request(request) + and owner == _PYTHON_NATIVE_TRANSCEIVER_OWNER + and self._transceiver + ) + + def _has_any_transfer_owner(self, request): + return self._has_exact_request(request) + + def _is_final_transfer_owner(self, request, owner=None): + assert request is self._request + if owner is None: + assert self._anonymous > 0 + else: + assert owner == _PYTHON_NATIVE_TRANSCEIVER_OWNER + assert self._transceiver + return int(self._transceiver) + self._anonymous == 1 + + def _requests_in_transfer(self): + if self._has_any_owner(): + return {self._request.py_request_id: self._request} + return {} + + def _requests_with_owner(self, owner): + if owner == _PYTHON_NATIVE_TRANSCEIVER_OWNER and self._transceiver: + return {self._request.py_request_id: self._request} + return {} + + def _end_transfer(self, request, owner=None): + assert request is self._request + if owner is None: + assert self._anonymous > 0 + self._anonymous -= 1 + if self._events is not None: + self._events.connector_owner_release(request) + else: + assert owner == _PYTHON_NATIVE_TRANSCEIVER_OWNER + assert self._transceiver + self._transceiver = False + if self._events is not None: + self._events.transceiver_owner_release(request) + return not self._has_any_owner() + + +def _initialize_bare_executor_lifecycle_state(executor): + """Install constructor-owned state needed by unbound lifecycle methods.""" + executor.enable_attention_dp = False + executor.dist = SimpleNamespace(rank=1, world_size=1) + executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._disagg_timed_out_gen_cancelled_ids = set() + executor._pending_native_context_completions = {} + executor._pending_transfer_terminals = {} + return executor + + +def _make_error_executor(request, *, fatal): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.active_requests = [request] + executor._deferred_transfer_terminations = {} + executor._error_budget = Mock(budget=1.0) + executor._error_budget.consume.return_value = fatal + executor._fatal_error = None + executor.is_shutdown = False + executor.waiting_queue = [] + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.enable_attention_dp = False + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = True + executor.force_terminate_ctx_for_partial_reuse = False + executor._pending_connector_completions = {} + executor._terminated_transfer_requests = {} + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = Queue() + return executor + + +@pytest.mark.parametrize("fatal", [False, True]) +def test_error_response_precedes_transfer_resource_termination(fatal): + request = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ) + request.is_generation_only_request.return_value = True + executor = _make_error_executor(request, fatal=fatal) + executor.kv_cache_transceiver.cancel_request.side_effect = [False, True] + + if fatal: + executor._handle_errors("transfer failure") + else: + executor._handle_errors("transfer failure", requests=[request], charge_budget=False) + + assert request.state == LlmRequestState.GENERATION_COMPLETE + assert executor.active_requests == [] + executor._enqueue_responses.assert_called_once() + executor._terminate_request.assert_not_called() + assert executor._deferred_transfer_terminations == {id(request): request} + + # This hook runs at the top of every executor iteration. The second + # cancellation attempt proves physical drain and releases the request. + executor._handle_disagg_cache_errors_synced() + + assert executor.kv_cache_transceiver.cancel_request.call_args_list == [ + call(request), + call(request), + ] + executor._terminate_request.assert_called_once_with(request) + assert executor._deferred_transfer_terminations == {} + + +def test_error_without_transceiver_does_not_defer_stale_transfer_state(): + request = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ) + request.is_generation_only_request.return_value = True + executor = _make_error_executor(request, fatal=False) + executor.kv_cache_transceiver = None + + executor._handle_errors("transfer failure", requests=[request], charge_budget=False) + + executor._terminate_request.assert_called_once_with(request) + assert executor._deferred_transfer_terminations == {} + + +def test_cpp_transceiver_keeps_existing_error_termination_contract(): + request = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ) + request.is_generation_only_request.return_value = True + executor = _make_error_executor(request, fatal=False) + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + + executor._handle_errors("transfer failure", requests=[request], charge_budget=False) + + executor.kv_cache_transceiver.cancel_request.assert_not_called() + executor._terminate_request.assert_called_once_with(request) + assert executor._deferred_transfer_terminations == {} + + +def test_receive_setup_error_still_uses_physical_drain_gate(): + request = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_TRANS_ERROR, + ) + request.is_generation_only_request.return_value = True + executor = _make_error_executor(request, fatal=False) + executor.kv_cache_transceiver.cancel_request.return_value = False + + executor._handle_errors("receive setup failed", requests=[request], charge_budget=False) + + executor._terminate_request.assert_not_called() + assert executor._deferred_transfer_terminations == {id(request): request} + + +def test_error_skips_pending_admission_but_terminates_sibling(): + pending = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + ) + sibling = Mock( + py_request_id=8, + py_client_id=None, + state=LlmRequestState.GENERATION_COMPLETE, + ) + executor = _make_error_executor(pending, fatal=False) + executor.active_requests = [pending, sibling] + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_pending_admission.side_effect = ( + lambda request: request is pending + ) + + executor._handle_errors( + "admission failed", + requests=[pending, sibling], + charge_budget=False, + ) + + executor._terminate_request.assert_called_once_with(sibling) + assert executor.active_requests == [] + + +def test_error_skips_pending_connector_completion_but_terminates_sibling(): + pending = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.GENERATION_COMPLETE, + ) + sibling = Mock( + py_request_id=8, + py_client_id=None, + state=LlmRequestState.GENERATION_COMPLETE, + ) + executor = _make_error_executor(pending, fatal=False) + executor.active_requests = [pending, sibling] + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor._pending_connector_completions = {id(pending): SimpleNamespace(request=pending)} + + executor._handle_errors( + "connector cleanup failed", + requests=[pending, sibling], + charge_budget=False, + ) + + executor._terminate_request.assert_called_once_with(sibling) + assert executor.active_requests == [] + + +def test_context_error_releases_async_manager_only_after_physical_drain(): + request = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + ) + request.is_generation_only_request.return_value = False + executor = _make_error_executor(request, fatal=False) + executor.kv_cache_transceiver.cancel_request.side_effect = [False, True] + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor.async_transfer_manager.has_transfer_owner.return_value = True + + def end_transfer(req, *, owner): + assert owner == _PYTHON_NATIVE_TRANSCEIVER_OWNER + if req.state != LlmRequestState.DISAGG_TRANS_ERROR: + req.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + return True + + executor.async_transfer_manager.end_transfer.side_effect = end_transfer + executor.force_terminate_ctx_for_partial_reuse = False + + executor._handle_errors("context transfer failed", requests=[request], charge_budget=False) + + executor.async_transfer_manager.end_transfer.assert_not_called() + executor._terminate_request.assert_not_called() + + executor._handle_disagg_cache_errors_synced() + + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor._terminate_request.assert_called_once_with(request) + assert request.state == LlmRequestState.DISAGG_TRANS_ERROR + assert executor._deferred_transfer_terminations == {} + + +def test_context_error_preserves_error_state_for_remaining_transfer_owner(): + request = Mock( + py_request_id=7, + state=LlmRequestState.GENERATION_COMPLETE, + ) + request.is_generation_only_request.return_value = False + executor = _make_error_executor(request, fatal=False) + executor._deferred_transfer_terminations = {id(request): request} + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor.async_transfer_manager.has_transfer_owner.return_value = True + executor.async_transfer_manager.end_transfer.return_value = False + executor.force_terminate_ctx_for_partial_reuse = False + + assert executor._finalize_deferred_transfer_termination(request) + + assert request.state == LlmRequestState.DISAGG_TRANS_ERROR + executor._terminate_request.assert_not_called() + assert executor._deferred_transfer_terminations == {} + + +def _make_native_launch_executor(request): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.kv_cache_manager = SimpleNamespace(release_index_slot=Mock()) + executor.kv_cache_transceiver = Mock( + requires_physical_drain_before_request_release=True, + kv_transfer_timeout_ms=None, + ) + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor.async_transfer_manager.has_transfer_owner.return_value = True + executor.async_transfer_manager.has_any_transfer_owner.return_value = True + executor.async_transfer_manager.end_transfer.return_value = True + executor.kv_connector_manager = None + executor._deferred_transfer_terminations = {} + executor._deferred_transfer_terminations_already_terminated = set() + executor._terminated_transfer_requests = {} + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = Mock() + return executor + + +def _make_native_launch_request(): + request = Mock( + py_request_id=7, + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + is_context_only_request=True, + is_context_finished=True, + is_finished_due_to_length=False, + is_finished_due_to_cancellation=False, + ) + request.is_generation_only_request.return_value = False + return request + + +def _make_native_completion_executor(request): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.kv_cache_transceiver = Mock(requires_physical_drain_before_request_release=True) + executor.kv_cache_transceiver.check_context_transfer_status.side_effect = [ + ([request.py_request_id], []), + ([], []), + ] + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + executor.active_requests = [request] + executor._deferred_transfer_terminations = {} + executor._terminated_transfer_requests = {} + executor.force_terminate_ctx_for_partial_reuse = False + executor._pending_transfer_responses = [] + executor._maybe_attach_ctx_usage = Mock() + executor._terminate_request = Mock() + executor._check_cache_transfer_errors = Mock() + return executor + + +def _make_native_completion_request(): + request = Mock( + py_request_id=7, + py_kv_transfer_timed_out=False, + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + cached_tokens=11, + ) + request.is_generation_only_request.return_value = False + request.create_response.return_value = Mock(result=SimpleNamespace()) + return request + + +@pytest.mark.parametrize("requires_physical_drain", [False, True]) +def test_transceiver_owner_mode_boundary(requires_physical_drain): + request = _make_native_launch_request() + request.py_kv_transfer_timed_out = False + executor = _make_native_launch_executor(request) + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = ( + requires_physical_drain + ) + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.active_requests = [] + + executor._send_kv_async([request]) + + expected_call = ( + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER) + if requires_physical_drain + else call(request) + ) + assert executor.async_transfer_manager.start_transfer.call_args == expected_call + executor.kv_cache_transceiver.respond_and_send_async.assert_called_once_with(request) + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + +@pytest.mark.parametrize("requires_physical_drain", [False, True]) +def test_context_status_owner_mode_boundary(requires_physical_drain): + request = Mock( + py_request_id=7, + py_kv_transfer_timed_out=False, + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + ) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.kv_cache_transceiver = Mock( + requires_physical_drain_before_request_release=requires_physical_drain + ) + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([7], []) + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor.async_transfer_manager.end_transfer.return_value = True + executor.active_requests = [] + executor._deferred_transfer_terminations = {} + executor._terminated_transfer_requests = {} + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = Mock() + executor._finalize_deferred_transfer_termination = Mock(return_value=False) + executor._end_transfer_and_maybe_terminate = Mock() + executor._check_cache_transfer_errors = Mock() + + executor._check_disagg_ctx_cache_transfer_status() + + if requires_physical_drain: + executor._end_transfer_and_maybe_terminate.assert_not_called() + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor._terminate_request.assert_called_once_with(request) + assert executor._pending_native_context_completions == {} + else: + executor._end_transfer_and_maybe_terminate.assert_called_once_with( + request, transfer_owner=None + ) + + +def test_native_completion_response_preparation_retries_exact_response(): + request = _make_native_completion_request() + response = request.create_response.return_value + executor = _make_native_completion_executor(request) + executor._maybe_attach_ctx_usage.side_effect = [ + RuntimeError("usage unavailable"), + None, + ] + + executor._check_disagg_ctx_cache_transfer_status() + + progress = executor._pending_native_context_completions[id(request)] + assert progress.request is request + terminal = executor._pending_transfer_terminals[id(request)] + assert terminal.response is response + assert executor._pending_transfer_responses == [] + executor.async_transfer_manager.end_transfer.assert_not_called() + + executor._check_disagg_ctx_cache_transfer_status() + + assert executor._pending_native_context_completions == {} + request.create_response.assert_called_once_with(False, 1) + assert executor._pending_transfer_responses == [(7, response)] + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor._terminate_request.assert_called_once_with(request) + + +def test_native_completion_unpin_failure_retries_named_owner_once_resolved(): + request = _make_native_completion_request() + executor = _make_native_completion_executor(request) + manager = executor.async_transfer_manager + release_attempts = 0 + + def release_owner(req, owner=None): + nonlocal release_attempts + release_attempts += 1 + if release_attempts == 1: + raise RuntimeError("unpin failed") + return manager._end_transfer(req, owner) + + manager.end_transfer.side_effect = release_owner + + executor._check_disagg_ctx_cache_transfer_status() + + progress = executor._pending_native_context_completions[id(request)] + assert executor._pending_transfer_terminals[id(request)].response_buffered + assert not progress.owner_released + executor._terminate_request.assert_not_called() + + executor._check_disagg_ctx_cache_transfer_status() + + assert executor._pending_native_context_completions == {} + assert manager.end_transfer.call_args_list == [ + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + ] + request.create_response.assert_called_once() + assert len(executor._pending_transfer_responses) == 1 + executor._terminate_request.assert_called_once_with(request) + + +def test_late_cancel_does_not_compete_with_consumed_native_completion(): + request = _make_native_completion_request() + request.is_child = False + request.parent_request_id = None + executor = _make_native_completion_executor(request) + manager = executor.async_transfer_manager + release_attempts = 0 + + def release_owner(req, owner=None): + nonlocal release_attempts + release_attempts += 1 + if release_attempts == 1: + raise RuntimeError("unpin failed") + return manager._end_transfer(req, owner) + + manager.end_transfer.side_effect = release_owner + executor._check_disagg_ctx_cache_transfer_status() + assert id(request) in executor._pending_native_context_completions + + executor.canceled_req_ids = [request.py_request_id] + executor.waiting_queue = Mock() + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [] + assert executor._pending_request_cancellations == {} + executor.kv_cache_transceiver.cancel_request.assert_not_called() + assert manager.end_transfer.call_count == 1 + request.finish_by_reason.assert_not_called() + + executor._check_disagg_ctx_cache_transfer_status() + + assert executor._pending_native_context_completions == {} + assert manager.end_transfer.call_args_list == [ + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + ] + assert len(executor._pending_transfer_responses) == 1 + executor._terminate_request.assert_called_once_with(request) + + +def test_native_status_hands_transport_drain_to_earlier_cancellation(): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + request.py_kv_transfer_timed_out = False + executor = _make_cancel_executor([request]) + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + executor.kv_cache_transceiver.cancel_request.return_value = False + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ( + [request.py_request_id], + [], + ) + executor._deferred_transfer_terminations = {} + executor._terminated_transfer_requests = {} + executor.force_terminate_ctx_for_partial_reuse = False + executor._check_cache_transfer_errors = Mock() + + executor._handle_canceled_requests() + progress = executor._pending_request_cancellations[id(request)] + assert not progress.transport_drained + + executor._check_disagg_ctx_cache_transfer_status() + + assert progress.transport_drained + assert executor._pending_native_context_completions == {} + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + executor._handle_canceled_requests() + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + request.finish_by_reason.assert_called_once() + assert executor._pending_request_cancellations == {} + assert executor.canceled_req_ids == [] + + +def test_error_does_not_override_buffered_native_success(): + request = _make_native_completion_request() + request.py_client_id = None + executor = _make_native_completion_executor(request) + manager = executor.async_transfer_manager + release_attempts = 0 + + def release_owner(req, owner=None): + nonlocal release_attempts + release_attempts += 1 + if release_attempts == 1: + raise RuntimeError("unpin failed") + return manager._end_transfer(req, owner) + + manager.end_transfer.side_effect = release_owner + executor._check_disagg_ctx_cache_transfer_status() + response = request.create_response.return_value + assert executor._pending_transfer_responses == [(7, response)] + + executor._error_budget = Mock(budget=1.0) + executor._error_budget.consume.return_value = False + executor._fatal_error = None + executor.is_shutdown = False + executor.waiting_queue = [] + executor.gather_all_responses = False + executor._enqueue_responses = Mock() + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = Queue() + + executor._handle_errors( + "unrelated executor error", + requests=[request], + charge_budget=False, + ) + + assert executor.active_requests == [request] + assert executor._pending_transfer_responses == [(7, response)] + executor._enqueue_responses.assert_called_once_with([]) + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + executor._check_disagg_ctx_cache_transfer_status() + + request.create_response.assert_called_once_with(False, 1) + assert executor._pending_native_context_completions == {} + assert executor.active_requests == [] + executor._terminate_request.assert_called_once_with(request) + + +def test_native_failure_emits_one_error_while_completion_ledger_retires(): + request = _make_native_completion_request() + request.py_client_id = None + request.is_context_only_request = True + request.state = LlmRequestState.DISAGG_TRANS_ERROR + executor = _make_native_completion_executor(request) + executor.kv_cache_transceiver.check_context_transfer_status.side_effect = [ + ([], [request.py_request_id]), + ([], []), + ] + executor.kv_cache_transceiver.cancel_request.return_value = True + manager = executor.async_transfer_manager + release_attempts = 0 + + def release_owner(req, owner=None): + nonlocal release_attempts + release_attempts += 1 + if release_attempts == 1: + raise RuntimeError("unpin failed") + return manager._end_transfer(req, owner) + + manager.end_transfer.side_effect = release_owner + executor._error_budget = Mock(budget=1.0) + executor._error_budget.consume.return_value = False + executor._fatal_error = None + executor.is_shutdown = False + executor.waiting_queue = [] + executor.gather_all_responses = False + executor._enqueue_responses = Mock() + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = Queue() + del executor._check_cache_transfer_errors + + executor._check_disagg_ctx_cache_transfer_status() + assert id(request) in executor._pending_native_context_completions + + executor._handle_errors( + "error observed while native cleanup is pending", + requests=[request], + charge_budget=False, + ) + + assert executor.active_requests == [] + executor._terminate_request.assert_not_called() + assert executor._enqueue_responses.call_count == 1 + assert executor._enqueue_responses.call_args.args[0][0][0] == 7 + + executor._check_disagg_ctx_cache_transfer_status() + + assert executor._pending_native_context_completions == {} + assert executor.active_requests == [] + executor._terminate_request.assert_called_once_with(request) + nonempty_response_calls = [ + args[0] for args, _kwargs in executor._enqueue_responses.call_args_list if args[0] + ] + assert len(nonempty_response_calls) == 1 + assert nonempty_response_calls[0][0][0] == request.py_request_id + assert nonempty_response_calls[0][0][1].error_msg == ( + "error observed while native cleanup is pending" + ) + request.create_response.assert_not_called() + + +def test_native_completion_termination_retry_does_not_release_twice(): + request = _make_native_completion_request() + executor = _make_native_completion_executor(request) + executor.resource_manager = Mock() + executor._prefetched_request_ids = Mock() + executor._prefetched_request_ids.discard.side_effect = [ + RuntimeError("post-release bookkeeping failed"), + None, + ] + executor.gather_all_responses = False + executor._request_resource_termination_progress = {} + executor._terminate_request = executor._do_terminate_request + + executor._check_disagg_ctx_cache_transfer_status() + + progress = executor._pending_native_context_completions[id(request)] + assert progress.owner_released + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor.resource_manager.free_resources.assert_called_once_with(request) + + executor._check_disagg_ctx_cache_transfer_status() + + assert executor._pending_native_context_completions == {} + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor.resource_manager.free_resources.assert_called_once_with(request) + assert executor._request_resource_termination_progress == {} + request.create_response.assert_called_once() + + +def test_native_launch_failure_before_session_retires_exact_owner(): + request = _make_native_launch_request() + executor = _make_native_launch_executor(request) + executor.kv_cache_transceiver.respond_and_send_async.side_effect = RuntimeError( + "launch failed before session" + ) + executor.kv_cache_transceiver.cancel_request.return_value = True + + with pytest.raises(RuntimeError, match="before session"): + executor._send_kv_async([request]) + + executor.async_transfer_manager.start_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor._terminate_request.assert_called_once_with(request) + assert executor._deferred_transfer_terminations == {} + assert request.state == LlmRequestState.DISAGG_TRANS_ERROR + + +def test_native_launch_failure_after_drained_session_retires_exact_owner(): + request = _make_native_launch_request() + executor = _make_native_launch_executor(request) + session_created = Mock() + + def launch_then_fail(req): + session_created(req) + raise RuntimeError("launch failed after session") + + executor.kv_cache_transceiver.respond_and_send_async.side_effect = launch_then_fail + executor.kv_cache_transceiver.cancel_request.return_value = True + + with pytest.raises(RuntimeError, match="after session"): + executor._send_kv_async([request]) + + session_created.assert_called_once_with(request) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor._terminate_request.assert_called_once_with(request) + assert executor._deferred_transfer_terminations == {} + + +def test_native_launch_failure_with_ambiguous_tasks_retains_exact_request(): + request = _make_native_launch_request() + executor = _make_native_launch_executor(request) + executor.kv_cache_transceiver.respond_and_send_async.side_effect = RuntimeError( + "launch outcome ambiguous" + ) + executor.kv_cache_transceiver.cancel_request.return_value = False + + with pytest.raises(RuntimeError, match="ambiguous"): + executor._send_kv_async([request]) + + assert executor._deferred_transfer_terminations == {id(request): request} + executor.async_transfer_manager.end_transfer.assert_not_called() + executor._terminate_request.assert_not_called() + assert request.state == LlmRequestState.DISAGG_TRANS_ERROR + + +def _make_cancel_request(request_id, state, *, generation_only): + request = Mock( + py_request_id=request_id, + parent_request_id=None, + is_child=False, + state=state, + py_decoding_iter=3, + is_finished=False, + is_finished_due_to_cancellation=False, + ) + request.is_generation_only_request.return_value = generation_only + + def finish_by_reason(_reason): + request.is_finished = True + request.is_finished_due_to_cancellation = True + request.state = LlmRequestState.GENERATION_COMPLETE + + request.finish_by_reason.side_effect = finish_by_reason + return request + + +def _make_cancel_executor(requests): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.canceled_req_ids = [request.py_request_id for request in requests] + executor.waiting_queue = Mock() + executor.active_requests = list(requests) + executor.kv_cache_transceiver = Mock(requires_physical_drain_before_request_release=True) + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_pending_admission.return_value = False + executor.async_transfer_manager.has_transfer_owner.return_value = False + executor.async_transfer_manager.has_any_transfer_owner.return_value = False + executor.async_transfer_manager.requests_in_transfer.return_value = {} + executor._terminated_transfer_requests = {} + return executor + + +def test_user_cancel_retires_native_context_owner_after_transport_drain(): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + executor = _make_cancel_executor([request]) + executor.async_transfer_manager.has_transfer_owner.return_value = True + events = [] + + def cancel_request(req): + events.append(("transport", req)) + return True + + def end_transfer(req, *, owner): + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + events.append(("owner", req, owner)) + return True + + def finish_by_reason(reason): + events.append(("finish", request, reason)) + request.is_finished = True + request.is_finished_due_to_cancellation = True + request.state = LlmRequestState.GENERATION_COMPLETE + + executor.kv_cache_transceiver.cancel_request.side_effect = cancel_request + executor.async_transfer_manager.end_transfer.side_effect = end_transfer + request.finish_by_reason.side_effect = finish_by_reason + + executor._handle_canceled_requests() + + assert [event[0] for event in events] == ["transport", "owner", "finish"] + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + assert request.state == LlmRequestState.GENERATION_COMPLETE + assert request.is_finished_due_to_cancellation + assert executor.canceled_req_ids == [] + + +def test_user_cancel_retries_native_owner_release_and_continues_siblings(): + retry_request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + sibling_request = _make_cancel_request( + 8, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + generation_only=True, + ) + executor = _make_cancel_executor([retry_request, sibling_request]) + executor.async_transfer_manager.has_transfer_owner.side_effect = ( + lambda request, _owner: request is retry_request + ) + executor.async_transfer_manager.end_transfer.side_effect = [ + RuntimeError("unpin failed"), + True, + ] + + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [7] + assert retry_request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + retry_request.finish_by_reason.assert_not_called() + sibling_request.finish_by_reason.assert_called_once() + + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [] + assert retry_request.state == LlmRequestState.GENERATION_COMPLETE + retry_request.finish_by_reason.assert_called_once() + assert executor.kv_cache_transceiver.cancel_request.call_args_list == [ + call(retry_request), + call(sibling_request), + ] + assert executor.async_transfer_manager.end_transfer.call_args_list == [ + call(retry_request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + call(retry_request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + ] + + +def test_user_cancel_retries_transport_error_and_continues_siblings(): + retry_request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_TRANS_ERROR, + generation_only=False, + ) + sibling_request = _make_cancel_request( + 8, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + generation_only=True, + ) + executor = _make_cancel_executor([retry_request, sibling_request]) + executor.async_transfer_manager.has_transfer_owner.side_effect = ( + lambda request, _owner: request is retry_request + ) + executor.kv_cache_transceiver.cancel_request.side_effect = [ + RuntimeError("transport failed"), + True, + True, + ] + executor.async_transfer_manager.end_transfer.return_value = True + + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [7] + assert retry_request.state == LlmRequestState.DISAGG_TRANS_ERROR + retry_request.finish_by_reason.assert_not_called() + sibling_request.finish_by_reason.assert_called_once() + executor.async_transfer_manager.end_transfer.assert_not_called() + + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [] + retry_request.finish_by_reason.assert_called_once() + executor.async_transfer_manager.end_transfer.assert_called_once_with( + retry_request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + + +def test_failed_cancellation_keeps_first_outcome_across_late_error(): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + request.py_client_id = None + executor = _make_cancel_executor([request]) + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + executor.kv_cache_transceiver.cancel_request.side_effect = [False, True] + executor._deferred_transfer_terminations = {} + executor._error_budget = Mock() + executor._error_budget.consume.return_value = False + executor._fatal_error = None + executor.is_shutdown = False + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.enable_attention_dp = False + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + executor.force_terminate_ctx_for_partial_reuse = False + + executor._handle_canceled_requests() + + progress = executor._pending_request_cancellations[id(request)] + assert not progress.transport_drained + executor._terminate_request.assert_not_called() + + executor._handle_errors( + "request failed during cancellation", + requests=[request], + charge_budget=False, + ) + + assert executor.active_requests == [request] + executor._terminate_request.assert_not_called() + + executor._handle_canceled_requests() + + assert executor._pending_request_cancellations == {} + assert executor.async_transfer_manager.end_transfer.call_args_list == [ + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER) + ] + + _configure_finished_response(executor, request) + executor._handle_responses() + + executor._terminate_request.assert_called_once_with(request) + + +def test_native_owner_release_retry_survives_response_processing(): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + request.is_finished = True + request.is_attention_dp_dummy = False + executor = _make_cancel_executor([request]) + executor.async_transfer_manager.has_transfer_owner.return_value = True + executor.async_transfer_manager.end_transfer.side_effect = [ + RuntimeError("unpin failed"), + True, + ] + executor.perf_manager = Mock() + executor.perf_manager.get_timestamp.return_value = 0 + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.enable_attention_dp = False + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [7] + assert executor._pending_request_cancellations[id(request)].request is request + + # A normal response pass must neither emit the pre-cancellation success nor + # drop the only active-list root before the owner release can retry. + assert executor._handle_responses() == [] + assert executor.active_requests == [request] + request.create_response.assert_not_called() + executor._terminate_request.assert_not_called() + + executor._handle_canceled_requests() + + assert executor.canceled_req_ids == [] + assert executor._pending_request_cancellations == {} + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor.async_transfer_manager.end_transfer.call_count == 2 + request.finish_by_reason.assert_called_once() + + +def _make_connector_completion_executor(request): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.kv_connector_manager = Mock() + executor.kv_cache_transceiver = Mock(requires_physical_drain_before_request_release=True) + executor.active_requests = [request] + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_pending_admission.return_value = False + executor.async_transfer_manager.has_any_transfer_owner.return_value = False + executor.async_transfer_manager.is_final_transfer_owner.return_value = True + executor.async_transfer_manager.requests_in_transfer.return_value = {} + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = Mock() + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor._maybe_attach_ctx_usage = Mock() + executor._pending_transfer_responses = [] + executor._pending_connector_completions = {} + executor._terminated_transfer_requests = {} + return executor + + +def _configure_finished_response(executor, request): + request.is_attention_dp_dummy = False + request.py_kv_transfer_timed_out = False + request.py_draft_tokens = [] + request.py_num_accepted_draft_tokens = 0 + request.py_per_pos_drafted = [] + request.py_per_pos_accepted = [] + request.return_perf_metrics = False + request.cached_tokens = 0 + request.is_disagg_context_complete_state = False + request.is_disagg_context_transmission_state = False + response = Mock(result=SimpleNamespace()) + request.create_response.return_value = response + executor.perf_manager = Mock() + executor.perf_manager.get_timestamp.return_value = 0 + executor.iter_counter = 1 + executor.stream_interval = 1 + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.enable_attention_dp = False + executor._maybe_attach_ctx_usage = Mock() + executor._enqueue_responses = Mock() + + +def _configure_connector_poll(executor, request): + executor.kv_connector_manager = _make_one_shot_connector_manager(request) + executor._pending_connector_completions = {} + executor._pending_transfer_responses = [] + executor._terminated_transfer_requests = {} + executor.force_terminate_ctx_for_partial_reuse = False + + +def _make_one_shot_connector_manager(request, *, empty_polls=0): + manager = Mock() + poll_count = 0 + + def get_finished(): + nonlocal poll_count + poll_count += 1 + if poll_count == empty_polls + 1: + return [request] + return [] + + manager.get_finished.side_effect = get_finished + return manager + + +@pytest.mark.parametrize("first_provider", ["native", "connector"]) +def test_mixed_provider_completion_publishes_and_terminates_once(first_provider): + request = _make_native_completion_request() + executor = _make_native_completion_executor(request) + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True, anonymous=1) + executor.kv_connector_manager = _make_one_shot_connector_manager(request) + executor._enqueue_responses = Mock() + executor._terminate_request.side_effect = executor._mark_transfer_terminal_teardown_complete + + if first_provider == "connector": + executor._kv_connector_terminate_requests() + # Connector success is not authoritative while native may still fail. + assert executor._pending_transfer_terminals == {} + request.create_response.assert_not_called() + + executor._check_disagg_ctx_cache_transfer_status() + + if first_provider == "native": + executor._kv_connector_terminate_requests() + + response = request.create_response.return_value + request.create_response.assert_called_once_with(False, 1) + assert executor._pending_transfer_responses == [(7, response)] + assert executor.async_transfer_manager.end_transfer.call_count == 2 + executor._terminate_request.assert_called_once_with(request) + + executor._flush_pending_transfer_responses() + + executor._enqueue_responses.assert_called_once_with([(7, response)]) + assert executor._pending_transfer_responses == [] + assert executor._pending_transfer_terminals == {} + + +def test_connector_first_does_not_mask_later_native_failure(): + request = _make_native_completion_request() + request.py_client_id = None + request.is_context_only_request = True + executor = _make_native_completion_executor(request) + executor.kv_cache_transceiver.check_context_transfer_status.side_effect = [ + ([], [request.py_request_id]), + ([], []), + ] + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True, anonymous=1) + executor.kv_connector_manager = _make_one_shot_connector_manager(request) + executor._enqueue_responses = Mock() + + executor._kv_connector_terminate_requests() + + assert executor._pending_transfer_terminals == {} + request.state = LlmRequestState.DISAGG_TRANS_ERROR + executor._check_disagg_ctx_cache_transfer_status() + + terminal = executor._pending_transfer_terminals[id(request)] + assert terminal.outcome.value == "failed" + request.create_response.assert_not_called() + + executor._error_budget = Mock(budget=1.0) + executor._error_budget.consume.return_value = False + executor._fatal_error = None + executor.is_shutdown = False + executor.waiting_queue = [] + executor.gather_all_responses = False + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = Queue() + executor._handle_errors("native transfer failed", requests=[request], charge_budget=False) + + responses = executor._enqueue_responses.call_args.args[0] + assert len(responses) == 1 + assert responses[0][1].error_msg == "native transfer failed" + executor._terminate_request.assert_called_once_with(request) + assert executor.async_transfer_manager.end_transfer.call_count == 2 + + +def test_connector_first_does_not_mask_later_cpp_failure(): + request = _make_native_completion_request() + request.py_client_id = None + request.is_context_only_request = True + executor = _make_native_completion_executor(request) + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor.kv_cache_transceiver.check_context_transfer_status.side_effect = [ + ([], [request.py_request_id]), + ([], []), + ] + # The connector and legacy C++ transceiver are both anonymous owners. + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=2) + executor.kv_connector_manager = _make_one_shot_connector_manager(request) + executor._enqueue_responses = Mock() + + executor._kv_connector_terminate_requests() + + assert executor._pending_transfer_terminals == {} + request.create_response.assert_not_called() + request.state = LlmRequestState.DISAGG_TRANS_ERROR + executor._check_disagg_ctx_cache_transfer_status() + + terminal = executor._pending_transfer_terminals[id(request)] + assert terminal.outcome.value == "failed" + request.create_response.assert_not_called() + + executor._error_budget = Mock(budget=1.0) + executor._error_budget.consume.return_value = False + executor._fatal_error = None + executor.is_shutdown = False + executor.waiting_queue = [] + executor.gather_all_responses = False + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = Queue() + executor._handle_errors("C++ transfer failed", requests=[request], charge_budget=False) + + responses = executor._enqueue_responses.call_args.args[0] + assert len(responses) == 1 + assert responses[0][1].error_msg == "C++ transfer failed" + request.create_response.assert_not_called() + executor._terminate_request.assert_called_once_with(request) + assert executor.async_transfer_manager.end_transfer.call_count == 2 + + +@pytest.mark.parametrize("late_event", ["error", "cancel"]) +def test_native_success_wins_late_sibling_window_event(late_event): + request = _make_native_completion_request() + request.py_client_id = None + request.is_child = False + request.parent_request_id = None + executor = _make_native_completion_executor(request) + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True, anonymous=1) + executor.kv_connector_manager = _make_one_shot_connector_manager(request) + executor._enqueue_responses = Mock() + + executor._check_disagg_ctx_cache_transfer_status() + terminal = executor._pending_transfer_terminals[id(request)] + assert terminal.outcome.value == "succeeded" + + if late_event == "error": + executor._error_budget = Mock(budget=1.0) + executor._error_budget.consume.return_value = False + executor._fatal_error = None + executor.is_shutdown = False + executor.waiting_queue = [] + executor.gather_all_responses = False + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = Queue() + executor._handle_errors("late sibling error", requests=[request], charge_budget=False) + else: + executor.canceled_req_ids = [request.py_request_id] + executor.waiting_queue = Mock() + executor._handle_canceled_requests() + request.finish_by_reason.assert_not_called() + + assert executor._pending_transfer_terminals[id(request)] is terminal + assert executor.active_requests == [request] + + executor._kv_connector_terminate_requests() + + request.create_response.assert_called_once() + executor._terminate_request.assert_called_once_with(request) + assert executor.async_transfer_manager.end_transfer.call_count == 2 + + +def test_connector_response_publication_failure_is_not_replayed(): + request = _make_native_completion_request() + response = request.create_response.return_value + executor = _make_connector_completion_executor(request) + executor._terminate_request.side_effect = executor._mark_transfer_terminal_teardown_complete + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + executor.kv_connector_manager.get_finished.side_effect = [[request], []] + executor._maybe_attach_ctx_usage.side_effect = [ + RuntimeError("usage unavailable"), + None, + ] + + executor._kv_connector_terminate_requests() + + terminal = executor._pending_transfer_terminals[id(request)] + assert terminal.response is response + assert not terminal.response_buffered + executor.async_transfer_manager.end_transfer.assert_not_called() + + executor._kv_connector_terminate_requests() + + request.create_response.assert_called_once_with(False, 0) + assert executor._pending_transfer_responses == [(7, response)] + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor._terminate_request.assert_called_once_with(request) + assert executor._pending_transfer_terminals[id(request)] is terminal + + executor._enqueue_responses = Mock(side_effect=RuntimeError("enqueue outcome ambiguous")) + with pytest.raises(RuntimeError, match="enqueue outcome ambiguous"): + executor._flush_pending_transfer_responses() + + assert executor._pending_transfer_responses == [(7, response)] + assert not terminal.response_published + assert terminal.response_publication_in_doubt + + with pytest.raises(RuntimeError, match="refusing to replay"): + executor._flush_pending_transfer_responses() + + executor._enqueue_responses.assert_called_once_with([(7, response)]) + executor._discard_pending_transfer_responses_after_shutdown() + assert executor._pending_transfer_responses == [] + assert executor._pending_transfer_terminals == {} + + +def test_connector_response_creation_failure_is_not_replayed(): + request = _make_native_completion_request() + executor = _make_connector_completion_executor(request) + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + executor.kv_connector_manager.get_finished.side_effect = [[request], []] + serialized_side_effects = [] + + def mutate_then_raise(*_args): + serialized_side_effects.append("consumed") + raise RuntimeError("response creation outcome ambiguous") + + request.create_response.side_effect = mutate_then_raise + + executor._kv_connector_terminate_requests() + + terminal = executor._pending_transfer_terminals[id(request)] + assert terminal.response_creation_in_doubt + assert serialized_side_effects == ["consumed"] + executor.async_transfer_manager.end_transfer.assert_not_called() + + executor._kv_connector_terminate_requests() + + assert serialized_side_effects == ["consumed"] + request.create_response.assert_called_once_with(False, 0) + executor.async_transfer_manager.end_transfer.assert_not_called() + assert executor._pending_connector_completions[id(request)].request is request + + +@pytest.mark.parametrize("transceiver_kind", ["connector_only", "capability_false"]) +def test_connector_owned_cancel_waits_for_connector_completion(transceiver_kind): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + executor = _make_cancel_executor([request]) + if transceiver_kind == "connector_only": + executor.kv_cache_transceiver = None + else: + executor.kv_cache_transceiver = Mock(requires_physical_drain_before_request_release=False) + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + executor._terminate_request = Mock() + _configure_finished_response(executor, request) + _configure_connector_poll(executor, request) + + executor._handle_canceled_requests() + finished_requests = executor._handle_responses() + + assert finished_requests == [] + assert executor.active_requests == [request] + request.create_response.assert_not_called() + executor._terminate_request.assert_not_called() + executor._enqueue_responses.assert_called_once_with([]) + executor._enqueue_responses.reset_mock() + + executor._kv_connector_terminate_requests() + + response = request.create_response.return_value + assert executor._pending_transfer_responses == [(7, response)] + executor._terminate_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + + executor._flush_pending_transfer_responses() + + executor._enqueue_responses.assert_called_once_with([(7, response)]) + request.create_response.assert_called_once_with(False, 0) + + +def test_capability_false_cancel_keeps_anonymous_status_owner_after_connector(): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + executor = _make_cancel_executor([request]) + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor.kv_cache_transceiver.cancel_request.return_value = True + # Default C++ transceiver and connector retain the legacy anonymous count. + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=2) + executor._terminate_request = Mock() + _configure_finished_response(executor, request) + _configure_connector_poll(executor, request) + + executor._handle_canceled_requests() + finished_requests = executor._handle_responses() + + assert finished_requests == [] + assert executor.active_requests == [request] + executor.async_transfer_manager.end_transfer.assert_not_called() + executor._terminate_request.assert_not_called() + executor._enqueue_responses.assert_called_once_with([]) + executor._enqueue_responses.reset_mock() + + executor._kv_connector_terminate_requests() + + response = request.create_response.return_value + assert executor._pending_transfer_responses == [(7, response)] + assert executor.active_requests == [request] + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor._terminate_request.assert_not_called() + + executor._flush_pending_transfer_responses() + + executor._enqueue_responses.assert_called_once_with([(7, response)]) + + executor._end_transfer_and_maybe_terminate(request) + + assert executor.async_transfer_manager.end_transfer.call_args_list == [ + call(request), + call(request), + ] + executor._terminate_request.assert_called_once_with(request) + + +@pytest.mark.parametrize("transceiver_kind", ["connector_only", "capability_false"]) +def test_connector_owned_error_waits_for_connector_completion(transceiver_kind): + request = Mock( + py_request_id=7, + py_client_id=None, + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + ) + request.is_generation_only_request.return_value = False + executor = _make_error_executor(request, fatal=False) + if transceiver_kind == "connector_only": + executor.kv_cache_transceiver = None + else: + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + _configure_connector_poll(executor, request) + + executor._handle_errors( + "connector-owned request failed", + requests=[request], + charge_budget=False, + ) + + assert executor.active_requests == [] + executor._terminate_request.assert_not_called() + + executor._kv_connector_terminate_requests() + + executor._terminate_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + + +def test_connector_unpin_failure_retries_without_replaying_response(): + request = _make_cancel_request( + 7, + LlmRequestState.GENERATION_COMPLETE, + generation_only=False, + ) + request.is_finished = True + request.is_finished_due_to_cancellation = True + request.cached_tokens = 0 + request.create_response.return_value = Mock(result=SimpleNamespace()) + executor = _make_connector_completion_executor(request) + executor.kv_connector_manager.get_finished.side_effect = [[request], []] + executor.async_transfer_manager.end_transfer.side_effect = [ + RuntimeError("unpin failed"), + True, + ] + + executor._kv_connector_terminate_requests() + + assert executor._pending_connector_completions[id(request)].request is request + assert request.state == LlmRequestState.GENERATION_COMPLETE + request.create_response.assert_called_once() + executor._terminate_request.assert_not_called() + + executor._kv_connector_terminate_requests() + + assert executor._pending_connector_completions == {} + assert executor.active_requests == [] + request.create_response.assert_called_once() + assert executor.async_transfer_manager.end_transfer.call_count == 2 + executor._terminate_request.assert_called_once_with(request) + assert request.state == LlmRequestState.GENERATION_COMPLETE + assert executor.kv_connector_manager.get_finished.call_args_list == [ + call(), + call(), + ] + + +def test_post_unpin_termination_retry_does_not_release_resources_twice(): + request = _make_cancel_request( + 7, + LlmRequestState.GENERATION_COMPLETE, + generation_only=False, + ) + request.is_finished = True + request.is_finished_due_to_cancellation = True + request.cached_tokens = 0 + request.create_response.return_value = Mock(result=SimpleNamespace()) + executor = _make_connector_completion_executor(request) + executor.kv_connector_manager.get_finished.side_effect = [[request], []] + executor.async_transfer_manager.end_transfer.return_value = True + executor.async_transfer_manager.has_transfer_owner.return_value = False + executor.resource_manager = Mock() + executor._prefetched_request_ids = Mock() + executor._prefetched_request_ids.discard.side_effect = [ + RuntimeError("post-release bookkeeping failed"), + None, + ] + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=1, world_size=1) + executor._request_resource_termination_progress = {} + executor._terminate_request = executor._do_terminate_request + + executor._kv_connector_terminate_requests() + + assert executor._pending_connector_completions[id(request)].owner_released + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor.resource_manager.free_resources.assert_called_once_with(request) + + executor._kv_connector_terminate_requests() + + assert executor._pending_connector_completions == {} + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor.resource_manager.free_resources.assert_called_once_with(request) + assert executor._request_resource_termination_progress == {} + request.create_response.assert_called_once() + + +@pytest.mark.parametrize("force_partial_reuse", [False, True]) +def test_cancelled_context_connector_owner_terminates_exactly_once(force_partial_reuse): + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + executor = _make_cancel_executor([request]) + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True, anonymous=1) + executor.force_terminate_ctx_for_partial_reuse = force_partial_reuse + executor._terminate_request = Mock(side_effect=executor._record_terminated_transfer_request) + + executor._handle_canceled_requests() + + request.is_attention_dp_dummy = False + request.py_kv_transfer_timed_out = False + request.py_draft_tokens = [] + request.py_num_accepted_draft_tokens = 0 + request.py_per_pos_drafted = [] + request.py_per_pos_accepted = [] + request.return_perf_metrics = False + request.cached_tokens = 0 + request.is_disagg_context_complete_state = False + request.is_disagg_context_transmission_state = False + response = Mock(result=SimpleNamespace()) + request.create_response.return_value = response + executor.perf_manager = Mock() + executor.perf_manager.get_timestamp.return_value = 0 + executor.iter_counter = 1 + executor.stream_interval = 1 + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.enable_attention_dp = False + executor._maybe_attach_ctx_usage = Mock() + executor._enqueue_responses = Mock() + + finished_requests = executor._handle_responses() + assert executor.active_requests == [] + assert finished_requests == [request] + request.create_response.assert_called_once_with(False, 0) + executor._enqueue_responses.assert_called_once_with([(7, response)]) + if force_partial_reuse: + executor._terminate_request.assert_called_once_with(request) + else: + executor._terminate_request.assert_not_called() + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_called_once_with(request) + assert request.state == LlmRequestState.GENERATION_COMPLETE + assert request.is_finished_due_to_cancellation + assert executor.async_transfer_manager.end_transfer.call_args_list == [ + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER), + call(request), + ] + assert executor._terminated_transfer_requests == {} + + +@pytest.mark.parametrize( + ("state", "generation_only"), + [ + (LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, True), + (LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, False), + ], +) +def test_user_cancel_without_native_context_owner_is_unchanged(state, generation_only): + request = _make_cancel_request( + 7, + state, + generation_only=generation_only, + ) + executor = _make_cancel_executor([request]) + + executor._handle_canceled_requests() + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_not_called() + request.finish_by_reason.assert_called_once() + assert request.state == LlmRequestState.GENERATION_COMPLETE + assert executor.canceled_req_ids == [] + + +def test_successful_request_termination_records_live_transfer_owner(): + request = Mock(py_request_id=7) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.resource_manager = Mock() + executor._prefetched_request_ids = {7} + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=1, world_size=1) + executor.kv_cache_transceiver = SimpleNamespace( + requires_physical_drain_before_request_release=True + ) + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_pending_admission.return_value = False + executor.async_transfer_manager.has_any_transfer_owner.return_value = True + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor._terminated_transfer_requests = {} + + executor._do_terminate_request(request) + + executor.resource_manager.free_resources.assert_called_once_with(request) + assert executor._prefetched_request_ids == set() + assert executor._terminated_transfer_requests == {id(request): request} + + +def test_failed_request_termination_is_not_recorded_as_complete(): + request = Mock(py_request_id=7) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.resource_manager = Mock() + executor.resource_manager.free_resources.side_effect = RuntimeError("free failed") + executor._prefetched_request_ids = {7} + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=1, world_size=1) + executor.kv_cache_transceiver = SimpleNamespace( + requires_physical_drain_before_request_release=True + ) + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_pending_admission.return_value = False + executor.async_transfer_manager.has_any_transfer_owner.return_value = True + executor.async_transfer_manager.requests_in_transfer.return_value = {7: request} + executor._terminated_transfer_requests = {} + + with pytest.raises(RuntimeError, match="free failed"): + executor._do_terminate_request(request) + + assert executor._prefetched_request_ids == {7} + assert executor._terminated_transfer_requests == {} + + +def test_request_termination_retry_does_not_release_resources_twice(): + request = Mock(py_request_id=7) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.resource_manager = Mock() + executor._prefetched_request_ids = Mock() + executor._prefetched_request_ids.discard.side_effect = [ + RuntimeError("post-release bookkeeping failed"), + None, + ] + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=1, world_size=1) + executor.kv_cache_transceiver = None + executor._request_resource_termination_progress = {} + + with pytest.raises(RuntimeError, match="bookkeeping failed"): + executor._do_terminate_request(request) + + assert executor._request_resource_termination_progress == {id(request): request} + + executor._do_terminate_request(request) + + executor.resource_manager.free_resources.assert_called_once_with(request) + assert executor._request_resource_termination_progress == {} + + +def test_request_termination_rejects_pending_transfer_admission(): + request = Mock(py_request_id=7) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_pending_admission.return_value = True + executor.resource_manager = Mock() + + with pytest.raises(RuntimeError, match="transfer admission is incomplete"): + executor._do_terminate_request(request) + + executor.resource_manager.free_resources.assert_not_called() + + +def _make_shutdown_executor(events, request): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.executor_request_queue = Mock() + executor.shutdown_event = Mock() + executor.hang_detector = Mock() + executor.hang_detector.detected.return_value = False + executor.worker_thread = Mock() + executor.dist = SimpleNamespace(pp_size=1, rank=1, world_size=1) + executor._shutdown_sleep_wakeup_listeners = Mock() + executor.worker_started = True + executor.model_engine = SimpleNamespace() + executor.draft_model_engine = None + executor.kv_cache_transceiver = SimpleNamespace( + shutdown=events.transceiver_shutdown, + cancel_request=events.cancel_request, + requires_physical_drain_before_request_release=True, + ) + manager = SimpleNamespace(shutdown=events.manager_shutdown) + executor.resource_manager = SimpleNamespace(resource_managers={"test": manager}) + executor._deferred_transfer_terminations = {id(request): request} + executor._terminated_transfer_requests = {} + executor._pending_transfer_responses = [] + executor._terminate_request = events.terminate_request + executor.virtual_memory_pools = None + executor.sampler = object() + executor.dwdp_manager = None + return executor + + +def test_shutdown_discards_post_loop_native_response_without_collective(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = _make_native_completion_request() + request.is_dummy_request = False + executor = _make_shutdown_executor(events, request) + events.terminate_request.side_effect = executor._mark_transfer_terminal_teardown_complete + executor.active_requests = [request] + executor._deferred_transfer_terminations = {} + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + executor.force_terminate_ctx_for_partial_reuse = False + executor._maybe_attach_ctx_usage = Mock() + executor._enqueue_responses = Mock() + executor._ingest_native_context_completions( + [request.py_request_id], [], executor.async_transfer_manager.requests_in_transfer() + ) + events.transceiver_shutdown.return_value = True + + executor.shutdown() + + request.create_response.assert_called_once_with(False, 1) + executor._enqueue_responses.assert_not_called() + assert executor._pending_transfer_responses == [] + assert executor._pending_native_context_completions == {} + events.terminate_request.assert_called_once_with(request) + events.manager_shutdown.assert_called_once_with() + + +def test_capability_false_shutdown_does_not_call_legacy_transceiver(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + request.is_generation_only_request.return_value = False + executor = _make_shutdown_executor(events, request) + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor.dist = SimpleNamespace(pp_size=1, rank=0, world_size=2, allreduce=Mock()) + executor._deferred_transfer_terminations = {} + executor.async_transfer_manager = _ExactTransferManager(request) + + executor.shutdown() + + events.transceiver_shutdown.assert_not_called() + events.cancel_request.assert_not_called() + executor.dist.allreduce.assert_not_called() + events.manager_shutdown.assert_called_once_with() + + +def test_capability_false_shutdown_preserves_legacy_owner_agnostic_teardown(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + request.is_generation_only_request.return_value = False + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + executor.dist = SimpleNamespace(pp_size=1, rank=0, world_size=2, allreduce=Mock()) + + executor.shutdown() + + events.transceiver_shutdown.assert_not_called() + executor.async_transfer_manager.begin_shutdown.assert_not_called() + executor.async_transfer_manager.end_transfer.assert_not_called() + events.terminate_request.assert_not_called() + executor.dist.allreduce.assert_not_called() + events.manager_shutdown.assert_called_once_with() + + +def test_shutdown_waits_for_final_connector_owner_after_native_cancel(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.time.sleep", lambda _: None) + events = Mock() + request = _make_cancel_request( + 7, + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + generation_only=False, + ) + request.is_dummy_request = False + request.cached_tokens = 0 + request.create_response.return_value = Mock(result=SimpleNamespace()) + executor = _make_shutdown_executor(events, request) + events.terminate_request.side_effect = executor._mark_transfer_terminal_teardown_complete + executor.dist = SimpleNamespace(pp_size=1, rank=0, world_size=1) + executor._deferred_transfer_terminations = {} + executor.active_requests = [request] + executor.canceled_req_ids = [request.py_request_id] + executor.waiting_queue = Mock() + executor.async_transfer_manager = _ExactTransferManager( + request, transceiver=True, anonymous=1, events=events + ) + executor.kv_connector_manager = _make_one_shot_connector_manager(request) + executor._pending_connector_completions = {} + executor._pending_transfer_responses = [] + executor._maybe_attach_ctx_usage = Mock() + executor.force_terminate_ctx_for_partial_reuse = False + + events.cancel_request.return_value = False + executor._handle_canceled_requests() + assert executor._pending_request_cancellations[id(request)].request is request + + events.reset_mock() + events.transceiver_shutdown.return_value = True + events.cancel_request.return_value = True + + executor.shutdown() + + assert events.mock_calls == [ + call.transceiver_shutdown(), + call.cancel_request(request), + call.transceiver_owner_release(request), + call.connector_owner_release(request), + call.terminate_request(request), + call.manager_shutdown(), + ] + assert executor._pending_request_cancellations == {} + assert executor._pending_connector_completions == {} + assert executor._terminated_transfer_requests == {} + + +def test_shutdown_releases_deferred_request_only_after_transceiver_drain(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7) + executor = _make_shutdown_executor(events, request) + events.transceiver_shutdown.return_value = True + events.cancel_request.return_value = True + + executor.shutdown() + + assert events.mock_calls == [ + call.transceiver_shutdown(), + call.cancel_request(request), + call.terminate_request(request), + call.manager_shutdown(), + ] + assert executor._deferred_transfer_terminations == {} + + +def test_non_drained_shutdown_preserves_deferred_request_and_managers(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7) + executor = _make_shutdown_executor(events, request) + events.transceiver_shutdown.return_value = False + + with pytest.raises(RuntimeError, match="still owns active transfer targets"): + executor.shutdown() + + assert events.mock_calls == [call.transceiver_shutdown()] + assert executor._deferred_transfer_terminations == {id(request): request} + + +def test_in_doubt_resource_release_vetoes_manager_shutdown(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + executor.resource_manager.has_in_doubt_resource_releases = Mock(return_value=True) + + with pytest.raises(RuntimeError, match="release outcome is in doubt"): + executor.shutdown() + + events.manager_shutdown.assert_not_called() + + +def test_pending_resource_release_vetoes_manager_shutdown(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + executor.resource_manager.has_pending_resource_releases = Mock(return_value=True) + + with pytest.raises(RuntimeError, match="release is still pending"): + executor.shutdown() + + events.manager_shutdown.assert_not_called() + + +def test_shutdown_retires_native_owner_but_vetoes_connector_owner(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7) + request.is_generation_only_request.return_value = False + executor = _make_shutdown_executor(events, request) + events.transceiver_shutdown.return_value = True + events.cancel_request.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True, anonymous=1) + executor.force_terminate_ctx_for_partial_reuse = False + + with pytest.raises(RuntimeError, match="Asynchronous transfer ownership is still active"): + executor.shutdown() + + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + executor.async_transfer_manager.has_any_inflight_requests.assert_called_once_with() + assert executor._deferred_transfer_terminations == {} + events.terminate_request.assert_not_called() + events.manager_shutdown.assert_not_called() + + +def test_shutdown_drain_polls_connector_after_initial_empty_sample(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + sleep = Mock() + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.time.sleep", sleep) + events = Mock() + request = SimpleNamespace( + py_request_id=7, + is_finished_due_to_cancellation=False, + state=LlmRequestState.GENERATION_COMPLETE, + ) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + executor.active_requests = [] + executor.force_terminate_ctx_for_partial_reuse = False + events.transceiver_shutdown.return_value = True + executor.kv_connector_manager = _make_one_shot_connector_manager(request, empty_polls=1) + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + + executor.shutdown() + + assert ( + executor.kv_connector_manager.get_finished.call_count + == _SHUTDOWN_CONNECTOR_COMPLETION_POLLS + ) + sleep.assert_called_once_with(_SHUTDOWN_CONNECTOR_COMPLETION_POLL_INTERVAL_S) + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor.async_transfer_manager.has_any_inflight_requests.assert_called_with() + assert executor.async_transfer_manager.has_any_inflight_requests.call_count > 1 + events.terminate_request.assert_called_once_with(request) + events.manager_shutdown.assert_called_once_with() + + +def test_shutdown_connector_poll_error_finishes_uniform_drain_without_replay(): + request = SimpleNamespace( + py_request_id=7, + is_finished_due_to_cancellation=False, + state=LlmRequestState.GENERATION_COMPLETE, + ) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.dist = SimpleNamespace(world_size=2, allreduce=Mock(return_value=1)) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + executor._terminate_request_after_worker_shutdown = Mock() + executor._terminated_transfer_requests = {} + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.side_effect = [ + [request], + RuntimeError("local connector poll failed"), + ] + [[]] * (_SHUTDOWN_CONNECTOR_COMPLETION_POLLS - 2) + + with pytest.raises(RuntimeError, match="local connector poll failed"): + executor._drain_connector_completions_after_shutdown() + + assert ( + executor.kv_connector_manager.get_finished.call_count + == _SHUTDOWN_CONNECTOR_COMPLETION_POLLS + ) + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor._terminate_request_after_worker_shutdown.assert_called_once_with(request) + executor.dist.allreduce.assert_called_once() + assert executor.dist.allreduce.call_args.args[0] == 1 + + +def test_shutdown_connector_peer_error_raises_before_later_collectives(): + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.dist = SimpleNamespace(world_size=2, allreduce=Mock(return_value=1)) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = None + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.return_value = [] + + with pytest.raises(RuntimeError, match="another rank"): + executor._drain_connector_completions_after_shutdown() + + assert ( + executor.kv_connector_manager.get_finished.call_count + == _SHUTDOWN_CONNECTOR_COMPLETION_POLLS + ) + executor.dist.allreduce.assert_called_once() + assert executor.dist.allreduce.call_args.args[0] == 0 + + +@pytest.mark.parametrize("local_pending", [True, False]) +def test_shutdown_final_readiness_vote_precedes_manager_teardown(monkeypatch, local_pending): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request) + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.return_value = [] + executor.dist = SimpleNamespace( + pp_size=1, + rank=0, + world_size=2, + allreduce=Mock(side_effect=[0, 1]), + ) + executor.resource_manager.has_in_doubt_resource_releases = Mock(return_value=False) + executor.resource_manager.has_pending_resource_releases = Mock(return_value=local_pending) + + error = "release is still pending" if local_pending else "another rank" + with pytest.raises(RuntimeError, match=error): + executor.shutdown() + + assert [entry.args[0] for entry in executor.dist.allreduce.call_args_list] == [ + 0, + int(local_pending), + ] + events.manager_shutdown.assert_not_called() + + +@pytest.mark.parametrize("local_manager_failure", [True, False]) +def test_shutdown_post_manager_vote_precedes_engine_deletion(monkeypatch, local_manager_failure): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request) + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.return_value = [] + executor.dist = SimpleNamespace( + pp_size=1, + rank=0, + world_size=2, + allreduce=Mock(side_effect=[0, 0, 1]), + ) + if local_manager_failure: + events.manager_shutdown.side_effect = RuntimeError("local manager shutdown failed") + + error = "local manager shutdown failed" if local_manager_failure else "another rank" + with pytest.raises(RuntimeError, match=error): + executor.shutdown() + + assert [entry.args[0] for entry in executor.dist.allreduce.call_args_list] == [ + 0, + 0, + int(local_manager_failure), + ] + events.manager_shutdown.assert_called_once_with() + assert hasattr(executor, "model_engine") + progress = executor._resource_manager_shutdown_progress + assert len(progress) == 1 + assert progress[0].in_doubt is local_manager_failure + assert progress[0].completed is not local_manager_failure + + +def test_ambiguous_manager_shutdown_is_not_replayed(): + manager = Mock() + manager.shutdown.side_effect = RuntimeError("manager shutdown failed") + executor = object.__new__(PyExecutor) + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.resource_manager = SimpleNamespace(resource_managers={"test": manager}) + + for _attempt in range(2): + with pytest.raises(RuntimeError, match="manager shutdown"): + executor._shutdown_resource_managers_rank_uniform() + + manager.shutdown.assert_called_once_with() + + +def test_ambiguous_manager_shutdown_blocks_later_dependencies() -> None: + first_manager = Mock() + first_manager.shutdown.side_effect = RuntimeError("manager shutdown failed") + later_manager = Mock() + executor = object.__new__(PyExecutor) + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.resource_manager = SimpleNamespace( + resource_managers={ + "dependent": first_manager, + "kv_cache": later_manager, + } + ) + + for _attempt in range(2): + with pytest.raises(RuntimeError, match="manager shutdown"): + executor._shutdown_resource_managers_rank_uniform() + + first_manager.shutdown.assert_called_once_with() + later_manager.shutdown.assert_not_called() + + +def test_shutdown_discards_ownerless_unbuffered_terminal_before_teardown( + monkeypatch, +): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = SimpleNamespace(py_request_id=7, is_dummy_request=False) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request) + executor.active_requests = [request] + executor._claim_transfer_terminal( + request, _TransferTerminalOutcome.FAILED, "native_transceiver" + ) + events.terminate_request.side_effect = ( + lambda owner: executor._mark_transfer_terminal_teardown_complete(owner) + ) + + executor.shutdown() + + events.terminate_request.assert_called_once_with(request) + assert executor.active_requests == [] + assert executor._pending_transfer_terminals == {} + events.manager_shutdown.assert_called_once_with() + + +def test_shutdown_discard_retires_already_torn_down_terminal() -> None: + request = SimpleNamespace(py_request_id=7) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) + executor.async_transfer_manager = None + terminal = executor._claim_transfer_terminal( + request, _TransferTerminalOutcome.FAILED, "native_transceiver" + ) + terminal.teardown_complete = True + + executor._discard_unpublishable_transfer_terminals_after_shutdown() + + assert terminal.response_discarded + assert executor._pending_transfer_terminals == {} + + +def test_shutdown_native_failure_still_enters_connector_consensus(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + executor.async_transfer_manager = _ExactTransferManager(request) + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.return_value = [] + executor.dist = SimpleNamespace( + pp_size=1, + rank=0, + world_size=2, + allreduce=Mock(return_value=1), + ) + events.transceiver_shutdown.return_value = False + + with pytest.raises(RuntimeError, match="still owns active transfer targets"): + executor.shutdown() + + assert ( + executor.kv_connector_manager.get_finished.call_count + == _SHUTDOWN_CONNECTOR_COMPLETION_POLLS + ) + executor.dist.allreduce.assert_called_once() + assert executor.dist.allreduce.call_args.args[0] == 1 + events.manager_shutdown.assert_not_called() + + +@pytest.mark.parametrize("failure_site", ["cuda_graph", "cuda_synchronize"]) +def test_shutdown_local_cleanup_failure_still_enters_connector_consensus(monkeypatch, failure_site): + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request) + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.return_value = [] + executor.dist = SimpleNamespace( + pp_size=1, + rank=0, + world_size=2, + allreduce=Mock(return_value=1), + ) + + if failure_site == "cuda_graph": + executor.model_engine = SimpleNamespace( + _release_cuda_graphs=Mock(side_effect=RuntimeError("CUDA graph cleanup failed")) + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + expected_error = "CUDA graph cleanup failed" + else: + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: True, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.synchronize", + Mock(side_effect=RuntimeError("CUDA synchronize failed")), + ) + expected_error = "CUDA synchronize failed" + + with pytest.raises(RuntimeError, match=expected_error): + executor.shutdown() + + assert ( + executor.kv_connector_manager.get_finished.call_count + == _SHUTDOWN_CONNECTOR_COMPLETION_POLLS + ) + executor.dist.allreduce.assert_called_once() + assert executor.dist.allreduce.call_args.args[0] == 1 + events.manager_shutdown.assert_not_called() + + +def test_shutdown_retries_and_retires_unpolled_native_owner(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + request.is_generation_only_request.return_value = False + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.side_effect = [False, True] + events.cancel_request.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + executor.force_terminate_ctx_for_partial_reuse = False + + with pytest.raises(RuntimeError, match="still owns active transfer targets"): + executor.shutdown() + + assert executor._deferred_transfer_terminations == {id(request): request} + executor.async_transfer_manager.end_transfer.assert_not_called() + events.manager_shutdown.assert_not_called() + + executor.shutdown() + + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + events.terminate_request.assert_called_once_with(request) + events.manager_shutdown.assert_called_once_with() + assert executor._deferred_transfer_terminations == {} + + +def test_shutdown_does_not_reterminate_partial_reuse_request(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + request.is_generation_only_request.return_value = False + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + executor.force_terminate_ctx_for_partial_reuse = True + executor._terminated_transfer_requests = {id(request): request} + events.transceiver_shutdown.return_value = True + events.cancel_request.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + + executor.shutdown() + + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER + ) + events.terminate_request.assert_not_called() + events.manager_shutdown.assert_called_once_with() + assert executor._deferred_transfer_terminations == {} + assert executor._deferred_transfer_terminations_already_terminated == set() + assert executor._terminated_transfer_requests == {} + + +def test_shutdown_finalizes_pp_request_after_executor_loop_stops(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + request.is_generation_only_request.return_value = False + executor = _make_shutdown_executor(events, request) + executor._deferred_transfer_terminations = {} + executor._terminate_request = Mock() + events.transceiver_shutdown.return_value = True + events.cancel_request.return_value = True + executor.async_transfer_manager = _ExactTransferManager(request, transceiver=True) + executor.force_terminate_ctx_for_partial_reuse = False + termination_handler = object.__new__(DisaggPPTerminationHandler) + termination_handler._terminator_func = events.terminate_request + termination_handler._pending_termination = {7: request} + executor._disagg_pp_termination_handler = termination_handler + + executor.shutdown() + + executor._terminate_request.assert_not_called() + events.terminate_request.assert_called_once_with(request) + assert termination_handler._pending_termination == {} + events.manager_shutdown.assert_called_once_with() + + +def test_shutdown_finalizes_pending_only_pp_request(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + events = Mock() + request = Mock(py_request_id=7, is_dummy_request=False) + executor = _make_shutdown_executor(events, request) + executor.dist = SimpleNamespace(pp_size=2) + executor.executed_batch_queue = Mock() + executor.broadcast_sample_state_handler = Mock() + executor._deferred_transfer_terminations = {} + events.transceiver_shutdown.return_value = True + termination_handler = DisaggPPTerminationHandler(executor.dist, events.terminate_request) + termination_handler._pending_termination = {7: request} + termination_handler._send_handle = Mock() + executor._disagg_pp_termination_handler = termination_handler + + executor.shutdown() + + assert events.mock_calls == [ + call.transceiver_shutdown(), + call.terminate_request(request), + call.manager_shutdown(), + ] + assert termination_handler._pending_termination == {} + termination_handler._send_handle.wait.assert_not_called() + + +def test_pending_pp_termination_failure_retains_exact_request(): + request = Mock(py_request_id=7) + terminator = Mock(side_effect=[RuntimeError("release failed"), None]) + handler = DisaggPPTerminationHandler(SimpleNamespace(), terminator) + handler._pending_termination = {7: request} + + with pytest.raises(RuntimeError, match="request 7: release failed"): + handler.terminate_all_after_shutdown() + + assert handler._pending_termination == {7: request} + handler.terminate_all_after_shutdown() + assert handler._pending_termination == {} + assert terminator.call_count == 2 + + +def test_direct_post_shutdown_termination_failure_keeps_pending_owner(): + request = Mock(py_request_id=7) + terminator = Mock(side_effect=RuntimeError("release failed")) + handler = DisaggPPTerminationHandler(SimpleNamespace(), terminator) + handler._pending_termination = {7: request} + + with pytest.raises(RuntimeError, match="release failed"): + handler.terminate_after_shutdown(request) + + assert handler._pending_termination == {7: request} diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 8faadb0147b4..1808ba934625 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -19,6 +19,7 @@ """ import queue +import threading from types import SimpleNamespace import numpy as np @@ -26,15 +27,23 @@ from tensorrt_llm._torch.disaggregation.native.bounce import config as bcfg -# core (the BounceTransport contract + the TransferContext state machine) is PURE (no CUDA / NIXL +# core (the BounceTransport contract + RecvBounceContext state machine) is PURE (no CUDA / NIXL # imports), so it is always importable on CPU. from tensorrt_llm._torch.disaggregation.native.bounce import core as bcore +from tensorrt_llm._torch.disaggregation.native.receive_lifecycle import ( + LifecycleAction, + PhysicalState, + RecvTransferRegistry, + WriterMode, + WriterResult, +) # transport/transfer pull CUDA-binding + torch deps at import; skip gracefully when those are # absent (CPU-only env). Catch only ImportError so a genuine bug in the module still fails CI # instead of being silently turned into a skip. try: from tensorrt_llm._torch.disaggregation.native.bounce import buffer as bbuf + from tensorrt_llm._torch.disaggregation.native.bounce import gather_scatter as bgs from tensorrt_llm._torch.disaggregation.native.bounce import impl as btr _HAVE_TRANSPORT = True @@ -44,6 +53,13 @@ _MIB = 1024 * 1024 +def test_transfer_context_compatibility_alias(): + from tensorrt_llm._torch.disaggregation.native import bounce + + assert bounce.TransferContext is bounce.RecvBounceContext + assert "TransferContext" in bounce.__all__ + + # --------------------------------------------------------------------------- # # config.py — sizing + OOM guard (pure math, always runnable) # --------------------------------------------------------------------------- # @@ -186,36 +202,50 @@ def test_make_kv_result_msg_uses_binary_frame(result_name): # --------------------------------------------------------------------------- # -# fan-in safety gate — equal total//num_writers split only for uniform TP-by-head +# fan-in safety gate — equal split only within one PP stage # --------------------------------------------------------------------------- # def test_fanin_bounce_safe_gate(): - """Restrict multi-writer equal-split bounce to uniform TP-by-head. - - PP (overlap_pp_size>1 -> unequal per-writer sizes) and duplicate_head_factor>1 - (MLA / duplicate TP heads -> some ranks don't send KV yet count in - expected_transfers) must fall back to the per-fragment path. - """ + """Require one PP stage and no TP head duplication in either direction.""" tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") safe = tfr.Receiver._fanin_bounce_safe - def ov(dup, pp): - return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp) + def ov(dup, pp, peer_dup=1): + return SimpleNamespace( + duplicate_head_factor=dup, + peer_duplicate_head_factor=peer_dup, + overlap_pp_size=pp, + ) - def ri(lpp): - return SimpleNamespace(layer_num_per_pp=lpp) + # A single PP stage is safe only when neither peer duplicates KV heads. + assert safe(ov(1, 1)) is True + assert safe(ov(1, 0)) is False + assert safe(ov(2, 1)) is False # duplicate heads / MLA -> some don't send + assert safe(ov(1, 1, peer_dup=2)) is False # reciprocal duplication + # Every PP fan-in falls back regardless of peer-stage metadata. + assert safe(ov(1, 4)) is False + assert safe(ov(2, 4)) is False - # single PP stage (overlap_pp_size <= 1): only duplicate_head_factor matters - assert safe(ov(1, 1), ri([24])) is True - assert safe(ov(1, 0), ri([24])) is True - assert safe(ov(2, 1), ri([24])) is False # duplicate heads / MLA -> some don't send - # EVEN PP fan-in (equal layers per overlapping stage) -> allowed - assert safe(ov(1, 4), ri([20, 20, 20, 20])) is True - # UNEVEN PP fan-in -> per-writer sizes differ -> fall back - assert safe(ov(1, 4), ri([20, 20, 20, 19])) is False - # incomplete per-stage info (single element for a multi-stage fan-in) -> conservative fall back - assert safe(ov(1, 4), ri([20])) is False - # duplicate heads blocks even an otherwise-even PP split - assert safe(ov(2, 4), ri([20, 20, 20, 20])) is False + +def test_fanin_bounce_rejects_uniform_peer_pp_with_mismatched_boundaries(): + """Uniform peer stages can intersect one local PP stage by unequal extents.""" + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + + # Local PP rank 0 owns layers [0, 30), while globally uniform peer stages + # own [0, 20) and [20, 40). Their local intersections are 20 and 10 layers, + # so an equal two-writer bounce split would authorize the wrong boundary. + local_start, local_end = 0, 30 + peer_boundaries = ((0, 20), (20, 40)) + intersections = [ + max(0, min(local_end, end) - max(local_start, start)) for start, end in peer_boundaries + ] + assert intersections == [20, 10] + + overlap = SimpleNamespace( + duplicate_head_factor=1, + peer_duplicate_head_factor=1, + overlap_pp_size=2, + ) + assert tfr.Receiver._fanin_bounce_safe(overlap) is False # --------------------------------------------------------------------------- # @@ -225,15 +255,32 @@ def ri(lpp): class TestNoBounce: def test_noop_behaviour(self): nb = btr.NoBounceTransport() + factory_calls = [] assert nb.enabled is False - assert nb.reserve(SimpleNamespace()) is False + assert ( + nb.reserve( + SimpleNamespace(), + destination_intervals_factory=lambda: factory_calls.append(True) or [], + ) + is False + ) + assert factory_calls == [] assert nb.build_request(SimpleNamespace()) is None assert nb.writer_base(("r", 0), 1) is None assert nb.is_bounced(("r", 0)) is False + assert nb.mark_writer_exposed(("r", 0), 1) is False + nb.record_no_access(("r", 0), 1) + nb.mark_logical_failure(("r", 0)) + nb.mark_protocol_conflict(("r", 0)) + nb.mark_backend_quiesced() + assert nb.retry_settlements() + assert nb.retry_settlements("r") + assert nb.retry_settlements(("r", 0)) nb.record_result(("r", 0), 1) # no-op, must not raise nb.record_failure(("r", 0), 1) # no-op, must not raise nb.release_idle_reservation(("r", 0)) # no-op, must not raise nb.release_send(0) + nb.quarantine_send(0) nb.close() def test_create_bounce_none_cfg(self): @@ -249,6 +296,33 @@ def test_implements_bounce_transport_abc(self): assert isinstance(btr.NoBounceTransport(), bcore.BounceTransport) +# --------------------------------------------------------------------------- # +# page-table sizing — every attention pool view contributes its gathered extent +# --------------------------------------------------------------------------- # +@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") +def test_block_bytes_per_group_sums_every_physical_pool_view(): + from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup + + layer_group = AttentionLayerGroup( + pool_group_idx=0, + pool_views=[ + SimpleNamespace(pool_idx=0), + SimpleNamespace(pool_idx=1), + # Sender descriptor construction iterates views, so a repeated + # physical pool must contribute its full gathered extent again. + SimpleNamespace(pool_idx=0), + ], + ) + page_table = SimpleNamespace( + layer_groups=[layer_group], + pool_groups=[ + SimpleNamespace(pools=[SimpleNamespace(slot_bytes=96), SimpleNamespace(slot_bytes=32)]) + ], + ) + + assert btr.block_bytes_per_group(page_table) == [224] + + # --------------------------------------------------------------------------- # # Transport TP fan-in logic — reserve gate / writer_base / accumulate ordering # (GPU allocators, streams and the scatter worker are mocked out) @@ -262,6 +336,8 @@ def __init__(self, capacity_bytes, phys_chunk_size, name="kv_bounce"): self.next_id = 0 self.released = [] self.quarantined = [] + self.active = set() + self.closed = False @property def capacity(self): @@ -272,30 +348,55 @@ def reserve(self, size, timeout=None): return None sid = self.next_id self.next_id += 1 + self.active.add(sid) return sid, self.base def release(self, slot_id): + self.active.discard(slot_id) self.released.append(slot_id) - def quarantine(self, slot_id, grace_s): + def quarantine(self, slot_id): + self.active.discard(slot_id) self.quarantined.append(slot_id) - def reclaim_expired(self): - return 0 + @property + def has_outstanding(self): + return bool(self.active or self.quarantined) def reg_descs(self): return [] + def close(self): + if self.has_outstanding: + raise RuntimeError("live fake slots") + self.closed = True + + +class _DescriptorAlloc(_FakeAlloc): + """Fake allocator with a unique registration descriptor.""" + + def __init__(self, capacity_bytes, phys_chunk_size, name="kv_bounce"): + super().__init__(capacity_bytes, phys_chunk_size, name=name) + self.desc = object() + + def reg_descs(self): + return self.desc + def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_blocks=1): monkeypatch.setattr(btr, "SlotAllocator", _FakeAlloc) monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: 0) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_destroy_stream", + lambda self, attr_name: setattr(self, attr_name, None), + ) monkeypatch.setattr( btr.VmmBounceTransport, "_start_scatter_worker", lambda self, name: setattr(self, "_scatter_q", queue.Queue()), ) - agent = SimpleNamespace(register_memory=lambda d: None) + agent = SimpleNamespace(register_memory=lambda d: None, deregister_memory=lambda d: None) return btr.VmmBounceTransport( agent, device_id=0, @@ -308,147 +409,865 @@ def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bl def _recv_req(block_counts, rid=1, slice_id=0): return SimpleNamespace( - block_ids_per_layer_groups=[SimpleNamespace(size=n) for n in block_counts], + block_ids_per_layer_groups=[np.arange(n, dtype=np.int64) for n in block_counts], unique_rid=rid, slice_id=slice_id, bounce_dst_base=None, ) +def _write_meta(): + return SimpleNamespace( + src_ptrs=np.array([0x1000], dtype=np.int64), + dst_ptrs=np.array([0x2000], dtype=np.int64), + sizes=np.array([8], dtype=np.int64), + bounce_dst_base=0x3000, + dst_device_id=0, + peer_name="peer", + ) + + +def _run_mock_scatter_worker(transport, monkeypatch) -> None: + transport._scatter_ready = threading.Event() + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace( + cudaSetDevice=lambda device_id: ("ok",), + cudaStreamSynchronize=lambda stream: ("ok",), + ), + ) + monkeypatch.setattr(btr, "CUASSERT", lambda result: result[1:]) + monkeypatch.setattr(btr, "scatter_contiguous", lambda *args, **kwargs: None) + transport._scatter_q.put(None) + transport._scatter_loop() + + @pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") -class TestFanInReserve: - def test_reserve_stamps_base_and_per_writer(self, monkeypatch): +class TestTransportRollback: + def _construct(self, agent, *, capacity=1024): + return btr.VmmBounceTransport( + agent, + device_id=0, + capacity_bytes=capacity, + phys_chunk_size=512, + block_bytes_per_group=[8], + min_blocks=1, + ) + + def test_constructor_rolls_back_first_allocator_when_second_fails(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + if allocators: + raise RuntimeError("recv allocator failed") + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + with pytest.raises(RuntimeError, match="recv allocator failed"): + self._construct(SimpleNamespace()) + assert allocators[0].closed + + def test_constructor_rolls_back_partial_registration(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + deregistered = [] + + class Agent: + def __init__(self): + self.register_count = 0 + + def register_memory(self, desc): + self.register_count += 1 + if self.register_count == 2: + raise RuntimeError("second registration failed") + + def deregister_memory(self, desc): + deregistered.append(desc) + + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + with pytest.raises(RuntimeError, match="second registration failed"): + self._construct(Agent()) + # The failing registration may have mutated the backend before raising, + # so rollback must conservatively deregister both attempted descriptors. + assert deregistered == [allocator.desc for allocator in allocators] + assert all(allocator.closed for allocator in allocators) + + def test_incomplete_constructor_rollback_retains_retry_owner(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + allow_deregister = False + + class Agent: + def __init__(self): + self.register_count = 0 + + def register_memory(self, desc): + self.register_count += 1 + if self.register_count == 2: + raise RuntimeError("second registration failed") + + def deregister_memory(self, desc): + if not allow_deregister: + raise RuntimeError("rollback deregistration failed") + + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + with pytest.raises(btr.IncompleteBounceInitializationError) as raised: + self._construct(Agent()) + + owner = raised.value.owner + assert owner in btr._INCOMPLETE_TRANSPORTS + assert not any(allocator.closed for allocator in allocators) + + allow_deregister = True + owner.retry_initialization_rollback() + assert owner not in btr._INCOMPLETE_TRANSPORTS + assert all(allocator.closed for allocator in allocators) + + def test_cleanup_pending_transport_retries_retained_owner(self): + calls = [] + owner = SimpleNamespace(retry_initialization_rollback=lambda: calls.append(True)) + transport = btr.CleanupPendingBounceTransport(owner) + + transport.close() + transport.close() + + assert calls == [True] + + def test_create_bounce_preserves_incomplete_cleanup_owner(self, monkeypatch): + calls = [] + owner = SimpleNamespace(retry_initialization_rollback=lambda: calls.append(True)) + error = btr.IncompleteBounceInitializationError(owner, RuntimeError("setup failed")) + monkeypatch.setattr(btr, "block_bytes_per_group", lambda page_table: []) + monkeypatch.setattr( + btr.VmmBounceTransport, + "from_config", + lambda *args, **kwargs: (_ for _ in ()).throw(error), + ) + + transport = btr.create_bounce(object(), object(), device_id=0, page_table=object()) + assert isinstance(transport, btr.CleanupPendingBounceTransport) + transport.close() + assert calls == [True] + + def test_constructor_rolls_back_registrations_when_send_stream_fails(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + deregistered = [] + agent = SimpleNamespace( + register_memory=lambda desc: None, + deregister_memory=lambda desc: deregistered.append(desc), + ) + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_new_stream", + lambda self: (_ for _ in ()).throw(RuntimeError("stream failed")), + ) + with pytest.raises(RuntimeError, match="stream failed"): + self._construct(agent) + assert deregistered == [allocator.desc for allocator in allocators] + assert all(allocator.closed for allocator in allocators) + + def test_constructor_rolls_back_streams_when_thread_start_fails(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + class FailingThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + raise RuntimeError("thread start failed") + + def is_alive(self): + return False + + streams = iter((11, 22)) + destroyed = [] + + def destroy_stream(transport, attr_name): + stream = getattr(transport, attr_name) + if stream is not None: + destroyed.append(stream) + setattr(transport, attr_name, None) + + deregistered = [] + agent = SimpleNamespace( + register_memory=lambda desc: None, + deregister_memory=lambda desc: deregistered.append(desc), + ) + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: next(streams)) + monkeypatch.setattr(btr.VmmBounceTransport, "_destroy_stream", destroy_stream) + monkeypatch.setattr(btr.threading, "Thread", FailingThread) + with pytest.raises(RuntimeError, match="thread start failed"): + self._construct(agent) + assert destroyed == [22, 11] + assert deregistered == [allocator.desc for allocator in allocators] + assert all(allocator.closed for allocator in allocators) + + def test_constructor_rolls_back_when_scatter_worker_device_setup_fails(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + streams = iter((11, 22)) + destroyed = [] + deregistered = [] + agent = SimpleNamespace( + register_memory=lambda desc: None, + deregister_memory=lambda desc: deregistered.append(desc), + ) + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: next(streams)) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_destroy_stream", + lambda owner, attr_name: ( + destroyed.append(getattr(owner, attr_name)), + setattr(owner, attr_name, None), + ), + ) + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace(cudaSetDevice=lambda device_id: ("error",)), + ) + monkeypatch.setattr( + btr, + "CUASSERT", + lambda result: (_ for _ in ()).throw(RuntimeError("set device failed")), + ) + + with pytest.raises(RuntimeError, match="scatter worker failed to initialize"): + self._construct(agent) + + assert destroyed == [22, 11] + assert deregistered == [allocator.desc for allocator in allocators] + assert all(allocator.closed for allocator in allocators) + + @pytest.mark.parametrize("failure_stage", ["launch", "wait", "make_write"]) + def test_send_slot_rolls_back_on_build_failure(self, monkeypatch, failure_stage): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + monkeypatch.setattr( + transport, + "_rollback_send_slot", + lambda slot_id: (transport._send_alloc.release(slot_id), True)[1], + ) + if failure_stage == "launch": + monkeypatch.setattr( + transport, + "_launch_gather", + lambda *args: (_ for _ in ()).throw(RuntimeError("launch failed")), + ) + else: + monkeypatch.setattr(transport, "_launch_gather", lambda *args: 17) + if failure_stage == "wait": + monkeypatch.setattr( + transport, + "_wait_gather", + lambda *args: (_ for _ in ()).throw(RuntimeError("wait failed")), + ) + else: + monkeypatch.setattr(transport, "_wait_gather", lambda *args: None) + monkeypatch.setattr( + transport, + "_make_write", + lambda *args: (_ for _ in ()).throw(RuntimeError("make write failed")), + ) + + with pytest.raises(RuntimeError): + transport.build_request(_write_meta()) + assert transport._send_alloc.released == [0] + assert not transport._send_alloc.has_outstanding + + def test_failed_send_slot_releases_after_stream_fence(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + slot_id, _ = transport._send_alloc.reserve(8) + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace(cudaStreamSynchronize=lambda stream: ("ok",)), + ) + monkeypatch.setattr(btr, "CUASSERT", lambda result: result[1:]) + + assert transport._rollback_send_slot(slot_id) is True + + assert transport._send_alloc.released == [slot_id] + assert transport._send_alloc.quarantined == [] + + def test_failed_send_slot_quarantines_when_stream_fence_fails(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + slot_id, _ = transport._send_alloc.reserve(8) + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace(cudaStreamSynchronize=lambda stream: ("error",)), + ) + + def cuassert(result): + raise RuntimeError("stream fence failed") + + monkeypatch.setattr(btr, "CUASSERT", cuassert) + + assert transport._rollback_send_slot(slot_id) is False + + assert transport._send_alloc.released == [] + assert transport._send_alloc.quarantined == [slot_id] + monkeypatch.setattr( + transport, + "_launch_gather", + lambda *args: pytest.fail("an unhealthy send stream must not launch more gathers"), + ) + assert transport.build_request(_write_meta()) is None + + def test_build_raises_distinct_error_when_gather_source_remains_in_doubt(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + monkeypatch.setattr( + transport, + "_launch_gather", + lambda *args: (_ for _ in ()).throw(RuntimeError("launch failed")), + ) + monkeypatch.setattr(transport, "_rollback_send_slot", lambda slot_id: False) + + with pytest.raises(bcore.GatherSourceInDoubtError, match="without a positive CUDA fence"): + transport.build_request(_write_meta()) + + def test_quarantine_failure_still_reports_gather_source_in_doubt(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + monkeypatch.setattr( + transport, + "_launch_gather", + lambda *args: (_ for _ in ()).throw(RuntimeError("launch failed")), + ) + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace(cudaStreamSynchronize=lambda stream: ("error",)), + ) + monkeypatch.setattr( + btr, + "CUASSERT", + lambda result: (_ for _ in ()).throw(RuntimeError("stream fence failed")), + ) + monkeypatch.setattr( + transport._send_alloc, + "quarantine", + lambda slot_id: (_ for _ in ()).throw(RuntimeError("quarantine failed")), + ) + + with pytest.raises(bcore.GatherSourceInDoubtError, match="without a positive CUDA fence"): + transport.build_request(_write_meta()) + + assert not transport._send_stream_healthy + assert transport._send_alloc.active == {0} + + def test_ambiguous_nixl_slot_is_quarantined_from_reuse(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + slot_id, _ = transport._send_alloc.reserve(8) + + transport.quarantine_send(slot_id) + transport.release_send(slot_id) + + assert transport._send_alloc.released == [slot_id] + assert transport._send_alloc.quarantined == [slot_id] + assert transport._send_alloc.has_outstanding + with pytest.raises(RuntimeError, match="outstanding arena slots"): + transport.close() + + def test_event_record_failure_destroys_created_event(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + destroyed = [] + monkeypatch.setattr(btr, "gather_contiguous", lambda *args, **kwargs: None) + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace( + cudaEventCreate=lambda: ("ok", 17), + cudaEventRecord=lambda event, stream: ("error",), + cudaEventDestroy=lambda event: (destroyed.append(event), ("ok",))[1], + ), + ) + + def cuassert(result): + if result[0] == "error": + raise RuntimeError("event record failed") + return result[1:] + + monkeypatch.setattr(btr, "CUASSERT", cuassert) + with pytest.raises(RuntimeError, match="event record failed"): + transport._launch_gather(0x1000, _write_meta(), 8) + assert destroyed == [17] + + def test_event_wait_failure_still_destroys_event(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + transport._pending_events.append(17) + destroyed = [] + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace( + cudaEventSynchronize=lambda event: ("error",), + cudaEventDestroy=lambda event: (destroyed.append(event), ("ok",))[1], + ), + ) + + def cuassert(result): + if result[0] == "error": + raise RuntimeError("event wait failed") + return result[1:] + + monkeypatch.setattr(btr, "CUASSERT", cuassert) + with pytest.raises(RuntimeError, match="event wait failed"): + transport._wait_gather(17) + assert destroyed == [17] + + def test_close_retries_failed_event_destruction(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + transport._pending_events.append(17) + fail_once = {"value": True} + destroyed = [] + + def destroy_event(event): + destroyed.append(event) + if fail_once["value"]: + fail_once["value"] = False + return ("error",) + return ("ok",) + + monkeypatch.setattr(btr, "cudart", SimpleNamespace(cudaEventDestroy=destroy_event)) + + def cuassert(result): + if result[0] == "error": + raise RuntimeError("event destroy failed") + return result[1:] + + monkeypatch.setattr(btr, "CUASSERT", cuassert) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_destroy_stream", + lambda owner, attr_name: setattr(owner, attr_name, None), + ) + + with pytest.raises(RuntimeError, match="failed to destroy 1 CUDA event"): + transport.close() + assert transport._pending_events == [17] + assert not transport._send_alloc.closed + + transport.close() + assert transport._pending_events == [] + assert destroyed == [17, 17] + + def test_destroy_stream_evicts_only_its_device_metadata(self, monkeypatch): + transport = object.__new__(btr.VmmBounceTransport) + transport._device_id = 1 + transport._send_stream = 17 + operations = [] + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace( + cudaStreamDestroy=lambda stream: operations.append(("destroy", stream)) or ("ok",) + ), + ) + monkeypatch.setattr(btr, "CUASSERT", lambda result: result[1:]) + monkeypatch.setattr( + btr, + "release_meta_buffers", + lambda stream, device_id: operations.append(("release_metadata", stream, device_id)), + ) + + transport._destroy_stream("_send_stream") + + assert transport._send_stream is None + assert operations == [("release_metadata", 17, 1), ("destroy", 17)] + + def test_metadata_cache_keys_include_device_and_can_be_evicted(self): + dev0 = SimpleNamespace(index=0) + dev1 = SimpleNamespace(index=1) + assert bgs._meta_buffer_key(17, dev0) != bgs._meta_buffer_key(17, dev1) + bgs._meta_buffers[(0, 17)] = (object(), object(), 1) + bgs._meta_buffers[(1, 17)] = (object(), object(), 1) + try: + bgs.release_meta_buffers(17, 0) + assert (0, 17) not in bgs._meta_buffers + assert (1, 17) in bgs._meta_buffers + finally: + bgs._meta_buffers.pop((0, 17), None) + bgs._meta_buffers.pop((1, 17), None) + + def test_gathers_hold_stream_metadata_until_completion(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + first_waiting = threading.Event() + allow_first_completion = threading.Event() + second_started = threading.Event() + launches = [] + results = [] + + def launch(*args): + event = len(launches) + 1 + launches.append(event) + return event + + def wait(event): + if event == 1: + first_waiting.set() + assert allow_first_completion.wait(timeout=1) + + monkeypatch.setattr(transport, "_launch_gather", launch) + monkeypatch.setattr(transport, "_wait_gather", wait) + monkeypatch.setattr(transport, "_make_write", lambda *args: object()) + + first = threading.Thread( + target=lambda: results.append(transport.build_request(_write_meta())) + ) + + def run_second(): + second_started.set() + results.append(transport.build_request(_write_meta())) + + second = threading.Thread(target=run_second) + first.start() + assert first_waiting.wait(timeout=1) + second.start() + assert second_started.wait(timeout=1) + # The second launch cannot refill the shared stream metadata yet. + assert launches == [1] + allow_first_completion.set() + first.join(timeout=1) + second.join(timeout=1) + + assert not first.is_alive() and not second.is_alive() + assert launches == [1, 2] + assert len(results) == 2 + + def test_close_retries_only_failed_deregistrations(self, monkeypatch): + allocators = [] + + def allocator_factory(*args, **kwargs): + allocator = _DescriptorAlloc(*args, **kwargs) + allocators.append(allocator) + return allocator + + calls = [] + fail_once = {"value": True} + + class Agent: + def register_memory(self, desc): + pass + + def deregister_memory(self, desc): + calls.append(desc) + if desc is allocators[1].desc and fail_once["value"]: + fail_once["value"] = False + raise RuntimeError("deregister failed") + + monkeypatch.setattr(btr, "SlotAllocator", allocator_factory) + monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: 0) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_start_scatter_worker", + lambda self, name: setattr(self, "_scatter_q", queue.Queue()), + ) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_destroy_stream", + lambda self, attr_name: setattr(self, attr_name, None), + ) + transport = self._construct(Agent()) + + with pytest.raises(RuntimeError, match="failed to deregister"): + transport.close() + assert transport._registered_descs == [allocators[1].desc] + assert not any(allocator.closed for allocator in allocators) + + transport.close() + transport.close() + assert calls.count(allocators[0].desc) == 1 + assert calls.count(allocators[1].desc) == 2 + assert all(allocator.closed for allocator in allocators) + + def test_close_retries_only_failed_stream_destruction(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + transport._scatter_stream = 22 + calls = [] + failed_once = {"value": True} + + def destroy_stream(owner, attr_name): + stream = getattr(owner, attr_name) + if stream is None: + return + calls.append(stream) + if stream == 22 and failed_once["value"]: + failed_once["value"] = False + raise RuntimeError("stream destroy failed") + setattr(owner, attr_name, None) + + monkeypatch.setattr(btr.VmmBounceTransport, "_destroy_stream", destroy_stream) + with pytest.raises(RuntimeError, match="failed to destroy"): + transport.close() + assert transport._scatter_stream == 22 + assert transport._send_stream is None + assert not transport._send_alloc.closed + + transport.close() + transport.close() + assert calls.count(22) == 2 + assert calls.count(0) == 1 + + def test_concurrent_close_retires_resources_once(self, monkeypatch): + transport = _make_transport(monkeypatch, block_bytes_per_group=[8]) + entered = threading.Event() + release = threading.Event() + second_started = threading.Event() + stop_calls = [] + + def stop_worker(owner): + stop_calls.append(True) + entered.set() + assert release.wait(timeout=1) + + monkeypatch.setattr(btr.VmmBounceTransport, "_stop_scatter_worker", stop_worker) + monkeypatch.setattr( + btr.VmmBounceTransport, + "_destroy_stream", + lambda owner, attr_name: setattr(owner, attr_name, None), + ) + errors = [] + + def close(second=False): + if second: + second_started.set() + try: + transport.close() + except Exception as error: + errors.append(error) + + first = threading.Thread(target=close) + second = threading.Thread(target=lambda: close(second=True)) + first.start() + assert entered.wait(timeout=1) + second.start() + assert second_started.wait(timeout=1) + release.set() + first.join(timeout=1) + second.join(timeout=1) + + assert not first.is_alive() and not second.is_alive() + assert errors == [] + assert len(stop_calls) == 1 + + +@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") +class TestReceiveReserve: + def test_reserve_stamps_single_writer_base(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) # total = 2 * 100 = 200 - assert t.reserve(req, num_writers=2) is True + assert t.reserve(req, writer_ranks=[7]) is True assert req.bounce_dst_base == 0x100000 - # writer i lands at base + i * (total // num_writers); per_writer = 100. - assert t.writer_base((req.unique_rid, req.slice_id), 0) == 0x100000 - assert t.writer_base((req.unique_rid, req.slice_id), 1) == 0x100000 + 100 - - def test_reserve_uneven_fanin_falls_back(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[3]) - req = _recv_req([1]) # total = 3, not divisible by 2 - assert t.reserve(req, num_writers=2) is False + assert t.writer_base((req.unique_rid, req.slice_id), 7) == 0x100000 + assert t.writer_base((req.unique_rid, req.slice_id), 9) is None + + @pytest.mark.parametrize("slot_bytes", [[3], [100], [100, 100], [100, 200]]) + def test_reserve_rejects_every_multiwriter_plan(self, monkeypatch, slot_bytes): + t = _make_transport(monkeypatch, block_bytes_per_group=slot_bytes) + req = _recv_req([2] * len(slot_bytes)) + assert t.reserve(req, writer_ranks=[7, 3]) is False assert req.bounce_dst_base is None + assert not t._recv_alloc.has_outstanding - def test_reserve_heterogeneous_fanin_falls_back(self, monkeypatch): - # Two groups with different per-block sizes: the equal split would overrun a sub-region, so - # fall back (even though the total is divisible). + def test_reserve_heterogeneous_single_writer_ok(self, monkeypatch): + # A single writer has no split, so heterogeneous slot bytes are fine. t = _make_transport(monkeypatch, block_bytes_per_group=[100, 200]) - req = _recv_req([2, 2]) # total = 2*100 + 2*200 = 600 - assert t.reserve(req, num_writers=2) is False + assert t.reserve(_recv_req([2, 2]), writer_ranks=[7]) is True + + def test_reserve_single_writer_ok(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[3]) + req = _recv_req([1]) # total = 3, one writer -> no even-split requirement + assert t.reserve(req, writer_ranks=[3]) is True + + def test_reserve_counts_only_valid_block_ids(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([0]) + # Prefix/SWA tables may retain BAD_PAGE_INDEX entries for evicted + # blocks. Only the two extractable slots contribute bytes. + req.block_ids_per_layer_groups = [np.array([4, -1, 9, -1], dtype=np.int64)] + + assert t.reserve(req, writer_ranks=[3]) is True + ctx = t._reserved_map[(req.unique_rid, req.slice_id)] + assert ctx.per_writer_bytes == 200 + + @pytest.mark.parametrize( + "remaining_blocks", + [ + np.array([8, 9], dtype=np.int64), + np.array([31], dtype=np.int64), + ], + ids=["prefix-trimmed", "swa-trimmed"], + ) + def test_reserve_sizes_only_the_remaining_trimmed_suffix(self, monkeypatch, remaining_blocks): + t = _make_transport(monkeypatch, block_bytes_per_group=[64]) + req = _recv_req([0]) + req.block_ids_per_layer_groups = [remaining_blocks] + + assert t.reserve(req, writer_ranks=[3]) is True + ctx = t._reserved_map[(req.unique_rid, req.slice_id)] + assert ctx.per_writer_bytes == remaining_blocks.size * 64 + + def test_reserve_binds_trusted_destination_intervals(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[8]) + req = _recv_req([1]) + assert t.reserve( + req, + writer_ranks=[3], + destination_intervals=[(0x2000, 8)], + ) + ctx = t._reserved_map[(req.unique_rid, req.slice_id)] + assert ctx._destination_intervals == ((0x2000, 0x2008),) + + def test_reserve_rejects_aliased_destination_union_before_publication(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[200]) + req = _recv_req([1]) + + assert not t.reserve( + req, + writer_ranks=[3], + destination_intervals=[(0x2000, 100), (0x2000, 100)], + ) + assert req.bounce_dst_base is None + assert t._reserved_map == {} + assert t._recv_alloc.released == [0] + assert not t._recv_alloc.has_outstanding - def test_reserve_uniform_multigroup_fanin_ok(self, monkeypatch): - # Uniform slot bytes across present groups -> even byte split -> bounce allowed. - t = _make_transport(monkeypatch, block_bytes_per_group=[100, 100]) - assert t.reserve(_recv_req([2, 2]), num_writers=2) is True + def test_destination_intervals_factory_runs_only_after_allocator_admission(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[8]) + req = _recv_req([1]) + calls = [] + t._recv_alloc.reserve = lambda size, timeout=None: None - def test_reserve_heterogeneous_single_writer_ok(self, monkeypatch): - # num_writers==1 has no split, so heterogeneous slot bytes are fine. - t = _make_transport(monkeypatch, block_bytes_per_group=[100, 200]) - assert t.reserve(_recv_req([2, 2]), num_writers=1) is True + assert not t.reserve( + req, + writer_ranks=[3], + destination_intervals_factory=lambda: calls.append(True) or [(0x2000, 8)], + ) + assert calls == [] - def test_reserve_single_writer_ok(self, monkeypatch): + def test_destination_intervals_factory_canonicalizes_without_retaining_input(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[8]) + req = _recv_req([1]) + raw_intervals = {(0x2004, 4), (0x2000, 4)} + calls = [] + + def build_intervals(): + calls.append(t._pending_reservations) + return raw_intervals + + assert t.reserve( + req, + writer_ranks=[3], + destination_intervals_factory=build_intervals, + ) + ctx = t._reserved_map[(req.unique_rid, req.slice_id)] + assert calls == [1] + assert "destination_intervals" not in ctx.__dict__ + assert ctx._destination_intervals == ((0x2000, 0x2008),) + + raw_intervals.add((0x3000, 8)) + assert ctx._destination_intervals == ((0x2000, 0x2008),) + + def test_invalid_lazy_destination_intervals_release_reserved_slot(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[8]) + req = _recv_req([1]) + assert not t.reserve( + req, + writer_ranks=[3], + destination_intervals_factory=lambda: [], + ) + assert t._recv_alloc.released == [0] + assert not t._recv_alloc.has_outstanding + + def test_destination_intervals_factory_error_releases_reserved_slot(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[8]) + req = _recv_req([1]) + + def fail_to_build_intervals(): + raise RuntimeError("manifest construction failed") + + with pytest.raises(RuntimeError, match="manifest construction failed"): + t.reserve( + req, + writer_ranks=[3], + destination_intervals_factory=fail_to_build_intervals, + ) + assert t._recv_alloc.released == [0] + assert not t._recv_alloc.has_outstanding + + def test_invalid_destination_intervals_release_reserved_slot(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[8]) + req = _recv_req([1]) + assert not t.reserve(req, writer_ranks=[3], destination_intervals=[(0, 8)]) + assert t._recv_alloc.released == [0] + assert not t._recv_alloc.has_outstanding + + @pytest.mark.parametrize("writer_ranks", [[], [3, 3], [-1], [True], 1]) + def test_reserve_rejects_invalid_exact_writer_plan(self, monkeypatch, writer_ranks): t = _make_transport(monkeypatch, block_bytes_per_group=[3]) - req = _recv_req([1]) # total = 3, num_writers=1 -> no even-split requirement - assert t.reserve(req, num_writers=1) is True + req = _recv_req([1]) + assert t.reserve(req, writer_ranks=writer_ranks) is False + assert req.bounce_dst_base is None def test_reserve_too_small_falls_back(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_blocks=96) - assert t.reserve(_recv_req([4]), num_writers=1) is False # 4 < 96 blocks + assert t.reserve(_recv_req([4]), writer_ranks=[3]) is False # 4 < 96 blocks def test_reserve_unknown_slot_size_falls_back(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) # only 1 group known - assert t.reserve(_recv_req([2, 2]), num_writers=1) is False # 2nd group unknown + assert t.reserve(_recv_req([2, 2]), writer_ranks=[3]) is False # 2nd group unknown def test_reserve_oversize_falls_back(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[1000], capacity=500) - assert t.reserve(_recv_req([2]), num_writers=1) is False # total 2000 > cap 500 - - def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - # writer for the HIGHER src_base reports first; scatter must reorder by src_base. - t.record_result( - rid_slice, - 7, - np.array([20], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=200, - ) - assert t._scatter_q.empty() # only 1 of 2 writers terminal -> no scatter - assert not t._recv_alloc.released # region NOT freed while a writer is still pending - t.record_result( - rid_slice, - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=100, - ) - ctx, descs = t._scatter_q.get_nowait() - # each tail carries its OWN src_base; sorted (100 before 200) so the scatter is deterministic. - assert [t[0] for t in descs] == [100, 200] # per-writer src_base preserved - assert [list(t[1]) for t in descs] == [[10], [20]] # dst_ptrs - assert [list(t[2]) for t in descs] == [[8], [8]] # sizes - - def test_fanin_fallback_writer_leaves_survivor_at_its_own_base(self, monkeypatch): - # Regression: if one fan-in writer falls back to in-place (SUCCESS, empty tail) while a - # sibling bounces, the survivor must be scattered from ITS OWN src_base, not packed to 0. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - t.record_result( - rid_slice, 7, None, None - ) # writer 0 fell back to in-place: SUCCESS, no tail - assert t._scatter_q.empty() # only 1 of 2 writers terminal - t.record_result( - rid_slice, - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=100, - ) # writer 1 bounced to base+100 - ctx, descs = t._scatter_q.get_nowait() - assert [t[0] for t in descs] == [100] # only the survivor, read from base+100 (NOT 0) - assert [list(t[1]) for t in descs] == [[10]] - - def test_fanin_failed_then_success_releases_only_after_both(self, monkeypatch): - # A FAILED writer must not free the shared region until every writer is terminal. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - t.record_failure(rid_slice, 7) # first writer fails - assert not t._recv_alloc.released # region held while a sibling may still be in flight - assert t.is_bounced(rid_slice) is True - t.record_result( - rid_slice, 3, np.array([10], dtype=np.int64), np.array([8], dtype=np.int64), src_base=0 - ) - # all terminal, >=1 FAILED -> no scatter, release (both drained), region freed. - assert t._scatter_q.empty() - assert t._recv_alloc.released - assert not t._recv_alloc.quarantined # FAILED has drained -> release, NOT quarantine - assert t.is_bounced(rid_slice) is False + assert t.reserve(_recv_req([2]), writer_ranks=[3]) is False # total 2000 > cap 500 def test_on_done_fires_after_scatter_lands(self, monkeypatch): # The completion cb rides the context and fires only when the worker records the scatter as # done -> the gen never observes completion before the KV is scattered into place. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True + assert t.reserve(req, writer_ranks=[3]) is True rid_slice = (req.unique_rid, req.slice_id) calls = [] t.record_result( rid_slice, 3, np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=0, + np.array([200], dtype=np.int64), + src_base=0x100000, on_done=lambda ok: calls.append(ok), ) ctx, descs = t._scatter_q.get_nowait() @@ -464,7 +1283,7 @@ def test_empty_acc_fires_on_done_inline_and_releases(self, monkeypatch): # complete -> on_done(True) inline + slot released, nothing queued. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True + assert t.reserve(req, writer_ranks=[3]) is True rid_slice = (req.unique_rid, req.slice_id) calls = [] t.record_result(rid_slice, 3, None, None, on_done=lambda ok: calls.append(ok)) @@ -487,18 +1306,6 @@ def test_missing_key_is_dropped(self, monkeypatch): ) assert calls == [] # no-op, no callback - def test_duplicate_writer_is_ignored(self, monkeypatch): - # A duplicate SUCCESS from the same peer_rank must not double-count toward all-terminal. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - arr = (np.array([10], dtype=np.int64), np.array([8], dtype=np.int64)) - t.record_result(rid_slice, 7, *arr, src_base=0) - t.record_result(rid_slice, 7, *arr, src_base=0) # duplicate of the SAME writer - assert t._scatter_q.empty() # still only 1 distinct writer -> not all terminal - assert not t._recv_alloc.released - def test_scatter_write_result_non_bounce_fires_on_done(self): # Non-bounced path completes inline (the in-place WRITE already landed the KV). calls = [] @@ -511,7 +1318,7 @@ def test_release_idle_reservation_frees_slot_and_is_idempotent(self, monkeypatch # INIT-cancel immediate release (nothing published) must free the recv slot; idempotent. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True + assert t.reserve(req, writer_ranks=[3]) is True rid_slice = (req.unique_rid, req.slice_id) assert t.is_bounced(rid_slice) is True t.release_idle_reservation(rid_slice) @@ -519,19 +1326,536 @@ def test_release_idle_reservation_frees_slot_and_is_idempotent(self, monkeypatch assert t._recv_alloc.released # slot freed t.release_idle_reservation(rid_slice) # already gone -> no-op, must not raise - def test_orphan_reservation_quarantines_and_is_idempotent(self, monkeypatch): - # Giving up on an in-flight reservation must quarantine the region (a write may still land), - # not release or leak it; a second give-up is a no-op. + def test_orphan_reservation_waits_for_backend_quiescence(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True + assert t.reserve(req, writer_ranks=[7]) rid_slice = (req.unique_rid, req.slice_id) + t.orphan_reservation(rid_slice) - assert t._recv_alloc.quarantined == [0] # quarantined, not released + + assert t.is_bounced(rid_slice) + assert t._recv_alloc.released == [] + assert t._recv_alloc.quarantined == [] + + t.mark_backend_quiesced(rid_slice) + + assert not t.is_bounced(rid_slice) + assert t._recv_alloc.released == [0] + + def test_failure_callback_fires_only_after_exposed_writer_drains(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[7]) + key = (req.unique_rid, req.slice_id) + assert t.mark_writer_exposed(key, 7) + calls = [] + t.mark_logical_failure(key, on_done=lambda ok: calls.append(ok)) + assert calls == [] + assert t.is_bounced(key) + t.record_failure(key, 7) + assert calls == [False] + assert not t.is_bounced(key) + + def test_failed_settlement_ack_retries_without_double_release(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[7]) + key = (req.unique_rid, req.slice_id) + calls = [] + + def acknowledge(ok): + calls.append(ok) + if len(calls) == 1: + raise RuntimeError("registry temporarily unavailable") + + t.record_result(key, 7, None, None, on_done=acknowledge) + assert calls == [True] + assert t._recv_alloc.released == [0] + assert t.is_bounced(key) + + # Any later observation, including a duplicate result, retries the + # durable acknowledgement without releasing the slot a second time. + t.record_result(key, 7, None, None) + assert calls == [True, True] + assert t._recv_alloc.released == [0] + assert not t.is_bounced(key) + + def test_settlement_retry_replays_registry_update_until_session_ack(self, monkeypatch): + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + from tensorrt_llm._torch.disaggregation.base.transfer import KVSlice, SessionArgsBase + from tensorrt_llm.disaggregated_params import DisaggregatedParams + + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2], rid=101) + assert t.reserve(req, writer_ranks=[7]) + key = (req.unique_rid, req.slice_id) + + registry = RecvTransferRegistry() + assert registry.prepare(key, {7}, has_bounce_slot=True).accepted + assert registry.begin_publication(key, 7).publication_allowed + assert registry.mark_published(key, 7).accepted + ready = registry.record_result(key, 7, WriterResult.SUCCESS, WriterMode.BOUNCE) + assert ready.actions == (LifecycleAction.START_BOUNCE_SCATTER,) + + params = DisaggregatedParams( + disagg_request_id=key[0], + ctx_request_id=key[0], + ctx_dp_rank=0, + ) + task = tfr.KVRecvTask(key[0], KVSlice(), key[1], params, aux_slot=None) + task.status = tfr.TaskStatus.TRANSFERRING + task._perf_timer = None + + receiver = object.__new__(tfr.Receiver) + receiver._shutdown = True + receiver._recv_registry = registry + receiver._sessions_lock = threading.Lock() + receiver._bounce_lifecycle_delivery_lock = threading.Lock() + receiver._pending_bounce_lifecycle_deliveries = {} + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0) + ) + + session = object.__new__(tfr.RxSession) + session._closed = True + session.request_id = key[0] + session._base_args = SessionArgsBase(params) + session.lock = threading.Lock() + session._kv_tasks = [task] + session._receiver = receiver + receiver._sessions = {key[0]: session} + + delivered_updates = [] + process_update = tfr.RxSession.process_lifecycle_update.__get__(session, tfr.RxSession) + + def fail_once(update, *, peer_rank=None): + delivered_updates.append(update) + if len(delivered_updates) == 1: + raise RuntimeError("consumer temporarily unavailable") + process_update(update, peer_rank=peer_rank) + + session.process_lifecycle_update = fail_once + t.record_result( + key, + 7, + np.array([0x2000], dtype=np.int64), + np.array([200], dtype=np.int64), + src_base=0x100000, + on_done=lambda succeeded: receiver._finish_bounce(key, succeeded, 7), + ) + bounce_context, _descs = t._scatter_q.get_nowait() + t._apply(bounce_context.rid_slice, lambda context: context.finish_scatter(True)) + + assert task.status is tfr.TaskStatus.TRANSFERRING + assert registry.context_snapshot(key).physical_state is PhysicalState.DRAINED + assert t._recv_alloc.released == [0] + assert t.is_bounced(key) + + assert t.retry_settlements() + assert task.status is tfr.TaskStatus.TRANSFERRED + assert delivered_updates[0] is delivered_updates[1] + assert receiver._pending_bounce_lifecycle_deliveries == {} + assert t._recv_alloc.released == [0] + + def test_concurrent_settlement_retry_does_not_duplicate_callback(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[7]) + key = (req.unique_rid, req.slice_id) + callback_entered = threading.Event() + allow_callback = threading.Event() + calls = [] + + def acknowledge(ok): + calls.append(ok) + callback_entered.set() + assert allow_callback.wait(timeout=1) + + worker = threading.Thread( + target=lambda: t.record_result(key, 7, None, None, on_done=acknowledge) + ) + worker.start() + assert callback_entered.wait(timeout=1) + assert not t.retry_settlements() + allow_callback.set() + worker.join(timeout=1) + + assert not worker.is_alive() + assert calls == [True] + assert t._recv_alloc.released == [0] + assert t.retry_settlements() + + def test_scatter_worker_continues_after_physical_settlement_failure(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + first_req = _recv_req([1], rid=1) + second_req = _recv_req([1], rid=2) + assert t.reserve(first_req, writer_ranks=[7]) + assert t.reserve(second_req, writer_ranks=[7]) + first_key = (first_req.unique_rid, first_req.slice_id) + second_key = (second_req.unique_rid, second_req.slice_id) + for key, dst_ptr in ((first_key, 0x2000), (second_key, 0x3000)): + t.record_result( + key, + 7, + np.array([dst_ptr], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=0x100000, + ) + + release = t._recv_alloc.release + failed_once = False + + def fail_first_release(slot_id): + nonlocal failed_once + if slot_id == 0 and not failed_once: + failed_once = True + raise RuntimeError("allocator temporarily unavailable") + release(slot_id) + + t._recv_alloc.release = fail_first_release + _run_mock_scatter_worker(t, monkeypatch) + + assert t.is_bounced(first_key) + assert not t.is_bounced(second_key) + assert t._recv_alloc.released == [1] + assert t._scatter_worker_error is None + assert t._scatter_stream_healthy + assert t._accepting_reservations + assert t.retry_settlements(first_key) + assert t._recv_alloc.released == [1, 0] + + def test_scoped_settlement_retry_ignores_unrelated_request(self, monkeypatch): + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + first_req = _recv_req([1], rid=1) + second_req = _recv_req([1], rid=2) + assert t.reserve(first_req, writer_ranks=[7]) + assert t.reserve(second_req, writer_ranks=[7]) + first_key = (first_req.unique_rid, first_req.slice_id) + second_key = (second_req.unique_rid, second_req.slice_id) + second_calls = [] + + def reject_first(_succeeded): + raise RuntimeError("request A consumer unavailable") + + def accept_second_on_retry(succeeded): + second_calls.append(succeeded) + if len(second_calls) == 1: + raise RuntimeError("request B consumer temporarily unavailable") + + t.record_result(first_key, 7, None, None, on_done=reject_first) + t.record_result(second_key, 7, None, None, on_done=accept_second_on_retry) + assert t.is_bounced(first_key) + assert t.is_bounced(second_key) + + assert t.retry_settlements(second_key) + assert not t.is_bounced(second_key) + assert t.is_bounced(first_key) + + # Request-scoped drain must not inherit request A's retry failure. + session = SimpleNamespace( + disagg_request_id=second_req.unique_rid, + _receiver=SimpleNamespace( + _bounce=t, + _recv_registry=SimpleNamespace(is_request_drained=lambda rid: rid == 2), + ), + has_untracked_receive_activity=lambda: False, + ) + assert tfr.RxSession.resources_drained(session) + assert not t.retry_settlements(first_req.unique_rid) + assert not t.retry_settlements() + assert second_calls == [True, True] + + def test_unexpected_scatter_worker_failure_closes_admission(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + first_req = _recv_req([1], rid=1) + second_req = _recv_req([1], rid=2) + assert t.reserve(first_req, writer_ranks=[7]) + assert t.reserve(second_req, writer_ranks=[7]) + first_key = (first_req.unique_rid, first_req.slice_id) + second_key = (second_req.unique_rid, second_req.slice_id) + first_logical_failures, second_calls = [], [] + for key, dst_ptr, on_done, on_logical_failure in ( + (first_key, 0x2000, None, lambda: first_logical_failures.append(True)), + (second_key, 0x3000, lambda ok: second_calls.append(ok), None), + ): + t.record_result( + key, + 7, + np.array([dst_ptr], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=0x100000, + on_done=on_done, + on_logical_failure=on_logical_failure, + ) + + apply = t._apply + failed_once = False + + def fail_first_state_update(rid_slice, mutate): + nonlocal failed_once + if not failed_once: + failed_once = True + raise RuntimeError("unexpected state-machine failure") + apply(rid_slice, mutate) + + t._apply = fail_first_state_update + _run_mock_scatter_worker(t, monkeypatch) + + assert isinstance(t._scatter_worker_error, RuntimeError) + assert not t._scatter_stream_healthy + assert not t._accepting_reservations + assert t.is_bounced(first_key) + first = t._reserved_map[first_key] + assert first.logical_failed + assert first.scatter_state is bcore.ScatterState.FAILED + assert first_logical_failures == [True] + assert not t.is_bounced(second_key) + assert second_calls == [False] + assert t._recv_alloc.released == [1] + assert not t.reserve(_recv_req([1], rid=3), writer_ranks=[7]) + + def test_scatter_failure_suppresses_later_unlaunched_queue_entries(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + first_req = _recv_req([1], rid=1) + second_req = _recv_req([1], rid=2) + assert t.reserve(first_req, writer_ranks=[7]) + assert t.reserve(second_req, writer_ranks=[7]) + first_key = (first_req.unique_rid, first_req.slice_id) + second_key = (second_req.unique_rid, second_req.slice_id) + first_calls, second_calls, first_logical_failures = [], [], [] + t.record_result( + first_key, + 7, + np.array([0x2000], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=0x100000, + on_done=lambda ok: first_calls.append(ok), + on_logical_failure=lambda: first_logical_failures.append(True), + ) + t.record_result( + second_key, + 7, + np.array([0x3000], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=0x100000, + on_done=lambda ok: second_calls.append(ok), + ) + t._scatter_ready = threading.Event() + monkeypatch.setattr( + btr, + "cudart", + SimpleNamespace(cudaSetDevice=lambda device_id: ("ok",)), + ) + monkeypatch.setattr(btr, "CUASSERT", lambda result: result[1:]) + monkeypatch.setattr( + btr, + "scatter_contiguous", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("scatter failed")), + ) + + t._scatter_loop() + + assert t.is_bounced(first_key) # the launched CUDA work remains ambiguous + assert first_calls == [] + assert first_logical_failures == [True] + assert not t.is_bounced(second_key) # this entry never launched and settles as failed + assert second_calls == [False] + assert t._recv_alloc.released == [1] + + def test_scatter_failure_fails_request_but_retains_physical_ownership(self, monkeypatch): + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + from tensorrt_llm._torch.disaggregation.base.transfer import KVSlice, SessionArgsBase + from tensorrt_llm._torch.disaggregation.native.receive_lifecycle import LogicalState + from tensorrt_llm.disaggregated_params import DisaggregatedParams + + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([1], rid=101) + assert t.reserve(req, writer_ranks=[7]) + key = (req.unique_rid, req.slice_id) + + registry = RecvTransferRegistry() + assert registry.prepare(key, {7}, has_bounce_slot=True).accepted + assert registry.begin_publication(key, 7).publication_allowed + assert registry.mark_published(key, 7).accepted + registry.record_result(key, 7, WriterResult.SUCCESS, WriterMode.BOUNCE) + + params = DisaggregatedParams( + disagg_request_id=key[0], + ctx_request_id=key[0], + ctx_dp_rank=0, + ) + task = tfr.KVRecvTask(key[0], KVSlice(), key[1], params, aux_slot=None) + task.status = tfr.TaskStatus.TRANSFERRING + task.lifecycle_managed = True + task._perf_timer = None + + receiver = object.__new__(tfr.Receiver) + receiver._shutdown = True + receiver._recv_registry = registry + receiver._bounce = t + receiver._sessions_lock = threading.Lock() + receiver._bounce_lifecycle_delivery_lock = threading.Lock() + receiver._pending_bounce_lifecycle_deliveries = {} + receiver._pending_bounce_logical_failure_deliveries = {} + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0) + ) + + session = object.__new__(tfr.RxSession) + session._closed = True + session.request_id = key[0] + session._base_args = SessionArgsBase(params) + session.lock = threading.Lock() + session._kv_tasks = [task] + session._receiver = receiver + session._need_aux = False + session._terminal_status = None + session._active_receive_dispatches = 0 + receiver._sessions = {key[0]: session} + + t.record_result( + key, + 7, + np.array([0x2000], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=0x100000, + on_done=lambda succeeded: receiver._finish_bounce(key, succeeded, 7), + on_logical_failure=lambda: receiver._fail_bounce_logically(key), + ) + bounce_context, _descs = t._scatter_q.get_nowait() + t._apply(bounce_context.rid_slice, lambda context: context.finish_scatter(False)) + + snapshot = registry.context_snapshot(key) + assert snapshot.logical_state is LogicalState.FAILED + assert snapshot.physical_state is not PhysicalState.DRAINED + assert task.status is tfr.TaskStatus.ERROR + assert not tfr.RxSession.resources_drained(session) + assert t.is_bounced(key) assert t._recv_alloc.released == [] - assert t.is_bounced(rid_slice) is False # settled and removed from the live map - t.orphan_reservation(rid_slice) # already gone -> no-op, must not raise - assert t._recv_alloc.quarantined == [0] + + def test_backend_quiescence_callback_fires_after_slot_settlement(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[7]) + key = (req.unique_rid, req.slice_id) + assert t.mark_writer_exposed(key, 7) + t.mark_logical_failure(key) + calls = [] + + t.mark_backend_quiesced(key, on_done=lambda ok: calls.append(ok)) + + assert calls == [False] + assert not t.is_bounced(key) + assert t._recv_alloc.released == [0] + + def test_close_refuses_live_receive_context(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[3]) + with pytest.raises(RuntimeError, match="live receive contexts"): + t.close() + assert not t.reserve(_recv_req([2], rid=2), writer_ranks=[4]) + + def test_close_finishes_already_queued_scatter(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[3]) + key = (req.unique_rid, req.slice_id) + t.record_result( + key, + 3, + np.array([10], dtype=np.int64), + np.array([200], dtype=np.int64), + src_base=0x100000, + ) + + def worker(): + while True: + item = t._scatter_q.get() + try: + if item is None: + return + ctx, _descs = item + t._apply(ctx.rid_slice, lambda c: c.finish_scatter(True)) + finally: + t._scatter_q.task_done() + + t._scatter_thread = threading.Thread(target=worker) + t._scatter_thread.start() + t.close() + assert not t.is_bounced(key) + assert t._recv_alloc.released == [0] + + def test_close_retries_settlement_created_while_draining_scatter(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([2]) + assert t.reserve(req, writer_ranks=[3]) + key = (req.unique_rid, req.slice_id) + callback_calls = [] + + def acknowledge(ok): + callback_calls.append(ok) + if len(callback_calls) == 1: + raise RuntimeError("registry temporarily unavailable") + + t.record_result( + key, + 3, + np.array([10], dtype=np.int64), + np.array([200], dtype=np.int64), + src_base=0x100000, + on_done=acknowledge, + ) + + def worker(): + while True: + item = t._scatter_q.get() + try: + if item is None: + return + ctx, _descs = item + t._apply(ctx.rid_slice, lambda c: c.finish_scatter(True)) + finally: + t._scatter_q.task_done() + + t._scatter_thread = threading.Thread(target=worker) + t._scatter_thread.start() + t.close() + + assert callback_calls == [True, True] + assert t._pending_settlements == {} + assert t._recv_alloc.released == [0] + + def test_close_drains_scatter_queue_before_destroying_arenas(self): + t = object.__new__(btr.VmmBounceTransport) + t._init_recv_state() + t._send_alloc = _FakeAlloc(1024, 512) + t._recv_alloc = _FakeAlloc(1024, 512) + t._reg_descs = [object()] + t._agent = SimpleNamespace( + deregister_memory=lambda d: pytest.fail("must not use a shut-down agent") + ) + t._scatter_q = queue.Queue() + processed = [] + + def worker(): + while True: + item = t._scatter_q.get() + try: + if item is None: + return + processed.append(item) + finally: + t._scatter_q.task_done() + + t._scatter_thread = threading.Thread(target=worker) + t._scatter_thread.start() + t._scatter_q.put("accepted-before-close") + t.close() + assert processed == ["accepted-before-close"] + assert t._send_alloc.closed and t._recv_alloc.closed # --------------------------------------------------------------------------- # @@ -567,32 +1891,103 @@ def test_blocks_until_release_when_full(self, monkeypatch): a.reserve(512) assert a.reserve(512, timeout=0.05) is None # full -> times out + def test_close_atomically_rejects_racing_reservations(self, monkeypatch): + close_started = threading.Event() + allow_close = threading.Event() + + class BlockingBuffer: + size = 512 + base_ptr = 0x1000 + + @staticmethod + def reg_descs(): + return [] + + @staticmethod + def close(): + close_started.set() + assert allow_close.wait(timeout=2) + + monkeypatch.setattr(bbuf, "Buffer", lambda *a, **k: BlockingBuffer()) + allocator = btr.SlotAllocator(512, 512) + close_thread = threading.Thread(target=allocator.close) + close_thread.start() + assert close_started.wait(timeout=1) + + try: + assert allocator.reserve(512, timeout=0.1) is None + finally: + allow_close.set() + close_thread.join(timeout=2) + + assert not close_thread.is_alive() + + def test_failed_buffer_close_keeps_admission_closed_and_is_retryable(self, monkeypatch): + class FailOnceBuffer: + size = 512 + base_ptr = 0x1000 + + def __init__(self): + self.close_calls = 0 + + @staticmethod + def reg_descs(): + return [] + + def close(self): + self.close_calls += 1 + if self.close_calls == 1: + raise RuntimeError("injected buffer close failure") + + buffer = FailOnceBuffer() + monkeypatch.setattr(bbuf, "Buffer", lambda *a, **k: buffer) + allocator = btr.SlotAllocator(512, 512) + + with pytest.raises(RuntimeError, match="injected buffer close failure"): + allocator.close() + assert allocator.reserve(512, timeout=0.1) is None + + allocator.close() + assert buffer.close_calls == 2 + # --------------------------------------------------------------------------- # -# core.TransferContext — the pure drain-before-release state machine. +# core.RecvBounceContext — the pure drain-before-release state machine. # No CUDA / NIXL / allocator, so these run on any CPU (no skipif). # --------------------------------------------------------------------------- # class TestLifecycle: - def _ctx(self, num_writers, per_writer_bytes=100, base_addr=0x1000): - return bcore.TransferContext( + def _ctx( + self, + writer_ranks=(3,), + per_writer_bytes=100, + base_addr=0x1000, + destination_intervals=None, + ): + return bcore.RecvBounceContext( rid_slice=(1, 0), slot_id=0, base_addr=base_addr, per_writer_bytes=per_writer_bytes, - num_writers=num_writers, + writer_ranks=tuple(writer_ranks), + destination_intervals=destination_intervals, ) def _dst(self, v=10): - return dict(dst_ptrs=np.array([v], dtype=np.int64), sizes=np.array([8], dtype=np.int64)) + return dict( + dst_ptrs=np.array([v], dtype=np.int64), + sizes=np.array([100], dtype=np.int64), + ) def test_writer_base_layout(self): - c = self._ctx(3, per_writer_bytes=0x64, base_addr=0x1000) - assert [c.writer_base(i) for i in range(3)] == [0x1000, 0x1064, 0x10C8] + c = self._ctx((7, 3, 11), per_writer_bytes=0x64, base_addr=0x1000) + assert [c.writer_base(rank) for rank in (7, 3, 11)] == [0x1000, 0x1064, 0x10C8] + with pytest.raises(KeyError): + c.writer_base(9) def test_single_writer_success_scatters_then_releases(self): - c = self._ctx(1) + c = self._ctx() assert not c.ready_to_scatter() and not c.ready_to_settle() - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.record_writer_result(3, succeeded=True, src_base=0x1000, **self._dst()) assert c.ready_to_scatter() c.begin_scatter() assert not c.ready_to_settle() # scatter not landed yet @@ -604,18 +1999,20 @@ def test_single_writer_success_scatters_then_releases(self): assert c.settle() is None # idempotent one-shot def test_fanin_holds_until_all_terminal(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) + c = self._ctx((7, 3)) + c.record_writer_result(7, succeeded=True, src_base=0x1000, **self._dst()) assert not c.ready_to_scatter() # 1/2 writers assert not c.ready_to_settle() # drain-before-release - c.record_writer_result(3, succeeded=True, src_base=100, **self._dst()) + c.record_writer_result(3, succeeded=True, src_base=0x1000 + 100, **self._dst(110)) assert c.ready_to_scatter() # all success -> scatter def test_fanin_failed_then_success_releases(self): - c = self._ctx(2) + c = self._ctx((7, 3)) + assert c.mark_writer_exposed(7) + assert c.mark_writer_exposed(3) c.record_writer_result(7, succeeded=False) assert not c.ready_to_settle() # a sibling is still pending -> hold - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.record_writer_result(3, succeeded=True, src_base=0x1000 + 100, **self._dst()) assert not c.ready_to_scatter() # >=1 FAILED -> skip scatter assert c.ready_to_settle() ret = c.settle() @@ -623,58 +2020,226 @@ def test_fanin_failed_then_success_releases(self): assert ret.success is False assert c.state is bcore.TransferState.FAILED - def test_orphan_quarantines(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) - c.mark_orphaned() # the other writer is in-doubt - assert c.ready_to_settle() - ret = c.settle() - assert ret.disposition is bcore.Disposition.QUARANTINE and ret.success is False + def test_logical_failure_retains_exposed_writer_and_suppresses_scatter(self): + c = self._ctx((7, 3)) + assert c.mark_writer_exposed(7) + c.mark_logical_failure() # rank 3 was never exposed, rank 7 remains in doubt + assert c.pending_exposed_writers == (7,) assert c.state is bcore.TransferState.QUARANTINED + assert not c.ready_to_settle() + c.record_writer_result(7, succeeded=True, src_base=0x1000, **self._dst()) + assert not c.ready_to_scatter() + ret = c.settle() + assert ret.disposition is bcore.Disposition.RELEASE and ret.success is False + + def test_backend_quiescence_retires_in_doubt_writer(self): + c = self._ctx((7, 3)) + assert c.mark_writer_exposed(7) + c.mark_logical_failure() + assert not c.ready_to_settle() + c.mark_backend_quiesced() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_failure_before_publication_releases_immediately(self): + c = self._ctx((7, 3)) + c.mark_logical_failure() + assert c.pending_exposed_writers == () + assert c.ready_to_settle() + assert c.settle().success is False + + def test_protocol_conflict_before_publication_requires_backend_quiescence(self): + c = self._ctx((7, 3)) + + c.mark_protocol_conflict() - def test_empty_tail_success_releases_without_scatter(self): - c = self._ctx(1) + assert c.pending_exposed_writers == () + assert not c.ready_to_settle() + c.mark_backend_quiesced() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_absent_tail_direct_fallback_succeeds_without_scatter(self): + c = self._ctx() c.record_writer_result(3, succeeded=True) # no dst tail assert not c.ready_to_scatter() assert c.ready_to_settle() - assert c.settle().disposition is bcore.Disposition.RELEASE + settlement = c.settle() + assert settlement.disposition is bcore.Disposition.RELEASE + assert settlement.success is True + + @pytest.mark.parametrize("sizes", ([101], [0], [-1])) + def test_invalid_scatter_extent_fails_closed(self, sizes): + c = self._ctx(per_writer_bytes=100) + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([10], dtype=np.int64), + sizes=np.array(sizes, dtype=np.int64), + ) + assert not c.ready_to_scatter() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_empty_scatter_tail_fails_closed(self): + c = self._ctx() + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + ) + assert not c.ready_to_scatter() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_overlapping_destinations_within_writer_fail_closed(self): + c = self._ctx() + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([0x2000, 0x2030], dtype=np.int64), + sizes=np.array([60, 40], dtype=np.int64), + ) + assert not c.ready_to_scatter() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_overlapping_destinations_across_writers_fail_closed(self): + c = self._ctx((7, 3)) + c.record_writer_result( + 7, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([0x2000], dtype=np.int64), + sizes=np.array([100], dtype=np.int64), + ) + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000 + 100, + dst_ptrs=np.array([0x2030], dtype=np.int64), + sizes=np.array([100], dtype=np.int64), + ) + assert c.logical_failed + assert not c.ready_to_scatter() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_scatter_destinations_must_stay_inside_trusted_intervals(self): + c = self._ctx( + per_writer_bytes=0x100, + destination_intervals=[(0x2000, 0x100)], + ) + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([0x20F0], dtype=np.int64), + sizes=np.array([0x100], dtype=np.int64), + ) + assert not c.ready_to_scatter() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_scatter_destinations_reject_truncated_in_range_intervals(self): + c = self._ctx( + per_writer_bytes=0x100, + destination_intervals=[(0x2000, 0x100)], + ) + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([0x2000, 0x2080], dtype=np.int64), + sizes=np.array([0x80, 0x20], dtype=np.int64), + ) + assert not c.ready_to_scatter() + assert c.ready_to_settle() + assert c.settle().success is False + + def test_scatter_destinations_accept_exact_heterogeneous_fragment_shapes(self): + c = self._ctx( + per_writer_bytes=0x100, + destination_intervals=[(0x2000, 0x100)], + ) + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000, + dst_ptrs=np.array([0x2000, 0x2040], dtype=np.int64), + sizes=np.array([0x40, 0xC0], dtype=np.int64), + ) + assert c.ready_to_scatter() + + @pytest.mark.parametrize( + "destination_intervals", + [[], [(0, 8)], [((1 << 64) - 1, 2)], [(0x2000, -1)]], + ) + def test_invalid_trusted_destination_intervals_are_rejected(self, destination_intervals): + with pytest.raises(ValueError, match="destination interval"): + self._ctx(destination_intervals=destination_intervals) + + def test_unlaunched_scatter_can_be_suppressed_and_settled(self): + c = self._ctx() + c.record_writer_result(3, succeeded=True, src_base=0x1000, **self._dst()) + c.begin_scatter() + c.suppress_scatter() + assert c.ready_to_settle() + assert c.settle().success is False def test_writers_locked_after_scatter_drops_late_writer(self): - c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c = self._ctx() + c.record_writer_result(3, succeeded=True, src_base=0x1000, **self._dst()) c.begin_scatter() # SCATTERING -> frozen c.record_writer_result(9, succeeded=False) # a late / reordered report assert 9 not in c._writer_ok # dropped, cannot re-arm the state def test_duplicate_writer_dedup(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) + c = self._ctx((7, 3)) + c.record_writer_result(7, succeeded=True, src_base=0x1000, **self._dst()) c.record_writer_result(7, succeeded=False) # same rank again -> ignored assert c._writer_ok[7] is True assert not c.ready_to_settle() # still only 1 distinct writer of 2 - def test_scatter_failure_releases_as_failed(self): - c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + def test_scatter_failure_retains_ownership_without_a_positive_fence(self): + c = self._ctx() + c.record_writer_result(3, succeeded=True, src_base=0x1000, **self._dst()) c.begin_scatter() c.finish_scatter(False) # scatter kernel failed - ret = c.settle() - assert ret.disposition is bcore.Disposition.RELEASE and ret.success is False - - def test_orphan_after_scatter_is_ignored(self): - # once SCATTERING, all writers already reported SUCCESS -> nothing is in doubt, so a late - # orphan (e.g. a racing cancel) must NOT downgrade a clean transfer to quarantine. - c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + assert not c.ready_to_settle() + assert c.settle() is None + assert c.scatter_state is bcore.ScatterState.FAILED + assert c.logical_failed + + def test_logical_failure_during_scatter_waits_and_reports_failure(self): + c = self._ctx() + c.record_writer_result(3, succeeded=True, src_base=0x1000, **self._dst()) c.begin_scatter() - c.mark_orphaned() # no-op after SCATTERING + c.mark_logical_failure() + assert not c.ready_to_settle() c.finish_scatter(True) - assert c.settle().disposition is bcore.Disposition.RELEASE + assert c.settle().success is False + + def test_unexpected_rank_does_not_count(self): + c = self._ctx((7, 3)) + assert not c.mark_writer_exposed(9) + assert not c.record_writer_result(9, succeeded=False) + assert c.pending_exposed_writers == () + + def test_no_access_transition_can_complete_direct_success(self): + c = self._ctx((7, 3)) + assert c.mark_writer_no_access(7, succeeded=True) + c.record_writer_result(3, succeeded=True) + assert c.ready_to_settle() + assert c.settle().success is True # --------------------------------------------------------------------------- # -# SlotAllocator quarantine — orphaned regions are held out of reuse for a grace -# period, then reclaimed off a timer (never handed out while a stray write could land) +# SlotAllocator quarantine — in-doubt regions are held until explicit quiescence evidence. # --------------------------------------------------------------------------- # @pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.buffer import needs CUDA bindings") class TestQuarantine: @@ -692,23 +2257,27 @@ def test_quarantined_region_is_not_reused(self, monkeypatch): a = self._alloc(monkeypatch, cap=1024) # two 512-byte slots a.reserve(512) # [0, 512) stays live sb, _ = a.reserve(512) # [512, 1024) - a.quarantine(sb, grace_s=float("inf")) # hold the tail out of reuse + a.quarantine(sb) # hold the tail out of reuse indefinitely assert a.quarantined_bytes == 512 # live [0,512) + quarantined [512,1024) => arena full => next reserve can't fit. assert a.reserve(512, timeout=0.05) is None - def test_reclaim_returns_expired_to_free(self, monkeypatch): + def test_explicit_quiescence_release_returns_region_to_free(self, monkeypatch): a = self._alloc(monkeypatch, cap=512) # single slot s, _ = a.reserve(512) - a.quarantine(s, grace_s=0.0) # deadline already in the past - assert a.reserve(512, timeout=0.05) is None # still held until reclaimed - assert a.reclaim_expired() == 1 + a.quarantine(s) + assert a.reserve(512, timeout=0.05) is None + assert a.release_quarantined(s) assert a.quarantined_bytes == 0 assert a.reserve(512, timeout=0.5) is not None # now reusable - def test_inf_grace_never_reclaims(self, monkeypatch): + def test_close_refuses_live_or_quarantined_slot(self, monkeypatch): a = self._alloc(monkeypatch, cap=512) s, _ = a.reserve(512) - a.quarantine(s, grace_s=float("inf")) # close-only reclaim - assert a.reclaim_expired() == 0 - assert a.quarantined_bytes == 512 + with pytest.raises(RuntimeError, match="live or quarantined"): + a.close() + a.quarantine(s) + with pytest.raises(RuntimeError, match="live or quarantined"): + a.close() + assert a.release_quarantined(s) + a.close() diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 76a0d8db25b3..58c8000a6666 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.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. + """Test KV Transfer with KVCacheManager (V1) and KVCacheManagerV2 (V2).""" import os @@ -27,7 +42,11 @@ TokenRange, WaitResult, ) -from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker, TransferWorkerConfig +from tensorrt_llm._torch.disaggregation.native.transfer import ( + TaskStatus, + TransferWorker, + TransferWorkerConfig, +) from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType @@ -36,6 +55,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp +from tensorrt_llm.disaggregated_params import DisaggScheduleStyle from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger @@ -626,6 +646,9 @@ def add_and_verify_request( valid_gen_transfer_workers = gen_transfer_workers unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + schedule_style = ( + DisaggScheduleStyle.CONTEXT_FIRST if send_first else DisaggScheduleStyle.GENERATION_FIRST + ) ctx_request = LlmRequest( request_id=ctx_request_id, max_new_tokens=1, @@ -636,7 +659,10 @@ def add_and_verify_request( is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, ) - ctx_request.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + ctx_request.py_disaggregated_params = DisaggregatedParams( + disagg_request_id=unique_rid, + schedule_style=schedule_style, + ) ctx_request.add_new_token(8 + ctx_request_id, 0) ctx_request.py_draft_tokens = [ @@ -676,6 +702,7 @@ def add_and_verify_request( ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, + schedule_style=schedule_style, ) # Add sequence to gen KV cache managers gen_kv_caches = [] @@ -1328,16 +1355,17 @@ def test_session_cancel_after_send(): page_table = ctx_transfer_worker._rank_info.page_table block_ids_per_groups = [np.array([], dtype=np.int64) for _ in page_table.layer_groups] kv_slice = KVSlice(is_last_slice=True, block_ids_per_layer_groups=block_ids_per_groups) - future = tx_session.send(kv_slice) + tx_session.send(kv_slice) + task = tx_session.kv_tasks[0] # No receiver registered yet; task is INIT. tx_session.cancel() assert tx_session.status == SessionStatus.CANCELLED assert tx_session.has_failed() - # Future for the cancelled INIT task must raise. - with pytest.raises(Exception): - future.result(timeout=5.0) + assert task.wait(timeout=5.0) + assert task.status == TaskStatus.ERROR + assert tx_session.wait_complete() == WaitResult.FAILED tx_session.close() finally: ctx_transfer_worker.shutdown() @@ -1347,7 +1375,7 @@ def test_session_cancel_after_send(): @pytest.mark.timeout(60) def test_session_cancel_twice(): - """cancel() called twice must be a no-op on the second call.""" + """Repeated cancel() calls stay safe and retry terminal settlement.""" setup = create_transfer_worker_setup( ctx_tp=1, ctx_pp=1, @@ -1388,14 +1416,14 @@ def test_session_cancel_twice(): try: tx_session = ctx_transfer_worker.create_tx_session(ctx_request) tx_session.cancel() - tx_session.cancel() # Second call must be a no-op + tx_session.cancel() # Replays any terminal result whose first send failed. assert tx_session.status == SessionStatus.CANCELLED assert not tx_session.has_transferring_tasks() tx_session.close() rx_session = gen_transfer_worker.create_rx_session(gen_request) rx_session.cancel() - rx_session.cancel() # Second call must be a no-op + rx_session.cancel() # Receiver cancellation remains idempotent. assert rx_session.status == SessionStatus.CANCELLED assert not rx_session.has_transferring_tasks() rx_session.close() @@ -1460,6 +1488,7 @@ def test_session_has_transferring_tasks_false(): kv_slice = KVSlice(is_last_slice=True, block_ids_per_layer_groups=block_ids_per_groups) tx_session.send(kv_slice) assert not tx_session.has_transferring_tasks() + tx_session.cancel() tx_session.close() # RxSession: no tasks yet diff --git a/tests/unittest/disaggregated/test_messenger.py b/tests/unittest/disaggregated/test_messenger.py index 172c2bf2686c..2c9c1e9ac03e 100644 --- a/tests/unittest/disaggregated/test_messenger.py +++ b/tests/unittest/disaggregated/test_messenger.py @@ -1,8 +1,25 @@ +# 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 socket import time import unittest +from threading import Event, Lock, Thread import pytest +import zmq from parameterized import parameterized from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message @@ -142,5 +159,436 @@ def test_zmq_messenger_double_start_listener(dynamic_endpoint): messenger.stop() +class _CloseTrackingSocket: + def __init__(self): + self.closed = False + self.close_calls = 0 + self.setsockopt_calls = 0 + self.send_calls = 0 + self.send_flags = [] + self.send_multipart_calls = 0 + + def send(self, _message, flags=0): + self.send_calls += 1 + self.send_flags.append(flags) + + def send_multipart(self, _message): + self.send_multipart_calls += 1 + + def setsockopt(self, _option, _value): + self.setsockopt_calls += 1 + + def close(self): + self.close_calls += 1 + self.closed = True + + +class _FailOnceSocket(_CloseTrackingSocket): + def __init__(self, operation): + super().__init__() + self._operation = operation + self._failed = False + + def setsockopt(self, option, value): + super().setsockopt(option, value) + if self._operation == "setsockopt" and not self._failed: + self._failed = True + raise RuntimeError("injected setsockopt failure") + + def close(self): + self.close_calls += 1 + if self._operation == "close" and not self._failed: + self._failed = True + raise RuntimeError("injected close failure") + self.closed = True + + +class _CloseTrackingContext: + def __init__(self): + self.closed = False + self.term_calls = 0 + + def term(self): + self.term_calls += 1 + self.closed = True + + +class _FailOnceContext(_CloseTrackingContext): + def term(self): + self.term_calls += 1 + if self.term_calls == 1: + raise RuntimeError("injected context termination failure") + self.closed = True + + +class _ControllableThread: + def __init__(self): + self.alive = True + self.join_calls = [] + + def is_alive(self): + return self.alive + + def join(self, timeout=None): + self.join_calls.append(timeout) + + +def _make_test_messenger(main_socket=None, listener_thread=None): + messenger = object.__new__(ZMQMessenger) + messenger._lock = Lock() + messenger._stop_lock = Lock() + messenger._socket_io_lock = Lock() + messenger._control_send_lock = Lock() + messenger._io_waiters_lock = Lock() + messenger._io_waiters = 0 + messenger._io_waiters_drained = Event() + messenger._io_waiters_drained.set() + messenger._closed = False + messenger._closing = False + messenger._stop_event = Event() + messenger._socket = main_socket or _CloseTrackingSocket() + messenger._internal_socket = _CloseTrackingSocket() + messenger._control_socket = _CloseTrackingSocket() + messenger._context = _CloseTrackingContext() + messenger._listener_thread = listener_thread + return messenger + + +def _wait_for_closing(messenger, timeout=1): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + with messenger._lock: + if messenger._closing: + return + time.sleep(0.001) + pytest.fail("messenger did not enter the closing state") + + +def test_stop_retries_without_closing_sockets_under_live_listener(): + messenger = _make_test_messenger(listener_thread=_ControllableThread()) + + with pytest.raises(RuntimeError, match="listener thread did not terminate"): + messenger.stop(timeout=0) + + assert messenger._stop_event.is_set() + assert messenger._closed is False + assert messenger._listener_thread.join_calls == [0] + assert messenger._internal_socket.send_calls == 1 + assert messenger._internal_socket.send_flags == [zmq.DONTWAIT] + assert messenger._socket.close_calls == 0 + assert messenger._internal_socket.close_calls == 0 + assert messenger._control_socket.close_calls == 0 + assert messenger._context.term_calls == 0 + + messenger._listener_thread.alive = False + messenger.stop(timeout=1) + messenger.stop(timeout=1) + + assert messenger._closed is True + assert messenger._socket.close_calls == 1 + assert messenger._internal_socket.close_calls == 1 + assert messenger._control_socket.close_calls == 1 + assert messenger._context.term_calls == 1 + + +@pytest.mark.parametrize( + ("operation", "error_match", "expected_close_calls"), + [ + ("setsockopt", "main socket linger configuration failed", 1), + ("close", "main socket close failed", 2), + ], +) +def test_stop_retries_failed_socket_teardown(operation, error_match, expected_close_calls): + main_socket = _FailOnceSocket(operation) + messenger = _make_test_messenger(main_socket=main_socket) + + with pytest.raises(RuntimeError, match=error_match): + messenger.stop() + + assert messenger._closing is True + assert messenger._closed is False + assert main_socket.closed is False + assert messenger._internal_socket.closed is True + assert messenger._control_socket.closed is True + assert messenger._context.term_calls == 0 + + messenger.stop() + + assert messenger._closed is True + assert main_socket.setsockopt_calls == 2 + assert main_socket.close_calls == expected_close_calls + # Sockets that closed on the first attempt are not touched by the retry. + assert messenger._internal_socket.setsockopt_calls == 1 + assert messenger._internal_socket.close_calls == 1 + assert messenger._control_socket.setsockopt_calls == 1 + assert messenger._control_socket.close_calls == 1 + assert messenger._context.term_calls == 1 + + +def test_stop_retries_failed_context_termination_without_reclosing_sockets(): + messenger = _make_test_messenger() + messenger._context = _FailOnceContext() + + with pytest.raises(RuntimeError, match="context termination failed"): + messenger.stop() + + assert messenger._closing is True + assert messenger._closed is False + assert messenger._socket.closed is True + assert messenger._internal_socket.closed is True + assert messenger._control_socket.closed is True + assert messenger._context.term_calls == 1 + + messenger.stop() + + assert messenger._closed is True + assert messenger._socket.close_calls == 1 + assert messenger._internal_socket.close_calls == 1 + assert messenger._control_socket.close_calls == 1 + assert messenger._context.term_calls == 2 + + +class _BlockingSendSocket(_CloseTrackingSocket): + def __init__(self): + super().__init__() + self.first_send_started = Event() + self.second_send_started = Event() + self.release_first_send = Event() + self.first_send_finished = Event() + self.close_overlapped_send = False + self._call_lock = Lock() + + def send_multipart(self, _message): + with self._call_lock: + self.send_multipart_calls += 1 + call_index = self.send_multipart_calls + if call_index == 1: + self.first_send_started.set() + assert self.release_first_send.wait(timeout=2) + self.first_send_finished.set() + else: + self.second_send_started.set() + + def close(self): + if self.first_send_started.is_set() and not self.first_send_finished.is_set(): + self.close_overlapped_send = True + super().close() + + +def test_main_socket_serializes_concurrent_sends(): + main_socket = _BlockingSendSocket() + messenger = _make_test_messenger(main_socket=main_socket) + errors = [] + + def send(message): + try: + messenger.send([message]) + except BaseException as error: + errors.append(error) + + first = Thread(target=send, args=(b"first",)) + second = Thread(target=send, args=(b"second",)) + first.start() + assert main_socket.first_send_started.wait(timeout=1) + second.start() + try: + assert not main_socket.second_send_started.wait(timeout=0.1) + finally: + main_socket.release_first_send.set() + first.join(timeout=2) + second.join(timeout=2) + + assert not first.is_alive() + assert not second.is_alive() + assert errors == [] + assert main_socket.second_send_started.is_set() + assert main_socket.send_multipart_calls == 2 + + +class _BlockingReceiveSocket(_CloseTrackingSocket): + def __init__(self): + super().__init__() + self.receive_started = Event() + self.release_receive = Event() + self.send_started = Event() + + def recv_multipart(self): + self.receive_started.set() + assert self.release_receive.wait(timeout=2) + return [b"response"] + + def send_multipart(self, _message): + super().send_multipart(_message) + self.send_started.set() + + +def test_main_socket_serializes_receive_with_send(): + main_socket = _BlockingReceiveSocket() + messenger = _make_test_messenger(main_socket=main_socket) + responses = [] + errors = [] + + def receive(): + try: + responses.append(messenger.receive()) + except BaseException as error: + errors.append(error) + + def send(): + try: + messenger.send([b"request"]) + except BaseException as error: + errors.append(error) + + receive_thread = Thread(target=receive) + send_thread = Thread(target=send) + receive_thread.start() + assert main_socket.receive_started.wait(timeout=1) + send_thread.start() + try: + assert not main_socket.send_started.wait(timeout=0.1) + finally: + main_socket.release_receive.set() + receive_thread.join(timeout=2) + send_thread.join(timeout=2) + + assert not receive_thread.is_alive() + assert not send_thread.is_alive() + assert errors == [] + assert responses == [[b"response"]] + assert main_socket.send_started.is_set() + + +def test_stop_waits_for_admitted_send_and_rejects_late_send(): + main_socket = _BlockingSendSocket() + messenger = _make_test_messenger(main_socket=main_socket) + errors = [] + + def run_send(message): + try: + messenger.send([message]) + except BaseException as error: + errors.append(error) + + first = Thread(target=run_send, args=(b"first",)) + stop_thread = Thread(target=messenger.stop) + late = Thread(target=run_send, args=(b"late",)) + first.start() + assert main_socket.first_send_started.wait(timeout=1) + stop_thread.start() + _wait_for_closing(messenger) + late.start() + try: + assert main_socket.close_calls == 0 + finally: + main_socket.release_first_send.set() + first.join(timeout=2) + stop_thread.join(timeout=2) + late.join(timeout=2) + + assert not first.is_alive() + assert not stop_thread.is_alive() + assert not late.is_alive() + assert len(errors) == 1 + assert isinstance(errors[0], RuntimeError) + assert str(errors[0]) == "ZMQMessenger is stopping or closed" + assert main_socket.send_multipart_calls == 1 + assert main_socket.close_calls == 1 + assert not main_socket.close_overlapped_send + assert messenger._closed + + +def test_stop_times_out_without_closing_under_blocked_main_socket_io(): + main_socket = _BlockingSendSocket() + messenger = _make_test_messenger(main_socket=main_socket) + send_errors = [] + + def run_send(): + try: + messenger.send([b"blocked"]) + except BaseException as error: + send_errors.append(error) + + send_thread = Thread(target=run_send) + send_thread.start() + assert main_socket.first_send_started.wait(timeout=1) + + started = time.monotonic() + with pytest.raises(RuntimeError, match="I/O did not quiesce"): + messenger.stop(timeout=0.05) + elapsed = time.monotonic() - started + + assert elapsed < 0.5 + assert main_socket.close_calls == 0 + assert not main_socket.close_overlapped_send + + main_socket.release_first_send.set() + send_thread.join(timeout=2) + assert not send_thread.is_alive() + assert send_errors == [] + + messenger.stop(timeout=1) + assert messenger._closed + assert main_socket.close_calls == 1 + + +def test_stop_does_not_deadlock_listener_callback_entering_send(): + main_socket = _CloseTrackingSocket() + callback_started = Event() + enter_send = Event() + send_rejected = Event() + messenger = _make_test_messenger(main_socket=main_socket) + + def listener_callback(): + callback_started.set() + assert enter_send.wait(timeout=2) + with pytest.raises(RuntimeError, match="stopping or closed"): + messenger.send([b"callback"]) + send_rejected.set() + + listener_thread = Thread(target=listener_callback) + messenger._listener_thread = listener_thread + listener_thread.start() + assert callback_started.wait(timeout=1) + + stop_thread = Thread(target=messenger.stop) + stop_thread.start() + _wait_for_closing(messenger) + enter_send.set() + stop_thread.join(timeout=2) + listener_thread.join(timeout=2) + + assert not stop_thread.is_alive() + assert not listener_thread.is_alive() + assert send_rejected.is_set() + assert main_socket.send_multipart_calls == 0 + assert main_socket.close_calls == 1 + assert messenger._closed + + +def test_stop_can_run_from_listener_thread(): + messenger = _make_test_messenger() + start_stop = Event() + errors = [] + + def stop_from_listener(): + assert start_stop.wait(timeout=1) + try: + messenger.stop(timeout=1) + except BaseException as error: + errors.append(error) + + listener_thread = Thread(target=stop_from_listener) + messenger._listener_thread = listener_thread + listener_thread.start() + start_stop.set() + listener_thread.join(timeout=2) + + assert not listener_thread.is_alive() + assert errors == [] + assert messenger._closed + assert messenger._socket.close_calls == 1 + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/disaggregated/test_receive_lifecycle.py b/tests/unittest/disaggregated/test_receive_lifecycle.py new file mode 100644 index 000000000000..c455726a89aa --- /dev/null +++ b/tests/unittest/disaggregated/test_receive_lifecycle.py @@ -0,0 +1,595 @@ +# 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. +"""CPU-only tests for native receiver transfer ownership.""" + +from __future__ import annotations + +import pytest + +from tensorrt_llm._torch.disaggregation.native.receive_lifecycle import ( + BounceState, + ExposureState, + LifecycleAction, + LogicalState, + PhysicalState, + RecvTransferRegistry, + WriterMode, + WriterResult, +) + +KEY = (101, 0) +WRITERS = frozenset({3, 7}) + + +def _actions(update) -> set[LifecycleAction]: + return set(update.actions) + + +def _prepare( + registry: RecvTransferRegistry, + *, + key=KEY, + writers=WRITERS, + bounce: bool = True, +): + # The initial containment implementation retains Python request ownership. + # Allocator-backed destination leases belong to a future phase. + update = registry.prepare( + key, + writers, + has_bounce_slot=bounce, + ) + assert update.accepted + return update + + +def _publish(registry: RecvTransferRegistry, rank: int, *, key=KEY) -> None: + update = registry.begin_publication(key, rank) + assert update.accepted and update.publication_allowed + assert registry.mark_published(key, rank).accepted + + +def _settle_bounce(registry: RecvTransferRegistry, *, key=KEY, succeeded: bool = True): + update = registry.finish_bounce_scatter(key, succeeded=succeeded) + assert not update.conflict + if succeeded: + assert update.accepted + return update + + +class TestAdmissionAndPublication: + def test_prepare_requires_an_exact_nonempty_writer_set(self) -> None: + registry = RecvTransferRegistry() + + with pytest.raises(ValueError, match="must not be empty"): + registry.prepare(KEY, set(), has_bounce_slot=False) + with pytest.raises(ValueError, match="non-negative integer"): + registry.prepare(KEY, {-1}, has_bounce_slot=False) + with pytest.raises(ValueError, match="non-negative integer"): + registry.prepare(KEY, {True}, has_bounce_slot=False) + + _prepare(registry) + snapshot = registry.context_snapshot(KEY) + assert snapshot is not None + assert snapshot.expected_writers == WRITERS + assert [writer.rank for writer in snapshot.writers] == [3, 7] + + def test_identical_prepare_is_idempotent_but_conflicting_prepare_is_rejected(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + + duplicate = registry.prepare( + KEY, + WRITERS, + has_bounce_slot=True, + ) + conflict = registry.prepare( + KEY, + {3}, + has_bounce_slot=True, + ) + + assert not duplicate.accepted and duplicate.duplicate and not duplicate.conflict + assert not conflict.accepted and conflict.conflict and not conflict.duplicate + assert registry.context_snapshot(KEY).expected_writers == WRITERS + + def test_cancel_before_publication_closes_gate_and_releases_once(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + + (cancelled,) = registry.cancel_request(KEY[0]) + + assert cancelled.logical_state is LogicalState.CANCELLED + assert cancelled.physical_state is PhysicalState.DRAINING + assert cancelled.writers_terminal + assert _actions(cancelled) == { + LifecycleAction.NOTIFY_CANCELLED, + LifecycleAction.RELEASE_BOUNCE, + } + settled = _settle_bounce(registry) + assert settled.physical_state is PhysicalState.DRAINED + assert _actions(settled) == {LifecycleAction.CONTEXT_DRAINED} + rejected = registry.begin_publication(KEY, 3) + assert not rejected.accepted and not rejected.publication_allowed + assert registry.cancel_request(KEY[0])[0].actions == () + + def test_begin_publication_is_the_atomic_cancellation_boundary(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + + authorized = registry.begin_publication(KEY, 3) + (cancelled,) = registry.cancel_request(KEY[0]) + + assert authorized.publication_allowed + assert cancelled.logical_state is LogicalState.CANCELLED + assert cancelled.physical_state is PhysicalState.IN_DOUBT + assert not cancelled.writers_terminal + assert LifecycleAction.RELEASE_BOUNCE not in cancelled.actions + snapshot = registry.context_snapshot(KEY) + writers = {writer.rank: writer for writer in snapshot.writers} + assert writers[3].exposure is ExposureState.POSSIBLY_EXPOSED + assert writers[7].exposure is ExposureState.NEVER_EXPOSED + + # The send may complete after cancellation because authorization won. + assert registry.mark_published(KEY, 3).accepted + terminal = registry.record_result(KEY, 3, WriterResult.FAILED, WriterMode.BOUNCE) + assert terminal.physical_state is PhysicalState.DRAINING + assert terminal.writers_terminal + assert LifecycleAction.RELEASE_BOUNCE in terminal.actions + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + + def test_definitive_nonpublication_unwinds_partial_fanout(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + assert registry.begin_publication(KEY, 3).publication_allowed + + terminal = registry.mark_never_published(KEY, 3) + + assert terminal.logical_state is LogicalState.FAILED + assert terminal.physical_state is PhysicalState.DRAINING + assert terminal.writers_terminal + snapshot = registry.context_snapshot(KEY) + assert all(writer.exposure is ExposureState.NEVER_EXPOSED for writer in snapshot.writers) + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + + def test_partial_publication_failure_retains_only_possibly_exposed_writer(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + assert registry.begin_publication(KEY, 3).publication_allowed + + failed = registry.fail_context(KEY, "second ZMQ publication failed") + + assert failed.logical_state is LogicalState.FAILED + assert failed.physical_state is PhysicalState.IN_DOUBT + assert not failed.writers_terminal + assert _actions(failed) == {LifecycleAction.NOTIFY_FAILURE} + snapshot = registry.context_snapshot(KEY) + writers = {writer.rank: writer for writer in snapshot.writers} + assert writers[3].exposure is ExposureState.POSSIBLY_EXPOSED + assert writers[7].exposure is ExposureState.NEVER_EXPOSED + + terminal = registry.record_result(KEY, 3, WriterResult.FAILED, WriterMode.BOUNCE) + assert terminal.physical_state is PhysicalState.DRAINING + assert terminal.writers_terminal + assert LifecycleAction.RELEASE_BOUNCE in terminal.actions + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + + def test_each_writer_can_cross_publication_boundary_only_once(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + + first = registry.begin_publication(KEY, 3) + duplicate = registry.begin_publication(KEY, 3) + + assert first.publication_allowed + assert not duplicate.publication_allowed + assert not duplicate.accepted and duplicate.duplicate + + def test_unexpected_writer_cannot_be_published(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + + update = registry.begin_publication(KEY, 9) + + assert not update.accepted and update.conflict + assert update.logical_state is LogicalState.FAILED + assert update.physical_state is PhysicalState.DRAINING + assert registry.context_snapshot(KEY).expected_writers == WRITERS + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + + +class TestWriterResults: + def test_first_failure_survives_consumer_detach_until_late_sibling_result(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + _publish(registry, 3) + _publish(registry, 7) + + failed = registry.record_result(KEY, 3, WriterResult.FAILED, WriterMode.BOUNCE) + + assert failed.logical_state is LogicalState.FAILED + assert failed.physical_state is PhysicalState.DRAINING + assert not failed.writers_terminal + assert _actions(failed) == {LifecycleAction.NOTIFY_FAILURE} + registry.detach_consumer(KEY[0]) + assert registry.context_snapshot(KEY) is not None + assert not registry.is_request_drained(KEY[0]) + + late = registry.record_result(KEY, 7, WriterResult.SUCCESS, WriterMode.BOUNCE) + + assert late.logical_state is LogicalState.FAILED + assert late.physical_state is PhysicalState.DRAINING + assert late.writers_terminal and not late.bounce_pending + assert LifecycleAction.START_BOUNCE_SCATTER not in late.actions + assert _actions(late) == {LifecycleAction.RELEASE_BOUNCE} + settled = _settle_bounce(registry) + assert settled.physical_state is PhysicalState.DRAINED + assert _actions(settled) == {LifecycleAction.CONTEXT_DRAINED} + + def test_unexpected_writer_does_not_satisfy_exact_writer_set(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, bounce=False) + _publish(registry, 3) + _publish(registry, 7) + + unexpected = registry.record_result(KEY, 9, WriterResult.SUCCESS, WriterMode.DIRECT) + + assert unexpected.conflict + assert unexpected.logical_state is LogicalState.FAILED + assert unexpected.physical_state is PhysicalState.IN_DOUBT + assert not unexpected.writers_terminal + snapshot = registry.context_snapshot(KEY) + assert all(writer.result is None for writer in snapshot.writers) + + registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + expected = registry.record_result(KEY, 7, WriterResult.SUCCESS, WriterMode.DIRECT) + assert expected.writers_terminal + assert expected.physical_state is PhysicalState.IN_DOUBT + + (quiesced,) = registry.mark_backend_quiesced() + assert quiesced.physical_state is PhysicalState.DRAINED + + def test_identical_result_is_idempotent(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=False) + _publish(registry, 3) + + first = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + duplicate = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + + assert first.logical_state is LogicalState.SUCCEEDED + assert LifecycleAction.NOTIFY_SUCCESS in first.actions + assert not duplicate.accepted and duplicate.duplicate + assert not duplicate.conflict and duplicate.actions == () + + def test_contradictory_result_fails_closed_until_backend_quiescence(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3, 7}, bounce=False) + _publish(registry, 3) + _publish(registry, 7) + registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + + conflict = registry.record_result(KEY, 3, WriterResult.FAILED, WriterMode.DIRECT) + + assert conflict.conflict + assert conflict.logical_state is LogicalState.FAILED + assert conflict.physical_state is PhysicalState.IN_DOUBT + registry.record_result(KEY, 7, WriterResult.SUCCESS, WriterMode.DIRECT) + assert not registry.is_request_drained(KEY[0]) + + registry.mark_backend_quiesced() + assert registry.is_request_drained(KEY[0]) + + def test_contradiction_after_single_writer_success_corrects_logical_outcome(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=False) + _publish(registry, 3) + + success = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + conflict = registry.record_result(KEY, 3, WriterResult.FAILED, WriterMode.DIRECT) + + assert LifecycleAction.NOTIFY_SUCCESS in success.actions + assert conflict.conflict + assert conflict.logical_state is LogicalState.FAILED + assert conflict.physical_state is PhysicalState.IN_DOUBT + assert conflict.actions == (LifecycleAction.NOTIFY_FAILURE,) + assert not registry.is_request_drained(KEY[0]) + + registry.mark_backend_quiesced() + assert registry.is_request_drained(KEY[0]) + + @pytest.mark.parametrize("writer_rank", [3, 9], ids=["expected", "unexpected"]) + def test_nonterminal_protocol_conflict_before_publication_requires_backend_fence( + self, writer_rank + ) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=False) + + conflict = registry.record_protocol_conflict(KEY, writer_rank, "unknown result status") + + assert conflict.conflict + assert conflict.logical_state is LogicalState.FAILED + assert conflict.physical_state is PhysicalState.IN_DOUBT + assert not registry.is_request_drained(KEY[0]) + + (quiesced,) = registry.mark_backend_quiesced() + assert quiesced.physical_state is PhysicalState.DRAINED + + def test_direct_target_rejects_bounce_result_mode(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=False) + _publish(registry, 3) + + conflict = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.BOUNCE) + + assert not conflict.accepted and conflict.conflict + assert conflict.logical_state is LogicalState.FAILED + assert conflict.physical_state is PhysicalState.IN_DOUBT + snapshot = registry.context_snapshot(KEY) + (writer,) = snapshot.writers + assert writer.mode is WriterMode.BOUNCE + assert writer.conflict + + (quiesced,) = registry.mark_backend_quiesced() + assert quiesced.physical_state is PhysicalState.DRAINED + + def test_bounce_offer_accepts_mixed_direct_and_bounce_writers(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, bounce=True) + _publish(registry, 3) + _publish(registry, 7) + + direct = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + ready = registry.record_result(KEY, 7, WriterResult.SUCCESS, WriterMode.BOUNCE) + + assert direct.accepted and not direct.conflict + assert direct.logical_state is LogicalState.PENDING + assert LifecycleAction.START_BOUNCE_SCATTER not in direct.actions + assert ready.accepted and ready.actions == (LifecycleAction.START_BOUNCE_SCATTER,) + assert ready.bounce_pending + snapshot = registry.context_snapshot(KEY) + assert snapshot is not None + modes = {writer.rank: writer.mode for writer in snapshot.writers} + assert modes == {3: WriterMode.DIRECT, 7: WriterMode.BOUNCE} + settled = _settle_bounce(registry) + assert settled.logical_state is LogicalState.SUCCEEDED + assert settled.physical_state is PhysicalState.DRAINED + + def test_bounce_offer_with_only_direct_writers_skips_scatter(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, bounce=True) + _publish(registry, 3) + _publish(registry, 7) + + registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + terminal = registry.record_result(KEY, 7, WriterResult.SUCCESS, WriterMode.DIRECT) + + assert terminal.logical_state is LogicalState.SUCCEEDED + assert terminal.physical_state is PhysicalState.DRAINING + assert terminal.writers_terminal + assert LifecycleAction.START_BOUNCE_SCATTER not in terminal.actions + assert _actions(terminal) == { + LifecycleAction.NOTIFY_SUCCESS, + LifecycleAction.RELEASE_BOUNCE, + } + settled = _settle_bounce(registry) + assert settled.logical_state is LogicalState.SUCCEEDED + assert settled.physical_state is PhysicalState.DRAINED + assert _actions(settled) == {LifecycleAction.CONTEXT_DRAINED} + + +class TestBounceSettlement: + def test_bounce_success_waits_for_scatter(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + _publish(registry, 3) + _publish(registry, 7) + + registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.BOUNCE) + ready = registry.record_result(KEY, 7, WriterResult.SUCCESS, WriterMode.BOUNCE) + + assert ready.logical_state is LogicalState.PENDING + assert ready.physical_state is PhysicalState.DRAINING + assert ready.writers_terminal and ready.bounce_pending + assert ready.actions == (LifecycleAction.START_BOUNCE_SCATTER,) + assert registry.context_snapshot(KEY).bounce_state is BounceState.SCATTERING + + settled = _settle_bounce(registry) + + assert settled.logical_state is LogicalState.SUCCEEDED + assert settled.physical_state is PhysicalState.DRAINED + assert not settled.bounce_pending + assert _actions(settled) == { + LifecycleAction.NOTIFY_SUCCESS, + LifecycleAction.CONTEXT_DRAINED, + } + + def test_scatter_failure_is_logical_failure_but_physically_drained(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=True) + _publish(registry, 3) + ready = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.BOUNCE) + assert ready.bounce_pending + + failed = _settle_bounce(registry, succeeded=False) + + assert failed.logical_state is LogicalState.FAILED + assert failed.physical_state is PhysicalState.DRAINED + assert LifecycleAction.NOTIFY_FAILURE in failed.actions + assert LifecycleAction.RELEASE_BOUNCE not in failed.actions + + def test_cancel_while_scatter_runs_waits_for_settlement(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=True) + _publish(registry, 3) + ready = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.BOUNCE) + assert ready.bounce_pending + + (cancelled,) = registry.cancel_request(KEY[0]) + assert cancelled.logical_state is LogicalState.CANCELLED + assert cancelled.physical_state is PhysicalState.IN_DOUBT + assert cancelled.bounce_pending + assert LifecycleAction.RELEASE_BOUNCE not in cancelled.actions + + settled = _settle_bounce(registry) + assert settled.logical_state is LogicalState.CANCELLED + assert settled.physical_state is PhysicalState.DRAINED + assert LifecycleAction.RELEASE_BOUNCE not in settled.actions + + +class TestTimeoutQuiescenceAndShutdown: + def test_timeout_retains_exposed_resources_until_late_results(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + _publish(registry, 3) + _publish(registry, 7) + + (timed_out,) = registry.timeout_request(KEY[0]) + + assert timed_out.logical_state is LogicalState.FAILED + assert timed_out.physical_state is PhysicalState.IN_DOUBT + assert not timed_out.writers_terminal + assert _actions(timed_out) == {LifecycleAction.NOTIFY_FAILURE} + + first = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.BOUNCE) + assert first.physical_state is PhysicalState.IN_DOUBT + late = registry.record_result(KEY, 7, WriterResult.FAILED, WriterMode.BOUNCE) + assert late.physical_state is PhysicalState.DRAINING + assert late.writers_terminal + assert LifecycleAction.START_BOUNCE_SCATTER not in late.actions + assert LifecycleAction.RELEASE_BOUNCE in late.actions + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + + def test_backend_quiescence_recovers_lost_result_without_claiming_success(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + _publish(registry, 3) + _publish(registry, 7) + registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.BOUNCE) + + (quiesced,) = registry.mark_backend_quiesced() + + assert quiesced.logical_state is LogicalState.FAILED + assert quiesced.physical_state is PhysicalState.DRAINING + assert quiesced.writers_terminal + assert LifecycleAction.START_BOUNCE_SCATTER not in quiesced.actions + assert LifecycleAction.RELEASE_BOUNCE in quiesced.actions + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + + def test_consumer_detach_suppresses_notification_not_retirement(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=False) + _publish(registry, 3) + registry.detach_consumer(KEY[0]) + + terminal = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + + assert terminal.logical_state is LogicalState.SUCCEEDED + assert LifecycleAction.NOTIFY_SUCCESS not in terminal.actions + assert LifecycleAction.CONTEXT_DRAINED in terminal.actions + + def test_request_drain_and_retirement_cover_every_slice(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry, key=(101, 0), writers={3}, bounce=False) + _prepare(registry, key=(101, 1), writers={3}, bounce=False) + + registry.cancel_request(101) + assert registry.is_request_drained(101) + assert registry.retire_request(101) == ((101, 0), (101, 1)) + assert registry.snapshot().contexts == () + + def test_atomic_retirement_keeps_every_slice_when_one_is_not_drained(self) -> None: + registry = RecvTransferRegistry() + first_key = (101, 0) + second_key = (101, 1) + _prepare(registry, key=first_key, writers={3}, bounce=False) + _prepare(registry, key=second_key, writers={3}, bounce=False) + _publish(registry, 3, key=first_key) + first = registry.record_result(first_key, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + assert first.physical_state is PhysicalState.DRAINED + + assert not registry.retire_request_if_drained(101) + assert registry.context_snapshot(first_key) is not None + assert registry.context_snapshot(second_key) is not None + + registry.cancel_request(101) + assert registry.retire_request_if_drained(101) + assert registry.context_snapshot(first_key) is None + assert registry.context_snapshot(second_key) is None + + def test_target_mode_tracks_offer_and_is_removed_with_request_index(self) -> None: + registry = RecvTransferRegistry() + direct_key = (101, 0) + bounce_key = (101, 1) + other_request_key = (202, 0) + _prepare(registry, key=direct_key, writers={3}, bounce=False) + _prepare(registry, key=bounce_key, writers={3}, bounce=True) + _prepare(registry, key=other_request_key, writers={7}, bounce=True) + + assert registry.target_mode(direct_key) is WriterMode.DIRECT + assert registry.target_mode(bounce_key) is WriterMode.BOUNCE + assert registry.target_mode((101, 99)) is None + + registry.cancel_request(101) + _settle_bounce(registry, key=bounce_key) + assert registry.retire_request_if_drained(101) + assert registry.target_mode(direct_key) is None + assert registry.target_mode(bounce_key) is None + assert registry.target_mode(other_request_key) is WriterMode.BOUNCE + + # Reusing the retired request ID must not retain stale index entries. + replacement_key = (101, 2) + _prepare(registry, key=replacement_key, writers={11}, bounce=False) + assert registry.target_mode(replacement_key) is WriterMode.DIRECT + updates = registry.cancel_request(101) + assert tuple(update.key for update in updates) == (replacement_key,) + + def test_shutdown_closes_admission_and_retains_possible_access(self) -> None: + registry = RecvTransferRegistry() + _prepare(registry) + _publish(registry, 3) + + (shutdown,) = registry.begin_shutdown() + + assert not registry.snapshot().accepting + assert shutdown.logical_state is LogicalState.CANCELLED + assert shutdown.physical_state is PhysicalState.IN_DOUBT + assert not registry.is_drained() + rejected = registry.prepare((202, 0), {3}, has_bounce_slot=False) + assert not rejected.accepted + assert "closed" in rejected.reason + + (quiesced,) = registry.mark_backend_quiesced() + assert quiesced.physical_state is PhysicalState.DRAINING + assert not registry.is_drained() + assert _settle_bounce(registry).physical_state is PhysicalState.DRAINED + assert registry.is_drained() + assert registry.begin_shutdown() == () + + +class TestInitialContainmentOwnership: + def test_terminal_slice_emits_only_initial_containment_actions(self) -> None: + """The future allocator-backed phase is not part of this lifecycle API.""" + registry = RecvTransferRegistry() + _prepare(registry, writers={3}, bounce=False) + _publish(registry, 3) + + terminal = registry.record_result(KEY, 3, WriterResult.SUCCESS, WriterMode.DIRECT) + + assert terminal.physical_state is PhysicalState.DRAINED + assert _actions(terminal) == { + LifecycleAction.NOTIFY_SUCCESS, + LifecycleAction.CONTEXT_DRAINED, + } diff --git a/tests/unittest/disaggregated/test_receive_ownership_wiring.py b/tests/unittest/disaggregated/test_receive_ownership_wiring.py new file mode 100644 index 000000000000..16dea6e6fb4c --- /dev/null +++ b/tests/unittest/disaggregated/test_receive_ownership_wiring.py @@ -0,0 +1,3252 @@ +# 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. +"""CPU-only wiring tests for native receive ownership.""" + +from __future__ import annotations + +import gc +import threading +import weakref +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base.transfer import ( + KVSlice, + SessionArgsBase, + SessionStatus, + WaitResult, +) +from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import IdentityMapper +from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap +from tensorrt_llm._torch.disaggregation.native.receive_lifecycle import ( + LogicalState, + PhysicalState, + RecvTransferRegistry, +) +from tensorrt_llm._torch.disaggregation.native.transfer import ( + _KV_RESULT_PREFIX, + _NON_DRAINED_TRANSFER_WORKERS, + AgentResult, + AuxSendTask, + KVRecvTask, + KVSendTask, + MessageType, + RankInfoServer, + Receiver, + RecvReqInfo, + RxSession, + Sender, + SendOperationState, + TaskStatus, + TransferSourceInDoubtError, + TransferWorker, + TxSession, + WriteMeta, + WriteMetaType, +) +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.page import ( + BUFFER_ENTRY_DTYPE, + AttentionLayerGroup, + KVCachePageTable, + PhysicalPool, + PhysicalPoolGroup, + PoolView, +) +from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle + +REQUEST_ID = 101 +WRITER_RANK = 3 + + +def _recv_info(*, writer_rank: int = WRITER_RANK, slice_id: int = 0) -> RecvReqInfo: + return RecvReqInfo( + sender_req_id=REQUEST_ID, + instance_name="receiver", + instance_rank=writer_rank, + block_ids_per_layer_groups=[], + unique_rid=REQUEST_ID, + aux_slot=7, + slice_id=slice_id, + ) + + +def _make_send_task_and_session(*, channel: str = "kv"): + params = _params() + session = SimpleNamespace( + lock=threading.Lock(), + kv_tasks=[], + aux_task=None, + status=SessionStatus.READY, + disagg_request_id=REQUEST_ID, + source_owner=object(), + _accepting_operations=True, + _closed=False, + ) + if channel == "kv": + task = KVSendTask(KVSlice(), params, 0, session=session) + session.kv_tasks.append(task) + else: + task = AuxSendTask(params, 7, session=session) + session.aux_task = task + task._unique_rid = REQUEST_ID + return task, session + + +def _make_owned_send_task_and_session(*, channel: str, expected_transfers: int): + sender = _make_sender_for_shutdown() + sender._shutdown_complete = True + sender.setup_session = Mock() + sender.cancel_session = Mock() + sender.retry_terminal_results = Mock() + sender._send_operation_message = Mock(return_value=True) + sender._enqueue_owned = Mock() + sender._instance_rank = 0 + sender._bounce = Mock() + + params = _params() + session = TxSession(REQUEST_ID, params, sender, source_owner=object()) + session.receiver_ready = True + if channel == "kv": + task = KVSendTask(KVSlice(), params, 0, session=session) + session.kv_tasks.append(task) + else: + completed_kv_task = KVSendTask(KVSlice(), params, 0, session=session) + assert completed_kv_task.complete() + session.kv_tasks.append(completed_kv_task) + task = AuxSendTask(params, 7, session=session) + session.aux_task = task + session._need_aux = True + task._unique_rid = REQUEST_ID + + write_meta = WriteMeta( + task=task, + expected_transfers=expected_transfers, + peer_name="receiver1", + peer_rank=1, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + slice_id=0 if channel == "kv" else None, + meta_type=WriteMetaType.KV if channel == "kv" else WriteMetaType.AUX, + session=session, + ) + if channel == "kv": + sender._build_kv_write_meta = Mock(return_value=write_meta) + else: + sender._build_aux_write_meta = Mock(return_value=write_meta) + return sender, session, task, write_meta + + +def _make_sender_for_shutdown() -> Sender: + sender = object.__new__(Sender) + sender._shutdown_attempt_lock = threading.Lock() + sender._operation_admission_lock = threading.RLock() + sender._dealer_admission_closed = False + sender._shutdown = False + sender._shutdown_complete = False + sender._listener_stopped = True + sender._shutdown_sentinels_sent = False + sender._messenger = Mock() + sender._send_task_queues = [] + sender._worker_threads = [] + sender._failed_thread_dealers = [] + sender._failed_thread_dealers_lock = threading.Lock() + sender._in_doubt_transfers = [] + sender._in_doubt_transfers_lock = threading.Lock() + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._pre_cancelled_rids = set() + sender._pre_session_terminal_results = {} + sender._pre_session_terminal_results_lock = threading.Lock() + sender._pre_session_terminal_retry_lock = threading.Lock() + sender._next_pre_session_terminal_retry_at = 0.0 + sender._peer_requests = {} + sender._peer_requests_timestamps = {} + sender._peer_requests_lock = threading.Lock() + sender._loaded_remote_agents = set() + sender._loaded_remote_agents_lock = threading.Lock() + sender._agent = Mock() + sender._dealers = {} + sender._dealers_lock = threading.Lock() + sender._instance_rank = 0 + return sender + + +class _FakeBounce: + enabled = False + + def __init__(self) -> None: + self.logical_failure_callbacks: list = [] + self.protocol_conflict_callbacks: list = [] + self.backend_quiesced_callbacks: list = [] + self.release_calls: list[tuple[int, int]] = [] + + def is_bounced(self, _key) -> bool: + return False + + def reserve(self, _receiver_req, _writer_ranks, **_kwargs) -> bool: + return False + + def release_idle_reservation(self, key) -> None: + self.release_calls.append(key) + + def mark_logical_failure(self, key, on_done) -> None: + self.logical_failure_callbacks.append((key, on_done)) + + def mark_protocol_conflict(self, key, on_done) -> None: + self.protocol_conflict_callbacks.append((key, on_done)) + + def mark_backend_quiesced(self, key, on_done) -> None: + self.backend_quiesced_callbacks.append((key, on_done)) + + def retry_settlements(self, _scope=None) -> bool: + return True + + +class _BlockingBounce(_FakeBounce): + enabled = True + + def __init__(self) -> None: + super().__init__() + self.reserve_started = threading.Event() + self.allow_reserve = threading.Event() + self._state_lock = threading.Lock() + self._reservation_active = False + self.released_active_reservations: list[tuple[int, int]] = [] + + def reserve(self, _receiver_req, _writer_ranks, **_kwargs) -> bool: + self.reserve_started.set() + assert self.allow_reserve.wait(timeout=2) + with self._state_lock: + self._reservation_active = True + return True + + def release_idle_reservation(self, key) -> None: + super().release_idle_reservation(key) + with self._state_lock: + if self._reservation_active: + self._reservation_active = False + self.released_active_reservations.append(key) + + +class _FailOnceLogicalFailureBounce(_FakeBounce): + def __init__(self) -> None: + super().__init__() + self.logical_failure_attempts: list[tuple[int, int]] = [] + + def mark_logical_failure(self, key, on_done) -> None: + self.logical_failure_attempts.append(key) + if len(self.logical_failure_attempts) == 1: + raise RuntimeError("injected cancellation failure") + super().mark_logical_failure(key, on_done) + + +def _params( + *, + request_id: int = REQUEST_ID, + ctx_dp_rank: int | None = 0, + schedule_style: DisaggScheduleStyle = DisaggScheduleStyle.CONTEXT_FIRST, +) -> DisaggregatedParams: + return DisaggregatedParams( + disagg_request_id=request_id, + ctx_request_id=request_id, + ctx_dp_rank=ctx_dp_rank, + schedule_style=schedule_style, + ) + + +def _make_receiver(*, bounce=None, writer_ranks=(WRITER_RANK,)) -> Receiver: + receiver = object.__new__(Receiver) + receiver._shutdown = True # Keep object.__new__ fixtures out of __del__ cleanup. + receiver._shutdown_started = False + receiver._shutdown_attempt_lock = threading.Lock() + receiver._dealer_admission_open = True + receiver._dealers = {} + receiver._dealers_lock = threading.Lock() + receiver._sessions = {} + receiver._sessions_lock = threading.Lock() + receiver._pre_cancelled_rids = set() + receiver._recv_registry = RecvTransferRegistry() + receiver._bounce_lifecycle_delivery_lock = threading.Lock() + receiver._pending_bounce_lifecycle_deliveries = {} + receiver._pending_bounce_logical_failure_deliveries = {} + receiver._bounce = bounce or _FakeBounce() + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0), + get_peer_overlap=lambda _peer_info, _dp_rank: PeerOverlap( + overlap_pp_size=1, + duplicate_head_factor=1, + ranks=list(writer_ranks), + ), + ) + peer_info = SimpleNamespace( + dp_size=1, + sender_endpoints=[f"tcp://sender-{rank}" for rank in range(max(writer_ranks) + 1)], + ) + receiver._get_sender_info = lambda _params: peer_info + receiver._build_recv_req_info = lambda task: RecvReqInfo( + sender_req_id=REQUEST_ID, + instance_name="receiver", + instance_rank=0, + block_ids_per_layer_groups=[], + unique_rid=task._unique_rid, + aux_slot=task._aux_slot, + slice_id=task.slice_id, + ) + receiver._request_sender_data = Mock() + receiver.send_cancel_to_senders = Mock() + # Most object.__new__ fixtures intentionally omit the real registrar's + # page-table/mapping metadata. Tests exercising bounce admission opt in to + # a proven single-writer map unless they explicitly test proof failure. + receiver._single_writer_bounce_exact = Mock(return_value=True) + return receiver + + +def test_destination_intervals_use_view_index_for_sparse_physical_pool() -> None: + receiver = object.__new__(Receiver) + page_table = KVCachePageTable( + tokens_per_block=32, + layer_groups=[ + AttentionLayerGroup( + pool_group_idx=0, + pool_views=[ + PoolView( + pool_idx=1, + buffer_entries=np.empty(0, dtype=BUFFER_ENTRY_DTYPE), + ) + ], + ) + ], + pool_groups=[ + PhysicalPoolGroup( + pools=[ + PhysicalPool(base_address=0x1000, slot_bytes=8, num_slots=16), + PhysicalPool(base_address=0x4000, slot_bytes=16, num_slots=16), + ] + ) + ], + ) + extractor = KVRegionExtractorV1(page_table) + receiver._registrar = SimpleNamespace(self_extractor=extractor) + task = SimpleNamespace( + _kv_slice=SimpleNamespace( + block_ids_per_layer_groups=[[7]], + mamba_state_index=None, + ) + ) + + intervals = receiver._destination_intervals(task) + + assert intervals == {(0x4070, 16)} + + +def _make_legacy_session(cohorts) -> tuple[RxSession, KVRecvTask]: + params = _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST) + task = KVRecvTask(REQUEST_ID, KVSlice(), 0, params, aux_slot=None) + task.set_valid_writer_cohorts(cohorts) + for rank in {rank for cohort in cohorts for rank in cohort}: + task.mark_writer_exposed(rank) + task.close_publication() + receiver = _make_receiver() + session = object.__new__(RxSession) + session._closed = True # Keep object.__new__ fixtures out of __del__ cleanup. + session._base_args = SessionArgsBase(params) + session._receiver = receiver + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.aux_slot = 7 + session._kv_tasks = [task] + session._need_aux = True + session._aux_results = {} + session._aux_exposed_writer_ranks = set(task._exposed_writer_ranks) + session._aux_publication_closed = True + session._aux_result_conflict = False + session._aux_drained = False + session._aux_status = TaskStatus.INIT + session._terminal_status = None + session._exception = None + session._sender_endpoints = set() + session._selected_writer_cohort = None + return session, task + + +@pytest.mark.parametrize("channel", ["kv", "aux"]) +@pytest.mark.parametrize( + "cohorts,results,expected_history,expected_conflict,expected_status", + [ + ( + ((1, 2),), + ((1, AgentResult.SUCCESS), (2, AgentResult.SUCCESS)), + (False, True), + False, + TaskStatus.TRANSFERRED, + ), + ( + ((1, 2), (3, 4)), + ( + (3, AgentResult.SUCCESS), + (4, AgentResult.SUCCESS), + ), + (False, True), + False, + TaskStatus.TRANSFERRED, + ), + ( + ((1, 2),), + ((9, AgentResult.SUCCESS),), + (False,), + True, + TaskStatus.ERROR, + ), + ( + ((1, 2), (3, 4)), + ((1, AgentResult.SUCCESS), (3, AgentResult.SUCCESS)), + (False, False), + True, + TaskStatus.ERROR, + ), + ( + ((1, 2), (3, 4)), + ( + (3, AgentResult.SUCCESS), + (1, AgentResult.SUCCESS), + (2, AgentResult.SUCCESS), + ), + (False, False, False), + True, + TaskStatus.ERROR, + ), + ( + ((1, 2),), + ( + (1, AgentResult.SUCCESS), + (1, AgentResult.SUCCESS), + (2, AgentResult.SUCCESS), + ), + (False, False, True), + False, + TaskStatus.TRANSFERRED, + ), + ( + ((1, 2),), + ((1, AgentResult.SUCCESS), (1, AgentResult.FAILED)), + (False, False), + True, + TaskStatus.ERROR, + ), + ( + ((1, 2),), + ((1, AgentResult.FAILED), (2, AgentResult.SUCCESS)), + (False, True), + False, + TaskStatus.ERROR, + ), + ], + ids=[ + "deterministic-valid", + "legacy-adp-selected-cohort", + "unexpected-writer", + "partial-results-across-adp-cohorts", + "complete-cohort-plus-partial-sibling", + "idempotent-duplicate", + "conflicting-duplicate", + "failure-then-late-sibling", + ], +) +def test_legacy_kv_and_aux_follow_selected_candidate_cohort( + channel, + cohorts, + results, + expected_history, + expected_conflict, + expected_status, +) -> None: + session, task = _make_legacy_session(cohorts) + drained_history = [] + + for rank, status in results: + if channel == "kv": + session.process_kv_agent_result(rank, 0, True, status) + drained_history.append(task.legacy_resources_drained) + else: + session.process_aux_agent_result(rank, status) + drained_history.append(session._aux_drained) + + assert tuple(drained_history) == expected_history + if channel == "kv": + assert task._legacy_result_conflict is expected_conflict + assert task.status is expected_status + else: + assert session._aux_result_conflict is expected_conflict + assert session._aux_status is expected_status + + +def test_aux_protocol_conflict_requires_full_exposed_ledger_or_backend_fence() -> None: + session, _task = _make_legacy_session(((1, 2),)) + + assert not session._mark_aux_writer_exposed_locked(9) + assert 9 not in session._aux_exposed_writer_ranks + + session.process_aux_protocol_conflict(9, "unknown result status") + session.process_aux_agent_result(1, AgentResult.SUCCESS) + session.process_aux_agent_result(2, AgentResult.SUCCESS) + + assert session._selected_writer_cohort == frozenset({1, 2}) + assert session._aux_results.keys() == {1, 2} + assert not session._aux_drained + + session.mark_backend_quiesced() + + assert session._aux_drained + assert session._aux_status is TaskStatus.ERROR + + +def test_aux_protocol_conflict_reopens_initially_drained_exposure_ledger() -> None: + session, _task = _make_legacy_session(((1, 2),)) + session._aux_exposed_writer_ranks.clear() + session._aux_drained = True + + session.process_aux_protocol_conflict(9, "unknown result status") + + assert session._aux_exposed_writer_ranks == {9} + assert not session._aux_drained + + session.mark_backend_quiesced() + assert session._aux_drained + + +def test_aux_protocol_conflict_before_first_receive_requires_backend_fence() -> None: + receiver = _make_receiver() + aux_buffer = SimpleNamespace( + alloc_slot=Mock(return_value=SimpleNamespace(id=7)), + free_slot=Mock(), + ) + session = RxSession( + REQUEST_ID, + _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + receiver, + aux_buffer=aux_buffer, + ) + + assert session._aux_drained + + session.process_aux_protocol_conflict(9, "unknown result status") + + assert session._aux_exposed_writer_ranks == {9} + assert not session._aux_drained + assert session.has_untracked_receive_activity() + + session.mark_backend_quiesced() + assert session._aux_drained + session.close() + aux_buffer.free_slot.assert_called_once_with(7) + + +def test_adp_cross_channel_cohort_conflict_reopens_aux_and_retains_kv() -> None: + session, task = _make_legacy_session(((1, 2), (3, 4))) + + session.process_aux_agent_result(1, AgentResult.SUCCESS) + session.process_aux_agent_result(2, AgentResult.SUCCESS) + assert session._selected_writer_cohort == frozenset({1, 2}) + assert session._aux_drained + + session.process_kv_agent_result(3, 0, True, AgentResult.SUCCESS) + + assert task._legacy_result_conflict + assert not task.legacy_resources_drained + assert session._aux_result_conflict + assert not session._aux_drained + + +def test_adp_selected_cohort_is_consistent_across_slices() -> None: + session, first = _make_legacy_session(((1, 2), (3, 4))) + session._need_aux = False + session._aux_drained = True + second = KVRecvTask(REQUEST_ID, KVSlice(), 1, first._params, aux_slot=None) + second.set_valid_writer_cohorts(((1, 2), (3, 4))) + for rank in (1, 2, 3, 4): + second.mark_writer_exposed(rank) + second.close_publication() + session._kv_tasks.append(second) + + session.process_kv_agent_result(1, 0, True, AgentResult.SUCCESS) + session.process_kv_agent_result(2, 0, True, AgentResult.SUCCESS) + assert first.status is TaskStatus.TRANSFERRED + + session.process_kv_agent_result(3, 1, True, AgentResult.SUCCESS) + + assert second._legacy_result_conflict + assert not second.legacy_resources_drained + assert second.status is TaskStatus.ERROR + + +@pytest.mark.parametrize("status_code", [0, 255], ids=["known", "unknown"]) +def test_negative_wire_slice_id_does_not_alias_last_task(status_code) -> None: + session, task = _make_legacy_session(((WRITER_RANK,),)) + receiver = session._receiver + receiver._sessions[REQUEST_ID] = session + + receiver._process_kv_agent_result( + b"sender", + [ + MessageType.KV_AGENT_RESULT, + _KV_RESULT_PREFIX.pack(WRITER_RANK, REQUEST_ID, -1, True, status_code), + ], + ) + + assert task._legacy_results == {} + assert not task._legacy_result_conflict + assert task.status is TaskStatus.INIT + assert session._aux_results == {} + + +def test_dispatch_configures_one_deterministic_writer_cohort() -> None: + receiver = _make_receiver(writer_ranks=(1, 2)) + session = RxSession(REQUEST_ID, _params(ctx_dp_rank=0), receiver) + session._closed = True + + session.receive(KVSlice()) + + task = session._kv_tasks[0] + assert task._valid_writer_cohorts == (frozenset({1, 2}),) + assert task.expected_transfers == 2 + assert task.lifecycle_managed + assert receiver._request_sender_data.call_count == 2 + + +def test_cancel_during_partial_publication_drains_only_exposed_aux_writer() -> None: + receiver = _make_receiver(writer_ranks=(1, 2)) + aux_buffer = SimpleNamespace( + alloc_slot=Mock(return_value=SimpleNamespace(id=7)), + free_slot=Mock(), + ) + session = RxSession( + REQUEST_ID, + _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + receiver, + aux_buffer=aux_buffer, + ) + session._closed = True + + receiver._request_sender_data.side_effect = lambda *_args: session.cancel() + + session.receive(KVSlice(is_last_slice=True)) + + assert receiver._request_sender_data.call_count == 1 + assert session._aux_exposed_writer_ranks == {1} + assert session._aux_results == {} + assert session._aux_publication_closed + assert not session._aux_drained + + session.process_aux_agent_result(1, AgentResult.FAILED) + + assert session._aux_results == {1: AgentResult.FAILED} + assert session._aux_drained + assert session._aux_status is TaskStatus.ERROR + + +def test_later_suppressed_slice_cannot_revoke_earlier_aux_exposure() -> None: + receiver = _make_receiver(writer_ranks=(1, 2)) + aux_buffer = SimpleNamespace( + alloc_slot=Mock(return_value=SimpleNamespace(id=7)), + free_slot=Mock(), + ) + session = RxSession( + REQUEST_ID, + _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + receiver, + aux_buffer=aux_buffer, + ) + session._closed = True + + session.receive(KVSlice(is_last_slice=False)) + assert session._aux_exposed_writer_ranks == {1, 2} + assert not session._aux_publication_closed + + receiver._request_sender_data.side_effect = lambda *_args: session.cancel() + session.receive(KVSlice(is_last_slice=True)) + + assert receiver._request_sender_data.call_count == 3 + assert session._aux_exposed_writer_ranks == {1, 2} + assert session._aux_results == {} + assert session._aux_publication_closed + + session.process_aux_agent_result(1, AgentResult.FAILED) + assert not session._aux_drained + + session.process_aux_agent_result(2, AgentResult.FAILED) + assert session._aux_drained + assert session._aux_status is TaskStatus.ERROR + + +def test_dispatch_configures_each_adp_group_as_a_candidate_cohort() -> None: + receiver = _make_receiver(writer_ranks=(1, 2, 3, 4)) + receiver._registrar.get_peer_overlap = lambda _peer_info, dp_rank: PeerOverlap( + overlap_pp_size=1, + duplicate_head_factor=1, + ranks=[1, 2] if dp_rank == 0 else [3, 4], + ) + peer_info = SimpleNamespace( + dp_size=2, + sender_endpoints=[f"tcp://sender-{rank}" for rank in range(5)], + ) + receiver._get_sender_info = lambda _params: peer_info + session = RxSession(REQUEST_ID, _params(ctx_dp_rank=None), receiver) + session._closed = True + + session.receive(KVSlice()) + + task = session._kv_tasks[0] + assert task._valid_writer_cohorts == ( + frozenset({1, 2}), + frozenset({3, 4}), + ) + assert task.expected_transfers == 2 + assert not task.lifecycle_managed + assert receiver._request_sender_data.call_count == 4 + + +def test_pre_cancelled_session_closes_receive_admission() -> None: + receiver = _make_receiver() + receiver._pre_cancelled_rids.add(REQUEST_ID) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + + assert session.status is SessionStatus.CANCELLED + + with pytest.raises(RuntimeError, match="new receive slices are not accepted"): + session.receive(KVSlice()) + + assert session._kv_tasks == [] + receiver._request_sender_data.assert_not_called() + assert REQUEST_ID in receiver._pre_cancelled_rids + + +def test_receiver_pre_cancel_setup_rolls_back_and_retains_tombstone_on_failure() -> None: + receiver = _make_receiver() + receiver._pre_cancelled_rids.add(REQUEST_ID) + session = SimpleNamespace( + disagg_request_id=REQUEST_ID, + cancel=Mock(side_effect=RuntimeError("injected pre-cancel failure")), + ) + + with pytest.raises(RuntimeError, match="pre-cancel failure"): + receiver.setup_session(session) + + assert REQUEST_ID not in receiver._sessions + assert REQUEST_ID in receiver._pre_cancelled_rids + + +def test_receiver_cancel_for_live_session_creates_durable_tombstone() -> None: + receiver = _make_receiver() + session = SimpleNamespace(disagg_request_id=REQUEST_ID, cancel=Mock()) + receiver._sessions[REQUEST_ID] = session + + receiver._handle_cancel_session([MessageType.CANCEL_SESSION, str(REQUEST_ID).encode("ascii")]) + + session.cancel.assert_called_once_with() + assert REQUEST_ID in receiver._pre_cancelled_rids + + +def test_local_rx_session_cancel_creates_durable_tombstone_and_seals_admission() -> None: + receiver = _make_receiver() + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + + session.cancel() + + assert REQUEST_ID in receiver._pre_cancelled_rids + with pytest.raises(RuntimeError, match="new receive slices are not accepted"): + session.receive(KVSlice()) + + +def test_local_rx_session_timeout_creates_durable_tombstone_and_seals_admission() -> None: + receiver = _make_receiver() + session = RxSession(REQUEST_ID, _params(), receiver, timeout_s=0) + session._closed = True + session.resources_drained = Mock(side_effect=(False, False, True)) + + session.wait_complete(blocking=True) + + assert REQUEST_ID in receiver._pre_cancelled_rids + with pytest.raises(RuntimeError, match="new receive slices are not accepted"): + session.receive(KVSlice()) + receiver.send_cancel_to_senders.assert_called_once_with(REQUEST_ID, set()) + + +def test_receiver_shutdown_closes_admission_and_duplicate_ids_are_rejected() -> None: + receiver = _make_receiver() + receiver.begin_shutdown() + + with pytest.raises(RuntimeError, match="shutting down"): + receiver.setup_session(SimpleNamespace(disagg_request_id=REQUEST_ID)) + + receiver = _make_receiver() + first = SimpleNamespace(disagg_request_id=REQUEST_ID) + receiver.setup_session(first) + with pytest.raises(RuntimeError, match="already registered"): + receiver.setup_session(SimpleNamespace(disagg_request_id=REQUEST_ID)) + + +def test_receiver_shutdown_retries_and_continues_after_session_cancel_failure() -> None: + receiver = _make_receiver() + first = SimpleNamespace( + disagg_request_id=REQUEST_ID, + cancel=Mock(side_effect=[RuntimeError("injected failure"), None]), + ) + second = SimpleNamespace(disagg_request_id=REQUEST_ID + 1, cancel=Mock()) + receiver._sessions = {REQUEST_ID: first, REQUEST_ID + 1: second} + + with pytest.raises(RuntimeError, match="encountered 1 error"): + receiver.begin_shutdown() + + assert receiver._shutdown_started + first.cancel.assert_called_once_with() + second.cancel.assert_called_once_with() + + receiver.begin_shutdown() + + assert first.cancel.call_count == 2 + assert second.cancel.call_count == 2 + + +def test_receiver_shutdown_is_serialized_and_closes_dealer_admission() -> None: + receiver = _make_receiver() + receiver._shutdown = False + receiver._listener_stopped = False + receiver._messenger = Mock() + dealer = Mock() + receiver._dealers = {"tcp://sender": dealer} + listener_stop_started = threading.Event() + allow_listener_stop = threading.Event() + + def stop_listener() -> None: + listener_stop_started.set() + assert allow_listener_stop.wait(timeout=2) + + receiver._messenger.stop.side_effect = stop_listener + results = [] + threads = [ + threading.Thread(target=lambda: results.append(receiver.shutdown())) for _ in range(2) + ] + threads[0].start() + assert listener_stop_started.wait(timeout=1) + threads[1].start() + + with pytest.raises(RuntimeError, match="dealer admission is closed"): + receiver._get_or_connect_dealer("tcp://late-sender") + + allow_listener_stop.set() + for thread in threads: + thread.join(timeout=2) + + assert all(not thread.is_alive() for thread in threads) + assert results == [True, True] + receiver._messenger.stop.assert_called_once_with() + dealer.stop.assert_called_once_with() + + +def test_receiver_shutdown_does_not_remove_replacement_dealer() -> None: + receiver = _make_receiver() + receiver._shutdown = False + receiver._listener_stopped = True + old_dealer = Mock() + replacement = Mock() + endpoint = "tcp://sender" + receiver._dealers = {endpoint: old_dealer} + old_dealer.stop.side_effect = lambda: receiver._dealers.__setitem__(endpoint, replacement) + + assert receiver.shutdown() is False + assert receiver._dealers[endpoint] is replacement + + assert receiver.shutdown() is True + replacement.stop.assert_called_once_with() + + +def test_rx_session_cancel_retries_failed_contexts_and_continues_siblings() -> None: + bounce = _FailOnceLogicalFailureBounce() + receiver = _make_receiver(bounce=bounce) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + session.receive(KVSlice(is_last_slice=False)) + session.receive(KVSlice(is_last_slice=True)) + + with pytest.raises(RuntimeError, match="cancellation encountered 1 error"): + session.cancel() + + assert bounce.logical_failure_attempts == [(REQUEST_ID, 0), (REQUEST_ID, 1)] + assert session.status is SessionStatus.CANCELLED + assert all(task.status is TaskStatus.ERROR for task in session._kv_tasks) + + session.cancel() + + assert bounce.logical_failure_attempts == [ + (REQUEST_ID, 0), + (REQUEST_ID, 1), + (REQUEST_ID, 0), + (REQUEST_ID, 1), + ] + + +def test_rx_session_cancel_retries_legacy_idle_release_after_task_failure() -> None: + bounce = _FakeBounce() + bounce.release_idle_reservation = Mock( + side_effect=[RuntimeError("injected idle-release failure"), None] + ) + receiver = _make_receiver(bounce=bounce) + session = RxSession( + REQUEST_ID, + _params( + ctx_dp_rank=None, + schedule_style=DisaggScheduleStyle.GENERATION_FIRST, + ), + receiver, + ) + session._closed = True + session.receive(KVSlice()) + + with pytest.raises(RuntimeError, match="injected idle-release failure"): + session.cancel() + + assert session._kv_tasks[0].status is TaskStatus.ERROR + + session.cancel() + + assert bounce.release_idle_reservation.call_count == 2 + + +def test_rx_session_serializes_slice_dispatch_before_closing_aux_publication() -> None: + receiver = _make_receiver() + aux_buffer = SimpleNamespace( + alloc_slot=Mock(return_value=SimpleNamespace(id=7)), + free_slot=Mock(), + ) + session = RxSession( + REQUEST_ID, + _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + receiver, + aux_buffer=aux_buffer, + ) + session._closed = True + aux_buffer.alloc_slot.assert_called_once_with() + assert session.aux_slot == 7 + first_send_started = threading.Event() + allow_first_send = threading.Event() + second_receive_started = threading.Event() + second_receive_finished = threading.Event() + call_lock = threading.Lock() + send_count = 0 + thread_errors = [] + + def request_sender_data(*_args) -> None: + nonlocal send_count + with call_lock: + send_count += 1 + current_send = send_count + if current_send == 1: + first_send_started.set() + assert allow_first_send.wait(timeout=2) + + receiver._request_sender_data.side_effect = request_sender_data + + def receive_first_slice() -> None: + try: + session.receive(KVSlice(is_last_slice=False)) + except Exception as e: + thread_errors.append(e) + + first_thread = threading.Thread(target=receive_first_slice) + + def receive_last_slice() -> None: + second_receive_started.set() + try: + session.receive(KVSlice(is_last_slice=True)) + except Exception as e: + thread_errors.append(e) + finally: + second_receive_finished.set() + + second_thread = threading.Thread(target=receive_last_slice) + first_thread.start() + assert first_send_started.wait(timeout=1) + second_thread.start() + assert second_receive_started.wait(timeout=1) + try: + assert not second_receive_finished.wait(timeout=0.05) + assert len(session._kv_tasks) == 1 + assert not session._aux_publication_closed + finally: + allow_first_send.set() + first_thread.join(timeout=2) + second_thread.join(timeout=2) + + assert not first_thread.is_alive() + assert not second_thread.is_alive() + assert not thread_errors + assert len(session._kv_tasks) == 2 + assert session._aux_publication_closed + assert session._aux_exposed_writer_ranks == {WRITER_RANK} + + with pytest.raises(RuntimeError, match="already admitted its last slice"): + session.receive(KVSlice()) + + +def test_rx_session_close_is_exactly_once_and_closes_receive_admission() -> None: + receiver = _make_receiver() + aux_buffer = SimpleNamespace( + alloc_slot=Mock(return_value=SimpleNamespace(id=7)), + free_slot=Mock(), + ) + session = RxSession( + REQUEST_ID, + _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + receiver, + aux_buffer=aux_buffer, + ) + session.mark_backend_quiesced() + start = threading.Barrier(3) + errors = [] + + def close() -> None: + start.wait() + try: + session.close() + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=close) for _ in range(2)] + for thread in threads: + thread.start() + start.wait() + for thread in threads: + thread.join(timeout=1) + + assert all(not thread.is_alive() for thread in threads) + assert not errors + aux_buffer.free_slot.assert_called_once_with(7) + assert REQUEST_ID not in receiver._sessions + + with pytest.raises(RuntimeError, match="new receive slices are not accepted"): + session.receive(KVSlice()) + + +def test_bounce_reservation_wait_is_outside_session_cancel_gate() -> None: + bounce = _BlockingBounce() + receiver = _make_receiver(bounce=bounce) + receiver._destination_intervals = Mock(return_value=()) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + receive_finished = threading.Event() + cancel_finished = threading.Event() + + def receive() -> None: + session.receive(KVSlice()) + receive_finished.set() + + def cancel() -> None: + session.cancel() + cancel_finished.set() + + receive_thread = threading.Thread(target=receive) + receive_thread.start() + assert bounce.reserve_started.wait(timeout=1) + + cancel_thread = threading.Thread(target=cancel) + cancel_thread.start() + try: + assert cancel_finished.wait(timeout=1), "cancel blocked behind bounce.reserve()" + finally: + bounce.allow_reserve.set() + receive_thread.join(timeout=2) + cancel_thread.join(timeout=2) + + assert receive_finished.is_set() + assert bounce.released_active_reservations == [(REQUEST_ID, 0)] + assert session._kv_tasks[0].status is TaskStatus.ERROR + receiver._request_sender_data.assert_not_called() + + +def test_receiver_shutdown_waits_for_admitted_receive_dispatch() -> None: + receiver = _make_receiver() + receiver._shutdown = False + receiver._listener_stopped = False + receiver._messenger = Mock() + discovery_started = threading.Event() + allow_discovery = threading.Event() + peer_info = SimpleNamespace( + dp_size=1, + sender_endpoints=[f"tcp://sender-{rank}" for rank in range(WRITER_RANK + 1)], + ) + + def get_sender_info(_params): + discovery_started.set() + allow_discovery.wait() + return peer_info + + receiver._get_sender_info = get_sender_info + session = RxSession(REQUEST_ID, _params(), receiver) + receive_errors = [] + + def receive() -> None: + try: + session.receive(KVSlice()) + except Exception as e: + receive_errors.append(e) + + receive_thread = threading.Thread( + target=receive, + name="blocked-receive", + daemon=True, + ) + shutdown_thread = None + receive_thread.start() + try: + assert discovery_started.wait(timeout=1) + + shutdown_results: list[bool] = [] + shutdown_errors: list[Exception] = [] + shutdown_finished = threading.Event() + + def shutdown() -> None: + try: + shutdown_results.append(receiver.shutdown()) + except Exception as e: + shutdown_errors.append(e) + finally: + shutdown_finished.set() + + shutdown_thread = threading.Thread( + target=shutdown, + name="receiver-shutdown", + daemon=True, + ) + shutdown_thread.start() + assert shutdown_finished.wait(timeout=5), ( + "receiver.shutdown() blocked behind admitted sender discovery" + ) + assert not shutdown_errors + assert shutdown_results == [False] + receiver._messenger.stop.assert_not_called() + assert not receiver.transfers_drained + with pytest.raises(RuntimeError, match="new receive slices are not accepted"): + session.receive(KVSlice()) + finally: + # Always unblock discovery so a failed deadlock assertion cannot leak + # either worker into the rest of the test process. + allow_discovery.set() + if shutdown_thread is not None: + shutdown_thread.join(timeout=2) + receive_thread.join(timeout=2) + + assert shutdown_thread is not None + assert not shutdown_thread.is_alive() + assert not receive_thread.is_alive() + assert not receive_errors + assert receiver.transfers_drained + assert receiver.shutdown() is True + receiver._messenger.stop.assert_called_once_with() + session.close() + + +def test_rx_session_close_waits_for_dispatch_and_prevents_later_receive() -> None: + receiver = _make_receiver() + send_started = threading.Event() + allow_send = threading.Event() + + def request_sender_data(*_args) -> None: + send_started.set() + assert allow_send.wait(timeout=2) + + receiver._request_sender_data.side_effect = request_sender_data + session = RxSession(REQUEST_ID, _params(), receiver) + receive_finished = threading.Event() + close_started = threading.Event() + close_finished = threading.Event() + close_errors = [] + + def receive() -> None: + session.receive(KVSlice()) + receive_finished.set() + + def close() -> None: + close_started.set() + try: + session.close() + except Exception as e: + close_errors.append(e) + finally: + close_finished.set() + + receive_thread = threading.Thread(target=receive) + receive_thread.start() + assert send_started.wait(timeout=1) + + close_thread = threading.Thread(target=close) + close_thread.start() + assert close_started.wait(timeout=1) + try: + assert not close_finished.wait(timeout=0.05) + assert REQUEST_ID in receiver._sessions + finally: + allow_send.set() + receive_thread.join(timeout=2) + close_thread.join(timeout=2) + + assert receive_finished.is_set() + assert close_finished.is_set() + assert len(close_errors) == 1 + assert "receive resources are not drained" in str(close_errors[0]) + assert REQUEST_ID in receiver._sessions + with pytest.raises(RuntimeError, match="new receive slices are not accepted"): + session.receive(KVSlice()) + session._closed = True + + +def test_mamba_receive_uses_direct_path_until_bound_layout_includes_state_bytes() -> None: + bounce = _FakeBounce() + bounce.enabled = True + bounce.reserve = Mock(return_value=True) + receiver = _make_receiver(bounce=bounce) + receiver._destination_intervals = Mock(return_value=()) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + + session.receive(KVSlice(mamba_state_index=0)) + + bounce.reserve.assert_not_called() + receiver._destination_intervals.assert_not_called() + assert receiver._recv_registry.target_mode((REQUEST_ID, 0)) is not None + + +def test_standalone_receive_dispatch_error_closes_publication_transactionally() -> None: + receiver = _make_receiver() + receiver.dispatch_task = Mock(side_effect=RuntimeError("prepare failed")) + destination_owner = object() + session = RxSession( + REQUEST_ID, + _params(), + receiver, + destination_owner=destination_owner, + ) + + with pytest.raises(RuntimeError, match="prepare failed"): + session.receive(KVSlice()) + + task = session._kv_tasks[0] + assert task.status is TaskStatus.ERROR + assert task._publication_closed + assert session._active_receive_dispatches == 0 + assert session.status is SessionStatus.ERROR + assert session.resources_drained() + session.close() + assert REQUEST_ID not in receiver._sessions + assert session._destination_owner is None + + +def test_standalone_receive_dispatch_error_retains_ambiguous_target_owner() -> None: + receiver = _make_receiver() + destination_owner = object() + session = RxSession( + REQUEST_ID, + _params(), + receiver, + destination_owner=destination_owner, + ) + key = (REQUEST_ID, 0) + + def fail_after_publication_started(task) -> None: + task.set_valid_writer_cohorts(((WRITER_RANK,),)) + task.lifecycle_managed = True + task.status = TaskStatus.TRANSFERRING + assert receiver._recv_registry.prepare(key, (WRITER_RANK,), has_bounce_slot=False).accepted + assert receiver._recv_registry.begin_publication(key, WRITER_RANK).publication_allowed + task.mark_writer_exposed(WRITER_RANK) + raise RuntimeError("publication outcome is ambiguous") + + receiver.dispatch_task = fail_after_publication_started + + with pytest.raises(RuntimeError, match="publication outcome is ambiguous"): + session.receive(KVSlice()) + + snapshot = receiver._recv_registry.context_snapshot(key) + assert snapshot.logical_state is LogicalState.CANCELLED + assert snapshot.physical_state is PhysicalState.IN_DOUBT + assert session._kv_tasks[0]._publication_closed + assert not session.resources_drained() + with pytest.raises(RuntimeError, match="not drained"): + session.close() + assert receiver._sessions[REQUEST_ID] is session + assert session._destination_owner is destination_owner + session._closed = True + + +def test_pp_fanin_with_uniform_peer_stages_uses_direct_path() -> None: + bounce = _FakeBounce() + bounce.enabled = True + bounce.reserve = Mock(return_value=True) + receiver = _make_receiver(bounce=bounce, writer_ranks=(0, 1)) + receiver._registrar.get_peer_overlap = lambda _peer_info, _dp_rank: PeerOverlap( + overlap_pp_size=2, + duplicate_head_factor=1, + peer_duplicate_head_factor=1, + ranks=[0, 1], + ) + receiver._get_sender_info = lambda _params: SimpleNamespace( + dp_size=1, + layer_num_per_pp=[20, 20], + sender_endpoints=["tcp://sender-0", "tcp://sender-1"], + ) + receiver._destination_intervals = Mock(return_value=()) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + + # A local [0, 30) stage intersects the uniform peer stages [0, 20) and + # [20, 40) by 20 and 10 layers. Peer uniformity therefore cannot authorize + # an equal two-writer bounce subdivision. + session.receive(KVSlice()) + + bounce.reserve.assert_not_called() + receiver._destination_intervals.assert_not_called() + assert receiver._request_sender_data.call_count == 2 + + +def test_tp_fanin_uses_direct_path_without_byte_accurate_writer_extents() -> None: + bounce = _FakeBounce() + bounce.enabled = True + bounce.reserve = Mock(return_value=True) + receiver = _make_receiver(bounce=bounce, writer_ranks=(0, 1)) + receiver._destination_intervals = Mock(return_value=()) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + + session.receive(KVSlice()) + + bounce.reserve.assert_not_called() + receiver._destination_intervals.assert_not_called() + assert receiver._request_sender_data.call_count == 2 + + +def test_partial_layer_mapping_uses_direct_path() -> None: + bounce = _FakeBounce() + bounce.enabled = True + bounce.reserve = Mock(return_value=True) + receiver = _make_receiver(bounce=bounce) + receiver._single_writer_bounce_exact = Receiver._single_writer_bounce_exact.__get__( + receiver, Receiver + ) + receiver._registrar.self_rank_info.layer_num_per_pp = [12] + receiver._get_sender_info = lambda _params: SimpleNamespace( + dp_size=1, + layer_num_per_pp=[10], + sender_endpoints=["tcp://sender-0"] * (WRITER_RANK + 1), + ) + receiver._destination_intervals = Mock(return_value=()) + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + + session.receive(KVSlice()) + + bounce.reserve.assert_not_called() + receiver._destination_intervals.assert_not_called() + receiver._request_sender_data.assert_called_once() + + +def _bounce_proof_page_table(*slot_bytes: int, tokens_per_block: int = 32): + pool_views = [ + PoolView( + pool_idx=pool_idx, + buffer_entries=np.empty(0, dtype=BUFFER_ENTRY_DTYPE), + ) + for pool_idx in range(len(slot_bytes)) + ] + return KVCachePageTable( + tokens_per_block=tokens_per_block, + layer_groups=[AttentionLayerGroup(pool_group_idx=0, pool_views=pool_views)], + pool_groups=[ + PhysicalPoolGroup( + pools=[ + PhysicalPool( + base_address=0x1000 + pool_idx * 0x1000, + slot_bytes=size, + num_slots=16, + ) + for pool_idx, size in enumerate(slot_bytes) + ] + ) + ], + ) + + +def _single_writer_bounce_proof(local_page_table, peer_page_table, mapping) -> bool: + receiver = object.__new__(Receiver) + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(layer_num_per_pp=[1], page_table=local_page_table), + get_pool_mapping=Mock(return_value=mapping), + get_kv_map=Mock(return_value=IdentityMapper()), + ) + peer_info = SimpleNamespace(layer_num_per_pp=[1], page_table=peer_page_table) + return receiver._single_writer_bounce_exact(peer_info) + + +def test_single_writer_bounce_requires_byte_exact_pool_bijection() -> None: + local_page_table = _bounce_proof_page_table(64, 128) + peer_page_table = _bounce_proof_page_table(64, 128) + + assert _single_writer_bounce_proof( + local_page_table, + peer_page_table, + {(0, 0): (0, 0), (0, 1): (0, 1)}, + ) + assert not _single_writer_bounce_proof( + local_page_table, + peer_page_table, + {(0, 0): (0, 0), (0, 1): (0, 0)}, + ) + + +def test_single_writer_bounce_rejects_asymmetric_slot_bytes() -> None: + assert not _single_writer_bounce_proof( + _bounce_proof_page_table(64), + _bounce_proof_page_table(128), + {(0, 0): (0, 0)}, + ) + + +def test_single_writer_bounce_rejects_mismatched_block_geometry() -> None: + assert not _single_writer_bounce_proof( + _bounce_proof_page_table(64, tokens_per_block=32), + _bounce_proof_page_table(64, tokens_per_block=64), + {(0, 0): (0, 0)}, + ) + + +def test_receiver_clear_session_does_not_hold_map_lock_during_bounce_retry() -> None: + receiver = _make_receiver() + + class _RetryingSession: + def resources_drained(self) -> bool: + assert receiver._get_session(REQUEST_ID) is self + return True + + session = _RetryingSession() + receiver._sessions[REQUEST_ID] = session + errors = [] + + def clear() -> None: + try: + receiver.clear_session(REQUEST_ID) + except Exception as e: + errors.append(e) + + thread = threading.Thread(target=clear) + thread.start() + thread.join(timeout=1) + + assert not thread.is_alive(), "clear_session deadlocked through resources_drained()" + assert not errors + assert REQUEST_ID not in receiver._sessions + + +@pytest.mark.parametrize("has_bounce", [False, True], ids=["direct", "bounce"]) +def test_unknown_managed_wire_status_remains_undrained(has_bounce) -> None: + bounce = _FakeBounce() + receiver = _make_receiver(bounce=bounce) + key = (REQUEST_ID, 0) + assert receiver._recv_registry.prepare( + key, + {WRITER_RANK}, + has_bounce_slot=has_bounce, + ).accepted + assert receiver._recv_registry.begin_publication(key, WRITER_RANK).publication_allowed + assert receiver._recv_registry.mark_published(key, WRITER_RANK).accepted + message = [ + MessageType.KV_AGENT_RESULT, + _KV_RESULT_PREFIX.pack(WRITER_RANK, REQUEST_ID, 0, True, 255), + ] + + receiver._process_kv_agent_result(b"sender", message) + + snapshot = receiver._recv_registry.context_snapshot(key) + assert snapshot is not None + assert snapshot.logical_state is LogicalState.FAILED + assert snapshot.physical_state is PhysicalState.IN_DOUBT + assert snapshot.writers[0].result is None + assert not receiver._recv_registry.is_request_drained(REQUEST_ID) + assert len(bounce.protocol_conflict_callbacks) == int(has_bounce) + + +def test_post_success_contradiction_corrects_receiver_session_outcome() -> None: + receiver = _make_receiver() + session = RxSession(REQUEST_ID, _params(), receiver) + session._closed = True + session.receive(KVSlice()) + message_prefix = (WRITER_RANK, REQUEST_ID, 0, True) + + receiver._process_kv_agent_result( + b"sender", + [MessageType.KV_AGENT_RESULT, _KV_RESULT_PREFIX.pack(*message_prefix, 0)], + ) + assert session._kv_tasks[0].status is TaskStatus.TRANSFERRED + + receiver._process_kv_agent_result( + b"sender", + [MessageType.KV_AGENT_RESULT, _KV_RESULT_PREFIX.pack(*message_prefix, 1)], + ) + + assert session._kv_tasks[0].status is TaskStatus.ERROR + assert session.status is SessionStatus.ERROR + assert not receiver._recv_registry.is_request_drained(REQUEST_ID) + + +@pytest.mark.parametrize("peer_rank", [WRITER_RANK, WRITER_RANK + 1]) +@pytest.mark.parametrize("has_bounce", [False, True], ids=["direct", "bounce"]) +def test_unknown_managed_wire_status_before_publication_requires_backend_fence( + has_bounce, peer_rank +) -> None: + bounce = _FakeBounce() + receiver = _make_receiver(bounce=bounce) + key = (REQUEST_ID, 0) + assert receiver._recv_registry.prepare( + key, + {WRITER_RANK}, + has_bounce_slot=has_bounce, + ).accepted + + receiver._process_kv_agent_result( + b"sender", + [ + MessageType.KV_AGENT_RESULT, + _KV_RESULT_PREFIX.pack(peer_rank, REQUEST_ID, 0, True, 255), + ], + ) + + snapshot = receiver._recv_registry.context_snapshot(key) + assert snapshot.logical_state is LogicalState.FAILED + assert snapshot.physical_state is PhysicalState.IN_DOUBT + assert not receiver._recv_registry.is_request_drained(REQUEST_ID) + assert len(bounce.protocol_conflict_callbacks) == int(has_bounce) + + +def test_unknown_aux_status_fails_logically_without_retiring_writer() -> None: + session, task = _make_legacy_session(((WRITER_RANK,),)) + receiver = session._receiver + receiver._sessions[REQUEST_ID] = session + + receiver._process_aux_agent_result( + b"sender", + [ + MessageType.AUX_AGENT_RESULT, + str(WRITER_RANK).encode("ascii"), + str(REQUEST_ID).encode("ascii"), + b"NEWER_STATUS", + ], + ) + + assert session._aux_status is TaskStatus.ERROR + assert session._aux_result_conflict + assert not session._aux_drained + assert session._aux_results == {} + assert not task.legacy_resources_drained + + session.process_aux_agent_result(WRITER_RANK, AgentResult.FAILED) + + assert session._aux_drained + assert session._aux_status is TaskStatus.ERROR + + +def test_sender_cancel_before_session_replays_failure_for_stored_and_late_requests() -> None: + info = _recv_info(writer_rank=0) + sender = object.__new__(Sender) + sender._instance_rank = WRITER_RANK + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._peer_requests = {REQUEST_ID: {(0, 0): info}} + sender._peer_requests_lock = threading.Lock() + sender._pre_cancelled_rids = set() + sender._send_failed_result_to_receiver = Mock() + sender._send_aux_failed_result_to_receiver = Mock() + + sender._handle_cancel_session([MessageType.CANCEL_SESSION, str(REQUEST_ID).encode("ascii")]) + + assert REQUEST_ID in sender._pre_cancelled_rids + sender._send_failed_result_to_receiver.assert_called_once_with(info) + sender._send_aux_failed_result_to_receiver.assert_called_once_with(info) + + sender._send_failed_result_to_receiver.reset_mock() + sender._send_aux_failed_result_to_receiver.reset_mock() + sender._save_peer_req_info = Mock() + sender._respond_with_kv(b"receiver", [MessageType.REQUEST_DATA, info.to_bytes()]) + sender._save_peer_req_info.assert_called_once() + sender._send_failed_result_to_receiver.assert_called_once() + sender._send_aux_failed_result_to_receiver.assert_called_once() + + sender._shutdown = False + late_session = Mock(disagg_request_id=REQUEST_ID) + sender.setup_session(late_session) + late_session.cancel.assert_called_once_with() + assert REQUEST_ID in sender._pre_cancelled_rids + + +def test_sender_shutdown_late_request_retains_kv_and_aux_no_access_results() -> None: + sender = _make_sender_for_shutdown() + sender._shutdown = True + sender._instance_rank = WRITER_RANK + sender._save_peer_req_info = Mock() + sender._send_operation_message = Mock(return_value=False) + info = _recv_info(writer_rank=0) + + sender._respond_with_kv(b"receiver", [MessageType.REQUEST_DATA, info.to_bytes()]) + + assert REQUEST_ID in sender._pre_cancelled_rids + assert len(sender._pre_session_terminal_results) == 2 + assert {key[1] for key in sender._pre_session_terminal_results} == { + WriteMetaType.KV, + WriteMetaType.AUX, + } + assert sender._has_pending_pre_session_terminal_results() + sender._shutdown_complete = True + + +def test_sender_shutdown_late_request_cannot_admit_session_source_access() -> None: + sender = _make_sender_for_shutdown() + sender._save_peer_req_info = Mock() + sender._send_operation_message = Mock(return_value=False) + session = TxSession( + request_id=REQUEST_ID, + params=_params(), + sender=sender, + source_owner=object(), + ) + kv_task = KVSendTask(KVSlice(), _params(), 0, session=session) + aux_task = AuxSendTask(_params(), 7, session=session) + session.kv_tasks.append(kv_task) + session.aux_task = aux_task + sender._shutdown = True + info = _recv_info() + key = (info.instance_name, info.instance_rank) + + sender._respond_with_kv(b"receiver", [MessageType.REQUEST_DATA, info.to_bytes()]) + + for task in (kv_task, aux_task): + snapshot = task.operation_snapshot(key) + assert snapshot is not None + assert snapshot[0] is SendOperationState.TERMINAL + assert not snapshot[2] + assert not task.source_access_active + + assert REQUEST_ID in sender._pre_cancelled_rids + kv_task.mark_operation_result_delivered(key) + aux_task.mark_operation_result_delivered(key) + kv_task.fail(RuntimeError("shutdown")) + aux_task.fail(RuntimeError("shutdown")) + session.close() + sender._shutdown_complete = True + + +def test_sender_strongly_retains_registered_session_until_explicit_clear() -> None: + class Session: + disagg_request_id = REQUEST_ID + + sender = _make_sender_for_shutdown() + session = Session() + session_ref = weakref.ref(session) + + sender.setup_session(session) + del session + gc.collect() + + assert session_ref() is sender._sessions[REQUEST_ID] + + sender.clear_session(REQUEST_ID) + gc.collect() + assert session_ref() is None + sender._shutdown_complete = True + + +def test_tx_session_setup_failure_rolls_back_sender_and_source_owner() -> None: + class SourceOwner: + pass + + sender = _make_sender_for_shutdown() + info = _recv_info() + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): info}} + sender._registrar = Mock() + sender._registrar.get_peer_rank_info.side_effect = RuntimeError("peer lookup failed") + source_owner = SourceOwner() + source_owner_ref = weakref.ref(source_owner) + + with pytest.raises(RuntimeError, match="peer lookup failed") as exc_info: + TxSession( + request_id=REQUEST_ID, + params=_params(), + sender=sender, + source_owner=source_owner, + ) + + del source_owner + gc.collect() + assert sender._sessions == {} + assert source_owner_ref() is None + assert str(exc_info.value) == "peer lookup failed" + sender._shutdown_complete = True + + +def test_rx_session_setup_failure_rolls_back_destination_owner() -> None: + class DestinationOwner: + pass + + receiver = _make_receiver() + receiver.begin_shutdown() + destination_owner = DestinationOwner() + destination_owner_ref = weakref.ref(destination_owner) + + with pytest.raises(RuntimeError, match="shutting down") as exc_info: + RxSession( + request_id=REQUEST_ID, + params=_params(), + receiver=receiver, + destination_owner=destination_owner, + ) + + del destination_owner + gc.collect() + assert destination_owner_ref() is None + assert "shutting down" in str(exc_info.value) + + +@pytest.mark.parametrize( + "session_type", + [ + pytest.param(TxSession, id="tx"), + pytest.param(RxSession, id="rx"), + ], +) +@pytest.mark.parametrize("failure_stage", ["base-init", "aux-allocation"]) +def test_early_session_constructor_failure_releases_allocator_owner( + session_type, failure_stage, monkeypatch +) -> None: + class AllocatorOwner: + pass + + if session_type is TxSession: + endpoint_name = "sender" + endpoint = _make_sender_for_shutdown() + owner_name = "source_owner" + else: + endpoint_name = "receiver" + endpoint = _make_receiver() + owner_name = "destination_owner" + + aux_buffer = SimpleNamespace( + alloc_slot=Mock(), + free_slot=Mock(), + ) + error_message = f"{failure_stage} failed" + base_type = session_type.__mro__[1] + original_base_init = base_type.__init__ + failed_instances = [] + if failure_stage == "base-init": + + def capture_and_fail_base_init(session, *_args, **_kwargs) -> None: + failed_instances.append(session) + raise RuntimeError(error_message) + + monkeypatch.setattr(base_type, "__init__", capture_and_fail_base_init) + else: + + def capture_base_init(session, *args, **kwargs) -> None: + failed_instances.append(session) + original_base_init(session, *args, **kwargs) + + monkeypatch.setattr(base_type, "__init__", capture_base_init) + aux_buffer.alloc_slot.side_effect = RuntimeError(error_message) + + owner = AllocatorOwner() + owner_ref = weakref.ref(owner) + with pytest.raises(RuntimeError, match=error_message) as exc_info: + session_type( + request_id=REQUEST_ID, + params=_params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + aux_buffer=aux_buffer, + **{endpoint_name: endpoint, owner_name: owner}, + ) + + assert len(failed_instances) == 1 + failed_session = failed_instances.pop() + # Constructor rollback already retired every resource it could have + # acquired. Explicit close, and therefore the later __del__, must stop at + # that marker without touching terminal fields that may not exist yet. + failed_session.close() + failed_session_ref = weakref.ref(failed_session) + observed_error = str(exc_info.value) + del exc_info + del failed_session + del owner + gc.collect() + assert owner_ref() is None + assert failed_session_ref() is None + assert observed_error == error_message + if failure_stage == "base-init": + aux_buffer.alloc_slot.assert_not_called() + else: + aux_buffer.alloc_slot.assert_called_once_with() + aux_buffer.free_slot.assert_not_called() + if session_type is TxSession: + endpoint._shutdown_complete = True + + +def test_pre_session_terminal_result_is_retried_by_sender_shutdown() -> None: + sender = _make_sender_for_shutdown() + info = _recv_info() + info.aux_slot = None + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): info}} + sender._send_operation_message = Mock(side_effect=[False, True]) + + sender._handle_cancel_session([MessageType.CANCEL_SESSION, str(REQUEST_ID).encode("ascii")]) + + assert sender._has_pending_pre_session_terminal_results() + assert sender.shutdown() + assert not sender._has_pending_pre_session_terminal_results() + assert sender._send_operation_message.call_count == 2 + + +def test_pre_session_terminal_result_is_retried_by_runtime_progress() -> None: + sender = _make_sender_for_shutdown() + info = _recv_info() + info.aux_slot = None + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): info}} + sender._send_operation_message = Mock(side_effect=[False, True]) + + sender._handle_cancel_session([MessageType.CANCEL_SESSION, str(REQUEST_ID).encode("ascii")]) + + assert sender._has_pending_pre_session_terminal_results() + sender.sweep_stale_req_infos() + assert not sender._has_pending_pre_session_terminal_results() + assert sender._send_operation_message.call_count == 2 + sender._shutdown_complete = True + + +def test_sender_claims_one_exact_operation_during_request_send_race() -> None: + task, session = _make_send_task_and_session() + info = _recv_info() + sender = object.__new__(Sender) + sender._shutdown_complete = True + sender._instance_rank = 0 + write_meta = WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver3", + peer_rank=WRITER_RANK, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + dst_device_id=0, + slice_id=0, + is_last_slice=True, + session=session, + ) + sender._build_kv_write_meta = Mock(return_value=write_meta) + sender._enqueue_owned = Mock() + sender._send_operation_message = Mock(return_value=True) + + start = threading.Barrier(3) + results: list = [] + + def dispatch() -> None: + start.wait() + results.append(sender._dispatch_operation(task, info)) + + threads = [threading.Thread(target=dispatch) for _ in range(2)] + for thread in threads: + thread.start() + start.wait() + for thread in threads: + thread.join(timeout=1) + + assert all(not thread.is_alive() for thread in threads) + assert results == [None, None] + sender._build_kv_write_meta.assert_called_once_with(task, info) + sender._enqueue_owned.assert_called_once_with(write_meta) + assert task.source_access_active + + key = (info.instance_name, info.instance_rank) + terminal_message = (b"terminal",) + task.cache_terminal_message(key, terminal_message) + task.mark_operation_result_delivered(key) + task.finish_source_access() + + sender._dispatch_operation(task, info) + + sender._send_operation_message.assert_called_once_with(info, terminal_message) + sender._enqueue_owned.assert_called_once_with(write_meta) + + +@pytest.mark.parametrize("channel", ["kv", "aux"]) +def test_sender_settles_every_writer_after_descriptor_build_failures(channel) -> None: + task, _session = _make_send_task_and_session(channel=channel) + infos = {rank: _recv_info(writer_rank=rank) for rank in (1, 2)} + sender = object.__new__(Sender) + sender._shutdown_complete = True + sender._instance_rank = 0 + sender._send_operation_message = Mock(return_value=True) + sender._enqueue_owned = Mock() + build = Mock( + side_effect=[RuntimeError("first descriptor failure"), RuntimeError("second failure")] + ) + if channel == "kv": + sender._build_kv_write_meta = build + else: + sender._build_aux_write_meta = build + + with pytest.raises(RuntimeError, match="first descriptor failure"): + sender.dispatch_task(task, infos) + + assert build.call_count == 2 + assert sender._send_operation_message.call_count == 2 + assert not task.source_access_active + for info in infos.values(): + snapshot = task.operation_snapshot((info.instance_name, info.instance_rank)) + assert snapshot is not None + assert snapshot[0] is SendOperationState.TERMINAL + assert snapshot[2] + + +@pytest.mark.parametrize("channel", ["kv", "aux"]) +def test_sender_failure_is_monotonic_while_sibling_source_access_drains(channel) -> None: + task, session = _make_send_task_and_session(channel=channel) + first_info = _recv_info(writer_rank=1) + second_info = _recv_info(writer_rank=2) + first_meta = WriteMeta( + task=task, + expected_transfers=2, + peer_name="receiver1", + peer_rank=1, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + slice_id=0 if channel == "kv" else None, + session=session, + ) + sender = object.__new__(Sender) + sender._shutdown_complete = True + sender._instance_rank = 0 + sender._send_operation_message = Mock(return_value=True) + sender._enqueue_owned = Mock() + build = Mock(side_effect=[first_meta, RuntimeError("second descriptor failed")]) + if channel == "kv": + sender._build_kv_write_meta = build + else: + sender._build_aux_write_meta = build + + assert sender._dispatch_operation(task, first_info) is None + assert task.mark_transferring() + assert isinstance(sender._dispatch_operation(task, second_info), RuntimeError) + + assert task.status is TaskStatus.ERROR + assert task.source_access_active + assert not task.complete() + assert task.status is TaskStatus.ERROR + + task.finish_source_access() + first_meta.source_access_enrolled = False + assert not task.source_access_active + + +@pytest.mark.parametrize("channel", ["kv", "aux"]) +def test_cancelled_partial_fanout_drains_and_late_writer_gets_no_access(channel) -> None: + sender, session, task, first_meta = _make_owned_send_task_and_session( + channel=channel, expected_transfers=2 + ) + first_info = _recv_info(writer_rank=1) + late_info = _recv_info(writer_rank=2) + + assert sender._dispatch_operation(task, first_info) is None + assert task.mark_transferring() + session.cancel() + + assert session.status is SessionStatus.CANCELLED + assert task.status is TaskStatus.ERROR + assert task.source_access_active + assert session.has_transferring_tasks() + + first_key = (first_info.instance_name, first_info.instance_rank) + task.cache_terminal_message(first_key, (b"first-writer-terminal",)) + task.mark_operation_result_delivered(first_key) + task.finish_source_access() + first_meta.source_access_enrolled = False + if channel == "kv": + task.transferred_count = 1 + else: + task._transfer_count = 1 + + assert not session.has_transferring_tasks() + assert sender._dispatch_operation(task, late_info) is None + late_key = (late_info.instance_name, late_info.instance_rank) + late_snapshot = task.operation_snapshot(late_key) + assert late_snapshot is not None + assert late_snapshot[0] is SendOperationState.TERMINAL + assert late_snapshot[2] + assert not task.source_access_active + assert task.status is TaskStatus.ERROR + assert session.status is SessionStatus.CANCELLED + assert session.wait_complete(blocking=False) is WaitResult.FAILED + assert not session.has_transferring_tasks() + + +@pytest.mark.parametrize("channel", ["kv", "aux"]) +@pytest.mark.parametrize("transfer_succeeded", [True, False]) +def test_active_writer_completion_after_cancel_preserves_cancelled( + channel, transfer_succeeded, monkeypatch +) -> None: + sender, session, task, write_meta = _make_owned_send_task_and_session( + channel=channel, expected_transfers=1 + ) + info = _recv_info(writer_rank=1) + write_meta.src_ptrs = np.array([1], dtype=np.int64) + write_meta.dst_ptrs = np.array([2], dtype=np.int64) + write_meta.sizes = np.array([4], dtype=np.int64) + assert sender._dispatch_operation(task, info) is None + + transfer_started = threading.Event() + release_transfer = threading.Event() + transfer_status = Mock() + + def wait_for_transfer() -> bool: + transfer_started.set() + assert release_transfer.wait(timeout=1) + return transfer_succeeded + + transfer_status.wait.side_effect = wait_for_transfer + sender._agent.submit_transfer_requests.return_value = transfer_status + sender._device_id = 0 + sender._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="sender", instance_rank=0) + ) + dealer = Mock() + sender._get_or_connect_thread_dealer = Mock(return_value=dealer) + if channel == "kv": + monkeypatch.setattr( + "tensorrt_llm._torch.disaggregation.native.bounce.build_send_request", + lambda _bounce, _write_meta, _factory: (object(), None), + ) + else: + monkeypatch.setattr( + Sender, "_make_agent_request", staticmethod(lambda *_args, **_kwargs: object()) + ) + + errors = [] + + def deliver() -> None: + try: + if channel == "kv": + sender._deliver_kv_to_agent(write_meta) + else: + sender._deliver_aux_to_agent(write_meta) + except Exception as error: + errors.append(error) + + thread = threading.Thread(target=deliver) + thread.start() + assert transfer_started.wait(timeout=1) + assert task.status is TaskStatus.TRANSFERRING + + session.cancel() + assert session.status is SessionStatus.CANCELLED + assert task.status is TaskStatus.ERROR + assert session.has_transferring_tasks() + + release_transfer.set() + thread.join(timeout=1) + assert not thread.is_alive() + assert errors == [] + assert session.status is SessionStatus.CANCELLED + assert task.status is TaskStatus.ERROR + assert not task.source_access_active + assert not task.has_pending_result_delivery + assert not session.has_transferring_tasks() + + +@pytest.mark.parametrize("first_terminal", ["cancel", "error"]) +def test_tx_session_first_terminal_failure_wins_and_settlement_retries( + first_terminal, +) -> None: + sender, session, task, _write_meta = _make_owned_send_task_and_session( + channel="kv", expected_transfers=1 + ) + + if first_terminal == "cancel": + session.cancel() + first_exception = session.exception + session.set_exception("late worker failure") + session.cancel() + + assert session.status is SessionStatus.CANCELLED + assert session.exception is first_exception is None + else: + session.set_exception("first worker failure") + first_exception = session.exception + session.cancel() + session.set_exception("late worker failure") + + assert session.status is SessionStatus.ERROR + assert session.exception is first_exception + assert str(session.exception).endswith(": first worker failure") + + assert task.status is TaskStatus.ERROR + assert sender.cancel_session.call_count == 3 + + +def test_local_cancel_installs_tombstone_and_retries_settlement() -> None: + info = _recv_info() + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session.aux_task = None + session._unbound_terminal_results = {} + session._terminal_retry_lock = threading.Lock() + session._next_terminal_retry_at = 0.0 + session._terminal_status = SessionStatus.CANCELLED + session._closed = True + session._sender = None + sender = object.__new__(Sender) + sender._shutdown_complete = False + sender._instance_rank = 0 + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): info}} + sender._peer_requests_lock = threading.Lock() + sender._pre_cancelled_rids = set() + sender._send_operation_message = Mock(return_value=False) + sender.send_cancel_to_receivers = Mock() + + sender.cancel_session(session) + sender.cancel_session(session) + + assert REQUEST_ID in sender._pre_cancelled_rids + assert sender._send_operation_message.call_count == 4 + assert sender.send_cancel_to_receivers.call_count == 2 + assert len(session._unbound_terminal_results) == 2 + assert session.has_transferring_tasks() + + sender._send_operation_message.return_value = True + session._sender = sender + + assert session.has_failed() + assert not session.has_transferring_tasks() + sender._shutdown_complete = True + + +def test_tx_session_repeated_cancel_retries_sender_settlement() -> None: + sender = Mock() + sender._operation_admission_lock = threading.RLock() + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session.aux_task = None + session._terminal_status = None + session._accepting_operations = True + session._closed = True + session._sender = sender + + session.cancel() + session.cancel() + + assert session.status is SessionStatus.CANCELLED + assert not session._accepting_operations + assert sender.cancel_session.call_count == 2 + sender.cancel_session.assert_called_with(session) + + +def test_tx_session_error_installs_sender_tombstone() -> None: + sender = Mock() + sender._operation_admission_lock = threading.RLock() + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session.aux_task = None + session._terminal_status = None + session._accepting_operations = True + session._closed = True + session._sender = sender + + session.set_exception("transfer failed") + + assert session.status is SessionStatus.ERROR + assert not session._accepting_operations + sender.cancel_session.assert_called_once_with(session) + + +def test_sender_req_info_returns_a_stable_snapshot() -> None: + info = _recv_info() + sender = object.__new__(Sender) + sender._shutdown_complete = True + sender._instance_rank = 0 + sender._peer_requests_lock = threading.Lock() + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): info}} + + snapshot = sender._get_req_info(REQUEST_ID) + assert snapshot is not None + snapshot.clear() + + assert sender._peer_requests[REQUEST_ID] == {(WRITER_RANK, 0): info} + + +def test_external_backend_fence_coordinates_registry_and_bounce_settlement() -> None: + bounce = _FakeBounce() + receiver = _make_receiver(bounce=bounce) + key = (REQUEST_ID, 0) + assert receiver._recv_registry.prepare( + key, + {WRITER_RANK}, + has_bounce_slot=True, + ).accepted + assert receiver._recv_registry.begin_publication(key, WRITER_RANK).publication_allowed + assert receiver._recv_registry.mark_published(key, WRITER_RANK).accepted + + receiver.mark_backend_quiesced() + + snapshot = receiver._recv_registry.context_snapshot(key) + assert snapshot is not None + assert snapshot.physical_state is PhysicalState.DRAINING + assert not receiver._recv_registry.is_request_drained(REQUEST_ID) + assert len(bounce.backend_quiesced_callbacks) == 1 + + callback_key, callback = bounce.backend_quiesced_callbacks[0] + assert callback_key == key + callback(True) + + assert receiver._recv_registry.is_request_drained(REQUEST_ID) + assert receiver._recv_registry.context_snapshot(key).physical_state is PhysicalState.DRAINED + + +def test_sender_retries_listener_stop_after_first_failure() -> None: + sender = object.__new__(Sender) + sender._instance_rank = 0 + sender._shutdown = False + sender._shutdown_complete = False + sender._listener_stopped = False + sender._shutdown_sentinels_sent = False + sender._messenger = Mock() + sender._messenger.stop.side_effect = [RuntimeError("busy"), None] + sender._send_task_queues = [Mock()] + sender._worker_threads = [] + sender._loaded_remote_agents = set() + sender._loaded_remote_agents_lock = threading.Lock() + sender._agent = Mock() + sender._dealers = {} + + assert sender.shutdown() is False + sender._send_task_queues[0].put.assert_not_called() + + assert sender.shutdown() is True + assert sender._messenger.stop.call_count == 2 + sender._send_task_queues[0].put.assert_called_once_with(None) + + +def test_sender_shutdown_continues_siblings_after_cancel_failure() -> None: + sender = _make_sender_for_shutdown() + sender._listener_stopped = False + sender.retry_terminal_results = Mock() + first = Mock(disagg_request_id=REQUEST_ID) + first.cancel.side_effect = [RuntimeError("injected cancel failure"), None] + first.has_transferring_tasks.return_value = False + second = Mock(disagg_request_id=REQUEST_ID + 1) + second.has_transferring_tasks.return_value = False + sender._sessions = {REQUEST_ID: first, REQUEST_ID + 1: second} + + assert sender.shutdown() is False + + first.cancel.assert_called_once_with() + first.close.assert_not_called() + second.cancel.assert_called_once_with() + second.close.assert_called_once_with() + sender._messenger.stop.assert_called_once_with() + assert sender._listener_stopped + assert sender._sessions == {REQUEST_ID: first} + + assert sender.shutdown() is True + assert first.cancel.call_count == 2 + first.close.assert_called_once_with() + assert sender._sessions == {} + + +def test_sender_shutdown_continues_siblings_after_close_failure() -> None: + sender = _make_sender_for_shutdown() + sender.retry_terminal_results = Mock() + first = Mock(disagg_request_id=REQUEST_ID) + first.has_transferring_tasks.return_value = False + first.close.side_effect = [RuntimeError("injected close failure"), None] + second = Mock(disagg_request_id=REQUEST_ID + 1) + second.has_transferring_tasks.return_value = False + sender._sessions = {REQUEST_ID: first, REQUEST_ID + 1: second} + + assert sender.shutdown() is False + + first.close.assert_called_once_with() + second.close.assert_called_once_with() + assert sender._sessions == {REQUEST_ID: first} + + assert sender.shutdown() is True + assert first.close.call_count == 2 + assert sender._sessions == {} + + +def test_sender_shutdown_forces_cached_terminal_result_progress() -> None: + sender = _make_sender_for_shutdown() + info = _recv_info() + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): info}} + sender._send_operation_message = Mock(return_value=True) + + params = _params() + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(params) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session.aux_task = None + session._unbound_terminal_results = {} + session._terminal_retry_lock = threading.Lock() + session._next_terminal_retry_at = float("inf") + session._sender = sender + session._need_aux = False + session._terminal_status = SessionStatus.CANCELLED + session._closed = False + session._accepting_operations = False + session._close_lock = threading.Lock() + session._aux_buffer = None + session.aux_slot = None + source_owner = object() + session._source_owner = source_owner + session.cancel = Mock() + + task = KVSendTask(KVSlice(), params, 0, session=session) + session.kv_tasks.append(task) + key = (info.instance_name, info.instance_rank) + admitted, state, _, _ = task.admit_operation( + key, + allow_source_access=True, + no_access_message=(b"unused",), + ) + assert admitted and state is SendOperationState.PENDING + terminal_message = (b"terminal",) + task.cache_terminal_message(key, terminal_message) + task.finish_source_access() + task.fail(RuntimeError("logical transfer failure")) + sender._sessions[REQUEST_ID] = session + + assert sender.shutdown() + + session.cancel.assert_called_once_with() + sender._send_operation_message.assert_called_once_with(info, terminal_message) + assert not task.has_pending_result_delivery + assert sender._sessions == {} + assert session._source_owner is None + assert session._closed + + +def test_sender_shutdown_is_serialized_and_closes_late_control_admission() -> None: + sender = _make_sender_for_shutdown() + task_queue = Mock() + sender._send_task_queues = [task_queue] + sender._loaded_remote_agents = {"peer0"} + dealer = Mock() + sender._dealers = {"tcp://receiver": dealer} + invalidation_started = threading.Event() + release_invalidation = threading.Event() + + def invalidate(_agent_name: str) -> None: + invalidation_started.set() + assert release_invalidation.wait(timeout=1) + + sender._agent.invalidate_remote_agent.side_effect = invalidate + results: list[bool] = [] + threads = [threading.Thread(target=lambda: results.append(sender.shutdown())) for _ in range(2)] + + threads[0].start() + assert invalidation_started.wait(timeout=1) + threads[1].start() + sender._handle_cancel_session([MessageType.CANCEL_SESSION, b"999"]) + assert 999 not in sender._pre_cancelled_rids + assert sender._pre_session_terminal_results == {} + + release_invalidation.set() + for thread in threads: + thread.join(timeout=1) + + assert all(not thread.is_alive() for thread in threads) + assert results == [True, True] + task_queue.put.assert_called_once_with(None) + sender._agent.invalidate_remote_agent.assert_called_once_with("peer0") + dealer.stop.assert_called_once_with() + + +def test_tx_send_admission_finishes_before_sender_shutdown_sentinel() -> None: + sender = object.__new__(Sender) + sender._instance_rank = 0 + sender._operation_admission_lock = threading.RLock() + sender._shutdown = False + sender._shutdown_complete = False + sender._listener_stopped = True + sender._shutdown_sentinels_sent = False + task_queue = Mock() + sender._send_task_queues = [task_queue] + sender._num_threads = 1 + sender._worker_threads = [] + sender._failed_thread_dealers = [] + sender._failed_thread_dealers_lock = threading.Lock() + sender._in_doubt_transfers = [] + sender._in_doubt_transfers_lock = threading.Lock() + sender._loaded_remote_agents = set() + sender._loaded_remote_agents_lock = threading.Lock() + sender._agent = Mock() + sender._dealers = {} + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): _recv_info()}} + sender._peer_requests_lock = threading.Lock() + sender._sessions = {} + sender._sessions_lock = threading.Lock() + + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session.aux_task = None + session._need_aux = False + session._terminal_status = None + session._accepting_operations = True + session._closed = False + session._unbound_terminal_results = {} + session._source_owner = object() + session._sender = sender + sender._sessions[REQUEST_ID] = session + + build_started = threading.Event() + release_build = threading.Event() + + def build(task, _info): + build_started.set() + assert release_build.wait(timeout=1) + return WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver3", + peer_rank=WRITER_RANK, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + slice_id=0, + session=session, + ) + + sender._build_kv_write_meta = build + send_errors: list[Exception] = [] + shutdown_results: list[bool] = [] + + def send() -> None: + try: + session.send(KVSlice()) + except Exception as e: + send_errors.append(e) + + send_thread = threading.Thread(target=send) + send_thread.start() + assert build_started.wait(timeout=1) + + shutdown_thread = threading.Thread(target=lambda: shutdown_results.append(sender.shutdown())) + shutdown_thread.start() + task_queue.put.assert_not_called() + + release_build.set() + send_thread.join(timeout=1) + shutdown_thread.join(timeout=1) + + assert not send_thread.is_alive() + assert not shutdown_thread.is_alive() + assert not send_errors + assert shutdown_results == [False] + assert len(task_queue.put.call_args_list) == 2 + assert isinstance(task_queue.put.call_args_list[0].args[0], WriteMeta) + assert task_queue.put.call_args_list[1].args[0] is None + + session.kv_tasks[0].finish_source_access() + session._closed = True + sender._shutdown_complete = True + + +def test_sender_retries_failed_worker_local_dealer_close() -> None: + sender = object.__new__(Sender) + sender._instance_rank = 0 + sender._failed_thread_dealers = [] + sender._failed_thread_dealers_lock = threading.Lock() + dealer = Mock() + dealer.stop.side_effect = [ + RuntimeError("worker close failed"), + RuntimeError("first retry failed"), + None, + ] + + dealers = {"tcp://receiver": dealer} + sender._close_thread_dealers(0, dealers) + + assert dealers == {} + assert sender._failed_thread_dealers == [dealer] + + sender._operation_admission_lock = threading.RLock() + sender._shutdown = False + sender._shutdown_complete = False + sender._listener_stopped = True + sender._shutdown_sentinels_sent = False + sender._send_task_queues = [] + sender._worker_threads = [] + sender._in_doubt_transfers = [] + sender._in_doubt_transfers_lock = threading.Lock() + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._loaded_remote_agents = set() + sender._loaded_remote_agents_lock = threading.Lock() + sender._agent = Mock() + sender._dealers = {} + + assert sender.shutdown() is False + assert sender._failed_thread_dealers == [dealer] + + assert sender.shutdown() is True + assert sender._failed_thread_dealers == [] + assert dealer.stop.call_count == 3 + + +def test_sender_shutdown_retains_agent_for_in_doubt_transfer_context() -> None: + sender = object.__new__(Sender) + sender._instance_rank = 0 + sender._shutdown = False + sender._shutdown_complete = False + sender._listener_stopped = True + sender._shutdown_sentinels_sent = False + sender._send_task_queues = [] + sender._worker_threads = [] + sender._in_doubt_transfers = [object()] + sender._in_doubt_transfers_lock = threading.Lock() + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._loaded_remote_agents = {"peer0"} + sender._loaded_remote_agents_lock = threading.Lock() + sender._agent = Mock() + sender._dealers = {} + + assert sender.shutdown() is False + sender._agent.invalidate_remote_agent.assert_not_called() + + +def test_rank_info_server_retries_listener_stop_after_first_failure() -> None: + server = object.__new__(RankInfoServer) + server._shutdown = False + server._shutdown_started = False + server._messenger = Mock() + server._messenger.stop.side_effect = [RuntimeError("busy"), None] + + assert server.shutdown() is False + assert not server._shutdown + + assert server.shutdown() is True + assert server._messenger.stop.call_count == 2 + + +@pytest.mark.parametrize("factory", ["create_tx_session", "create_rx_session"]) +def test_transfer_worker_rejects_new_sessions_after_shutdown_starts(factory) -> None: + worker = object.__new__(TransferWorker) + worker._shutdown = True # Keep object.__new__ fixtures out of __del__ cleanup. + worker._shutdown_started = True + worker._session_admission_lock = threading.Lock() + + with pytest.raises(RuntimeError, match="shutting down"): + getattr(worker, factory)(SimpleNamespace()) + + +def test_transfer_worker_rx_session_retains_destination_owner_until_close() -> None: + class _DestinationOwner: + pass + + owner = _DestinationOwner() + owner.py_disaggregated_params = _params() + owner.py_request_id = REQUEST_ID + owner.prompt_len = 8 + owner.py_beam_width = 1 + owner_ref = weakref.ref(owner) + receiver = _make_receiver() + worker = object.__new__(TransferWorker) + worker._shutdown = True # Keep object.__new__ fixture out of __del__ cleanup. + worker._shutdown_started = False + worker._session_admission_lock = threading.Lock() + worker._receiver = receiver + worker._aux_buffer = None + worker._config = SimpleNamespace(rx_timeout_s=None) + + session = worker.create_rx_session(owner) + del owner + gc.collect() + + assert owner_ref() is not None + assert session._destination_owner is owner_ref() + + session.close() + gc.collect() + + assert owner_ref() is None + + +def test_transfer_worker_retains_owner_before_nested_shutdown_failure() -> None: + worker = object.__new__(TransferWorker) + worker._shutdown = False + worker._agent_shutdown_failed = False + worker._shutdown_started = False + worker._session_admission_lock = threading.Lock() + worker._rank_info_server = Mock() + worker._rank_info_server.shutdown.side_effect = [RuntimeError("listener failed"), True] + worker._receiver = None + worker._sender = None + worker._bounce = None + worker._agent = None + + try: + assert worker.shutdown() is False + assert worker in _NON_DRAINED_TRANSFER_WORKERS + + assert worker.shutdown() is True + assert worker not in _NON_DRAINED_TRANSFER_WORKERS + finally: + _NON_DRAINED_TRANSFER_WORKERS.discard(worker) + + +def test_transfer_worker_shutdown_attempts_are_serialized() -> None: + worker = object.__new__(TransferWorker) + worker._shutdown_attempt_lock = threading.Lock() + worker._shutdown = False + worker._agent_shutdown_failed = False + worker._shutdown_started = False + worker._session_admission_lock = threading.Lock() + worker._rank_info_server = Mock() + worker._receiver = None + worker._sender = None + worker._bounce = None + agent = Mock() + worker._agent = agent + descriptor = object() + worker._registered_mem = [descriptor] + shutdown_started = threading.Event() + release_shutdown = threading.Event() + + def stop_rank_info_server() -> bool: + shutdown_started.set() + assert release_shutdown.wait(timeout=1) + return True + + worker._rank_info_server.shutdown.side_effect = stop_rank_info_server + results: list[bool] = [] + threads = [threading.Thread(target=lambda: results.append(worker.shutdown())) for _ in range(2)] + + try: + threads[0].start() + assert shutdown_started.wait(timeout=1) + threads[1].start() + release_shutdown.set() + for thread in threads: + thread.join(timeout=1) + + assert all(not thread.is_alive() for thread in threads) + assert results == [True, True] + worker._rank_info_server.shutdown.assert_called_once_with() + agent.deregister_memory.assert_called_once_with(descriptor) + agent.shutdown.assert_called_once_with() + finally: + _NON_DRAINED_TRANSFER_WORKERS.discard(worker) + + +def test_tx_session_retains_source_owner_while_local_access_is_ambiguous() -> None: + task = SimpleNamespace(status=TaskStatus.TRANSFERRING) + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session._terminal_status = SessionStatus.CANCELLED + session.kv_tasks = [task] + session._need_aux = False + session.aux_task = None + session._closed = False + session._aux_buffer = None + session.aux_slot = None + session._sender = None + + assert not session.has_failed() + assert session.wait_complete(blocking=False) is None + with pytest.raises(RuntimeError, match="source work is pending or active"): + session.close() + + task.status = TaskStatus.TRANSFERRED + assert session.has_failed() + assert session.wait_complete(blocking=False) is WaitResult.FAILED + session.close() + + +def test_tx_session_includes_aux_source_access_in_drain_gate() -> None: + kv_task = SimpleNamespace(status=TaskStatus.TRANSFERRED) + aux_task = SimpleNamespace(status=TaskStatus.TRANSFERRING) + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session._terminal_status = SessionStatus.CANCELLED + session.kv_tasks = [kv_task] + session._need_aux = True + session.aux_task = aux_task + session._closed = False + session._aux_buffer = None + session.aux_slot = None + session._sender = None + + assert session.has_transferring_tasks() + assert not session.has_failed() + assert session.wait_complete(blocking=False) is None + with pytest.raises(RuntimeError, match="source work is pending or active"): + session.close() + + aux_task.status = TaskStatus.TRANSFERRED + assert not session.has_transferring_tasks() + assert session.wait_complete(blocking=False) is WaitResult.FAILED + session.close() + + +def test_tx_session_close_is_serialized_across_concurrent_callers() -> None: + sender = SimpleNamespace( + _operation_admission_lock=threading.RLock(), + clear_session=Mock(), + ) + aux_buffer = Mock() + free_started = threading.Event() + release_free = threading.Event() + + def free_slot(_slot: int) -> None: + free_started.set() + assert release_free.wait(timeout=1) + + aux_buffer.free_slot.side_effect = free_slot + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session._close_lock = threading.Lock() + session._terminal_status = SessionStatus.CANCELLED + session.kv_tasks = [] + session.aux_task = None + session._unbound_terminal_results = {} + session._closed = False + session._accepting_operations = False + session._aux_buffer = aux_buffer + session.aux_slot = 7 + session._source_owner = object() + session._sender = sender + errors: list[Exception] = [] + + def close() -> None: + try: + session.close() + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=close) for _ in range(2)] + threads[0].start() + assert free_started.wait(timeout=1) + threads[1].start() + release_free.set() + for thread in threads: + thread.join(timeout=1) + + assert all(not thread.is_alive() for thread in threads) + assert errors == [] + aux_buffer.free_slot.assert_called_once_with(7) + sender.clear_session.assert_called_once_with(REQUEST_ID) + assert session.aux_slot is None + assert session._source_owner is None + assert session._closed + + +def test_aux_pack_holds_sender_admission_until_fill_finishes() -> None: + class ObservableRLock: + def __init__(self) -> None: + self._lock = threading.RLock() + self.attempted = threading.Event() + + def __enter__(self): + self.attempted.set() + self._lock.acquire() + return self + + def __exit__(self, _exc_type, _exc_value, _traceback) -> None: + self._lock.release() + + sender = _make_sender_for_shutdown() + admission_lock = ObservableRLock() + sender._operation_admission_lock = admission_lock + aux_buffer = Mock() + aux_buffer.alloc_slot.return_value = SimpleNamespace(id=7) + fill_started = threading.Event() + release_fill = threading.Event() + + def fill_slot(_slot: int, _request) -> None: + fill_started.set() + assert release_fill.wait(timeout=1) + + aux_buffer.fill_slot.side_effect = fill_slot + session = TxSession( + request_id=REQUEST_ID, + params=_params(), + sender=sender, + aux_buffer=aux_buffer, + source_owner=object(), + ) + admission_lock.attempted.clear() + pack_errors: list[Exception] = [] + shutdown_results: list[bool] = [] + + def pack() -> None: + try: + session.pack_aux(SimpleNamespace()) + except Exception as e: + pack_errors.append(e) + + pack_thread = threading.Thread(target=pack) + shutdown_thread = threading.Thread(target=lambda: shutdown_results.append(sender.shutdown())) + pack_thread.start() + assert fill_started.wait(timeout=1) + admission_lock.attempted.clear() + shutdown_thread.start() + assert admission_lock.attempted.wait(timeout=1) + + assert not sender._shutdown + aux_buffer.free_slot.assert_not_called() + + release_fill.set() + pack_thread.join(timeout=1) + shutdown_thread.join(timeout=1) + + assert not pack_thread.is_alive() + assert not shutdown_thread.is_alive() + assert pack_errors == [] + assert shutdown_results == [True] + aux_buffer.fill_slot.assert_called_once() + aux_buffer.free_slot.assert_called_once_with(7) + assert session._closed + + +def test_aux_send_requires_packed_data_and_freezes_slot_contents() -> None: + sender = _make_sender_for_shutdown() + aux_buffer = Mock() + aux_buffer.alloc_slot.return_value = SimpleNamespace(id=7) + session = TxSession( + request_id=REQUEST_ID, + params=_params(), + sender=sender, + aux_buffer=aux_buffer, + source_owner=object(), + ) + + with pytest.raises(RuntimeError, match="must be packed before send"): + session.send_aux() + + request = SimpleNamespace() + session.pack_aux(request) + task = session.send_aux() + + with pytest.raises(RuntimeError, match="frozen after send"): + session.pack_aux(request) + + assert session.aux_task is task + aux_buffer.fill_slot.assert_called_once_with(7, request) + assert sender.shutdown() + + +def test_task_failure_does_not_hide_a_sibling_source_operation() -> None: + task = KVSendTask(KVSlice(), _params(), 0) + task.begin_source_access() + task.begin_source_access() + task.fail(RuntimeError("one writer failed")) + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(_params()) + session.request_id = REQUEST_ID + session._terminal_status = SessionStatus.ERROR + session.kv_tasks = [task] + session._need_aux = False + session.aux_task = None + + task.finish_source_access() + + assert task.status is TaskStatus.ERROR + assert session.has_transferring_tasks() + assert not session.has_failed() + assert session.wait_complete(blocking=False) is None + + task.finish_source_access() + assert not session.has_transferring_tasks() + assert session.has_failed() + + +def test_ambiguous_agent_exception_retains_context_when_quarantine_fails() -> None: + params = _params() + task = KVSendTask(KVSlice(), params, 0) + source_owner = object() + session = SimpleNamespace( + kv_tasks=[task], + lock=threading.Lock(), + status=SessionStatus.READY, + source_owner=source_owner, + ) + bounce = Mock() + bounce.build_request.return_value = (object(), 11) + bounce.quarantine_send.side_effect = RuntimeError("quarantine failed") + sender = object.__new__(Sender) + sender._sessions_lock = threading.Lock() + sender._get_session = lambda _rid: session + sender._bounce = bounce + sender._agent = Mock() + sender._agent.submit_transfer_requests.side_effect = RuntimeError("submit crossed backend") + sender._device_id = 0 + sender._instance_rank = 0 + sender._in_doubt_transfers = [] + sender._in_doubt_transfers_lock = threading.Lock() + write_meta = WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver0", + peer_rank=0, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([100], dtype=np.int64), + dst_ptrs=np.array([200], dtype=np.int64), + sizes=np.array([8], dtype=np.int64), + dst_device_id=0, + slice_id=0, + is_last_slice=True, + bounce_dst_base=300, + session=session, + ) + + with pytest.raises(TransferSourceInDoubtError): + sender._deliver_kv_to_agent(write_meta) + + assert task.status is TaskStatus.TRANSFERRING + assert task.source_access_active + bounce.quarantine_send.assert_called_once_with(11) + bounce.release_send.assert_not_called() + assert len(sender._in_doubt_transfers) == 1 + context = sender._in_doubt_transfers[0] + assert context.session is session + assert context.source_owner is source_owner + assert context.transfer_request is bounce.build_request.return_value[0] + assert context.bounce_slot_id == 11 + + +def test_ambiguous_gather_rollback_retains_sender_source_context() -> None: + from tensorrt_llm._torch.disaggregation.native.bounce import GatherSourceInDoubtError + + params = _params() + task = KVSendTask(KVSlice(), params, 0) + source_owner = object() + session = SimpleNamespace( + kv_tasks=[task], + lock=threading.Lock(), + status=SessionStatus.READY, + source_owner=source_owner, + ) + sender = object.__new__(Sender) + sender._sessions_lock = threading.Lock() + sender._get_session = lambda _rid: session + sender._bounce = Mock() + sender._bounce.build_request.side_effect = GatherSourceInDoubtError( + "gather rollback has no positive fence" + ) + sender._agent = Mock() + sender._device_id = 0 + sender._instance_rank = 0 + sender._in_doubt_transfers = [] + sender._in_doubt_transfers_lock = threading.Lock() + write_meta = WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver0", + peer_rank=0, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([100], dtype=np.int64), + dst_ptrs=np.array([200], dtype=np.int64), + sizes=np.array([8], dtype=np.int64), + dst_device_id=0, + slice_id=0, + is_last_slice=True, + bounce_dst_base=300, + session=session, + ) + + with pytest.raises(GatherSourceInDoubtError): + sender._deliver_kv_to_agent(write_meta) + + assert task.source_access_active + sender._agent.submit_transfer_requests.assert_not_called() + assert len(sender._in_doubt_transfers) == 1 + context = sender._in_doubt_transfers[0] + assert context.session is session + assert context.source_owner is source_owner + assert context.transfer_request is None + + +def test_fenced_pre_submit_failure_reports_terminal_writer_result() -> None: + params = _params() + task = KVSendTask(KVSlice(), params, 0) + session = SimpleNamespace( + kv_tasks=[task], + lock=threading.Lock(), + status=SessionStatus.READY, + source_owner=object(), + ) + sender = object.__new__(Sender) + sender._sessions_lock = threading.Lock() + sender._get_session = lambda _rid: session + sender._bounce = Mock() + sender._bounce.build_request.side_effect = RuntimeError("descriptor build failed") + sender._agent = Mock() + sender._device_id = 0 + sender._instance_rank = 0 + dealer = Mock() + sender._get_or_connect_thread_dealer = Mock(return_value=dealer) + write_meta = WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver0", + peer_rank=0, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([100], dtype=np.int64), + dst_ptrs=np.array([200], dtype=np.int64), + sizes=np.array([8], dtype=np.int64), + dst_device_id=0, + slice_id=0, + is_last_slice=True, + bounce_dst_base=300, + session=session, + ) + + sender._deliver_kv_to_agent(write_meta) + + assert task.status is TaskStatus.ERROR + assert not task.source_access_active + sender._agent.submit_transfer_requests.assert_not_called() + dealer.send.assert_called_once() + result = dealer.send.call_args.args[0] + assert result[0] == MessageType.KV_AGENT_RESULT + assert _KV_RESULT_PREFIX.unpack(result[1])[-1] == 1 + + +def test_tx_poll_retries_cached_terminal_result_after_worker_send_failure() -> None: + params = _params() + sender = object.__new__(Sender) + sender._shutdown_complete = True + sender._instance_rank = 0 + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): _recv_info()}} + sender._peer_requests_lock = threading.Lock() + sender._send_operation_message = Mock(return_value=True) + sender._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="sender", instance_rank=0) + ) + sender._device_id = 0 + dealer = Mock() + dealer.send.side_effect = RuntimeError("transient send failure") + sender._get_or_connect_thread_dealer = Mock(return_value=dealer) + + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(params) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session.aux_task = None + session._need_aux = False + session._terminal_status = None + session.receiver_ready = True + session._accepting_operations = True + session._unbound_terminal_results = {} + session._terminal_retry_lock = threading.Lock() + session._next_terminal_retry_at = 0.0 + session._source_owner = object() + session._sender = sender + session._closed = False + + task = KVSendTask(KVSlice(), params, 0, session=session) + session.kv_tasks.append(task) + key = ("receiver", WRITER_RANK) + admitted, state, _message, _delivered = task.admit_operation( + key, + allow_source_access=True, + no_access_message=(b"unused",), + ) + assert admitted + assert state is SendOperationState.PENDING + write_meta = WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver3", + peer_rank=WRITER_RANK, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + dst_device_id=0, + slice_id=0, + is_last_slice=True, + session=session, + source_access_enrolled=True, + operation_key=key, + ) + + sender._deliver_kv_to_agent(write_meta) + + assert not task.source_access_active + assert task.has_pending_result_delivery + assert session.has_transferring_tasks() + assert task.status is TaskStatus.TRANSFERRED + + assert not session.has_failed() + sender._send_operation_message.assert_called_once() + assert not task.has_pending_result_delivery + + +def test_aux_poll_retries_success_after_worker_result_send_failure() -> None: + params = _params() + sender = object.__new__(Sender) + sender._shutdown_complete = True + sender._instance_rank = 0 + sender._peer_requests = {REQUEST_ID: {(WRITER_RANK, 0): _recv_info()}} + sender._peer_requests_lock = threading.Lock() + sender._send_operation_message = Mock(return_value=True) + sender._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="sender", instance_rank=0) + ) + sender._device_id = 0 + dealer = Mock() + dealer.send.side_effect = RuntimeError("transient send failure") + sender._get_or_connect_thread_dealer = Mock(return_value=dealer) + + session = object.__new__(TxSession) + session._base_args = SessionArgsBase(params) + session.request_id = REQUEST_ID + session.lock = threading.Lock() + session.kv_tasks = [] + session._need_aux = True + session._terminal_status = None + session.receiver_ready = True + session._accepting_operations = True + session._unbound_terminal_results = {} + session._terminal_retry_lock = threading.Lock() + session._next_terminal_retry_at = 0.0 + session._source_owner = object() + session._sender = sender + session._closed = False + + task = AuxSendTask(params, 7, session=session) + session.aux_task = task + key = ("receiver", WRITER_RANK) + admitted, state, _message, _delivered = task.admit_operation( + key, + allow_source_access=True, + no_access_message=(b"unused",), + ) + assert admitted + assert state is SendOperationState.PENDING + write_meta = WriteMeta( + task=task, + expected_transfers=1, + peer_name="receiver3", + peer_rank=WRITER_RANK, + peer_endpoint="tcp://receiver", + unique_rid=REQUEST_ID, + src_ptrs=np.array([], dtype=np.int64), + dst_ptrs=np.array([], dtype=np.int64), + sizes=np.array([], dtype=np.int64), + dst_device_id=0, + meta_type=WriteMetaType.AUX, + session=session, + source_access_enrolled=True, + operation_key=key, + ) + + sender._deliver_aux_to_agent(write_meta) + + assert not task.source_access_active + assert task.has_pending_result_delivery + assert session.has_transferring_tasks() + assert task.status is TaskStatus.TRANSFERRED + + assert not session.has_failed() + sender._send_operation_message.assert_called_once() + assert not task.has_pending_result_delivery + + +def test_async_send_enrolls_source_owner_before_launch_and_retains_on_error() -> None: + request = SimpleNamespace( + request_id=REQUEST_ID, + py_disaggregated_params=None, + state=None, + ) + session = Mock() + session.has_transferring_tasks.return_value = True + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_send_session = False + transceiver._send_reqs = {} + transceiver._send_sessions = {REQUEST_ID: session} + + def get_session(owner): + assert transceiver._send_reqs == {REQUEST_ID: owner} + return session + + def fail_send(_slice): + assert transceiver._send_reqs == {REQUEST_ID: request} + raise RuntimeError("launch failed") + + transceiver._get_or_create_send_session = get_session + transceiver._create_kv_slice = Mock(return_value=object()) + session.send.side_effect = fail_send + + with pytest.raises(RuntimeError, match="launch failed"): + transceiver.respond_and_send_async(request) + + assert transceiver._send_reqs == {REQUEST_ID: request} + assert transceiver._send_sessions == {REQUEST_ID: session} + session.cancel.assert_called_once_with() diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index cfad765968e2..d024b3aff143 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -12,32 +12,96 @@ # 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. -"""Bounded polling tests for KvCacheTransceiverV2 Tx sessions.""" +"""Bounded polling and shutdown ownership tests for the native transceiver.""" from __future__ import annotations +import gc +import threading +import weakref from dataclasses import dataclass from typing import Optional from unittest.mock import Mock +import pytest + from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult -from tensorrt_llm._torch.disaggregation.native.transfer import TaskStatus, TxSession -from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.disaggregation.native.transfer import ( + Receiver, + RxSession, + Sender, + TaskStatus, + TransferWorker, + TxSession, +) +from tensorrt_llm._torch.disaggregation.transceiver import ( + _NON_DRAINED_TRANSCEIVERS, + KvCacheTransceiverV2, +) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm.bindings import LlmRequestState @dataclass class _FakeRequest: state: Optional[LlmRequestState] = None + request_id: int = 0 + py_disaggregated_params: Optional[object] = None + + +class _RetirementRequest: + """Request with an optionally fallible terminal-state assignment.""" + + def __init__( + self, + rid: int, + *, + fail_state: Optional[LlmRequestState] = None, + ) -> None: + self.request_id = rid + self.py_request_id = rid + self.py_disaggregated_params = None + self.py_kv_cache_xfer_bytes = 64 + self.set_kv_cache_size = Mock() + self._state: Optional[LlmRequestState] = None + self._fail_state = fail_state + self._failure_pending = fail_state is not None + self.terminal_state_set_calls = 0 + + @property + def state(self) -> Optional[LlmRequestState]: + return self._state + + @state.setter + def state(self, value: LlmRequestState) -> None: + if value in ( + LlmRequestState.DISAGG_CONTEXT_COMPLETE, + LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE, + LlmRequestState.DISAGG_TRANS_ERROR, + ): + self.terminal_state_set_calls += 1 + if value == self._fail_state: + if self._failure_pending: + self._failure_pending = False + raise RuntimeError("terminal state update failed") + self._state = value class _FakeTransferWorker: - def __init__(self) -> None: + def __init__(self, shutdown_results: Optional[list[bool]] = None) -> None: self.sweep_count = 0 + self.shutdown_results = shutdown_results or [] + self.shutdown_count = 0 def sweep_stale_req_infos(self) -> None: self.sweep_count += 1 + def shutdown(self) -> bool: + self.shutdown_count += 1 + if self.shutdown_results: + return self.shutdown_results.pop(0) + return True + class _FakeSession: def __init__( @@ -48,13 +112,16 @@ def __init__( status: SessionStatus = SessionStatus.READY, is_completed: bool = False, has_failed: bool = False, + has_transferring_tasks: bool = False, ) -> None: self._rid = rid self._wait_result = wait_result self._status = status self._is_completed = is_completed self._has_failed = has_failed + self._has_transferring_tasks = has_transferring_tasks self.blocking_calls: list[bool] = [] + self.cancel_calls = 0 self.closed = False @property @@ -75,10 +142,39 @@ def is_completed(self) -> bool: def has_failed(self) -> bool: return self._has_failed + def cancel(self) -> None: + self.cancel_calls += 1 + + def has_transferring_tasks(self) -> bool: + return self._has_transferring_tasks + def close(self) -> None: self.closed = True +class _FakeWorkerThread: + def __init__(self, *, alive: bool) -> None: + self.alive = alive + self.join_calls: list[Optional[float]] = [] + + def join(self, timeout: Optional[float] = None) -> None: + self.join_calls.append(timeout) + + def is_alive(self) -> bool: + return self.alive + + +@dataclass +class _FakeReceiveTaskOwnership: + lifecycle_managed: bool + status: TaskStatus + publication_started: bool = True + + @property + def legacy_resources_drained(self) -> bool: + return self.lifecycle_managed or not self.publication_started + + class _FakeTask: def __init__(self, status: TaskStatus, wait_result: bool = True) -> None: self.status = status @@ -105,10 +201,10 @@ def _make_transceiver( transceiver._transfer_worker = _FakeTransferWorker() transceiver._ctx_consensus = lambda local_ids: list(local_ids) transceiver._ctx_consensus_outcome = ( - lambda _to_process, cancelled, failed, completed, timed_out: ( - cancelled, - failed, - completed, + lambda _to_process, cancelled, failed, completed, timed_out, cleanup_ready: ( + [rid for rid in cancelled if rid in cleanup_ready], + [rid for rid in failed if rid in cleanup_ready], + [rid for rid in completed if rid in cleanup_ready], timed_out, ) ) @@ -150,6 +246,47 @@ def test_context_transfer_status_bounded_poll_keeps_not_ready_session_queued() - assert transceiver._transfer_worker.sweep_count == 1 +def test_context_transfer_status_retains_cancelled_source_until_physically_terminal() -> None: + session = _FakeSession( + rid=19, + wait_result=None, + status=SessionStatus.CANCELLED, + has_failed=True, + has_transferring_tasks=True, + ) + transceiver = _make_transceiver({19: session}) + + completed, failed = transceiver.check_context_transfer_status(at_least_request_num=1) + + assert completed == [] + assert failed == [] + assert session.blocking_calls == [False] + assert not session.closed + assert 19 in transceiver._send_sessions + assert 19 in transceiver._send_reqs + + +def test_context_transfer_status_reports_drained_remote_cancellation_as_failed() -> None: + rid = 24 + session = _FakeSession( + rid=rid, + wait_result=WaitResult.FAILED, + status=SessionStatus.CANCELLED, + has_failed=True, + ) + req = _FakeRequest(request_id=rid) + transceiver = _make_transceiver({rid: session}, {rid: req}) + + completed, failed = transceiver.check_context_transfer_status(at_least_request_num=1) + + assert completed == [] + assert failed == [rid] + assert session.closed + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert rid not in transceiver._send_sessions + assert rid not in transceiver._send_reqs + + def test_context_transfer_status_block_all_uses_blocking_wait() -> None: session = _FakeSession(rid=12, wait_result=WaitResult.COMPLETED) req = _FakeRequest() @@ -189,6 +326,275 @@ def test_context_transfer_status_zero_budget_processes_task_level_failure() -> N assert 13 not in transceiver._send_reqs +def test_context_provisional_cleanup_retries_and_readvertises_without_extra_consensus() -> None: + rid = 70 + request = _RetirementRequest(rid) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + session.close = Mock(side_effect=[RuntimeError("close failed"), None]) + transceiver = _make_transceiver({rid: session}, {rid: request}) + transceiver._ctx_consensus = Mock(side_effect=lambda local_ids: list(local_ids)) + + assert transceiver.check_context_transfer_status(0, mark_complete=True) == ([], []) + assert request.state is None + assert transceiver._send_sessions == {rid: session} + assert transceiver._send_reqs == {rid: request} + assert set(transceiver._get_context_retirements()) == {rid} + + # The closed/provisional candidate remains in the normal ready + # advertisement. The existing outcome consensus is the only prepare + # barrier; no post-decision consensus round is added. + assert transceiver.check_context_transfer_status(0, mark_complete=True) == ([rid], []) + consensus_inputs = [call.args[0] for call in transceiver._ctx_consensus.call_args_list] + assert consensus_inputs == [[rid], [rid]] + assert transceiver._get_context_retirements() == {} + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + assert session.close.call_count == 2 + + +def test_context_commit_invariant_failure_latches_fail_stop() -> None: + rid = 76 + request = _RetirementRequest(rid, fail_state=LlmRequestState.DISAGG_CONTEXT_COMPLETE) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + transceiver = _make_transceiver({rid: session}, {rid: request}) + + with pytest.raises(RuntimeError, match="terminal state update failed"): + transceiver.check_context_transfer_status(0, mark_complete=True) + + assert transceiver._shutdown_started + assert transceiver._retirement_fault is not None + assert set(transceiver._get_context_retirements()) == {rid} + with pytest.raises(RuntimeError, match="retirement invariant failed"): + transceiver.check_context_transfer_status(0, mark_complete=True) + _NON_DRAINED_TRANSCEIVERS.discard(transceiver) + + +def test_pending_context_retirement_does_not_block_positive_wait_on_unrelated_request() -> None: + rid = 71 + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + session.close = Mock(side_effect=[RuntimeError("close failed"), None]) + request = _RetirementRequest(rid) + transceiver = _make_transceiver({rid: session}, {rid: request}) + transceiver._ctx_consensus = Mock(side_effect=lambda local_ids: list(local_ids)) + + assert transceiver.check_context_transfer_status(0, mark_complete=True) == ([], []) + + unrelated_rid = 72 + unrelated_session = _FakeSession(rid=unrelated_rid, wait_result=None) + transceiver._send_sessions[unrelated_rid] = unrelated_session + transceiver._send_reqs[unrelated_rid] = _FakeRequest(request_id=unrelated_rid) + + assert transceiver.check_context_transfer_status(1, mark_complete=True) == ([rid], []) + assert unrelated_session.blocking_calls == [] + assert unrelated_rid in transceiver._send_sessions + + +def test_provisional_context_owner_blocks_cancellation_and_rid_replacement() -> None: + rid = 77 + request = _RetirementRequest(rid) + replacement = _RetirementRequest(rid) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + session.close = Mock(side_effect=RuntimeError("close pending")) + transceiver = _make_transceiver({rid: session}, {rid: request}) + transceiver._wait_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + + assert transceiver.check_context_transfer_status(0) == ([], []) + assert set(transceiver._get_context_retirements()) == {rid} + assert not transceiver.cancel_request(request) + with pytest.raises(RuntimeError, match="pending terminal retirement"): + transceiver.respond_and_send_async(replacement) + + assert transceiver._send_sessions == {rid: session} + assert transceiver._send_reqs == {rid: request} + + +def test_active_status_snapshot_blocks_cancellation_retirement_until_enrollment() -> None: + rid = 76 + request = _RetirementRequest(rid) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + transceiver = _make_transceiver({rid: session}, {rid: request}) + transceiver._wait_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + + snapshot_taken = threading.Event() + release_consensus = threading.Event() + + def pause_after_snapshot(local_ids: list[int]) -> list[int]: + snapshot_taken.set() + assert release_consensus.wait(timeout=2) + return list(local_ids) + + transceiver._ctx_consensus = pause_after_snapshot + status_results: list[tuple[list[int], list[int]]] = [] + errors: list[BaseException] = [] + + def poll_status() -> None: + try: + status_results.append(transceiver.check_context_transfer_status(0)) + except BaseException as error: + errors.append(error) + + poll_thread = threading.Thread(target=poll_status) + poll_thread.start() + assert snapshot_taken.wait(timeout=1) + try: + assert transceiver.cancel_request(request) is False + assert session.cancel_calls == 1 + assert not session.closed + assert transceiver._send_sessions == {rid: session} + assert transceiver._send_reqs == {rid: request} + finally: + release_consensus.set() + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert status_results == [([rid], [])] + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + + +def test_prepare_context_cancellation_does_not_reverse_consensus_promotion() -> None: + rid = 75 + request = _FakeRequest(request_id=rid) + transceiver = _make_transceiver({}) + transceiver._wait_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._transfer_worker.has_all_peer_req_infos_for_send = Mock(return_value=True) + + consensus_started = threading.Event() + release_consensus = threading.Event() + + def pause_consensus(local_ids: list[int]) -> list[int]: + consensus_started.set() + assert release_consensus.wait(timeout=2) + return list(local_ids) + + transceiver._ctx_consensus = pause_consensus + errors: list[BaseException] = [] + + def prepare() -> None: + try: + transceiver.prepare_context_requests([request]) + except BaseException as error: + errors.append(error) + + prepare_thread = threading.Thread(target=prepare) + prepare_thread.start() + assert consensus_started.wait(timeout=1) + try: + assert transceiver.cancel_request(request) is False + assert transceiver._wait_reqs == {rid: request} + assert transceiver._get_cancelled_wait_reqs() == {rid: request} + finally: + release_consensus.set() + prepare_thread.join(timeout=2) + + assert not prepare_thread.is_alive() + assert errors == [] + assert request.state == LlmRequestState.CONTEXT_INIT + assert transceiver._wait_reqs == {} + assert transceiver._get_cancelled_wait_reqs() == {rid: request} + + assert transceiver.cancel_request(request) is True + assert transceiver._wait_reqs == {} + assert transceiver._get_cancelled_wait_reqs() == {} + + +def test_prepare_context_cancellation_before_wait_enrollment_blocks_admission() -> None: + rid = 73 + request = _FakeRequest(request_id=rid) + transceiver = _make_transceiver({}) + transceiver._wait_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._transfer_worker.has_all_peer_req_infos_for_send = Mock(return_value=True) + transceiver._ctx_consensus = Mock(side_effect=lambda local_ids: list(local_ids)) + + status_admitted = threading.Event() + release_enrollment = threading.Event() + prepare_impl = transceiver._prepare_context_requests_impl + + def pause_before_enrollment(requests: list[_FakeRequest]) -> None: + status_admitted.set() + assert release_enrollment.wait(timeout=2) + prepare_impl(requests) + + transceiver._prepare_context_requests_impl = pause_before_enrollment + errors: list[BaseException] = [] + + def prepare() -> None: + try: + transceiver.prepare_context_requests([request]) + except BaseException as error: + errors.append(error) + + prepare_thread = threading.Thread(target=prepare) + prepare_thread.start() + assert status_admitted.wait(timeout=1) + try: + assert transceiver.cancel_request(request) is False + assert transceiver._wait_reqs == {} + assert transceiver._get_cancelled_wait_reqs() == {rid: request} + finally: + release_enrollment.set() + prepare_thread.join(timeout=2) + + assert not prepare_thread.is_alive() + assert errors == [] + assert request.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + assert transceiver._wait_reqs == {rid: request} + assert transceiver._get_cancelled_wait_reqs() == {rid: request} + transceiver._ctx_consensus.assert_called_once_with([]) + + assert transceiver.cancel_request(request) is True + assert transceiver._wait_reqs == {} + assert transceiver._get_cancelled_wait_reqs() == {} + + +def test_respond_and_send_rejects_wait_ledger_owner_collision() -> None: + rid = 74 + original = _FakeRequest(request_id=rid) + replacement = _FakeRequest(request_id=rid) + transceiver = _make_transceiver({}) + transceiver._shutdown_started = False + transceiver._wait_reqs = {rid: original} + + with pytest.raises(RuntimeError, match="different live source owner"): + transceiver.respond_and_send_async(replacement) + with pytest.raises(RuntimeError, match="still waiting for scheduler promotion"): + transceiver.respond_and_send_async(original) + + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + assert transceiver._wait_reqs == {rid: original} + + def test_context_transfer_status_skips_consensus_when_never_sent() -> None: # A worker that never sends skips the ctx consensus even when TP sync would need it, but still # sweeps so nothing leaks. @@ -209,21 +615,15 @@ def test_context_transfer_status_skips_consensus_when_never_sent() -> None: assert transceiver._transfer_worker.sweep_count == 1 -def test_context_transfer_status_never_sent_no_sync_is_a_noop() -> None: - # With no tp/pp sync (e.g. attention_dp), a never-sent worker skips the consensus and the sweep, - # unchanged from before -- a true no-op, so the fix can't slow attention_dp workers. +def test_context_transfer_idle_fast_path_still_runs_sender_progress() -> None: transceiver = object.__new__(KvCacheTransceiverV2) transceiver._ever_had_send_session = False transceiver._ctx_need_tp_sync = False transceiver._ctx_need_pp_sync = False - transceiver._send_sessions = {} - transceiver._send_reqs = {} transceiver._transfer_worker = _FakeTransferWorker() - transceiver._ctx_consensus = Mock(side_effect=AssertionError("consensus must be skipped")) assert transceiver.check_context_transfer_status(at_least_request_num=0) == ([], []) - transceiver._ctx_consensus.assert_not_called() - assert transceiver._transfer_worker.sweep_count == 0 # matches the original early-out exactly + assert transceiver._transfer_worker.sweep_count == 1 def test_gen_transfer_status_enters_consensus_when_sync_required() -> None: @@ -235,7 +635,7 @@ def test_gen_transfer_status_enters_consensus_when_sync_required() -> None: transceiver._gen_consensus = Mock(return_value=[]) transceiver._build_to_process = Mock(return_value=[]) transceiver._gen_consensus_outcome = Mock(return_value=([], [], [])) - transceiver._close_failed_sessions = Mock() + transceiver._dist = Mock(rank=0) completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=0) @@ -243,32 +643,668 @@ def test_gen_transfer_status_enters_consensus_when_sync_required() -> None: assert failed == [] assert cancelled == [] transceiver._gen_consensus.assert_called_once_with([]) + transceiver._gen_consensus_outcome.assert_called_once_with([], [], [], [], []) + + +def test_gen_transfer_status_retains_logical_failure_until_physical_drain() -> None: + rid = 14 + session = _FakeSession( + rid=rid, + wait_result=None, + has_failed=True, + ) + req = _FakeRequest(request_id=rid) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: req} + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._gen_consensus_outcome = ( + lambda _to_process, cancelled, failed, completed, _cleanup_ready: ( + cancelled, + failed, + completed, + ) + ) + + completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=0) + + assert completed == [] + assert failed == [] + assert cancelled == [] + assert session.blocking_calls == [False] + assert not session.closed + assert req.state is None + assert transceiver._recv_sessions == {rid: session} + assert transceiver._recv_reqs == {rid: req} + + session._wait_result = WaitResult.FAILED + transceiver._dist = Mock(rank=0) + completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=0) + + assert completed == [] + assert failed == [rid] + assert cancelled == [] + assert session.closed + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert rid not in transceiver._recv_sessions + assert rid not in transceiver._recv_reqs + + +def test_gen_transfer_status_retains_cancelled_session_until_wait_is_terminal() -> None: + rid = 16 + session = _FakeSession( + rid=rid, + wait_result=None, + status=SessionStatus.CANCELLED, + has_failed=True, + ) + req = _FakeRequest(request_id=rid) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: req} + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._gen_consensus_outcome = ( + lambda _to_process, cancelled, failed, completed, _cleanup_ready: ( + cancelled, + failed, + completed, + ) + ) + + completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=0) + + assert completed == [] + assert failed == [] + assert cancelled == [] + assert not session.closed + assert transceiver._recv_sessions == {rid: session} + assert transceiver._recv_reqs == {rid: req} + + session._wait_result = WaitResult.FAILED + completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=0) + + assert completed == [] + assert failed == [] + assert cancelled == [req] + assert session.closed + assert rid not in transceiver._recv_sessions + assert rid not in transceiver._recv_reqs + + +@pytest.mark.parametrize( + "failure_stage", + ["kv_size", "aux", "history", "close"], +) +def test_generation_provisional_preparation_retries_each_fallible_phase(failure_stage: str) -> None: + rid = 73 + request = _RetirementRequest(rid) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + session.close = Mock() + if failure_stage == "close": + session.close.side_effect = [RuntimeError("close failed"), None] + + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver.kv_transfer_poll_interval_ms = 0 + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: request} + transceiver._gen_consensus = Mock(side_effect=lambda local_ids: list(local_ids)) + transceiver._gen_consensus_outcome = ( + lambda _to_process, cancelled, failed, completed, cleanup_ready: ( + [rid for rid in cancelled if rid in cleanup_ready], + [rid for rid in failed if rid in cleanup_ready], + [rid for rid in completed if rid in cleanup_ready], + ) + ) + transceiver._need_aux_transfer = Mock(return_value=True) + transceiver._apply_aux = Mock() + transceiver._assert_disagg_history_declared = Mock() + transceiver._dist = Mock(rank=0) + if failure_stage == "kv_size": + request.set_kv_cache_size.side_effect = [ + RuntimeError("kv size failed"), + None, + ] + elif failure_stage == "aux": + transceiver._apply_aux.side_effect = [RuntimeError("aux failed"), None] + elif failure_stage == "history": + transceiver._assert_disagg_history_declared.side_effect = [ + RuntimeError("history failed"), + None, + ] + + assert transceiver.check_gen_transfer_status(0) == ([], [], []) + progress = transceiver._get_generation_retirements()[rid] + failed_flag = { + "kv_size": "kv_size_set", + "aux": "aux_applied", + "history": "history_checked", + "close": "session_closed", + }[failure_stage] + assert not getattr(progress, failed_flag) + assert request.state is None + assert transceiver._recv_sessions == {rid: session} + assert transceiver._recv_reqs == {rid: request} + + assert transceiver.check_gen_transfer_status(0) == ([rid], [], []) + consensus_inputs = [call.args[0] for call in transceiver._gen_consensus.call_args_list] + assert consensus_inputs == [[rid], [rid]] + assert transceiver._get_generation_retirements() == {} + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + assert request.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + + expected_kv_size_calls = 2 if failure_stage == "kv_size" else 1 + assert request.set_kv_cache_size.call_count == expected_kv_size_calls + request.set_kv_cache_size.assert_called_with(64) + expected_aux_calls = 2 if failure_stage == "aux" else 1 + assert transceiver._apply_aux.call_count == expected_aux_calls + expected_history_calls = 2 if failure_stage == "history" else 1 + assert transceiver._assert_disagg_history_declared.call_count == expected_history_calls + expected_close_calls = 2 if failure_stage == "close" else 1 + assert session.close.call_count == expected_close_calls + assert request.terminal_state_set_calls == 1 + + +def test_generation_provisional_candidate_waits_for_peer_and_readvertises() -> None: + rid = 75 + request = _RetirementRequest(rid) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = True + transceiver.kv_transfer_poll_interval_ms = 0 + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: request} + transceiver._gen_consensus = Mock(side_effect=lambda local_ids: list(local_ids)) + # The local candidate is fully prepared on the first poll, but the peer is + # not cleanup-ready. On the next poll the peer catches up and the same + # existing outcome consensus admits the decision. + transceiver._gen_consensus_outcome = Mock(side_effect=[([], [], []), ([], [], [rid])]) + transceiver._need_aux_transfer = Mock(return_value=False) + transceiver._assert_disagg_history_declared = Mock() + transceiver._dist = Mock(rank=0) + + assert transceiver.check_gen_transfer_status(0) == ([], [], []) + assert request.state is None + assert transceiver._recv_sessions == {rid: session} + assert transceiver._recv_reqs == {rid: request} + assert set(transceiver._get_generation_retirements()) == {rid} + assert session.closed + assert session.blocking_calls == [False] + + assert transceiver.check_gen_transfer_status(0) == ([rid], [], []) + assert request.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + assert session.blocking_calls == [False] + assert [call.args[0] for call in transceiver._gen_consensus.call_args_list] == [[rid], [rid]] + + +def test_shutdown_retains_decided_retirement_until_local_commit() -> None: + rid = 74 + close_allowed = False + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + + def close_session() -> None: + if not close_allowed: + raise RuntimeError("close still pending") + session.closed = True + + session.close = Mock(side_effect=close_session) + request = _RetirementRequest(rid) + transceiver = _make_transceiver({rid: session}, {rid: request}) + transceiver._shutdown = False + transceiver._shutdown_started = False + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._wait_reqs = {} + + assert transceiver.check_context_transfer_status(0, mark_complete=True) == ([], []) + try: + assert transceiver.shutdown() is False + assert transceiver in _NON_DRAINED_TRANSCEIVERS + assert set(transceiver._get_context_retirements()) == {rid} + assert transceiver._send_sessions == {rid: session} + assert session.cancel_calls == 0 + + close_allowed = True + assert transceiver.shutdown() is True + assert transceiver not in _NON_DRAINED_TRANSCEIVERS + assert transceiver._get_context_retirements() == {} + assert transceiver._send_sessions == {} + assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + finally: + _NON_DRAINED_TRANSCEIVERS.discard(transceiver) + + +def test_gen_transfer_timeout_is_not_cleanup_ready() -> None: + rid = 22 + session = _FakeSession(rid=rid, wait_result=WaitResult.TIMEOUT) + req = _FakeRequest(request_id=rid) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: req} + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._gen_consensus_outcome = Mock(return_value=([], [], [])) + + completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=None) + + assert completed == [] + assert failed == [] + assert cancelled == [] + transceiver._gen_consensus_outcome.assert_called_once_with([rid], [], [], [], []) + assert not session.closed + assert transceiver._recv_sessions == {rid: session} + assert transceiver._recv_reqs == {rid: req} + + +def test_async_receive_enrolls_owner_before_allocation_and_retains_on_publication_error() -> None: + rid = 17 + req = _FakeRequest(request_id=rid) + session = Mock() + session.receive.side_effect = RuntimeError("publication failed") + session.has_transferring_tasks.return_value = True + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._ever_had_recv_session = False + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._kv_size_rank_factor = 1 + transceiver._slice_num_bytes = Mock(return_value=8) + + def create_rx_session(owner): + assert transceiver._recv_reqs == {rid: owner} + return session + + def create_kv_slice(owner): + assert transceiver._recv_reqs == {rid: owner} + assert transceiver._recv_sessions == {rid: session} + return object() + + transceiver._transfer_worker = Mock() + transceiver._transfer_worker.create_rx_session.side_effect = create_rx_session + transceiver._create_kv_slice = Mock(side_effect=create_kv_slice) + + with pytest.raises(RuntimeError, match="publication failed"): + transceiver.request_and_receive_async(req) + + assert transceiver._recv_reqs == {rid: req} + assert transceiver._recv_sessions == {rid: session} + + +def test_async_send_rejects_a_distinct_owner_for_a_live_request_id() -> None: + rid = 20 + original = _FakeRequest(request_id=rid) + replacement = _FakeRequest(request_id=rid) + session = _FakeSession( + rid=rid, + wait_result=None, + has_transferring_tasks=True, + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._shutdown_started = False + transceiver._send_sessions = {rid: session} + transceiver._send_reqs = {rid: original} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._ever_had_send_session = True + + with pytest.raises(RuntimeError, match="different live source owner"): + transceiver.respond_and_send_async(replacement) + + assert transceiver._send_reqs == {rid: original} + assert transceiver._send_sessions == {rid: session} + + +def test_context_status_poll_retains_owner_and_blocks_replacement() -> None: + rid = 21 + original = _FakeRequest(request_id=rid) + replacement = _FakeRequest(request_id=rid) + wait_started = threading.Event() + release_wait = threading.Event() + errors: list[BaseException] = [] + results: list[tuple[list[int], list[int]]] = [] + + original_session = Mock() + original_session.status = SessionStatus.READY + original_session.is_completed.return_value = False + original_session.has_failed.side_effect = release_wait.is_set + original_session.has_transferring_tasks.return_value = False + + def wait_complete(*, blocking: bool) -> WaitResult: + assert not blocking + wait_started.set() + assert release_wait.wait(timeout=2) + return WaitResult.FAILED + + original_session.wait_complete.side_effect = wait_complete + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._send_sessions = {rid: original_session} + transceiver._send_reqs = {rid: original} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._wait_reqs = {} + transceiver._ever_had_send_session = True + transceiver._ctx_need_tp_sync = False + transceiver._ctx_need_pp_sync = False + transceiver._transfer_worker = _FakeTransferWorker() + transceiver._ctx_consensus = lambda local_ids: list(local_ids) + transceiver._ctx_consensus_outcome = ( + lambda _to_process, cancelled, failed, completed, timed_out, _cleanup_ready: ( + cancelled, + failed, + completed, + timed_out, + ) + ) + + def poll_status() -> None: + try: + results.append(transceiver.check_context_transfer_status(at_least_request_num=1)) + except BaseException as error: + errors.append(error) + + poll_thread = threading.Thread(target=poll_status) + poll_thread.start() + assert wait_started.wait(timeout=1) + try: + assert transceiver.cancel_request(original) is False + assert transceiver._send_sessions == {rid: original_session} + assert transceiver._send_reqs == {rid: original} + with pytest.raises(RuntimeError, match="different live source owner"): + transceiver.respond_and_send_async(replacement) + finally: + release_wait.set() + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert results == [([], [rid])] + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + original_session.cancel.assert_called_once_with() + original_session.close.assert_called_once_with() + + +def test_gen_status_poll_retains_owner_and_blocks_replacement() -> None: + rid = 23 + original = _FakeRequest(request_id=rid) + replacement = _FakeRequest(request_id=rid) + wait_started = threading.Event() + release_wait = threading.Event() + errors: list[BaseException] = [] + results: list[tuple[list[int], list[int], list[_FakeRequest]]] = [] + + original_session = Mock() + original_session.status = SessionStatus.READY + original_session.is_completed.return_value = False + original_session.has_failed.side_effect = release_wait.is_set + original_session.has_transferring_tasks.return_value = False + + def wait_complete(*, blocking: bool) -> WaitResult: + assert not blocking + wait_started.set() + assert release_wait.wait(timeout=2) + return WaitResult.FAILED + + original_session.wait_complete.side_effect = wait_complete + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._recv_sessions = {rid: original_session} + transceiver._recv_reqs = {rid: original} + transceiver._wait_reqs = {} + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver._dist = Mock(rank=0) + transceiver.kv_transfer_poll_interval_ms = 1000 + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._gen_consensus_outcome = ( + lambda _to_process, cancelled, failed, completed, _cleanup_ready: ( + cancelled, + failed, + completed, + ) + ) + + def poll_status() -> None: + try: + results.append(transceiver.check_gen_transfer_status(at_least_request_num=1)) + except BaseException as error: + errors.append(error) + + poll_thread = threading.Thread(target=poll_status) + poll_thread.start() + assert wait_started.wait(timeout=1) + try: + assert transceiver.cancel_request(original) is False + assert transceiver._recv_sessions == {rid: original_session} + assert transceiver._recv_reqs == {rid: original} + with pytest.raises(RuntimeError, match="different live destination owner"): + transceiver.request_and_receive_async(replacement) + finally: + release_wait.set() + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert results == [([], [rid], [])] + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + original_session.cancel.assert_called_once_with() + original_session.close.assert_called_once_with() + + +def test_cancel_request_retains_receive_session_until_writers_drain() -> None: + rid = 15 + session = _FakeSession( + rid=rid, + wait_result=None, + has_transferring_tasks=True, + ) + req = _FakeRequest(request_id=rid) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._wait_reqs = {} + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: req} + + assert not transceiver.cancel_request(req) + assert session.cancel_calls == 1 + assert not session.closed + assert transceiver._recv_sessions == {rid: session} + assert transceiver._recv_reqs == {rid: req} + + session._has_transferring_tasks = False + + assert transceiver.cancel_request(req) + assert session.cancel_calls == 2 + assert session.closed + assert rid not in transceiver._recv_sessions + assert rid not in transceiver._recv_reqs + + +def test_cancel_request_continues_with_receive_after_send_cleanup_error() -> None: + rid = 25 + req = _FakeRequest(request_id=rid) + send_session = Mock() + send_session.cancel.side_effect = [RuntimeError("send cancel failed"), None] + send_session.has_transferring_tasks.return_value = False + recv_session = Mock() + recv_session.has_transferring_tasks.return_value = False + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._wait_reqs = {} + transceiver._send_sessions = {rid: send_session} + transceiver._send_reqs = {rid: req} + transceiver._recv_sessions = {rid: recv_session} + transceiver._recv_reqs = {rid: req} + + assert not transceiver.cancel_request(req) + + assert transceiver._send_sessions == {rid: send_session} + assert transceiver._send_reqs == {rid: req} + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + recv_session.cancel.assert_called_once_with() + recv_session.close.assert_called_once_with() + + assert transceiver.cancel_request(req) + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + assert send_session.cancel.call_count == 2 + send_session.close.assert_called_once_with() + + +def test_prepare_context_requests_rejects_admission_after_shutdown() -> None: + req = _FakeRequest(request_id=26) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = True + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._wait_reqs = {} + transceiver._transfer_worker = Mock() + + with pytest.raises(RuntimeError, match="shutting down"): + transceiver.prepare_context_requests([req]) + + assert transceiver._wait_reqs == {} + transceiver._transfer_worker.has_all_peer_req_infos_for_send.assert_not_called() + + +def test_prepare_context_requests_rejects_replacement_wait_owner() -> None: + rid = 27 + original = _FakeRequest(request_id=rid) + replacement = _FakeRequest(request_id=rid) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._wait_reqs = {rid: original} + transceiver._transfer_worker = Mock() + + with pytest.raises(RuntimeError, match="different live source owner"): + transceiver.prepare_context_requests([replacement]) + + assert transceiver._wait_reqs == {rid: original} + + +def test_prepare_context_requests_does_not_partially_enroll_on_owner_collision() -> None: + existing_rid = 28 + new_req = _FakeRequest(request_id=29) + original = _FakeRequest(request_id=existing_rid) + replacement = _FakeRequest(request_id=existing_rid) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._wait_reqs = {existing_rid: original} + transceiver._transfer_worker = Mock() + + with pytest.raises(RuntimeError, match="different live source owner"): + transceiver.prepare_context_requests([new_req, replacement]) + + assert transceiver._wait_reqs == {existing_rid: original} + assert new_req.state is None + + +def test_generation_timeout_keeps_request_active_until_cancel_drains() -> None: + request = Mock() + request.py_request_id = 18 + request.is_attention_dp_dummy = False + request.py_kv_transfer_timed_out = True + executor = object.__new__(PyExecutor) + executor.active_requests = [request] + executor.perf_manager = Mock() + executor.perf_manager.get_timestamp.return_value = 0.0 + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.cancel_request.return_value = False + executor._enqueue_responses = Mock() + executor.enable_attention_dp = False + executor.dist = Mock(world_size=1) + + finished = executor._handle_responses() + + assert finished == [] + assert executor.active_requests == [request] + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) def test_consensus_outcome_uses_single_batched_allgather() -> None: - # The cancelled/failed/completed id lists are exchanged with ONE allgather + # Outcomes and physical cleanup readiness are exchanged with ONE allgather # (packed as a list-of-lists) instead of three; verify a single call and that - # union (cancelled/failed) + intersection (completed) semantics are preserved. + # union (cancelled/failed) + intersection (completed/readiness) semantics are preserved. transceiver = object.__new__(KvCacheTransceiverV2) calls: list = [] def fake_allgather(payload): calls.append(payload) - # rank0 = this rank's [cancelled, failed, completed]; rank1 = a peer rank. - return [payload, [[], [99], [7, 8]]] + # rank0 = this rank's [cancelled, failed, completed, cleanup_ready]; + # rank1 = a peer rank. + return [payload, [[], [99], [7, 8], [1, 2, 7, 8, 99]]] to_process = [1, 2, 7, 8, 99] new_cancelled, new_failed, new_completed = transceiver._consensus_outcome( - to_process, [1], [2], [7], fake_allgather, True + to_process, + [1], + [2], + [7], + [1, 2, 7, 8, 99], + fake_allgather, + True, ) assert len(calls) == 1 # batched: a single allgather, not three - assert calls[0] == [[1], [2], [7]] + assert calls[0] == [[1], [2], [7], [1, 2, 7, 8, 99]] assert new_cancelled == [1] # union of cancelled across ranks assert new_failed == [2, 99] # union of failed across ranks assert new_completed == [7] # intersection only (8 is completed on the peer only) +def test_consensus_retains_remote_failure_until_every_rank_is_cleanup_ready() -> None: + transceiver = object.__new__(KvCacheTransceiverV2) + rid = 21 + + def peer_failed(payload): + return [payload, [[], [rid], [], [rid]]] + + cancelled, failed, completed = transceiver._consensus_outcome( + [rid], [], [], [], [], peer_failed, True + ) + assert (cancelled, failed, completed) == ([], [], []) + + cancelled, failed, completed = transceiver._consensus_outcome( + [rid], [], [], [], [rid], peer_failed, True + ) + assert (cancelled, failed, completed) == ([], [rid], []) + + def test_tx_session_wait_complete_defaults_to_blocking() -> None: task = _FakeTask(TaskStatus.INIT, wait_result=False) session = _make_tx_session([task]) @@ -290,6 +1326,8 @@ def test_tx_session_wait_complete_nonblocking_reports_later_task_error() -> None failed_task = _FakeTask(TaskStatus.ERROR) session = _make_tx_session([pending_task, failed_task]) + assert session.wait_complete(blocking=False) is None + pending_task.status = TaskStatus.TRANSFERRED assert session.wait_complete(blocking=False) == WaitResult.FAILED assert pending_task.wait_calls == [] assert failed_task.wait_calls == [] @@ -312,6 +1350,7 @@ def test_check_context_runs_consensus_after_a_send() -> None: transceiver.check_context_transfer_status(0) transceiver._ctx_consensus.assert_called_once() + transceiver._ctx_consensus_outcome.assert_called_once_with([], [], [], [], [], []) def test_prepare_context_requests_skips_consensus_when_nothing_waiting() -> None: @@ -323,3 +1362,643 @@ def test_prepare_context_requests_skips_consensus_when_nothing_waiting() -> None transceiver.prepare_context_requests([]) transceiver._ctx_consensus.assert_not_called() + + +def test_sender_shutdown_retains_remote_agents_while_worker_is_alive() -> None: + sender = object.__new__(Sender) + worker = _FakeWorkerThread(alive=True) + dealer = Mock() + sender._shutdown = False + sender._shutdown_started = False + sender._shutdown_complete = False + sender._shutdown_sentinels_sent = False + sender._messenger = Mock() + sender._send_task_queues = [Mock()] + sender._worker_threads = [worker] + sender._loaded_remote_agents = {"peer-agent"} + sender._loaded_remote_agents_lock = threading.Lock() + sender._agent = Mock() + sender._dealers = {"peer": dealer} + + assert sender.shutdown() is False + + assert sender._shutdown_complete is False + sender._agent.invalidate_remote_agent.assert_not_called() + dealer.stop.assert_not_called() + assert sender._loaded_remote_agents == {"peer-agent"} + + worker.alive = False + + assert sender.shutdown() is True + assert sender._shutdown is True + assert sender._shutdown_complete is True + sender._agent.invalidate_remote_agent.assert_called_once_with("peer-agent") + dealer.stop.assert_called_once_with() + + +def test_transfer_worker_does_not_treat_local_agent_shutdown_as_remote_quiescence() -> None: + worker = object.__new__(TransferWorker) + worker._shutdown = False + worker._shutdown_started = False + worker._shutdown_complete = False + worker._rank_info_server = Mock() + worker._sender = Mock() + worker._sender.shutdown.return_value = True + worker._receiver = Mock() + worker._receiver.transfers_drained = False + worker._agent = Mock() + worker._bounce = Mock() + worker._registered_mem = [Mock()] + + assert worker.shutdown() is False + + assert worker._shutdown is False + worker._receiver.begin_shutdown.assert_called_once_with() + worker._agent.shutdown.assert_not_called() + worker._agent.deregister_memory.assert_not_called() + worker._bounce.close.assert_not_called() + worker._receiver.shutdown.assert_not_called() + + +def _make_receiver_with_session(session: RxSession) -> Receiver: + receiver = object.__new__(Receiver) + receiver._shutdown = True # Keep object.__new__ fixtures out of __del__ cleanup. + receiver._recv_registry = Mock() + receiver._recv_registry.is_drained.return_value = True + receiver._sessions_lock = threading.Lock() + receiver._sessions = {41: weakref.ref(session)} + return receiver + + +def test_receiver_keeps_session_owner_alive_until_explicit_retirement() -> None: + receiver = object.__new__(Receiver) + receiver._shutdown = True # Keep object.__new__ fixtures out of __del__ cleanup. + receiver._sessions = {} + receiver._sessions_lock = threading.Lock() + receiver._pre_cancelled_rids = set() + session = _FakeSession(rid=40, wait_result=None) + + receiver.setup_session(session) + session_ref = weakref.ref(session) + del session + gc.collect() + + assert session_ref() is not None + assert receiver._sessions[40] is session_ref() + + +def _make_receive_ownership_session( + task: _FakeReceiveTaskOwnership, + *, + need_aux: bool, + aux_drained: bool, +) -> RxSession: + session = object.__new__(RxSession) + session._closed = True # Keep object.__new__ test fixtures out of __del__ cleanup. + session.lock = threading.Lock() + session._kv_tasks = [task] + session._need_aux = need_aux + session._aux_status = TaskStatus.TRANSFERRED if aux_drained else TaskStatus.TRANSFERRING + session._aux_drained = aux_drained + return session + + +def test_receiver_not_drained_while_legacy_receive_is_active() -> None: + session = _make_receive_ownership_session( + _FakeReceiveTaskOwnership( + lifecycle_managed=False, + status=TaskStatus.TRANSFERRING, + ), + need_aux=False, + aux_drained=True, + ) + receiver = _make_receiver_with_session(session) + + assert receiver.transfers_drained is False + + +def test_receiver_not_drained_after_published_legacy_failure() -> None: + session = _make_receive_ownership_session( + _FakeReceiveTaskOwnership( + lifecycle_managed=False, + status=TaskStatus.ERROR, + ), + need_aux=False, + aux_drained=True, + ) + receiver = _make_receiver_with_session(session) + + assert receiver.transfers_drained is False + + +def test_receiver_not_drained_while_aux_receive_is_active() -> None: + session = _make_receive_ownership_session( + _FakeReceiveTaskOwnership( + lifecycle_managed=True, + status=TaskStatus.TRANSFERRED, + ), + need_aux=True, + aux_drained=False, + ) + receiver = _make_receiver_with_session(session) + + assert receiver.transfers_drained is False + + +def test_receiver_not_drained_after_nonterminal_aux_failure() -> None: + session = _make_receive_ownership_session( + _FakeReceiveTaskOwnership( + lifecycle_managed=True, + status=TaskStatus.TRANSFERRED, + ), + need_aux=True, + aux_drained=False, + ) + session._aux_status = TaskStatus.ERROR + receiver = _make_receiver_with_session(session) + + assert receiver.transfers_drained is False + + +def test_transceiver_shutdown_retains_sessions_until_worker_drains() -> None: + send_session = _FakeSession(rid=51, wait_result=None) + recv_session = _FakeSession(rid=52, wait_result=None) + send_req = _FakeRequest(request_id=51) + recv_req = _FakeRequest(request_id=52) + wait_req = _FakeRequest(request_id=53) + worker = _FakeTransferWorker(shutdown_results=[False, True]) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._shutdown = False + transceiver._shutdown_started = False + transceiver._transfer_worker = worker + transceiver._send_sessions = {51: send_session} + transceiver._send_reqs = {51: send_req} + transceiver._recv_sessions = {52: recv_session} + transceiver._recv_reqs = {52: recv_req} + transceiver._wait_reqs = {53: wait_req} + + assert transceiver.shutdown() is False + assert transceiver._shutdown is False + assert worker.shutdown_count == 1 + assert not send_session.closed + assert not recv_session.closed + assert transceiver._send_sessions == {51: send_session} + assert transceiver._send_reqs == {51: send_req} + assert transceiver._recv_sessions == {52: recv_session} + assert transceiver._recv_reqs == {52: recv_req} + assert transceiver._wait_reqs == {53: wait_req} + + assert transceiver.shutdown() is True + assert transceiver._shutdown is True + assert worker.shutdown_count == 2 + assert send_session.closed + assert recv_session.closed + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + assert transceiver._wait_reqs == {} + + +@pytest.mark.parametrize("failure_stage", ["cancel", "worker", "close"]) +def test_transceiver_shutdown_retains_owner_when_cleanup_raises(failure_stage: str) -> None: + send_session = Mock() + recv_session = Mock() + worker = Mock() + worker.shutdown.return_value = True + if failure_stage == "cancel": + send_session.cancel.side_effect = [RuntimeError("cancel failed"), None] + elif failure_stage == "worker": + worker.shutdown.side_effect = [RuntimeError("worker failed"), True] + else: + send_session.close.side_effect = [RuntimeError("close failed"), None] + + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._shutdown = False + transceiver._shutdown_started = False + transceiver._transfer_worker = worker + transceiver._send_sessions = {51: send_session} + transceiver._send_reqs = {51: _FakeRequest(request_id=51)} + transceiver._recv_sessions = {52: recv_session} + transceiver._recv_reqs = {52: _FakeRequest(request_id=52)} + wait_req = _FakeRequest(request_id=53) + transceiver._wait_reqs = {53: wait_req} + + try: + assert transceiver.shutdown() is False + assert transceiver in _NON_DRAINED_TRANSCEIVERS + assert transceiver._send_sessions + if failure_stage == "worker": + assert transceiver._recv_sessions + recv_session.close.assert_not_called() + else: + assert transceiver._recv_sessions == {} + recv_session.cancel.assert_called_once_with() + recv_session.close.assert_called_once_with() + assert transceiver._wait_reqs == {53: wait_req} + + assert transceiver.shutdown() is True + assert transceiver not in _NON_DRAINED_TRANSCEIVERS + assert transceiver._send_sessions == {} + assert transceiver._recv_sessions == {} + assert transceiver._wait_reqs == {} + finally: + _NON_DRAINED_TRANSCEIVERS.discard(transceiver) + + +def test_shutdown_does_not_drain_provisional_owner_during_active_status_call() -> None: + rid = 78 + request = _RetirementRequest(rid) + session = _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + transceiver = _make_transceiver({rid: session}, {rid: request}) + transceiver._shutdown = False + transceiver._shutdown_started = False + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._wait_reqs = {} + + enrolled = threading.Event() + release_prepare = threading.Event() + original_prepare = transceiver._prepare_context_retirements + prepare_calls = 0 + + def pause_after_enrollment() -> None: + nonlocal prepare_calls + prepare_calls += 1 + if prepare_calls == 2: + enrolled.set() + assert release_prepare.wait(timeout=2) + original_prepare() + + transceiver._prepare_context_retirements = pause_after_enrollment + status_results: list[tuple[list[int], list[int]]] = [] + errors: list[BaseException] = [] + + def poll_status() -> None: + try: + status_results.append(transceiver.check_context_transfer_status(0, mark_complete=True)) + except BaseException as error: + errors.append(error) + + poll_thread = threading.Thread(target=poll_status) + poll_thread.start() + assert enrolled.wait(timeout=1) + try: + assert transceiver.shutdown() is False + assert set(transceiver._get_context_retirements()) == {rid} + assert transceiver._send_sessions == {rid: session} + assert transceiver._transfer_worker.shutdown_count == 0 + finally: + release_prepare.set() + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert status_results == [([rid], [])] + assert transceiver.shutdown() is True + assert transceiver._transfer_worker.shutdown_count == 1 + _NON_DRAINED_TRANSCEIVERS.discard(transceiver) + + +def test_blocking_generation_poll_shutdown_cancels_without_deadlock() -> None: + rid = 79 + request = _RetirementRequest(rid) + wait_started = threading.Event() + cancel_seen = threading.Event() + session = _FakeSession(rid=rid, wait_result=None) + + def wait_complete(*, blocking: bool) -> WaitResult: + assert blocking + wait_started.set() + assert cancel_seen.wait(timeout=2) + return WaitResult.FAILED + + def cancel_session() -> None: + session.cancel_calls += 1 + session._status = SessionStatus.CANCELLED + cancel_seen.set() + + session.wait_complete = wait_complete + session.cancel = cancel_session + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown = False + transceiver._shutdown_started = False + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver.kv_transfer_poll_interval_ms = 0 + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._recv_sessions = {rid: session} + transceiver._recv_reqs = {rid: request} + transceiver._wait_reqs = {} + transceiver._transfer_worker = _FakeTransferWorker() + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._gen_consensus_outcome = ( + lambda _to_process, cancelled, failed, completed, cleanup_ready: ( + [candidate for candidate in cancelled if candidate in cleanup_ready], + [candidate for candidate in failed if candidate in cleanup_ready], + [candidate for candidate in completed if candidate in cleanup_ready], + ) + ) + + status_results: list[tuple[list[int], list[int], list[_RetirementRequest]]] = [] + errors: list[BaseException] = [] + + def poll_status() -> None: + try: + status_results.append(transceiver.check_gen_transfer_status(None)) + except BaseException as error: + errors.append(error) + + poll_thread = threading.Thread(target=poll_status) + poll_thread.start() + assert wait_started.wait(timeout=1) + assert transceiver.shutdown() is False + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert status_results == [([], [], [request])] + assert transceiver._transfer_worker.shutdown_count == 0 + assert transceiver.shutdown() is True + assert transceiver._transfer_worker.shutdown_count == 1 + _NON_DRAINED_TRANSCEIVERS.discard(transceiver) + + +@pytest.mark.parametrize("direction", ["send", "receive"]) +def test_transceiver_shutdown_waits_for_async_enrollment_and_launch(direction) -> None: + rid = 61 if direction == "send" else 62 + req = _FakeRequest(request_id=rid) + launch_started = threading.Event() + release_launch = threading.Event() + shutdown_attempted = threading.Event() + shutdown_called = threading.Event() + errors: list[BaseException] = [] + session = Mock() + session.has_transferring_tasks.return_value = False + + def block_launch(_slice) -> None: + launch_started.set() + assert release_launch.wait(timeout=2) + + if direction == "send": + session.send.side_effect = block_launch + else: + session.receive.side_effect = block_launch + + worker = Mock() + + def shutdown_worker() -> bool: + shutdown_called.set() + return True + + worker.shutdown.side_effect = shutdown_worker + worker.create_rx_session.return_value = session + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._shutdown = False + transceiver._transfer_worker = worker + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._ever_had_send_session = False + transceiver._ever_had_recv_session = False + transceiver._create_kv_slice = Mock(return_value=object()) + transceiver._slice_num_bytes = Mock(return_value=8) + transceiver._kv_size_rank_factor = 1 + transceiver._finalize_send = Mock() + + if direction == "send": + + def get_send_session(_req): + transceiver._send_sessions[rid] = session + return session + + transceiver._get_or_create_send_session = get_send_session + launch = transceiver.respond_and_send_async + else: + launch = transceiver.request_and_receive_async + + def run_launch() -> None: + try: + launch(req) + except BaseException as error: + errors.append(error) + + def run_shutdown() -> None: + try: + shutdown_attempted.set() + assert transceiver.shutdown() + except BaseException as error: + errors.append(error) + + launch_thread = threading.Thread(target=run_launch) + shutdown_thread = threading.Thread(target=run_shutdown) + launch_thread.start() + assert launch_started.wait(timeout=1) + shutdown_thread.start() + try: + assert shutdown_attempted.wait(timeout=1) + assert not shutdown_called.wait(timeout=0.1) + finally: + release_launch.set() + launch_thread.join(timeout=2) + shutdown_thread.join(timeout=2) + + assert not launch_thread.is_alive() + assert not shutdown_thread.is_alive() + assert errors == [] + assert shutdown_called.is_set() + assert transceiver._shutdown + + +def test_sync_receive_releases_admission_gate_while_waiting_for_shutdown() -> None: + rid = 63 + req = _FakeRequest(request_id=rid) + wait_started = threading.Event() + cancel_seen = threading.Event() + shutdown_finished = threading.Event() + errors: list[BaseException] = [] + shutdown_results: list[bool] = [] + + session = Mock() + session.status = SessionStatus.CANCELLED + session.has_transferring_tasks.return_value = False + + def wait_complete(*, blocking: bool) -> WaitResult: + assert blocking + wait_started.set() + assert cancel_seen.wait(timeout=2) + return WaitResult.FAILED + + def cancel_session() -> None: + cancel_seen.set() + + session.wait_complete.side_effect = wait_complete + session.cancel.side_effect = cancel_session + + worker = Mock() + worker.create_rx_session.return_value = session + worker.shutdown.return_value = True + + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._shutdown = False + transceiver._transfer_worker = worker + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._create_kv_slice = Mock(return_value=object()) + + def run_receive() -> None: + try: + transceiver.request_and_receive_sync(req) + except BaseException as error: + errors.append(error) + + def run_shutdown() -> None: + try: + shutdown_results.append(transceiver.shutdown()) + except BaseException as error: + errors.append(error) + finally: + shutdown_finished.set() + + receive_thread = threading.Thread(target=run_receive) + shutdown_thread = threading.Thread(target=run_shutdown) + receive_thread.start() + assert wait_started.wait(timeout=1) + shutdown_thread.start() + try: + # shutdown() is the only path that signals cancel_seen. If the sync + # receive still held the admission gate while waiting, neither thread + # could make progress here. + assert shutdown_finished.wait(timeout=1) + finally: + cancel_seen.set() + receive_thread.join(timeout=2) + shutdown_thread.join(timeout=2) + + assert not receive_thread.is_alive() + assert not shutdown_thread.is_alive() + assert errors == [] + assert shutdown_results == [True] + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + session.close.assert_called_once_with() + + +@pytest.mark.parametrize("direction", ["send", "receive"]) +def test_cancel_request_waits_for_async_enrollment_and_launch(direction) -> None: + rid = 64 if direction == "send" else 65 + req = _FakeRequest(request_id=rid) + launch_started = threading.Event() + release_launch = threading.Event() + cancel_attempted = threading.Event() + cancel_called = threading.Event() + errors: list[BaseException] = [] + cancel_results: list[bool] = [] + + session = Mock() + session.has_transferring_tasks.return_value = False + session.cancel.side_effect = cancel_called.set + + def block_launch(_slice) -> None: + launch_started.set() + assert release_launch.wait(timeout=2) + + if direction == "send": + session.send.side_effect = block_launch + else: + session.receive.side_effect = block_launch + + worker = Mock() + worker.create_rx_session.return_value = session + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._session_admission_lock = threading.RLock() + transceiver._shutdown_started = False + transceiver._shutdown = False + transceiver._transfer_worker = worker + transceiver._send_sessions = {} + transceiver._send_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._wait_reqs = {} + transceiver._ever_had_send_session = False + transceiver._ever_had_recv_session = False + transceiver._create_kv_slice = Mock(return_value=object()) + transceiver._slice_num_bytes = Mock(return_value=8) + transceiver._kv_size_rank_factor = 1 + transceiver._finalize_send = Mock() + + if direction == "send": + + def get_send_session(_req): + transceiver._send_sessions[rid] = session + return session + + transceiver._get_or_create_send_session = get_send_session + launch = transceiver.respond_and_send_async + else: + launch = transceiver.request_and_receive_async + + def run_launch() -> None: + try: + launch(req) + except BaseException as error: + errors.append(error) + + def run_cancel() -> None: + try: + cancel_attempted.set() + cancel_results.append(transceiver.cancel_request(req)) + except BaseException as error: + errors.append(error) + + launch_thread = threading.Thread(target=run_launch) + cancel_thread = threading.Thread(target=run_cancel) + launch_thread.start() + assert launch_started.wait(timeout=1) + cancel_thread.start() + try: + assert cancel_attempted.wait(timeout=1) + assert not cancel_called.wait(timeout=0.1) + finally: + release_launch.set() + launch_thread.join(timeout=2) + cancel_thread.join(timeout=2) + + assert not launch_thread.is_alive() + assert not cancel_thread.is_alive() + assert errors == [] + assert cancel_results == [True] + assert cancel_called.is_set() + assert transceiver._send_sessions == {} + assert transceiver._send_reqs == {} + assert transceiver._recv_sessions == {} + assert transceiver._recv_reqs == {} + + +def test_transceiver_context_manager_surfaces_non_drained_shutdown() -> None: + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver.shutdown = Mock(return_value=False) + + with pytest.raises(RuntimeError, match="shutdown did not drain"): + transceiver.__exit__(None, None, None) + + +def test_transceiver_context_manager_does_not_mask_body_exception_on_shutdown_failure() -> None: + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver.shutdown = Mock(return_value=False) + + assert transceiver.__exit__(ValueError, ValueError("body failed"), None) is False