From e66e7bdbd415fa35208193a67806897a8dfc00b0 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:04:25 -0700 Subject: [PATCH 1/5] [None][docs] define Python KV transfer ownership contract Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../disagg-kv-transfer-ownership/README.md | 776 ++++++++++++++++++ .../implementation-plan.md | 291 +++++++ 2 files changed, 1067 insertions(+) create mode 100644 docs/design/disagg-kv-transfer-ownership/README.md create mode 100644 docs/design/disagg-kv-transfer-ownership/implementation-plan.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..9d3652835854 --- /dev/null +++ b/docs/design/disagg-kv-transfer-ownership/README.md @@ -0,0 +1,776 @@ +# 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-08 | +| **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 context per KV slice**, + with an exact per-writer ledger. A request-level handle aggregates slices for + user-visible completion. +- `TransferWorker` owns a registry that strongly retains transfer 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 + backend-wide quiescence 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` is not an implementation target. + +## 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 KV slice.** Each receive context contains a + per-writer ledger; a request-level handle aggregates one or more slice + contexts. +3. **The registry is the root owner.** Request or session removal cannot destroy + an active context. +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 generation-safe.** Results are matched by transfer, slice, + generation, and exact writer identity, not merely by 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 generation, writer, 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. + +## 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 +``` + +The feature is opt-in through `kv_cache_bounce_size_mb`; zero disables it. It is +constructed only by the Python native transceiver. The sender may still fall +back to the direct fragmented path when a transfer is not bounceable or no slot +is available. + +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. In this design, that class is treated as +an internal receive-bounce state object and should be renamed accordingly when +the general context is introduced. + +### Related work + +Status snapshot: 2026-07-08. + +| 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) | Open post-merge follow-up intended to exercise the Python bounce configuration. It is validation-adjacent, not ownership work. | +| [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. | +| [PRs #15238](https://github.com/NVIDIA/TensorRT-LLM/pull/15238) and [#15737](https://github.com/NVIDIA/TensorRT-LLM/pull/15737) | Open inputs to the C++ cancellation chain: gated NIXL in-flight cancellation/quiescence containment and sender liveness hardening. They are conceptual precedents, not Python-native dependencies. | +| [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. | +| [Chunked KV transfer design](../chunked-kv-transfer/README.md) and [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 +[`disagg-inflight-cancel-poison`](../disagg-inflight-cancel-poison/README.md) +design 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. + +## Problem Statement + +### Current ownership is fragmented + +The present 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. + +No object has both the information and authority to answer: **which asynchronous +accessors can still touch which allocation generation?** + +### Concrete remaining hazards + +#### Released bounce address can still be advertised + +The current +[`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. + +#### First fan-in failure can outlive its session + +The first failed writer currently +[`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. + +#### Timeout and quarantine do not prove retirement + +`mark_orphaned()` exists in the bounce state object but has no production +callers through the transport interface. Timeout, cancellation, partial +publication, and session destruction do not consistently transition the +physical owner into a drain or containment state. + +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. + +#### Shutdown order can destroy live memory + +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 the +[`BaseTransferAgent.shutdown()` quiescence contract](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. + +#### Count-based fan-in is not identity-safe + +The current 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. + +## 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** | Disabled (`kv_cache_bounce_size_mb == 0`) and enabled (`> 0`). | +| **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/ADP 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`. +- Unifying the Python bounce implementation with the independent C++ bounce + implementation in PR #15780. +- Redesigning rank consensus, scheduler cancellation, or request admission. +- 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 resource 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. | +| **Terminal evidence** | Positive evidence, defined by the backend contract, that a specific operation cannot perform later memory access. | +| **Quiescence** | A per-operation or backend-wide guarantee that no relevant asynchronous access can occur later. | +| **Quarantine** | Removal from reuse while quiescence is unknown. It does not become safe through elapsed time. | +| **Generation** | A nonce distinguishing successive uses of the same request/slice or arena slot identity. | +| **Consumer** | The request/session-facing observer that receives logical completion. It is not the physical owner. | + +## 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 — Register before publication.** All leases and the complete planned +writer-operation ledger, including candidate direct and bounce ranges, MUST be +installed in the registry before publication begins. + +**O2 — Mark possible access first.** Before the system attempts to send any +message containing an address, `publication_started()` MUST atomically set the +writer operation's exposure state to `MAY_ACCESS` and access state to +`POSSIBLE`. A 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 operation accounting.** A receive context MUST track every planned +writer-operation identity, generation, and candidate direct/bounce range. The +actual target mode is recorded after sender selection. Writer count alone is +insufficient. + +**O6 — Registry-first result routing.** Results 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. + +**O7 — Allocator-enforced KV leases.** Lease acquisition MUST atomically validate +and pin `(pool, block, allocation_generation)` before raw pointers are resolved +or descriptors are constructed. `free_resources()` atomically drops request +ownership and marks leased blocks pending-free. A 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 or advance another 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 substitutes for missing per-operation evidence) **and** all local +CUDA work that can touch it is complete. 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 the context lock.** Network operations, CUDA +synchronization, allocator callbacks, and consumer callbacks MUST execute +outside the context state lock. + +**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. + +## Ownership Model + +```mermaid +flowchart TB + TW["TransferWorker"] --> TR["TransferRegistry (strong context owner)"] + TR --> SC["SendTransferContext per slice"] + TR --> RC["RecvTransferContext per slice"] + RH["LlmRequest / session / request handle"] -. "consumer association" .-> SC + RH -. "consumer association" .-> RC + 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 physical context key is conceptually: + +```text +(transfer_id, slice_id, generation, sender_endpoint_epoch, receiver_endpoint_epoch) +``` + +A **slice** is the smallest address set that is published and physically retired +as one unit. One receive context contains one ledger entry per planned writer +operation; a writer that issues multiple NIXL operations for the slice has a +distinct operation ID for each. One request-level `TransferHandle` aggregates +all slice contexts needed to compute the logical request outcome. This supports +chunked and pipelined transfers without making a single request-wide context +serialize independent slices. + +### Component responsibilities + +| Component | Owns | Must not own or decide | +|---|---|---| +| `TransferRegistry` | Strong references to live contexts; identity lookup; shutdown drain accounting. | Request scheduling or KV allocation policy. | +| `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` | Request-facing aggregation and logical cancellation/notification. | Physical lease release. | +| `BounceTransport` | Arena mappings, registrations, allocators, streams, and lease issuance. | End-to-end transfer outcome. | +| KV manager | Underlying KV allocation and enforcement of active transfer leases. | Network completion inference. | +| NIXL agent | Submission/progress and documented quiescence evidence. | KV block reuse. | + +The current `bounce.TransferContext` should become an implementation detail such +as `RecvBounceContext` or be replaced by `RecvBounceLease` state. Expanding it +into the general owner would couple direct transfer, request notification, KV +allocation, and transport arena internals into the bounce package. + +### Resource retirement matrix + +| 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. | +| 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 `TransferHandle`. 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 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. + +## 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 every required writer to succeed and all required +scatter work to complete successfully. + +The first logical terminal outcome wins. 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. + +### 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`. + +**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 +``` + +The context and required leases MUST exist before deriving raw pointers or +constructing plans. The accessor enters `MAY_ACCESS` before CUDA launch, queue +publication, or NIXL submission. 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 safe ordering +is: + +1. Acquire all required KV and bounce leases. +2. Create the exact writer-operation plan and authorized candidate ranges. +3. Register the context. +4. Under the context lock, atomically set the selected writer operation's + exposure state to `MAY_ACCESS` and access state to `POSSIBLE`. +5. Release the lock and attempt to publish its addresses. +6. Record send/submission outcome without weakening `MAY_ACCESS` unless + non-delivery is proven. + +Cancellation before any writer enters `MAY_ACCESS` can retire immediately. +Cancellation after that boundary changes the logical outcome and starts drain; +it does not release memory. + +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 and result routing + +A result must identify: + +- protocol version/capability; +- transfer ID; +- slice ID; +- generation/attempt nonce; +- exact writer identity and writer-operation ID; +- sender and receiver endpoint/instance epochs, because ranks can be reused + after restart; +- actual direct/bounce mode; +- quiescence evidence and data outcome; +- validated scatter metadata when bounce succeeded. + +The registry validates identity before any request/session lookup. Duplicate +results are idempotent. Unexpected writer operations, contradictory duplicates, +and wrong generations 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. Generations MUST NOT be reused within an endpoint epoch. +Retired identities may leave bounded diagnostic tombstones; pruning a tombstone +does not authorize generation reuse or resource reclamation. + +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 backend-wide or driver-level teardown explicitly revokes +the relevant registrations and prevents later DMA. Merely 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. + +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 writer-operation 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 a +context or registry 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 TransferRegistry: + def prepare_send(self, plan: SendPlan) -> SendTransferContext: ... + def prepare_receive(self, plan: RecvPlan) -> RecvTransferContext: ... + def route_result(self, result: WriterOperationResult) -> None: ... + def cancel_consumer(self, transfer_id: TransferId, reason: str) -> None: ... + def begin_shutdown(self) -> None: ... + def drain(self, deadline: float | None) -> DrainResult: ... + + +class RecvTransferContext: + def publication_started( + self, operation: WriterOperationId + ) -> PublicationToken: ... + def mark_never_published( + self, operation: WriterOperationId, proof: NonDeliveryProof + ) -> None: ... + def record_result(self, result: WriterOperationResult) -> None: ... + def record_global_quiescence(self, evidence: QuiescenceEvidence) -> None: ... + def cancel_consumer(self, reason: str) -> None: ... + def detach_consumer(self) -> None: ... + + +class KVTransferLease: + def release(self) -> None: ... +``` + +Required properties: + +- Methods are thread-safe and idempotent. +- Registry insertion rejects a duplicate live identity. +- Resource release is private to the context. +- Request-facing handles cannot force physical retirement. +- `KVTransferLease` prevents allocator reuse even if request cleanup runs. +- `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. + +### 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 + +Shutdown follows this order: + +1. Stop admitting requests and publishing new addresses. +2. Mark remaining consumers failed/cancelled, but keep result listeners, + transfer workers, and CUDA completion workers alive. +3. Drain endpoint contexts until the shutdown deadline. +4. For network operations still in doubt, request backend-wide quiescence + through a + documented `quiesce()`/drain operation. +5. Synchronize remaining gather/scatter work. +6. Release transfer leases and remove retired contexts. +7. Stop listeners and completion workers. +8. Confirm that any non-KV registration owners, including auxiliary transfer + paths, have independently drained; then deregister memory and destroy VMM + mappings. +9. Destroy the backend agent. + +The backend likely needs two phases: quiesce while registrations are still +usable, followed by final destruction after deregistration. If the deadline +expires and quiescence cannot be proven, shutdown returns a non-drained result +and must not deregister, unmap, or reuse the affected memory. The worker remains +alive and retryable, and listeners/completion workers remain available for a +later drain attempt. 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 generation or mode 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 and result fields; or +2. reject the mixed-version session explicitly. + +Silent downgrade is not allowed when it would remove generation-safe identity, +exact writer 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. + +If implementation Phase 1 initially retains the existing wire identity, it must +prove and enforce that transfer IDs and tombstones cannot be reused during the +lifetime of any late result, including worker recreation. If that proof is not +available, generation and protocol negotiation move into Phase 1. One of those +two conditions is a hard Phase 1 exit criterion, not optional follow-up work. + +## 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..2cf199f0af79 --- /dev/null +++ b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md @@ -0,0 +1,291 @@ +# 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-08 | + +## 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. +- Context lookup and result processing are O(1) average; metadata is O(writers) + per slice. +- Locks are per context or sharded; no global lock is held across NIXL or CUDA + work. +- Executor polling remains bounded and non-blocking. +- 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. + +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. + +## Observability + +At minimum, expose counters or structured logs for: + +- active send and receive contexts; +- contexts logically terminal but physically draining; +- in-doubt context count, bytes, age, and reason; +- active source/destination KV leases; +- active send/receive bounce leases; +- direct, bounce, and fallback writers; +- duplicate, unexpected, contradictory, and stale-generation results; +- partial-publication failures; +- shutdown drain duration and non-drained context count; +- admission failures caused by safely retained capacity. + +Logs identify transfer ID, slice, generation, writer identity, path, logical +outcome, physical state, and held leases without logging raw data contents. + +## 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:** Medium-large; approximately 7–10 production files, 600–1,000 +production lines, and 700–1,200 test lines. Generation negotiation or a new V1 +lease binding can increase this estimate. + +- Add `TransferRegistry` and `RecvTransferContext` above the bounce package. +- Split logical outcome, exposure, access, target, and data state. +- Create the exact writer-operation plan and range authorization before + publication. +- Move the publication boundary before message send. +- Route all results through the registry before session lookup. +- Drain late sibling results after logical failure/session removal. +- Replace timed bounce reuse with indefinite in-doubt retention. +- Rename or narrow the current bounce-only `TransferContext`. +- Add a minimum allocator-enforced destination-KV lease before resolving or + publishing destination pointers; Phase 1 cannot claim direct-path safety + without it. +- Treat any failure lacking documented quiescence evidence as in doubt. +- Preselect the exact ADP writer cohort or exclude that topology from the first + rollout. +- Add safe interim shutdown behavior: no unmap if contexts remain in doubt. +- Meet the hard identity gate: add generation/versioning or prove and enforce + non-reuse across late results and worker recreation. + +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` +- 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 8–12 files, 500–1,000 production lines, and +600–1,000 test lines. A small C++/nanobind extension may be required. + +- 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. +- Add `SendTransferContext` and track gather, NIXL, send-bounce, descriptor, and + result-delivery state independently. +- Integrate with or generalize `AsyncTransferManager` rather than creating a + second request-retention mechanism. +- 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 — Two-phase shutdown and backend quiescence + +**Size:** Medium; approximately 300–600 production lines and 400–700 test lines. + +- Add backend quiesce/drain and reorder shutdown before + deregistration/unmapping. +- Keep non-drained workers retryable; prohibit destructor fallback unmapping. +- 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 the required lifecycle metrics. + +### Phase 4 — Protocol and liveness hardening + +**Size:** Medium; approximately 200–500 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 core effort is approximately 1,500–3,000 production lines plus substantial +concurrency and integration tests. The C++ `CacheTransceiver` remains untouched; +cross-language work is limited to shared KV-manager or NIXL binding contracts +that the Python runtime consumes. + +## Validation Plan + +### Deterministic unit and fault-injection tests + +- Cancel immediately before, during, and after address publication. +- 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-generation results. +- Gather launch, descriptor construction, NIXL submission, result-send, and + scatter exceptions. +- 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. +- Exactly-once release under concurrent cancel/result and reordered duplicates. +- Lease acquisition racing request `free_resources()`. +- Overlapping slice leases on the same KV allocation generation. +- Malformed, overflowing, misaligned, or out-of-range NIXL and scatter metadata. +- 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. +- Gather/scatter worker death or asynchronous CUDA failure. +- Sender GPU lease retirement while result delivery remains pending. + +Tests must use barriers/hooks rather than probabilistic sleeps for the critical +publication and teardown races. + +### Configuration matrix + +- bounce disabled and enabled; +- Python KV manager V2 and C++-backed V1; +- single writer and multi-writer TP/ADP fan-in; +- direct, bounce, and per-writer fallback; +- single slice and chunked/multi-slice transfer; +- PyTorch executor and AutoDeploy with Python transceiver; +- bounce-ineligible/non-attention layer-group fallback; +- normal completion, cancellation, timeout, peer failure, and shutdown. + +### Performance validation + +- Direct-path throughput and CPU overhead with ownership enabled. +- Bounce throughput and gather/scatter overlap. +- Context-registry lock contention under high concurrency. +- Source KV early-release timing for chunked transfer. +- Memory high-water mark with normal drain and deliberately in-doubt contexts. +- Shutdown drain latency. + +### 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 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 generation lease. +13. Every local CUDA/NIXL accessor is independently tracked through quiescence. +14. Phase 1 either provides generation-safe identity or enforces non-reuse + across worker recreation and all possible late results. + +## Risks and Mitigations + +| Risk | Mitigation | +|---|---| +| Lock-order deadlock across registry, session, allocator, CUDA, and callbacks | Per-context serialization; no blocking operations or callbacks under lifecycle locks; lock-order tests. | +| 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. | +| Stale result advances a new transfer | Generation and endpoint epoch, tombstones, exact writer validation, and negotiated protocol. | +| 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. | +| 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`. | + +## Alternatives Considered + +### Expand the current bounce `TransferContext` + +Rejected as the general architecture. That object is 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. Can `unique_rid` be reused across worker recreation or retry? If yes or not + provably no, generation/version work is a Phase 1 prerequisite. +4. What supplies a stable endpoint epoch across 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. Is result acknowledgement/retry required for the first production rollout, + or is observable indefinite retention acceptable for all Python-native + direct and bounce transfers? +9. 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. +10. Can every claimed TP/ADP topology identify its exact potential writer cohort + before address fan-out, or which cells must be gated initially? + +No answer to these questions may weaken the invariant by interpreting +uncertainty as retirement. From 4228850d29c1c2201da40f6d77e3fcc3932d745f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:39:19 -0700 Subject: [PATCH 2/5] [None][docs] harden KV transfer ownership contract Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docs/design/README.md | 13 + .../disagg-kv-transfer-ownership/README.md | 966 +++++++++++++++--- .../implementation-plan.md | 807 ++++++++++++++- 3 files changed, 1560 insertions(+), 226 deletions(-) create mode 100644 docs/design/README.md 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 index 9d3652835854..7c8621be4d48 100644 --- a/docs/design/disagg-kv-transfer-ownership/README.md +++ b/docs/design/disagg-kv-transfer-ownership/README.md @@ -1,3 +1,8 @@ + + # Python Native Disaggregated KV Transfer Ownership and Lifecycle | Field | Value | @@ -5,7 +10,7 @@ | **Owner** | Chien-Chun Hung | | **Status** | Draft design contract | | **Created** | 2026-07-08 | -| **Last updated** | 2026-07-08 | +| **Last updated** | 2026-07-13 | | **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) | @@ -29,11 +34,13 @@ quiescent. This document defines the target ownership contract for the Python native transceiver: -- The unit of physical ownership is an **endpoint-local context per KV slice**, - with an exact per-writer ledger. A request-level handle aggregates slices for - user-visible completion. -- `TransferWorker` owns a registry that strongly retains transfer contexts. - `LlmRequest` and session objects are associated consumers, not the root owners. +- 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. @@ -55,11 +62,13 @@ 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 KV slice.** Each receive context contains a - per-writer ledger; a request-level handle aggregates one or more slice - contexts. -3. **The registry is the root owner.** Request or session removal cannot destroy - an active context. +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. @@ -71,17 +80,29 @@ The following decisions are normative for the implementation plan. 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 generation-safe.** Results are matched by transfer, slice, - generation, and exact writer identity, not merely by writer count. +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 generation, writer, or mode fields - require capability/version negotiation or explicit mixed-version rejection - before address publication. +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 @@ -99,10 +120,16 @@ source KV pages -> CUDA scatter into destination KV pages ``` -The feature is opt-in through `kv_cache_bounce_size_mb`; zero disables it. It is -constructed only by the Python native transceiver. The sender may still fall -back to the direct fragmented path when a transfer is not bounceable or no slot -is available. +`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. The merged baseline has a bounce-configured integration test marked +as GB200/GB300-eligible and scheduled on GB200, but it does not assert that the +transport or transfer actually engaged bounce. PR #16116 proposes that positive +signal. Treat those MNNVL cells as candidates, not proven active-bounce coverage, +until an equivalent assertion lands. Other environments may also 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) @@ -114,28 +141,29 @@ the general context is introduced. ### Related work -Status snapshot: 2026-07-08. +Status snapshot: 2026-07-13. | 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) | Open post-merge follow-up intended to exercise the Python bounce configuration. It is validation-adjacent, not ownership work. | +| [PR #16116](https://github.com/NVIDIA/TensorRT-LLM/pull/16116) | Open post-merge follow-up adding positive active-bounce validation and partial bounce-only lifecycle containment: it proposes orphaning/quarantining in-flight receive reservations during cancellation/session close and propagating 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. | -| [PRs #15238](https://github.com/NVIDIA/TensorRT-LLM/pull/15238) and [#15737](https://github.com/NVIDIA/TensorRT-LLM/pull/15737) | Open inputs to the C++ cancellation chain: gated NIXL in-flight cancellation/quiescence containment and sender liveness hardening. They are conceptual precedents, not Python-native dependencies. | +| [PR #15238](https://github.com/NVIDIA/TensorRT-LLM/pull/15238) | Open input to the C++ cancellation chain for gated NIXL in-flight cancellation and quiescence containment. It is a 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. | -| [Chunked KV transfer design](../chunked-kv-transfer/README.md) and [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 #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 -[`disagg-inflight-cancel-poison`](../disagg-inflight-cancel-poison/README.md) -design 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. +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. ## Problem Statement @@ -186,10 +214,12 @@ a Python object-lifetime issue. #### Timeout and quarantine do not prove retirement -`mark_orphaned()` exists in the bounce state object but has no production -callers through the transport interface. Timeout, cancellation, partial -publication, and session destruction do not consistently transition the -physical owner into a drain or containment state. +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 @@ -225,10 +255,11 @@ completion, and a stale result can collide with a later reuse of the same | **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** | Disabled (`kv_cache_bounce_size_mb == 0`) and enabled (`> 0`). | +| **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** | GB200/GB300 MNNVL are candidate cells in the existing test marker, with the merged baseline scheduled on GB200; neither counts as active-bounce coverage without a positive transport-active and per-transfer engagement assertion. 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/ADP fan-in and multiple slices/chunks. | +| **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. | @@ -258,14 +289,23 @@ safety applies. | 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 resource generation, followed by exactly-once lease release. | +| **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** | A per-operation or backend-wide guarantee that no relevant asynchronous access can occur later. | | **Quarantine** | Removal from reuse while quiescence is unknown. It does not become safe through elapsed time. | -| **Generation** | A nonce distinguishing successive uses of the same request/slice or arena slot identity. | +| **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 @@ -290,15 +330,23 @@ The following are explicitly **not** terminal evidence: ### Required invariants -**O1 — Register before publication.** All leases and the complete planned -writer-operation ledger, including candidate direct and bounce ranges, MUST be -installed in the registry before publication begins. - -**O2 — Mark possible access first.** Before the system attempts to send any -message containing an address, `publication_started()` MUST atomically set the -writer operation's exposure state to `MAY_ACCESS` and access state to -`POSSIBLE`. A send error may revert to `NEVER_EXPOSED`/`QUIESCED` only when the -messaging layer positively guarantees non-delivery. +**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 @@ -309,26 +357,74 @@ predicate. It MAY trigger retirement when that predicate is already satisfied. be reported immediately, but physical retirement MUST wait for all possible accessors. -**O5 — Exact operation accounting.** A receive context MUST track every planned -writer-operation identity, generation, and candidate direct/bounce range. The -actual target mode is recorded after sender selection. Writer count alone is -insufficient. - -**O6 — Registry-first result routing.** Results 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. - -**O7 — Allocator-enforced KV leases.** Lease acquisition MUST atomically validate -and pin `(pool, block, allocation_generation)` before raw pointers are resolved -or descriptors are constructed. `free_resources()` atomically drops request -ownership and marks leased blocks pending-free. A 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. +**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 or advance another generation. +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 @@ -341,9 +437,11 @@ CUDA work that can touch it is complete. 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 the context lock.** Network operations, CUDA -synchronization, allocator callbacks, and consumer callbacks MUST execute -outside the context state lock. +**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 @@ -354,17 +452,213 @@ 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. +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. + +**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 (strong context owner)"] - TR --> SC["SendTransferContext per slice"] - TR --> RC["RecvTransferContext per slice"] - RH["LlmRequest / session / request handle"] -. "consumer association" .-> SC - RH -. "consumer association" .-> RC + 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"] @@ -380,30 +674,43 @@ flowchart TB ### Ownership granularity -The physical context key is conceptually: +The endpoint-local physical context key is conceptually: ```text -(transfer_id, slice_id, generation, sender_endpoint_epoch, receiver_endpoint_epoch) +(role, local_endpoint_epoch, admission_authority_epoch, + transfer_id, transfer_generation, local_context_id) ``` -A **slice** is the smallest address set that is published and physically retired -as one unit. One receive context contains one ledger entry per planned writer -operation; a writer that issues multiple NIXL operations for the slice has a -distinct operation ID for each. One request-level `TransferHandle` aggregates -all slice contexts needed to compute the logical request outcome. This supports -chunked and pipelined transfers without making a single request-wide context -serialize independent slices. +`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` | Strong references to live contexts; identity lookup; shutdown drain accounting. | Request scheduling or KV allocation policy. | +| `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` | Request-facing aggregation and logical cancellation/notification. | Physical lease release. | +| `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 and enforcement of active transfer leases. | Network completion inference. | +| 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 current `bounce.TransferContext` should become an implementation detail such @@ -411,8 +718,28 @@ as `RecvBounceContext` or be replaced by `RecvBounceLease` state. Expanding it into the general owner 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. | @@ -425,6 +752,8 @@ allocation, and transport arena internals into the bounce package. | 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 @@ -435,7 +764,9 @@ until all of its responsibilities are settled. ### Relationship to `LlmRequest` The request and transfer lifetimes largely overlap, so they should be associated -through `TransferHandle`. They must not be collapsed into one owner: +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; @@ -446,9 +777,78 @@ through `TransferHandle`. They must not be collapsed into one owner: Request cleanup does not perform a check-then-free against `TransferHandle`. Instead, `free_resources()` atomically drops request ownership in the KV manager; -leased 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. +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. + +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 @@ -466,12 +866,16 @@ 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 every required writer to succeed and all required -scatter work to complete successfully. +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. 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. +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 @@ -496,6 +900,12 @@ MAY_ACCESS -> NEVER_EXPOSED (only with positive non-delivery proof) 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. @@ -515,32 +925,76 @@ accessors with their own state: ```text PLANNED -> MAY_ACCESS -> QUIESCED +PLANNED -> QUIESCED (positive proof that launch/queue/submission did not occur) ``` -The context and required leases MUST exist before deriving raw pointers or -constructing plans. The accessor enters `MAY_ACCESS` before CUDA launch, queue -publication, or NIXL submission. 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. +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 safe ordering -is: - -1. Acquire all required KV and bounce leases. -2. Create the exact writer-operation plan and authorized candidate ranges. -3. Register the context. -4. Under the context lock, atomically set the selected writer operation's - exposure state to `MAY_ACCESS` and access state to `POSSIBLE`. -5. Release the lock and attempt to publish its addresses. -6. Record send/submission outcome without weakening `MAY_ACCESS` unless - non-delivery is proven. - -Cancellation before any writer enters `MAY_ACCESS` can retire immediately. -Cancellation after that boundary changes the logical outcome and starts drain; -it does not release memory. +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 @@ -572,32 +1026,62 @@ 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 and result routing +### Identity, control, and result routing -A result must identify: +Every wire message that can mutate lifecycle state, including cancellation, +must carry a common identity envelope with: - protocol version/capability; - transfer ID; -- slice ID; -- generation/attempt nonce; -- exact writer identity and writer-operation 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. -The registry validates identity before any request/session lookup. Duplicate -results are idempotent. Unexpected writer operations, contradictory duplicates, -and wrong generations 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. +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. Generations MUST NOT be reused within an endpoint epoch. -Retired identities may leave bounded diagnostic tombstones; pruning a tombstone -does not authorize generation reuse or resource reclamation. +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 @@ -629,7 +1113,8 @@ 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 writer-operation cohort through an address-free handshake; +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. @@ -640,8 +1125,8 @@ count is not an acceptable substitute. ### Threading and callbacks -State changes are serialized per context. The implementation must not hold a -context or registry lock while: +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; @@ -658,45 +1143,129 @@ but cannot interrupt physical finalization. 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 cancel_consumer(self, transfer_id: TransferId, reason: str) -> 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, operation: WriterOperationId + self, authorization: WriterAuthorizationId ) -> PublicationToken: ... def mark_never_published( - self, operation: WriterOperationId, proof: NonDeliveryProof + 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. -- Registry insertion rejects a duplicate live identity. -- Resource release is private to the context. +- `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 @@ -720,54 +1289,111 @@ per-operation cancel/recovery API may remain a later liveness improvement. ## Shutdown Contract +`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. + Shutdown follows this order: -1. Stop admitting requests and publishing new addresses. +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. Drain endpoint contexts until the shutdown deadline. -4. For network operations still in doubt, request backend-wide quiescence - through a - documented `quiesce()`/drain operation. -5. Synchronize remaining gather/scatter work. -6. Release transfer leases and remove retired contexts. -7. Stop listeners and completion workers. -8. Confirm that any non-KV registration owners, including auxiliary transfer - paths, have independently drained; then deregister memory and destroy VMM - mappings. -9. Destroy the backend agent. - -The backend likely needs two phases: quiesce while registrations are still -usable, followed by final destruction after deregistration. If the deadline -expires and quiescence cannot be proven, shutdown returns a non-drained result -and must not deregister, unmap, or reuse the affected memory. The worker remains -alive and retryable, and listeners/completion workers remain available for a -later drain attempt. Destructors MUST NOT perform fallback unmapping. The caller -may escalate process health through the adjacent cancellation/poison policy, but +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 or + backend-wide quiescence through a documented `quiesce()`/drain operation + before the deadline. +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 generation or mode fields with -another implicit same-version assumption. +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 and result fields; or + identity, result, and lifecycle-control fields; or 2. reject the mixed-version session explicitly. -Silent downgrade is not allowed when it would remove generation-safe identity, -exact writer 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. - -If implementation Phase 1 initially retains the existing wire identity, it must -prove and enforce that transfer IDs and tombstones cannot be reused during the -lifetime of any late result, including worker recreation. If that proof is not -available, generation and protocol negotiation move into Phase 1. One of those -two conditions is a hard Phase 1 exit criterion, not optional follow-up work. +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 diff --git a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md index 2cf199f0af79..612c133d28e0 100644 --- a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md +++ b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md @@ -1,3 +1,8 @@ + + # Implementation Plan for Python Native Disaggregated KV Transfer Ownership [< Back to ownership contract](README.md) @@ -7,44 +12,137 @@ | **Owner** | Chien-Chun Hung | | **Status** | Draft implementation plan | | **Created** | 2026-07-08 | -| **Last updated** | 2026-07-08 | +| **Last updated** | 2026-07-13 | ## 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. -- Context lookup and result processing are O(1) average; metadata is O(writers) - per slice. +- 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. +- 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. + +| 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 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; -- direct, bounce, and fallback writers; -- duplicate, unexpected, contradictory, and stale-generation results; +- 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. -Logs identify transfer ID, slice, generation, writer identity, path, logical -outcome, physical state, and held leases without logging raw data contents. +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 @@ -53,28 +151,123 @@ Size estimates are directional and include only production code unless noted. ### Phase 1 — Receiver safety foundation -**Size:** Medium-large; approximately 7–10 production files, 600–1,000 -production lines, and 700–1,200 test lines. Generation negotiation or a new V1 -lease binding can increase this estimate. - -- Add `TransferRegistry` and `RecvTransferContext` above the bounce package. +**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. -- Create the exact writer-operation plan and range authorization before - publication. +- 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. -- Route all results through the registry before session lookup. +- 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. - Rename or narrow the current bounce-only `TransferContext`. -- Add a minimum allocator-enforced destination-KV lease before resolving or - publishing destination pointers; Phase 1 cannot claim direct-path safety - without it. +- 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. +- 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. -- Add safe interim shutdown behavior: no unmap if contexts remain in doubt. -- Meet the hard identity gate: add generation/versioning or prove and enforce - non-reuse across late results and worker recreation. +- 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: @@ -82,6 +275,12 @@ Likely files: - `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 @@ -90,8 +289,10 @@ sender-and-receiver contract. ### Phase 2 — Complete KV leases and sender context -**Size:** Large; approximately 8–12 files, 500–1,000 production lines, and -600–1,000 test lines. A small C++/nanobind extension may be required. +**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. @@ -99,29 +300,78 @@ sender-and-receiver contract. - Acquire source leases before gather or direct submission. - Add `SendTransferContext` and track gather, NIXL, send-bounce, descriptor, and result-delivery state independently. +- 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 — Two-phase shutdown and backend quiescence - -**Size:** Medium; approximately 300–600 production lines and 400–700 test lines. - -- Add backend quiesce/drain and reorder shutdown before - deregistration/unmapping. -- Keep non-drained workers retryable; prohibit destructor fallback unmapping. +### 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 the required lifecycle metrics. +- 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 200–500 production lines plus protocol tests. +**Size:** Medium; approximately 300–700 production lines plus protocol tests. - Add result acknowledgement/retry or receiver query for bounded recovery from message loss. @@ -130,64 +380,363 @@ sender-and-receiver contract. - Consider per-operation abort only if the backend can guarantee no later memory access after acknowledgement. -The core effort is approximately 1,500–3,000 production lines plus substantial +The core effort is approximately 4,000–7,300 production lines plus substantial concurrency and integration tests. The C++ `CacheTransceiver` remains untouched; cross-language work is limited to shared KV-manager or NIXL binding contracts that the Python runtime consumes. +## 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-generation results. +- 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. -- Sender GPU lease retirement while result delivery remains pending. +- 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 disabled and enabled; +- bounce not requested, requested but factory-inactive, and arena-active; +- positive bounce-engagement evidence on candidate GB200/GB300 MNNVL fabric-VMM + cells; the merged config-enabled GB200 test is not proof, and + configured-but-inactive runs count as direct fallback; - Python KV manager V2 and C++-backed V1; -- single writer and multi-writer TP/ADP fan-in; +- 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 -- Direct-path throughput and CPU overhead with ownership enabled. -- Bounce throughput and gather/scatter overlap. -- Context-registry lock contention under high concurrency. -- Source KV early-release timing for chunked transfer. -- Memory high-water mark with normal drain and deliberately in-doubt contexts. -- Shutdown drain latency. +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 @@ -195,7 +744,7 @@ 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 generations. +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. @@ -205,26 +754,144 @@ The design is implemented only when all of the following are true: 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 generation lease. +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 either provides generation-safe identity or enforces non-reuse - across worker recreation and all possible late results. +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 registry, session, allocator, CUDA, and callbacks | Per-context serialization; no blocking operations or callbacks under lifecycle locks; lock-order tests. | +| 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. | -| Stale result advances a new transfer | Generation and endpoint epoch, tombstones, exact writer validation, and negotiated protocol. | +| 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 @@ -269,23 +936,51 @@ These do not change the safety contract, but they determine phase sequencing. 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. Can `unique_rid` be reused across worker recreation or retry? If yes or not - provably no, generation/version work is a Phase 1 prerequisite. -4. What supplies a stable endpoint epoch across rank/process restart? +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. Is result acknowledgement/retry required for the first production rollout, +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? -9. Which independent auxiliary-buffer paths can outlive their request/session? +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. -10. Can every claimed TP/ADP topology identify its exact potential writer cohort +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. From 4e22b07beb9d80bd27ccdbb2a7b2c58ad2c9e8ed Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:46:09 -0700 Subject: [PATCH 3/5] [None][fix] harden Python native KV transfer ownership Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/kvCacheManager.h | 2 + .../batch_manager/kvCacheManager.cpp | 34 +- .../batch_manager/kvCacheManagerTest.cpp | 22 +- .../disagg-kv-transfer-ownership/README.md | 275 +- .../implementation-plan.md | 51 +- .../disaggregation/native/bounce/__init__.py | 14 +- .../disaggregation/native/bounce/buffer.py | 86 +- .../disaggregation/native/bounce/core.py | 433 +- .../native/bounce/gather_scatter.py | 23 +- .../disaggregation/native/bounce/impl.py | 980 +++- .../_torch/disaggregation/native/messenger.py | 228 +- .../native/receive_lifecycle.py | 915 ++++ .../_torch/disaggregation/native/transfer.py | 4010 ++++++++++++++--- .../_torch/disaggregation/transceiver.py | 778 +++- .../_torch/pyexecutor/kv_cache_transceiver.py | 16 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 1237 ++++- .../_torch/pyexecutor/py_executor_creator.py | 33 +- .../_torch/pyexecutor/resource_manager.py | 167 +- .../integration/test_lists/test-db/l0_a10.yml | 5 + .../test_lists/test-db/l0_h100.yml | 3 + .../executor/test_async_transfer_manager.py | 347 +- .../executor/test_py_executor_teardown.py | 134 + .../test_py_executor_transfer_termination.py | 1487 ++++++ tests/unittest/disaggregated/test_bounce.py | 1580 ++++++- .../disaggregated/test_kv_transfer.py | 37 +- .../unittest/disaggregated/test_messenger.py | 448 ++ .../disaggregated/test_receive_lifecycle.py | 595 +++ .../test_receive_ownership_wiring.py | 2662 +++++++++++ .../test_transceiver_bounded_polling.py | 1123 ++++- 29 files changed, 16263 insertions(+), 1462 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/native/receive_lifecycle.py create mode 100644 tests/unittest/_torch/executor/test_py_executor_teardown.py create mode 100644 tests/unittest/_torch/executor/test_py_executor_transfer_termination.py create mode 100644 tests/unittest/disaggregated/test_receive_lifecycle.py create mode 100644 tests/unittest/disaggregated/test_receive_ownership_wiring.py diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 01adf276f878..4a3e205f0a6e 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -442,6 +442,8 @@ class KVCacheBlock : public std::enable_shared_from_this void decSchedulingRefCount(); + [[nodiscard]] SizeType32 getRefCount() const; + [[nodiscard]] bool hasRefs() const; [[nodiscard]] bool hasSchedulingRefs() const; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 266cdafc8d06..00dfe5469999 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; @@ -273,6 +274,11 @@ void KVCacheBlock::decSchedulingRefCount() mSchedulingRefCount--; } +SizeType32 KVCacheBlock::getRefCount() const +{ + return mRefCount; +} + bool KVCacheBlock::hasRefs() const { return mRefCount > 0; @@ -2993,6 +2999,15 @@ void WindowBlockManager::unpinBlocksById(std::vector const return; } + std::lock_guard lock(mLookupTree->getMutex()); + + // Validate the complete operation before changing any refcount. Repeated + // IDs represent multiple pin tokens, so aggregate their decrements and + // prove that the block has enough references for the entire batch. + std::unordered_map decrementCounts; + std::vector blocksToUnpin; + decrementCounts.reserve(blockIds.size()); + blocksToUnpin.reserve(blockIds.size()); for (auto const& blockId : blockIds) { TLLM_CHECK_WITH_INFO(blockId >= 0 && static_cast(blockId) < mAllBlocksById.size(), @@ -3000,11 +3015,20 @@ void WindowBlockManager::unpinBlocksById(std::vector const auto block = mAllBlocksById[blockId]; if (block && block->getBlockId() != KVCacheBlock::kCachedBlocksRootId) { - block->decRefCount(); - if (!block->hasRefs()) - { - mEvictionPolicy->releaseBlock(block); - } + auto const decrementCount = ++decrementCounts[blockId]; + TLLM_CHECK_WITH_INFO(block->getRefCount() >= decrementCount, + "Can't unpin block (id=%d) %d time(s) with refcount %d", static_cast(blockId), decrementCount, + block->getRefCount()); + 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..be355269c3de 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -4903,12 +4903,32 @@ 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()); 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_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), pinnedBlockIds.front()}), std::runtime_error); + + EXPECT_NO_THROW(kvCacheManager.unpinBlocksById(pinnedBlockIds)); auto const freeAfterUnpin = kvCacheManager.getNumFreeBlocks(); EXPECT_EQ(freeAfterUnpin, totalBlocks); } diff --git a/docs/design/disagg-kv-transfer-ownership/README.md b/docs/design/disagg-kv-transfer-ownership/README.md index 7c8621be4d48..d6bb646243e1 100644 --- a/docs/design/disagg-kv-transfer-ownership/README.md +++ b/docs/design/disagg-kv-transfer-ownership/README.md @@ -10,7 +10,7 @@ SPDX-License-Identifier: Apache-2.0 | **Owner** | Chien-Chun Hung | | **Status** | Draft design contract | | **Created** | 2026-07-08 | -| **Last updated** | 2026-07-13 | +| **Last updated** | 2026-07-14 | | **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) | @@ -47,14 +47,139 @@ transceiver: - 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 - backend-wide quiescence exists. The initial implementation chooses safety over - bounded reclamation. +- 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` is not an implementation target. +`CacheTransceiver` implementation, data path, and wire protocol are not +implementation targets. Shared `AsyncTransferManager`/`ResourceManager` +accounting may still fail closed while any registered owner or release is live, +but that accounting neither retires a C++ transport owner nor makes local C++ +object destruction evidence of remote/global quiescence. + +## 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. +Shared manager accounting may nevertheless veto teardown for a live owner or +release; the Python-native lifecycle does not retire a C++ transport owner. +Within the Python runtime it is not gated by `kv_cache_bounce_size_mb` or a +second ownership 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 exact rank-to-subrange mapping, evidence-based slot + retirement, explicit slot-settlement acknowledgement, no time-based + quarantine reuse, scatter-aware close, validated non-overlapping result + metadata, and lazy construction of destination manifests 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. + +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 the transactional `unpinBlocksById` operation. +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 also disabled for PP fan-in (`overlap_pp_size > 1`) in this +slice. Equal PP layer counts do not prove equal byte contributions when layer +groups have different block sizes, and a divisibility check after summing bytes +cannot prevent a writer from overflowing its assigned subrange. PP fan-in uses +the direct path until the receiver can authorize byte-accurate per-writer +extents and offsets before publishing any bounce address. + +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. + +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. They do not establish active-bounce engagement on fabric hardware, +attention-DP containment, mixed-version protocol liveness, or an +allocator-enforced KV lease. Those configurations remain outside a production +ownership rollout until their corresponding validation gates pass. ## Design Decisions @@ -135,9 +260,9 @@ 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. In this design, that class is treated as -an internal receive-bounce state object and should be renamed accordingly when -the general context is introduced. +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 @@ -165,11 +290,15 @@ 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. -## Problem Statement +## Pinned-Baseline Problem Statement and Residual Gaps -### Current ownership is fragmented +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. -The present path has several partial owners: +### 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 @@ -181,14 +310,17 @@ The present path has several partial owners: receive bounce contexts. - KV managers remain the authority that can return blocks to their pools. -No object has both the information and authority to answer: **which asynchronous -accessors can still touch which allocation generation?** +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 remaining hazards +### Concrete baseline hazards and containment status #### Released bounce address can still be advertised -The current +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 @@ -198,9 +330,14 @@ 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 -The first failed writer currently +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) @@ -212,6 +349,10 @@ 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 @@ -225,26 +366,40 @@ 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 -Sender workers use an unbounded transfer wait but are joined with a finite -deadline. The enclosing +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 the -[`BaseTransferAgent.shutdown()` quiescence contract](https://github.com/NVIDIA/TensorRT-LLM/blob/3bf37d25387c2b99ef6dba4555ed432bcedfcf4d/tensorrt_llm/_torch/disaggregation/base/agent.py#L85-L121) +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. +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. #### Count-based fan-in is not identity-safe -The current 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. +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 @@ -269,7 +424,8 @@ safety applies. ### Non-goals -- Replacing or porting the design into the separate C++ `CacheTransceiver`. +- 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 rank consensus, scheduler cancellation, or request admission. @@ -295,7 +451,7 @@ safety applies. | **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** | A per-operation or backend-wide guarantee that no relevant asynchronous access can occur later. | +| **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. | @@ -432,10 +588,12 @@ 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 substitutes for missing per-operation evidence) **and** all local -CUDA work that can touch it is complete. 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. +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 @@ -713,10 +871,12 @@ completion before a later chunk enrolls. | 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 current `bounce.TransferContext` should become an implementation detail such -as `RecvBounceContext` or be replaced by `RecvBounceLease` state. Expanding it -into the general owner would couple direct transfer, request notification, KV -allocation, and transport arena internals into the bounce package. +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 @@ -843,7 +1003,10 @@ 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. +not a runtime union inferred from a returned value. Shared teardown accounting +may still fail closed for a live owner or release, but this adapter neither +retires a C++ transport owner nor 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 @@ -1089,12 +1252,14 @@ 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 backend-wide or driver-level teardown explicitly revokes -the relevant registrations and prevents later DMA. Merely 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. +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: @@ -1102,7 +1267,7 @@ Bounded in-process reclamation requires one of: - 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. +- backend-wide quiescence backed by an external remote/global fence. Those are liveness improvements, not permission to weaken the baseline safety contract. @@ -1289,12 +1454,15 @@ per-operation cancel/recovery API may remain a later liveness improvement. ## Shutdown Contract -`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. +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. +For other transceivers, shared teardown accounting may veto destruction while a +generic owner or release remains live, but it cannot derive quiescence or retire +a C++ transport owner from local object destruction. Shutdown follows this order: @@ -1314,9 +1482,10 @@ Shutdown follows this order: 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 or - backend-wide quiescence through a documented `quiesce()`/drain operation - before the deadline. +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. diff --git a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md index 612c133d28e0..2456a2e2ace3 100644 --- a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md +++ b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md @@ -12,7 +12,7 @@ SPDX-License-Identifier: Apache-2.0 | **Owner** | Chien-Chun Hung | | **Status** | Draft implementation plan | | **Created** | 2026-07-08 | -| **Last updated** | 2026-07-13 | +| **Last updated** | 2026-07-14 | ## Performance and Capacity Contract @@ -52,6 +52,21 @@ 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. + | 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. | @@ -228,7 +243,9 @@ rollout-safe until the full Phase 1 gate passes. 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. -- Rename or narrow the current bounce-only `TransferContext`. +- 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, @@ -239,7 +256,9 @@ rollout-safe until the full Phase 1 gate passes. 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. + gate until Phase 2. Shared manager accounting may veto teardown for a live + generic owner or release, but the adapter must not retire a C++ transport + owner or treat local C++ 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, @@ -298,8 +317,10 @@ validation scope would otherwise become too broad. managers. - Support atomic pending-free and overlapping reference-counted slice leases. - Acquire source leases before gather or direct submission. -- Add `SendTransferContext` and track gather, NIXL, send-bounce, descriptor, and - result-delivery state independently. +- 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 @@ -381,9 +402,12 @@ adapters in `tensorrt_llm/_torch/pyexecutor/_util.py`. access after acknowledgement. The core effort is approximately 4,000–7,300 production lines plus substantial -concurrency and integration tests. The C++ `CacheTransceiver` remains untouched; -cross-language work is limited to shared KV-manager or NIXL binding contracts -that the Python runtime consumes. +concurrency and integration tests. 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. Shared `AsyncTransferManager`/`ResourceManager` accounting may observe +live generic ownership and fail closed, but it does not implement or retire the +C++ transport lifecycle. ## Rollout and Rollback @@ -895,12 +919,13 @@ The design is implemented only when all of the following are true: ## Alternatives Considered -### Expand the current bounce `TransferContext` +### Expand the PR #15618 bounce `TransferContext` -Rejected as the general architecture. That object is 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. +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 diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py b/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py index f014842d5254..13f0b69f8773 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, @@ -37,15 +45,17 @@ "Config", "Disposition", "FixedSizing", + "GatherSourceInDoubtError", "NoBounceTransport", "Plan", + "RecvBounceContext", "ScatterState", "Sizing", "SizingContext", "SlotAllocator", - "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..ec997d59bdae 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,157 @@ 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 - # 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 - def writer_base(self, writer_index: int) -> int: - """Where this fan-in writer writes in the region.""" + 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)) + 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, 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 +238,91 @@ 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 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 +333,101 @@ 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 + ) + 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 + 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 +435,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 +477,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,6 +512,32 @@ 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) -> bool: + """Retry durable physical-settlement acknowledgements; return whether all were delivered.""" + @abstractmethod def record_result( self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None @@ -218,9 +545,9 @@ def record_result( """Handle a writer's success; scatter and finalize once all writers reported.""" @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..1a68665d4322 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,297 @@ 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 + # 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 +438,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 +457,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,11 +520,28 @@ 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.""" + 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) + with self._reserved_map_lock: + if not self._accepting_reservations or not self._scatter_stream_healthy: + return self._skip_bounce("transport is closing") nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) if nblocks < self._min_blocks: return self._skip_bounce(f"{nblocks} blocks < min {self._min_blocks} (too small)") @@ -259,98 +579,196 @@ def reserve( 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) + + 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: + if settlement.disposition is Disposition.QUARANTINE: + self._recv_alloc.quarantine(settlement.slot_id) + else: + self._recv_alloc.release(settlement.slot_id) + except Exception: + with self._reserved_map_lock: + pending.physical_in_progress = False + raise + 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: - settlement.on_done(settlement.success) - except Exception as e: # never let the callback strand the arena + 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 + + def _retry_pending_settlements(self) -> bool: + with self._reserved_map_lock: + keys = tuple(self._pending_settlements) + for key in keys: + self._commit_pending_settlement(key) + with self._reserved_map_lock: + return not self._pending_settlements + + def retry_settlements(self) -> bool: + """Public non-blocking retry hook for the receiver's bounded poll path.""" + return self._retry_pending_settlements() def record_result( self, @@ -364,62 +782,228 @@ def record_result( """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) 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 _scatter_loop(self): + 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 + 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 + 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 +1019,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 +1045,52 @@ 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) -> bool: + return True + def record_result( self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=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 +1101,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 +1157,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 +1173,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..c2cfe28e9780 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" + ) + + # 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 - # Snapshot under lock as defense in depth. + 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,6 +2355,7 @@ def __init__( timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, beam_width: int = 1, + source_owner: Optional[LlmRequest] = None, ): super().__init__( sender, @@ -1180,18 +2366,40 @@ def __init__( self._sender: Sender # narrow base class type for Pylance self.request_id = request_id self._aux_buffer = aux_buffer + # Keep the allocator-level request owner alive until every queued, + # active, or in-doubt source operation has terminal evidence. + self._source_owner = source_owner 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) + try: + self._sender.setup_session(self) + except Exception: + 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._source_owner = None + self._closed = True + raise @property def disagg_request_id(self) -> int: @@ -1205,6 +2413,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 +2471,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 +2563,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): @@ -1275,25 +2582,52 @@ def cancel(self) -> None: 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(): + with self.lock: + self._accepting_operations = False + if self._terminal_status != SessionStatus.CANCELLED: + 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) + # 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 +2636,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 +2660,89 @@ 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}" - with self.lock: - self._exception = RuntimeError(msg) - self._terminal_status = SessionStatus.ERROR - for task in self.kv_tasks: - if not task.is_done: - task.fail(self._exception) - if self.aux_task is not None and not self.aux_task.is_done: - self.aux_task.fail(self._exception) + 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 self.lock: + self._exception = RuntimeError(msg) + self._terminal_status = SessionStatus.ERROR + self._accepting_operations = False + for task in self.kv_tasks: + if not task.is_done: + task.fail(self._exception) + if self.aux_task is not None and not self.aux_task.is_done: + self.aux_task.fail(self._exception) + # 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): - if getattr(self, "_closed", False): - return - self._closed = True + 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): + 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 +2771,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 +2814,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 +2911,29 @@ 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._shutdown_started = False + self._listener_stopped = False self._shutdown = False self._start_listener() @@ -1455,43 +2943,273 @@ 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 _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: self_ri = self._registrar.self_rank_info assert task._params.ctx_request_id is not None, ( @@ -1513,25 +3231,72 @@ 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: + """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 safe only when + the complete per-stage layer counts are uniform. ``reserve`` separately + rejects heterogeneous layer-group block sizes. + """ + if overlap.duplicate_head_factor != 1 or overlap.peer_duplicate_head_factor != 1: 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: + layer_num_per_pp = getattr(peer_ri, "layer_num_per_pp", None) + if ( + not layer_num_per_pp + or len(layer_num_per_pp) < overlap.overlap_pp_size + or len(set(layer_num_per_pp)) != 1 + ): return False return True + 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 logger.debug( @@ -1544,61 +3309,247 @@ 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) + task.set_valid_writer_cohorts(valid_writer_cohorts) # 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) + # _fanin_bounce_safe() (TP-by-head with no PP fan-in), and never under ADP broadcast (sender_dp_rank + # None), where the exact writer cohort is not selected before publication. + 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 (task.expected_transfers == 1 or self._fanin_bounce_safe(peer_overlap, 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 +3566,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,39 +3641,189 @@ 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() - def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): - if message[0] != MessageType.KV_AGENT_RESULT: - logger.error( - f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" + def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): + if message[0] != MessageType.KV_AGENT_RESULT: + logger.error( + 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]) + ) + 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, + ) + 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 - peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code = ( - _KV_RESULT_PREFIX.unpack(message[1]) - ) - from .bounce import decode_result_tail - 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 +3839,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 +3866,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,6 +3891,7 @@ def __init__( timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, beam_width: int = 1, + destination_owner: object | None = None, ): super().__init__( receiver, @@ -1772,17 +3901,47 @@ def __init__( 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_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._accepting_receives = True + self._active_receive_dispatches = 0 + self._last_slice_admitted = False + self._receiver_retired = False + # The current containment slice has no allocator-enforced destination + # lease. Retain the request that owns its KV allocation until receiver + # retirement proves every accessor and scatter is drained. + self._destination_owner = destination_owner self._terminal_status: Optional[SessionStatus] = None self._kv_tasks: list[KVRecvTask] = [] - self._aux_count = 0 + 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 + self._aux_drained = not self._need_aux self._aux_status: TaskStatus = TaskStatus.INIT self._sender_endpoints: set[str] = set() - self.lock = threading.Lock() - self._receiver.setup_session(self) + self._selected_writer_cohort: Optional[frozenset[int]] = None + try: + self._receiver.setup_session(self) + except Exception: + 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._destination_owner = None + self._closed = True + raise @property def disagg_request_id(self) -> int: @@ -1816,18 +3975,299 @@ 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}" ) - self._kv_tasks.append(task) - self._receiver.dispatch_task(task) + 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}" + ) + 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 _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._aux_exposed_writer_ranks.add(int(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: + self._aux_drained = True + 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 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. + with self._receive_lock: + with self.lock: + 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" + ) + 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) + 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._aux_exposed_writer_ranks.add(int(peer_rank)) + self._close_aux_publication_locked() + logger.error(detail) def process_kv_agent_result( self, @@ -1839,96 +4279,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 +4333,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._aux_exposed_writer_ranks.add(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._aux_exposed_writer_ranks.add(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 +4417,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(): + 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 +4648,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 +4702,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 +4714,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 +4757,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 +4811,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 +4835,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 +4950,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 +5062,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..c694afac21d1 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -1,3 +1,19 @@ +# 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 @@ -39,6 +55,11 @@ 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() + def _find_consensus_request_ids(request_ids_all_ranks, sync_size): frequency_map = defaultdict(int) @@ -54,6 +75,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 +86,9 @@ def __init__( kv_cache_manager: KVCacheManager, cache_transceiver_config: CacheTransceiverConfig, ): + self._session_admission_lock = threading.RLock() + self._shutdown_started = False + self._shutdown = False self._dist: Distributed = dist self._kv_cache_manager = kv_cache_manager self._mapping = mapping @@ -147,25 +175,105 @@ 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 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: set[int] = set() + for direction, sessions in ( + ("send", self._send_sessions), + ("receive", self._recv_sessions), + ): + for session in list(sessions.values()): + try: + # Cancellation closes future publication but leaves + # already-exposed resources owned until terminal + # results arrive. + session.cancel() + except Exception as e: + cancel_failed_session_ids.add(id(session)) + logger.error( + f"KvCacheTransceiverV2 shutdown failed to cancel {direction} " + f"session {session.disagg_request_id}: {e}" + ) + 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 + + 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 + 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: + return False + getattr(self, "_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 +452,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 +476,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): @@ -429,12 +587,29 @@ 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 _close_failed_sessions( + self, + sessions: dict, + reqs: dict, + failed: list, + expected_sessions: dict, + expected_reqs: dict, + ) -> list: + """Retire only the exact failed owners observed by one status poll.""" + retired = [] + with self._get_session_admission_lock(): + for rid in failed: + session = expected_sessions.get(rid) + req = expected_reqs.get(rid) + if sessions.get(rid) is not session or reqs.get(rid) is not req: + continue + session.close() + if sessions.get(rid) is session and reqs.get(rid) is req: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + reqs.pop(rid, None) + sessions.pop(rid, None) + retired.append(rid) + return retired def _apply_aux(self, session, req: LlmRequest): """Unpack aux tokens from session into request's context_phase_params.""" @@ -481,97 +656,232 @@ 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 + 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._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._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, ): - # 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. + with self._get_session_admission_lock(): + if getattr(self, "_shutdown_started", False): + return [], [] + self._transfer_worker.sweep_stale_req_infos() + if not self._ever_had_send_session: + return [], [] + send_sessions = dict(self._send_sessions) + 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) + local_completed, local_failed = self._collect_done(send_sessions, send_reqs) to_process = self._build_to_process( - self._send_sessions, + send_sessions, self._ctx_consensus(local_completed + local_failed), wait_num, block_all, ) - completed, timed_out, failed, cancelled = [], [], [], [] + completed, timed_out, failed, cancelled, cleanup_ready = [], [], [], [], [] for rid in to_process: - session = self._send_sessions[rid] + 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) + cleanup_ready.append(rid) elif result == WaitResult.COMPLETED: completed.append(rid) - elif result is None: - continue + cleanup_ready.append(rid) elif result == WaitResult.TIMEOUT: logger.warning( f"TxSession rid={session.disagg_request_id} timed out after {self._sender_future_timeout_ms}ms" @@ -580,52 +890,81 @@ def check_context_transfer_status( else: logger.warning(f"TxSession rid={session.disagg_request_id} failed") failed.append(rid) + cleanup_ready.append(rid) # 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() + retired_completed = [] + with self._get_session_admission_lock(): + for rid in completed: + session = send_sessions.get(rid) + req = send_reqs.get(rid) + if ( + self._send_sessions.get(rid) is not session + or self._send_reqs.get(rid) is not req + ): + continue + session.close() + if mark_complete: + req.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + if self._send_sessions.get(rid) is session and self._send_reqs.get(rid) is req: + self._send_reqs.pop(rid, None) + self._send_sessions.pop(rid, None) + retired_completed.append(rid) + # A remotely cancelled context transfer is terminal for its local + # AsyncTransferManager entry. Report it as an error so PyExecutor can + # unpin the request instead of silently losing the transceiver session. + terminal_failed = failed + cancelled + retired_failed = self._close_failed_sessions( + self._send_sessions, + self._send_reqs, + terminal_failed, + send_sessions, + send_reqs, + ) - return completed, failed + return retired_completed, retired_failed 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._get_session_admission_lock(): + if getattr(self, "_shutdown_started", False): + return [], [], [] + if not self._ever_had_recv_session and not self._gen_need_sync: + return [], [], [] + recv_sessions = dict(self._recv_sessions) + 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 need_progress = wait_num > 0 if need_progress: - self._poll_gen_sessions_for_poll_interval(wait_num) + self._poll_gen_sessions_for_poll_interval(wait_num, recv_sessions, recv_reqs) - local_completed, local_failed = self._collect_done(self._recv_sessions, self._recv_reqs) + local_completed, local_failed = self._collect_done(recv_sessions, recv_reqs) to_process = self._build_to_process( - self._recv_sessions, + recv_sessions, self._gen_consensus(local_completed + local_failed), 0 if need_progress else wait_num, block_all, ) - completed, failed, cancelled = [], [], [] + completed, failed, cancelled, cleanup_ready = [], [], [], [] for rid in to_process: - session = self._recv_sessions[rid] + 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 + cleanup_ready.append(rid) if session.status == SessionStatus.CANCELLED: # Session cancelled — either by local cancel_request() (user # cancel) or by a remote CANCEL_SESSION message (e.g. CTX @@ -636,57 +975,82 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): completed.append(rid) elif result == WaitResult.FAILED: failed.append(rid) - # else: None — KV done but aux still in flight; re-poll next cycle # 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 ) + retired_completed = [] 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: + with self._get_session_admission_lock(): + for rid in cancelled: + session = recv_sessions.get(rid) + req = recv_reqs.get(rid) + if ( + self._recv_sessions.get(rid) is not session + or self._recv_reqs.get(rid) is not req + ): + continue + session.close() + if self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req: + self._recv_reqs.pop(rid, None) + self._recv_sessions.pop(rid, None) + cancelled_reqs.append(req) + + for rid in completed: + session = recv_sessions.get(rid) + req = recv_reqs.get(rid) + if ( + self._recv_sessions.get(rid) is not session + or self._recv_reqs.get(rid) is not req + ): + continue + # 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) + session.close() + req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + if self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req: + self._recv_reqs.pop(rid, None) + self._recv_sessions.pop(rid, None) + retired_completed.append(rid) + retired_failed = self._close_failed_sessions( + self._recv_sessions, + self._recv_reqs, + failed, + recv_sessions, + recv_reqs, + ) + 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 len(self._recv_sessions) == 0 def _assert_disagg_history_declared(self, req: LlmRequest) -> None: """Verify the V2 scheduler pre-declared prompt_len as history. @@ -729,33 +1093,39 @@ 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(): + # 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: + self._wait_reqs.pop(rid, None) + + cleanup_pending = False + 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 + try: + session.cancel() + 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(). @@ -779,29 +1149,63 @@ def get_disaggregated_params(self) -> Dict[str, Any]: def prepare_context_requests(self, requests: List[LlmRequest]): # 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(): + if getattr(self, "_shutdown_started", False): + raise RuntimeError("KV cache transceiver is shutting down") + pending_admissions = {} + planned_wait_owners = dict(self._wait_reqs) + for req in requests: + rid = get_unique_rid(req) + assert rid is not None + 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._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: + 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..b6d5915b22f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -379,6 +379,37 @@ class BatchStatePP(BatchState): microbatch_id: int = -1 +_PYTHON_NATIVE_TRANSCEIVER_OWNER = "python_native_transceiver" + + +@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 + terminate_on_completion: 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 + response_prepared: bool = False + owner_released: bool = False + all_transfers_complete: Optional[bool] = None + active_request_removed: bool = False + request_terminated: bool = False + + class AsyncTransferManager: """ Handle asynchronous transfer of KV cache after a request has completed. @@ -393,21 +424,77 @@ 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 + def __init__(self, resource_manager: "ResourceManager", should_store_blocks: bool = True): @@ -424,10 +511,90 @@ 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._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 +604,53 @@ 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 + 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") - if self.should_store_blocks: - block_id = self.kv_cache_manager.store_blocks_for_reuse( - request, True) - else: - block_id = None + if existing_request is not None: + self._request_transfer_metadata[req_id].start_transfer(owner) + return - self._requests_in_transfer[req_id] = request - self._request_transfer_metadata[ - req_id] = self.RequestTransferMetadata(block_id) + 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._request_transfer_metadata[req_id].start_transfer() + 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,22 +659,38 @@ 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}") + existing_request = self._requests_in_transfer.get(req_id) + if existing_request is None: + logger.warning(f"Request {req_id} not found in transfer manager") return False + if existing_request is not request: + raise RuntimeError( + f"request {req_id} has a different active transfer owner") - if transfer_metadata.end_transfer(): - self._requests_in_transfer.pop(request.py_request_id) - self._request_transfer_metadata.pop(request.py_request_id) + transfer_metadata = self._request_transfer_metadata[req_id] + final_transfer = transfer_metadata.will_end_all_transfers(owner) - if self.should_store_blocks: - self.kv_cache_manager.unpin_blocks_by_id( - transfer_metadata.block_id) + # Keep the exact owner and request root enrolled until final unpin + # succeeds. A failed unpin is retryable without reconstructing which + # provider contribution or request allocation it belonged to. + if final_transfer and self.should_store_blocks: + self.kv_cache_manager.unpin_blocks_by_id( + transfer_metadata.pinned_block_ids) + + if transfer_metadata.end_transfer(owner): + self._requests_in_transfer.pop(req_id) + self._request_transfer_metadata.pop(req_id) # We don't want to overwrite any error state. if request.state != LlmRequestState.DISAGG_TRANS_ERROR: @@ -494,8 +700,30 @@ def end_transfer(self, request: LlmRequest) -> bool: return False + def has_transfer_owner(self, request: LlmRequest, owner: str) -> bool: + """Whether the exact request still has the named provider owner.""" + if self._requests_in_transfer.get(request.py_request_id) is not request: + return False + transfer_metadata = self._request_transfer_metadata.get( + request.py_request_id) + return (transfer_metadata is not None + and transfer_metadata.has_owner(owner)) + + def has_any_transfer_owner(self, request: LlmRequest) -> bool: + """Whether the exact request still has any enrolled provider owner.""" + return self._requests_in_transfer.get(request.py_request_id) is request + + def requests_with_owner(self, owner: str) -> Dict[int, LlmRequest]: + """Snapshot requests whose transfer count includes ``owner``.""" + return { + req_id: request + for req_id, request in self._requests_in_transfer.items() + if self._request_transfer_metadata[req_id].has_owner(owner) + } + 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) class PyExecutor: @@ -708,6 +936,35 @@ def __init__( # 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] = {} 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 +1332,94 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() - def _end_transfer_and_maybe_terminate(self, request: LlmRequest): + 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.""" + response = request.create_response(False, self.dist.rank) + if response: + response.result.cached_tokens = request.cached_tokens + self._maybe_attach_ctx_usage(request, response) + # With ADP, _enqueue_responses does a tp_gather collective. Calling + # it from a completion callback would deadlock because only the + # owning DP rank reaches that callback. + self._pending_transfer_responses.append( + (request.py_request_id, response)) + + 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 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) + if hasattr(self, "_terminated_transfer_requests"): + PyExecutor._forget_terminated_transfer_request_if_unowned( + 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._forget_terminated_transfer_request_if_unowned( + 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) + if hasattr(self, "_terminated_transfer_requests"): + PyExecutor._forget_terminated_transfer_request_if_unowned( + self, request) def _flush_pending_transfer_responses(self): """Enqueue buffered transfer-completion responses. @@ -1466,6 +1774,107 @@ def shutdown(self): # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): torch.cuda.synchronize() + transfer_manager = getattr(self, "async_transfer_manager", None) + if transfer_manager is not None: + transfer_manager.begin_shutdown() + is_python_native_transceiver = ( + self.kv_cache_transceiver is not None + and self._requires_physical_transfer_drain()) + 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: + defer_transceiver_owners() + transceiver_shutdown_result = self.kv_cache_transceiver.shutdown() + if transceiver_shutdown_result is False: + raise RuntimeError( + "KV cache transceiver still owns active transfer targets; " + "refusing to shut down resource managers") + if transceiver_shutdown_result is not True: + raise RuntimeError( + "Python-native KV cache transceiver did not prove physical " + "drain; 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: + drain_deferred_terminations( + terminate_after_worker_shutdown=True) + drain_pending_cancellations = getattr( + self, "_drain_pending_request_cancellations_after_shutdown", None) + if drain_pending_cancellations is not None: + drain_pending_cancellations() + if getattr(self, "_pending_request_cancellations", None): + raise RuntimeError( + "User cancellation ownership is still incomplete after " + "transceiver shutdown; refusing to shut down resource managers") + if getattr(self, "_deferred_transfer_terminations", None): + raise 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._kv_connector_terminate_requests( + terminate_after_worker_shutdown=True) + else: + drain_connector_completions = getattr( + self, "_drain_pending_connector_completions", None) + if drain_connector_completions is not None: + drain_connector_completions( + terminate_after_worker_shutdown=True) + if getattr(self, "_pending_connector_completions", None): + raise RuntimeError( + "Connector completion ownership is still incomplete after " + "executor shutdown; refusing to shut down resource managers") + 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") + 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") + termination_handler = getattr(self, "_disagg_pp_termination_handler", + None) + if termination_handler is not None: + 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") + + 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: + self._do_terminate_request(request) + if termination_progress: + raise RuntimeError( + "Request resource termination is still incomplete; refusing " + "to shut down resource managers") + + has_in_doubt_releases = getattr(self.resource_manager, + "has_in_doubt_resource_releases", None) + 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") + has_pending_releases = getattr(self.resource_manager, + "has_pending_resource_releases", None) + 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") for manager in self.resource_manager.resource_managers.values(): if manager: manager.shutdown() @@ -2837,7 +3246,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 +3792,108 @@ 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 _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", + {}) + 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 + 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 +4228,91 @@ 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)) + if not progress.response_prepared: + if progress.was_active: + PyExecutor._prepare_transfer_completion_response(self, request) + progress.response_prepared = True + + 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 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) + progress.request_terminated = True + self._forget_terminated_transfer_request_if_unowned(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 _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) + pending = self._get_pending_connector_completions() + for request in self.kv_connector_manager.get_finished(): + 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: + pending[request_key] = _ConnectorCompletionProgress( + request=request, + was_active=bool(self.kv_cache_transceiver + and request in self.active_requests), + ) + + self._drain_pending_connector_completions( + terminate_after_worker_shutdown=terminate_after_worker_shutdown) def _kv_connector_wait_for_save(self): if self.kv_connector_manager is not None: @@ -3872,6 +4462,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 +6637,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 +6686,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 @@ -6104,6 +6747,8 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): finished_requests, error_requests = self.kv_cache_transceiver.check_context_transfer_status( atLeastNum) + transfer_owner = (_PYTHON_NATIVE_TRANSCEIVER_OWNER + if self._requires_physical_transfer_drain() else None) completed_req_ids = set(finished_requests + error_requests) @@ -6119,7 +6764,9 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): request = requests_in_transfer[request_id] - self._end_transfer_and_maybe_terminate(request) + 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 +6779,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 @@ -6510,6 +7168,75 @@ def _handle_errors(self, failed_requests = (list(self.active_requests) if requests is None else requests) + pending_ownership_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) + if self._has_pending_ownership_transition(request): + pending_cancellations = getattr( + self, "_pending_request_cancellations", {}) + cancellation_progress = pending_cancellations.get(request_key) + if (cancellation_progress is not None + and cancellation_progress.request is request): + # _handle_errors removes the request from active_requests, + # so the cancellation ledger must either perform final + # teardown or hand it to a remaining provider owner. + cancellation_progress.terminate_on_completion = True + 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: req_id = request.py_request_id request.state = LlmRequestState.GENERATION_COMPLETE @@ -6517,6 +7244,10 @@ def _handle_errors(self, request_id=req_id, error_msg=error_msg, client_id=request.py_client_id) + 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 if requests is None: self.active_requests.clear() else: @@ -6526,7 +7257,15 @@ def _handle_errors(self, ] self._enqueue_responses(list(error_responses.items())) for request in failed_requests: - self._terminate_request(request) + request_key = id(request) + 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 +7282,37 @@ 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) + 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 +7330,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 +7353,216 @@ 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 + 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) + self._forget_terminated_transfer_request_if_unowned(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_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 and not request_was_already_terminated: + if terminate_after_worker_shutdown: + self._terminate_request_after_worker_shutdown(request) + else: + self._terminate_request(request) + deferred_terminations.pop(request_key) + already_terminated.discard(request_key) + self._forget_terminated_transfer_request_if_unowned(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 +7571,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 - - 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) - - # Clear list of requests marked for cancellation and add back those that failed to cancel. + 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: + pending[request_key] = _RequestCancellationProgress( + request=request, cancel_id=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 progress.terminate_on_completion: + try: + self._complete_request_cancellation_termination(progress) + except Exception as e: + logger.error( + f"Failed to terminate cancelled request " + f"{progress.cancel_id}; retaining cancellation for " + f"retry: {e}") + continue + if (completed and (not progress.terminate_on_completion + or progress.termination_handoff_complete) + and pending.get(request_key) is progress): + pending.pop(request_key) + + # 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) @@ -6753,7 +7760,14 @@ 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 + + # 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 +7787,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 @@ -6848,6 +7862,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: @@ -7032,6 +8052,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..a2f0e15c8423 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -333,6 +333,31 @@ 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() + 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 +1051,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..b470525cb523 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,41 @@ 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 + failure: Optional[str] = None + + +@dataclass +class _TransferResourceReleaseProgress: + request: LlmRequest + completed_managers: Set[ResourceManagerType] + in_doubt_manager: Optional[ResourceManagerType] = None + 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 +2635,132 @@ 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) + 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): + 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: + 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}") + + while progress.next_manager < len(progress.release_plan): + # Mark the individual manager attempt ambiguous before calling + # it. An exception may be raised after that manager has already + # released some state, so blindly invoking it again could + # double-free. Only a normal return proves this step complete. + progress.in_doubt_manager = progress.next_manager + try: + progress.release_plan[progress.next_manager].free_resources( + request) + except Exception as error: + progress.failure = f"{type(error).__name__}: {error}" + raise + progress.next_manager += 1 + progress.in_doubt_manager = None + progress.failure = None + + 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) + + 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: + 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 + try: + manager.free_resources(request) + except Exception as error: + progress.failure = f"{type(error).__name__}: {error}" + raise + progress.completed_managers.add(resource_mgr_type) + progress.in_doubt_manager = None + 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 + for progress in self._resource_release_progress.values()) + or any(progress.in_doubt_manager is not None 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..f9b40d04ae6a 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"); @@ -15,8 +15,10 @@ 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 @@ -34,27 +36,24 @@ def create_mock_resource_manager( 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 +68,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 +81,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 +108,142 @@ 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_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_preserves_owner_for_retry(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + kv_cache_manager.unpin_blocks_by_id.side_effect = [RuntimeError("unpin failed"), None] + 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"): + 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 manager.end_transfer(request, owner="python_native") + assert manager.requests_in_transfer() == {} + assert manager._request_transfer_metadata == {} + assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + assert kv_cache_manager.unpin_blocks_by_id.call_count == 2 + + +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_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 +256,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 +271,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 +290,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 +308,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_py_executor_teardown.py b/tests/unittest/_torch/executor/test_py_executor_teardown.py new file mode 100644 index 000000000000..10d0a6611453 --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_executor_teardown.py @@ -0,0 +1,134 @@ +# 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 + + 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.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 + + +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("terminal_shutdown_result", [True, None]) +def test_estimation_teardown_retries_only_after_transceiver_drains( + terminal_shutdown_result, +): + """Manager teardown must wait for either current or legacy drain success.""" + transceiver = Mock() + transceiver.shutdown.side_effect = [False, terminal_shutdown_result] + py_executor = SimpleNamespace(kv_cache_transceiver=transceiver) + kv_cache_creator = Mock() + resources = {} + + with pytest.raises(RuntimeError, match="still owns active transfer targets"): + 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 + + +@pytest.mark.parametrize("shutdown_result", [True, None]) +def test_estimation_teardown_rejects_inflight_transfer_owner(shutdown_result): + transceiver = Mock() + 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..63f097ea2464 --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py @@ -0,0 +1,1487 @@ +# 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, + DisaggPPTerminationHandler, + PyExecutor, +) + + +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.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 _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 _make_error_executor(request, *, fatal): + executor = 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 = 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 + + +@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) + executor = 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._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() + + expected_owner = _PYTHON_NATIVE_TRANSCEIVER_OWNER if requires_physical_drain else None + executor._end_transfer_and_maybe_terminate.assert_called_once_with( + request, transfer_owner=expected_owner + ) + + +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 = 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_then_error_terminates_once_after_retry(): + 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 progress.terminate_on_completion + assert executor.active_requests == [] + executor._terminate_request.assert_not_called() + + executor._handle_canceled_requests() + + executor._terminate_request.assert_called_once_with(request) + assert executor._pending_request_cancellations == {} + assert executor.async_transfer_manager.end_transfer.call_args_list == [ + call(request, owner=_PYTHON_NATIVE_TRANSCEIVER_OWNER) + ] + + +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 = 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.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 = Mock() + executor.kv_connector_manager.get_finished.return_value = [request] + executor._pending_connector_completions = {} + executor._pending_transfer_responses = [] + executor._terminated_transfer_requests = {} + executor.force_terminate_ctx_for_partial_reuse = False + + +@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 == [request] + 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_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 == [request] + assert executor.active_requests == [] + executor.async_transfer_manager.end_transfer.assert_not_called() + executor._terminate_request.assert_not_called() + + executor._kv_connector_terminate_requests() + + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor._terminate_request.assert_not_called() + + 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] + 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 = object.__new__(PyExecutor) + executor.resource_manager = Mock() + executor._prefetched_request_ids = {7} + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=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 = 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) + 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 = 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) + 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 = 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 = 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) + 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._terminate_request = events.terminate_request + executor.virtual_memory_pools = None + executor.sampler = object() + executor.dwdp_manager = None + return executor + + +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._deferred_transfer_terminations = {} + executor.async_transfer_manager = _ExactTransferManager(request) + + executor.shutdown() + + events.transceiver_shutdown.assert_not_called() + events.cancel_request.assert_not_called() + events.manager_shutdown.assert_called_once_with() + + +def test_capability_false_shutdown_does_not_claim_anonymous_owner_drain(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) + + with pytest.raises(RuntimeError, match="Asynchronous transfer ownership"): + executor.shutdown() + + events.transceiver_shutdown.assert_not_called() + executor.async_transfer_manager.end_transfer.assert_not_called() + events.terminate_request.assert_not_called() + events.manager_shutdown.assert_not_called() + + +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, + ) + 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) + 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 = Mock() + executor.kv_connector_manager.get_finished.return_value = [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_ingests_finished_connector_owner_before_veto(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.is_available", + lambda: False, + ) + 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 = Mock() + executor.kv_connector_manager.get_finished.return_value = [request] + executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) + + executor.shutdown() + + executor.kv_connector_manager.get_finished.assert_called_once_with() + executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor.async_transfer_manager.has_any_inflight_requests.assert_called_once_with() + events.terminate_request.assert_called_once_with(request) + events.manager_shutdown.assert_called_once_with() + + +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..90903997dc03 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 @@ -186,31 +195,31 @@ 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 for uniform TP-by-head and even PP # --------------------------------------------------------------------------- # 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 uniform PP layers 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) - # single PP stage (overlap_pp_size <= 1): only duplicate_head_factor matters + # A single PP stage is safe only when neither peer duplicates KV heads. 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, 1, peer_dup=2), ri([24])) is False # reciprocal duplication + # Equal PP layers are safe; reserve separately checks per-block byte sizes. assert safe(ov(1, 4), ri([20, 20, 20, 20])) is True - # UNEVEN PP fan-in -> per-writer sizes differ -> fall back + # Uneven PP fan-in also falls 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 @@ -225,15 +234,30 @@ 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() 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 +273,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 +313,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 +325,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, @@ -315,21 +393,673 @@ def _recv_req(block_counts, rid=1, slice_id=0): ) +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", + ) + + +@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") +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 TestFanInReserve: def test_reserve_stamps_base_and_per_writer(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, 3]) 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 + # Exact rank order determines layout; per_writer = 100. + assert t.writer_base((req.unique_rid, req.slice_id), 7) == 0x100000 + assert t.writer_base((req.unique_rid, req.slice_id), 3) == 0x100000 + 100 + assert t.writer_base((req.unique_rid, req.slice_id), 9) is None 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.reserve(req, writer_ranks=[7, 3]) is False assert req.bounce_dst_base is None def test_reserve_heterogeneous_fanin_falls_back(self, monkeypatch): @@ -337,48 +1067,136 @@ def test_reserve_heterogeneous_fanin_falls_back(self, monkeypatch): # fall back (even though the total is divisible). 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(req, writer_ranks=[7, 3]) is False assert req.bounce_dst_base is None 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 + assert t.reserve(_recv_req([2, 2]), writer_ranks=[7, 3]) is True def test_reserve_heterogeneous_single_writer_ok(self, monkeypatch): - # num_writers==1 has no split, so heterogeneous slot bytes are fine. + # A single writer 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 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, num_writers=1 -> no even-split requirement - assert t.reserve(req, num_writers=1) is True + req = _recv_req([1]) # total = 3, one writer -> no even-split requirement + assert t.reserve(req, writer_ranks=[3]) is True + + 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_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 + + assert not t.reserve( + req, + writer_ranks=[3], + destination_intervals_factory=lambda: calls.append(True) or [(0x2000, 8)], + ) + assert calls == [] + + 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 = {(0x2080, 0x80), (0x2000, 0x80)} + 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, 0x2100),) + + raw_intervals.add((0x3000, 8)) + assert ctx._destination_intervals == ((0x2000, 0x2100),) + + 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]) + 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 + assert t.reserve(_recv_req([2]), writer_ranks=[3]) is False # total 2000 > cap 500 - def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): + def test_fanin_scatters_in_exact_rank_plan_order(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True + assert t.reserve(req, writer_ranks=[3, 7]) is True rid_slice = (req.unique_rid, req.slice_id) - # writer for the HIGHER src_base reports first; scatter must reorder by src_base. + # Writer for the higher planned base reports first; scatter follows the rank plan. t.record_result( rid_slice, 7, np.array([20], dtype=np.int64), np.array([8], dtype=np.int64), - src_base=200, + src_base=0x100000 + 100, ) 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 @@ -387,11 +1205,11 @@ def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): 3, np.array([10], dtype=np.int64), np.array([8], dtype=np.int64), - src_base=100, + src_base=0x100000, ) 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 + # Each tail carries its own planned src_base in exact rank-plan order. + assert [t[0] for t in descs] == [0x100000, 0x100000 + 100] assert [list(t[1]) for t in descs] == [[10], [20]] # dst_ptrs assert [list(t[2]) for t in descs] == [[8], [8]] # sizes @@ -400,7 +1218,7 @@ def test_fanin_fallback_writer_leaves_survivor_at_its_own_base(self, monkeypatch # 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 + assert t.reserve(req, writer_ranks=[7, 3]) is True rid_slice = (req.unique_rid, req.slice_id) t.record_result( rid_slice, 7, None, None @@ -411,23 +1229,29 @@ def test_fanin_fallback_writer_leaves_survivor_at_its_own_base(self, monkeypatch 3, np.array([10], dtype=np.int64), np.array([8], dtype=np.int64), - src_base=100, + src_base=0x100000 + 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 [t[0] for t in descs] == [0x100000 + 100] 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 + assert t.reserve(req, writer_ranks=[7, 3]) is True rid_slice = (req.unique_rid, req.slice_id) + assert t.mark_writer_exposed(rid_slice, 7) + assert t.mark_writer_exposed(rid_slice, 3) 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 + rid_slice, + 3, + np.array([10], dtype=np.int64), + np.array([8], dtype=np.int64), + src_base=0x100000 + 100, ) # all terminal, >=1 FAILED -> no scatter, release (both drained), region freed. assert t._scatter_q.empty() @@ -440,7 +1264,7 @@ def test_on_done_fires_after_scatter_lands(self, monkeypatch): # 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( @@ -448,7 +1272,7 @@ def test_on_done_fires_after_scatter_lands(self, monkeypatch): 3, np.array([10], dtype=np.int64), np.array([8], dtype=np.int64), - src_base=0, + src_base=0x100000, on_done=lambda ok: calls.append(ok), ) ctx, descs = t._scatter_q.get_nowait() @@ -464,7 +1288,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)) @@ -491,11 +1315,11 @@ 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 + assert t.reserve(req, writer_ranks=[7, 3]) 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 + t.record_result(rid_slice, 7, *arr, src_base=0x100000) + t.record_result(rid_slice, 7, *arr, src_base=0x100000) # duplicate same writer assert t._scatter_q.empty() # still only 1 distinct writer -> not all terminal assert not t._recv_alloc.released @@ -511,7 +1335,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 +1343,332 @@ 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.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] + 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 + 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.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([8], 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_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 = [], [] + t.record_result( + first_key, + 7, + np.array([0x2000], dtype=np.int64), + np.array([8], dtype=np.int64), + src_base=0x100000, + on_done=lambda ok: first_calls.append(ok), + ) + t.record_result( + second_key, + 7, + np.array([0x3000], dtype=np.int64), + np.array([8], 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 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_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([8], 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([8], 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 +1704,100 @@ 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)) 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 +1809,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(20)) 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 +1830,209 @@ 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_empty_tail_success_releases_without_scatter(self): - c = self._ctx(1) + 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() + + 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, 0x2004], dtype=np.int64), + sizes=np.array([8, 8], 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([8], dtype=np.int64), + ) + c.record_writer_result( + 3, + succeeded=True, + src_base=0x1000 + 100, + dst_ptrs=np.array([0x2004], dtype=np.int64), + sizes=np.array([8], 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([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_complete_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 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 + assert not c.ready_to_settle() + assert c.settle() is None + assert c.scatter_state is bcore.ScatterState.FAILED - 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()) + 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 +2050,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..dee5931b3358 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 @@ -1328,16 +1347,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 +1367,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 +1408,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 +1480,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..77ebe01ad72d --- /dev/null +++ b/tests/unittest/disaggregated/test_receive_ownership_wiring.py @@ -0,0 +1,2662 @@ +# 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.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_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) -> 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._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() + 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_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() + session = RxSession( + REQUEST_ID, + _params(schedule_style=DisaggScheduleStyle.GENERATION_FIRST), + receiver, + ) + session._closed = True + 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() + assert allow_discovery.wait(timeout=2) + 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", + ) + receive_thread.start() + assert discovery_started.wait(timeout=1) + + assert receiver.shutdown() is 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()) + + allow_discovery.set() + receive_thread.join(timeout=2) + 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_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"): + 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 + sender._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 + + +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._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._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._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._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._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._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..a317d2f632c0 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -12,32 +12,58 @@ # 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 _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 +74,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 +104,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,7 +163,7 @@ 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: ( + lambda _to_process, cancelled, failed, completed, timed_out, _cleanup_ready: ( cancelled, failed, completed, @@ -150,6 +208,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() @@ -209,21 +308,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: @@ -245,30 +338,504 @@ def test_gen_transfer_status_enters_consensus_when_sync_required() -> None: transceiver._gen_consensus.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 + + +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_does_not_retire_replacement_owner() -> 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 + replacement_session = Mock() + + 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) + with transceiver._session_admission_lock: + transceiver._send_sessions[rid] = replacement_session + transceiver._send_reqs[rid] = replacement + finally: + release_wait.set() + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert results == [([], [])] + assert transceiver._send_sessions == {rid: replacement_session} + assert transceiver._send_reqs == {rid: replacement} + replacement_session.close.assert_not_called() + + +def test_gen_status_poll_does_not_retire_replacement_owner() -> 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 + replacement_session = Mock() + + 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.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) + with transceiver._session_admission_lock: + transceiver._recv_sessions[rid] = replacement_session + transceiver._recv_reqs[rid] = replacement + finally: + release_wait.set() + poll_thread.join(timeout=2) + + assert not poll_thread.is_alive() + assert errors == [] + assert results == [([], [], [])] + assert transceiver._recv_sessions == {rid: replacement_session} + assert transceiver._recv_reqs == {rid: replacement} + replacement_session.close.assert_not_called() + + +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 +857,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 == [] @@ -323,3 +892,519 @@ 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) + + +@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 From 6888a21bebce6ee615b2281598569a488ed10b32 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 21 Jul 2026 21:26:11 -0700 Subject: [PATCH 4/5] [None][fix] address KV transfer ownership review findings Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/kvCacheManager.h | 13 +- .../batch_manager/kvCacheManager.cpp | 38 +- .../batch_manager/kvCacheManagerTest.cpp | 54 +- .../disagg-kv-transfer-ownership/README.md | 140 +- .../implementation-plan.md | 82 +- .../disaggregation/native/bounce/__init__.py | 4 + .../disaggregation/native/bounce/core.py | 78 +- .../disaggregation/native/bounce/impl.py | 246 ++- .../_torch/disaggregation/native/transfer.py | 382 ++-- .../_torch/disaggregation/transceiver.py | 1014 ++++++++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 1724 ++++++++++++++--- .../_torch/pyexecutor/py_executor_creator.py | 7 + .../_torch/pyexecutor/resource_manager.py | 184 +- .../executor/test_async_transfer_manager.py | 221 ++- .../test_disagg_index_mapper_early_release.py | 5 +- .../executor/test_py_executor_teardown.py | 128 +- .../test_py_executor_transfer_termination.py | 1066 +++++++++- tests/unittest/disaggregated/test_bounce.py | 488 +++-- .../disaggregated/test_kv_transfer.py | 10 +- .../test_receive_ownership_wiring.py | 573 +++++- .../test_transceiver_bounded_polling.py | 647 ++++++- 21 files changed, 6159 insertions(+), 945 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 4a3e205f0a6e..4a134ac4f1d5 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -442,8 +442,6 @@ class KVCacheBlock : public std::enable_shared_from_this void decSchedulingRefCount(); - [[nodiscard]] SizeType32 getRefCount() const; - [[nodiscard]] bool hasRefs() const; [[nodiscard]] bool hasSchedulingRefs() const; @@ -1001,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. @@ -1258,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); @@ -1557,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); @@ -2216,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 00dfe5469999..336c7d526068 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -42,7 +42,7 @@ #include #include #include -#include +#include #include namespace tc = tensorrt_llm::common; @@ -274,11 +274,6 @@ void KVCacheBlock::decSchedulingRefCount() mSchedulingRefCount--; } -SizeType32 KVCacheBlock::getRefCount() const -{ - return mRefCount; -} - bool KVCacheBlock::hasRefs() const { return mRefCount > 0; @@ -2984,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(); + } } } @@ -3001,24 +3007,24 @@ void WindowBlockManager::unpinBlocksById(std::vector const std::lock_guard lock(mLookupTree->getMutex()); - // Validate the complete operation before changing any refcount. Repeated - // IDs represent multiple pin tokens, so aggregate their decrements and - // prove that the block has enough references for the entire batch. - std::unordered_map decrementCounts; + // 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; - decrementCounts.reserve(blockIds.size()); + 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) { - auto const decrementCount = ++decrementCounts[blockId]; - TLLM_CHECK_WITH_INFO(block->getRefCount() >= decrementCount, - "Can't unpin block (id=%d) %d time(s) with refcount %d", static_cast(blockId), decrementCount, - block->getRefCount()); + TLLM_CHECK_WITH_INFO( + block->hasRefs(), "Can't unpin block (id=%d) that is not allocated", static_cast(blockId)); blocksToUnpin.push_back(block); } } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index be355269c3de..a9746b52f360 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -4904,6 +4904,11 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById) 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(); @@ -4926,13 +4931,58 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById) }(); ASSERT_LT(freeBlockId, totalBlocks); EXPECT_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), freeBlockId}), std::runtime_error); - EXPECT_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), pinnedBlockIds.front()}), 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/disagg-kv-transfer-ownership/README.md b/docs/design/disagg-kv-transfer-ownership/README.md index d6bb646243e1..e1d091308f0d 100644 --- a/docs/design/disagg-kv-transfer-ownership/README.md +++ b/docs/design/disagg-kv-transfer-ownership/README.md @@ -10,7 +10,7 @@ SPDX-License-Identifier: Apache-2.0 | **Owner** | Chien-Chun Hung | | **Status** | Draft design contract | | **Created** | 2026-07-08 | -| **Last updated** | 2026-07-14 | +| **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) | @@ -56,9 +56,10 @@ 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 may still fail closed while any registered owner or release is live, -but that accounting neither retires a C++ transport owner nor makes local C++ -object destruction evidence of remote/global quiescence. +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 @@ -71,14 +72,15 @@ 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. -Shared manager accounting may nevertheless veto teardown for a live owner or -release; the Python-native lifecycle does not retire a C++ transport owner. -Within the Python runtime it is not gated by `kv_cache_bounce_size_mb` or a -second ownership 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. +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: @@ -100,11 +102,13 @@ Implemented in this slice: 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 exact rank-to-subrange mapping, evidence-based slot - retirement, explicit slot-settlement acknowledgement, no time-based - quarantine reuse, scatter-aware close, validated non-overlapping result - metadata, and lazy construction of destination manifests after slot - reservation succeeds; +- `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 @@ -117,7 +121,23 @@ Implemented in this slice: 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. + 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, @@ -130,7 +150,11 @@ 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 the transactional `unpinBlocksById` operation. +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 @@ -151,12 +175,16 @@ 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 also disabled for PP fan-in (`overlap_pp_size > 1`) in this -slice. Equal PP layer counts do not prove equal byte contributions when layer -groups have different block sizes, and a divisibility check after summing bytes -cannot prevent a writer from overflowing its assigned subrange. PP fan-in uses -the direct path until the receiver can authorize byte-accurate per-writer -extents and offsets before publishing any bounce address. +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 @@ -165,6 +193,16 @@ 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 @@ -176,10 +214,11 @@ 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. They do not establish active-bounce engagement on fabric hardware, -attention-DP containment, mixed-version protocol liveness, or an -allocator-enforced KV lease. Those configurations remain outside a production -ownership rollout until their corresponding validation gates pass. +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 @@ -249,12 +288,11 @@ source KV pages 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. The merged baseline has a bounce-configured integration test marked -as GB200/GB300-eligible and scheduled on GB200, but it does not assert that the -transport or transfer actually engaged bounce. PR #16116 proposes that positive -signal. Treat those MNNVL cells as candidates, not proven active-bounce coverage, -until an equivalent assertion lands. Other environments may also run only the -direct fallback. +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) @@ -266,17 +304,17 @@ object; the future general context remains separate. ### Related work -Status snapshot: 2026-07-13. +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) | Open post-merge follow-up adding positive active-bounce validation and partial bounce-only lifecycle containment: it proposes orphaning/quarantining in-flight receive reservations during cancellation/session close and propagating local gather/build failures. It does not provide direct-path or allocator-enforced KV ownership, generation-safe routing, or bounded safe reclamation. | +| [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) | Open input to the C++ cancellation chain for gated NIXL in-flight cancellation and quiescence containment. It is a conceptual precedent, not a Python-native dependency. | +| [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. | @@ -390,6 +428,10 @@ from submitting a later one-sided write to an address it already received. 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 @@ -411,7 +453,7 @@ generation and endpoint-epoch identity remain future work. | **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** | GB200/GB300 MNNVL are candidate cells in the existing test marker, with the merged baseline scheduled on GB200; neither counts as active-bounce coverage without a positive transport-active and per-transfer engagement assertion. Other hardware remains in direct-path scope and needs the same positive evidence before a bounce rollout. | +| **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. | @@ -428,7 +470,7 @@ safety applies. 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 rank consensus, scheduler cancellation, or request admission. +- 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. @@ -616,6 +658,11 @@ 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 @@ -1003,9 +1050,10 @@ 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. Shared teardown accounting -may still fail closed for a live owner or release, but this adapter neither -retires a C++ transport owner nor interprets local C++ object destruction as +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 @@ -1460,9 +1508,9 @@ the top-level owner of shutdown progress. Its `shutdown(deadline)` returns a 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. -For other transceivers, shared teardown accounting may veto destruction while a -generic owner or release remains live, but it cannot derive quiescence or retire -a C++ transport owner from local object destruction. +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: diff --git a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md index 2456a2e2ace3..1bdb82c8d4d4 100644 --- a/docs/design/disagg-kv-transfer-ownership/implementation-plan.md +++ b/docs/design/disagg-kv-transfer-ownership/implementation-plan.md @@ -12,7 +12,7 @@ SPDX-License-Identifier: Apache-2.0 | **Owner** | Chien-Chun Hung | | **Status** | Draft implementation plan | | **Created** | 2026-07-08 | -| **Last updated** | 2026-07-14 | +| **Last updated** | 2026-07-21 | ## Performance and Capacity Contract @@ -67,6 +67,20 @@ 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. | @@ -78,6 +92,34 @@ instrumentation contract and are not all exported by this first slice. | 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 @@ -256,9 +298,10 @@ rollout-safe until the full Phase 1 gate passes. 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. Shared manager accounting may veto teardown for a live - generic owner or release, but the adapter must not retire a C++ transport - owner or treat local C++ object destruction as quiescence evidence. + 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, @@ -401,13 +444,23 @@ adapters in `tensorrt_llm/_torch/pyexecutor/_util.py`. - Consider per-operation abort only if the backend can guarantee no later memory access after acknowledgement. -The core effort is approximately 4,000–7,300 production lines plus substantial -concurrency and integration tests. 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. Shared `AsyncTransferManager`/`ResourceManager` accounting may observe -live generic ownership and fail closed, but it does not implement or retire the -C++ transport lifecycle. +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 @@ -645,9 +698,10 @@ publication and teardown races. ### Configuration matrix - bounce not requested, requested but factory-inactive, and arena-active; -- positive bounce-engagement evidence on candidate GB200/GB300 MNNVL fabric-VMM - cells; the merged config-enabled GB200 test is not proof, and - configured-but-inactive runs count as direct fallback; +- 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; diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py b/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py index 13f0b69f8773..740913f002ca 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py @@ -39,6 +39,9 @@ scatter_write_result, ) +# Deprecated compatibility name retained for imports from the #15618 API. +TransferContext = RecvBounceContext + __all__ = [ "BounceTransport", "Buffer", @@ -53,6 +56,7 @@ "Sizing", "SizingContext", "SlotAllocator", + "TransferContext", "TransferState", "VmmBounceTransport", "WriterState", diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py index ec997d59bdae..4d6165d3443d 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py @@ -96,6 +96,7 @@ class RecvBounceContext: 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 physically terminal writer succeeded, keyed by its exact planned rank. _writer_ok: Dict[int, bool] = field(default_factory=dict) @@ -110,6 +111,8 @@ class RecvBounceContext: 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) @@ -148,6 +151,14 @@ def __post_init__(self, destination_intervals: Optional[Iterable[tuple[int, int] 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 @@ -255,6 +266,36 @@ 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. @@ -370,8 +411,13 @@ def record_writer_result( scatter_intervals = self._non_overlapping_destination_intervals( dst_values, size_values ) - valid_tail = scatter_intervals is not None and ( - sum(size_values) <= self.per_writer_bytes + # ``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 @@ -417,6 +463,10 @@ def finish_scatter(self, ok: bool) -> None: 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.""" @@ -535,14 +585,30 @@ 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) -> bool: - """Retry durable physical-settlement acknowledgements; return whether all were delivered.""" + 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, on_done=None) -> None: diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index 1a68665d4322..f8517dcbcaab 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -151,6 +151,7 @@ def __init__( 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. @@ -528,9 +529,13 @@ def reserve( 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.""" + """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") @@ -539,40 +544,28 @@ def reserve( "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") - nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) + 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 " @@ -694,6 +687,38 @@ def _apply(self, rid_slice: RidSlice, mutate: Callable[[RecvBounceContext], None 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 @@ -719,10 +744,15 @@ def _commit_pending_settlement(self, rid_slice: RidSlice) -> bool: self._recv_alloc.quarantine(settlement.slot_id) else: self._recv_alloc.release(settlement.slot_id) - except Exception: + except Exception as e: with self._reserved_map_lock: pending.physical_in_progress = False - raise + 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 @@ -758,17 +788,37 @@ def _commit_pending_settlement(self, rid_slice: RidSlice) -> bool: self._pending_settlements.pop(rid_slice, None) return True - def _retry_pending_settlements(self) -> bool: + @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(self._pending_settlements) + 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 self._pending_settlements + return not any( + self._settlement_in_scope(key, scope) for key in self._pending_settlements + ) - def retry_settlements(self) -> bool: - """Public non-blocking retry hook for the receiver's bounded poll path.""" - return self._retry_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, @@ -778,6 +828,7 @@ 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.""" @@ -788,6 +839,7 @@ def mut(ctx: RecvBounceContext) -> None: ) if accepted: ctx.set_on_done(on_done) + ctx.set_on_logical_failure(on_logical_failure) self._apply(rid_slice, mut) @@ -888,7 +940,49 @@ def _suppress_queued_scatters(self) -> None: finally: self._scatter_q.task_done() - def _scatter_loop(self): + 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: @@ -902,40 +996,45 @@ def _scatter_loop(self): if item is None: return # FIFO poison pill: every earlier scatter has finished ctx, descs = item - 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): + 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._scatter_stream_healthy = False - c.finish_scatter(ok) - - self._apply(ctx.rid_slice, finish) - if not ok: - self._suppress_queued_scatters() - return + self._suppress_queued_scatters() + return + except Exception: + self._fail_current_scatter(ctx) + raise finally: self._scatter_q.task_done() @@ -1062,11 +1161,18 @@ def mark_protocol_conflict(self, rid_slice, on_done=None) -> None: def mark_backend_quiesced(self, rid_slice=None, on_done=None) -> None: pass - def retry_settlements(self) -> bool: + 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 diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index c2cfe28e9780..f5533a33ceb5 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2357,47 +2357,52 @@ def __init__( 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 - self._aux_buffer = aux_buffer - # Keep the allocator-level request owner alive until every queued, - # active, or in-doubt source operation has terminal evidence. + # 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 - 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. + source_owner = None + self._aux_buffer = aux_buffer + 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._source_owner = None self._closed = True raise @@ -2577,23 +2582,19 @@ 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. """ admission_lock = getattr(self._sender, "_operation_admission_lock", None) with admission_lock if admission_lock is not None else nullcontext(): - with self.lock: - self._accepting_operations = False - if self._terminal_status != SessionStatus.CANCELLED: - 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) + 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) @@ -2676,6 +2677,30 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: return None return WaitResult.FAILED if terminal_failure else WaitResult.COMPLETED + 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._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(exception) + if self.aux_task is not None and not self.aux_task.is_done: + self.aux_task.fail(exception) + return True + def set_exception(self, reason: str = ""): msg = f"TxSession {self.disagg_request_id} exception" if reason: @@ -2683,15 +2708,11 @@ def set_exception(self, reason: str = ""): 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 self.lock: - self._exception = RuntimeError(msg) - self._terminal_status = SessionStatus.ERROR - self._accepting_operations = False - for task in self.kv_tasks: - if not task.is_done: - task.fail(self._exception) - if self.aux_task is not None and not self.aux_task.is_done: - self.aux_task.fail(self._exception) + 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: @@ -2702,6 +2723,11 @@ 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() @@ -2713,6 +2739,8 @@ def close(self): self._close_locked() def _close_locked(self): + if getattr(self, "_closed", False): + return self._retry_terminal_results() with self.lock: if getattr(self, "_closed", False): @@ -2932,6 +2960,9 @@ def __init__( 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 @@ -3210,6 +3241,36 @@ def _finish_bounce( 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, ( @@ -3230,25 +3291,94 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: ) @staticmethod - def _fanin_bounce_safe(overlap, peer_ri) -> bool: + 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 safe only when - the complete per-stage layer counts are uniform. ``reserve`` separately - rejects heterogeneous layer-group block sizes. + 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. """ - if overlap.duplicate_head_factor != 1 or overlap.peer_duplicate_head_factor != 1: + 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: - layer_num_per_pp = getattr(peer_ri, "layer_num_per_pp", None) + + 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 ( - not layer_num_per_pp - or len(layer_num_per_pp) < overlap.overlap_pp_size - or len(set(layer_num_per_pp)) != 1 + 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. @@ -3336,16 +3466,20 @@ def dispatch_task(self, task: KVRecvTask): # groups, but expected_transfers should reflect per-DP-group count since # only one DP group will actually process the context request. task.set_valid_writer_cohorts(valid_writer_cohorts) - # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: - # _fanin_bounce_safe() (TP-by-head with no PP fan-in), and never under ADP broadcast (sender_dp_rank - # None), where the exact writer cohort is not selected before publication. + # 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 (task.expected_transfers == 1 or self._fanin_bounce_safe(peer_overlap, peer_infos)) + and peer_overlap.overlap_pp_size == 1 + and task.expected_transfers == 1 + and self._single_writer_bounce_exact(peer_infos) ) session = self._get_session(task._unique_rid) if session is None: @@ -3785,6 +3919,7 @@ def on_done( 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) @@ -3893,53 +4028,58 @@ def __init__( 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 - # 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_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._accepting_receives = True - self._active_receive_dispatches = 0 - self._last_slice_admitted = False - self._receiver_retired = False - # The current containment slice has no allocator-enforced destination - # lease. Retain the request that owns its KV allocation until receiver - # retirement proves every accessor and scatter is drained. + # 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 - 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 - self._aux_drained = not self._need_aux - self._aux_status: TaskStatus = TaskStatus.INIT - self._sender_endpoints: set[str] = set() - self._selected_writer_cohort: Optional[frozenset[int]] = None + destination_owner = None + self._aux_buffer = aux_buffer + 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 + self._aux_drained = not self._need_aux + 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._destination_owner = None self._closed = True raise @@ -4173,6 +4313,28 @@ def receive(self, slice: KVSlice) -> None: 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 @@ -4459,7 +4621,7 @@ def seal_receive_admission(self) -> None: def resources_drained(self) -> bool: """Whether request cleanup can release every receive-side target.""" try: - if not self._receiver._bounce.retry_settlements(): + if not self._receiver._bounce.retry_settlements(self.disagg_request_id): return False except Exception as e: logger.error( diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index c694afac21d1..5e52de97db55 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -17,6 +17,8 @@ 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 @@ -61,6 +63,44 @@ _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) consensus = [] @@ -87,6 +127,13 @@ def __init__( 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 @@ -128,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. @@ -183,6 +240,86 @@ def _get_session_admission_lock(self): 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: @@ -196,23 +333,20 @@ def shutdown(self) -> bool: # 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: set[int] = set() - for direction, sessions in ( - ("send", self._send_sessions), - ("receive", self._recv_sessions), - ): - for session in list(sessions.values()): - try: - # Cancellation closes future publication but leaves - # already-exposed resources owned until terminal - # results arrive. - session.cancel() - except Exception as e: - cancel_failed_session_ids.add(id(session)) - logger.error( - f"KvCacheTransceiverV2 shutdown failed to cancel {direction} " - f"session {session.disagg_request_id}: {e}" - ) + 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: @@ -224,6 +358,18 @@ def shutdown(self) -> bool: 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), @@ -232,6 +378,13 @@ def shutdown(self) -> bool: 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: @@ -245,9 +398,15 @@ def shutdown(self) -> bool: sessions.pop(rid, None) requests.pop(rid, None) - if cancel_failed_session_ids or close_failed: + 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 @@ -578,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: @@ -587,29 +746,485 @@ def _build_to_process( to_process.append(rid) return to_process - def _close_failed_sessions( + 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, - failed: list, - expected_sessions: dict, - expected_reqs: dict, - ) -> list: - """Retire only the exact failed owners observed by one status poll.""" - retired = [] + *, + mark_complete: bool, + ) -> None: + """Persist exact owners before fallible context-side preparation.""" + + retirements = self._get_context_retirements() with self._get_session_admission_lock(): - for rid in failed: - session = expected_sessions.get(rid) - req = expected_reqs.get(rid) - if sessions.get(rid) is not session or reqs.get(rid) is not req: + 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 - session.close() - if sessions.get(rid) is session and reqs.get(rid) is req: - req.state = LlmRequestState.DISAGG_TRANS_ERROR - reqs.pop(rid, None) - sessions.pop(rid, None) - retired.append(rid) - return retired + 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.""" @@ -661,6 +1276,18 @@ def respond_and_send_async(self, req: LlmRequest): 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") @@ -755,6 +1382,8 @@ def request_and_receive_sync(self, req: LlmRequest): 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( @@ -796,6 +1425,8 @@ def request_and_receive_async(self, req: LlmRequest): 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( @@ -844,31 +1475,68 @@ def check_context_transfer_status( 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, ): # 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(): - if getattr(self, "_shutdown_started", False): - return [], [] self._transfer_worker.sweep_stale_req_infos() - if not self._ever_had_send_session: + pending_retirements = self._get_context_retirements() + if not self._ever_had_send_session and not pending_retirements: return [], [] - send_sessions = dict(self._send_sessions) + 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 + # 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( send_sessions, - self._ctx_consensus(local_completed + local_failed), + self._ctx_consensus(local_terminal_candidates), wait_num, block_all, ) - completed, timed_out, failed, cancelled, cleanup_ready = [], [], [], [], [] + observed_completed: list[int] = [] + observed_failed: list[int] = [] + observed_cancelled: list[int] = [] + timed_out: list[int] = [] for rid in to_process: + if rid in pending_rids: + continue session = send_sessions[rid] result = session.wait_complete(blocking=block_all) if result is None: @@ -877,11 +1545,9 @@ def check_context_transfer_status( # physical state. continue if session.status == SessionStatus.CANCELLED: - cancelled.append(rid) - cleanup_ready.append(rid) + observed_cancelled.append(rid) elif result == WaitResult.COMPLETED: - completed.append(rid) - cleanup_ready.append(rid) + 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" @@ -889,8 +1555,23 @@ def check_context_transfer_status( timed_out.append(rid) else: logger.warning(f"TxSession rid={session.disagg_request_id} failed") - failed.append(rid) - cleanup_ready.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( @@ -902,61 +1583,72 @@ def check_context_transfer_status( cleanup_ready, ) - retired_completed = [] - with self._get_session_admission_lock(): - for rid in completed: - session = send_sessions.get(rid) - req = send_reqs.get(rid) - if ( - self._send_sessions.get(rid) is not session - or self._send_reqs.get(rid) is not req - ): - continue - session.close() - if mark_complete: - req.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - if self._send_sessions.get(rid) is session and self._send_reqs.get(rid) is req: - self._send_reqs.pop(rid, None) - self._send_sessions.pop(rid, None) - retired_completed.append(rid) - # A remotely cancelled context transfer is terminal for its local - # AsyncTransferManager entry. Report it as an error so PyExecutor can - # unpin the request instead of silently losing the transceiver session. - terminal_failed = failed + cancelled - retired_failed = self._close_failed_sessions( - self._send_sessions, - self._send_reqs, - terminal_failed, - send_sessions, - send_reqs, + # 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), + ) ) - return retired_completed, retired_failed - def check_gen_transfer_status(self, at_least_request_num: Optional[int]): - with self._get_session_admission_lock(): - if getattr(self, "_shutdown_started", False): + with self._status_call() as admitted: + if not admitted: return [], [], [] - if not self._ever_had_recv_session and not self._gen_need_sync: + 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 [], [], [] - recv_sessions = dict(self._recv_sessions) + 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, 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( recv_sessions, - self._gen_consensus(local_completed + local_failed), + self._gen_consensus(local_terminal_candidates), 0 if need_progress else wait_num, block_all, ) - completed, failed, cancelled, cleanup_ready = [], [], [], [] + observed_completed: list[int] = [] + observed_failed: list[int] = [] + observed_cancelled: list[int] = [] for rid in to_process: + 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): @@ -964,65 +1656,40 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): # Neither a non-blocking miss nor a timeout proves that request # cleanup can no longer race a transfer accessor. continue - cleanup_ready.append(rid) 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) + 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, cleanup_ready ) - retired_completed = [] - cancelled_reqs = [] - with self._get_session_admission_lock(): - for rid in cancelled: - session = recv_sessions.get(rid) - req = recv_reqs.get(rid) - if ( - self._recv_sessions.get(rid) is not session - or self._recv_reqs.get(rid) is not req - ): - continue - session.close() - if self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req: - self._recv_reqs.pop(rid, None) - self._recv_sessions.pop(rid, None) - cancelled_reqs.append(req) - - for rid in completed: - session = recv_sessions.get(rid) - req = recv_reqs.get(rid) - if ( - self._recv_sessions.get(rid) is not session - or self._recv_reqs.get(rid) is not req - ): - continue - # 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) - session.close() - req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - if self._recv_sessions.get(rid) is session and self._recv_reqs.get(rid) is req: - self._recv_reqs.pop(rid, None) - self._recv_sessions.pop(rid, None) - retired_completed.append(rid) - retired_failed = self._close_failed_sessions( - self._recv_sessions, - self._recv_reqs, - failed, - recv_sessions, - recv_reqs, + retired_completed, retired_failed, cancelled_reqs = self._commit_generation_decisions( + ( + ("completed", completed), + ("failed", failed), + ("cancelled", cancelled), + ) ) if retired_failed: logger.warning( @@ -1050,7 +1717,7 @@ def _poll_gen_sessions_for_poll_interval( def check_gen_transfer_complete(self): with self._get_session_admission_lock(): - return len(self._recv_sessions) == 0 + 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. @@ -1095,12 +1762,48 @@ def cancel_request(self, req: LlmRequest) -> bool: rid = get_unique_rid(req) 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: - self._wait_reqs.pop(rid, None) - - cleanup_pending = False + 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), @@ -1108,8 +1811,23 @@ def cancel_request(self, req: LlmRequest) -> bool: 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 @@ -1147,16 +1865,34 @@ 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. with self._get_session_admission_lock(): - if getattr(self, "_shutdown_started", False): - raise RuntimeError("KV cache transceiver is shutting down") 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: @@ -1191,6 +1927,7 @@ def prepare_context_requests(self, requests: List[LlmRequest]): 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) } @@ -1204,6 +1941,11 @@ def prepare_context_requests(self, requests: List[LlmRequest]): 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) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b6d5915b22f0..f6c65990c046 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -380,6 +380,36 @@ class BatchStatePP(BatchState): _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 @@ -393,7 +423,6 @@ class _RequestCancellationProgress: all_transfers_complete: Optional[bool] = None logical_cancelled: bool = False decoding_iter_updated: bool = False - terminate_on_completion: bool = False termination_handoff_complete: bool = False @@ -403,13 +432,47 @@ class _ConnectorCompletionProgress: request: LlmRequest was_active: bool - response_prepared: bool = False 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. @@ -495,6 +558,18 @@ def in_doubt_reason(self) -> Optional[str]: 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): @@ -517,6 +592,8 @@ def __init__(self, # 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]: @@ -624,6 +701,15 @@ def start_transfer(self, request: LlmRequest, owner: Optional[str] = None): self._resume_first_transfer_admission(progress) return + 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") + existing_request = self._requests_in_transfer.get(req_id) if existing_request is not None and existing_request is not request: raise RuntimeError( @@ -670,60 +756,171 @@ def end_transfer(self, 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: - logger.warning(f"Request {req_id} not found in transfer manager") - return False + 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, + ) + self._transfer_completion_progress[req_id] = completion + + # 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} has a different active transfer owner") + 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.pinned_block_ids) + except Exception: + completion.unpin_in_doubt = True + raise + completion.unpin_complete = True - transfer_metadata = self._request_transfer_metadata[req_id] - final_transfer = transfer_metadata.will_end_all_transfers(owner) + if not completion.state_updated: + # We don't want to overwrite any error state. + if request.state not in (LlmRequestState.DISAGG_TRANS_ERROR, + LlmRequestState.DISAGG_CONTEXT_COMPLETE): + request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + completion.state_updated = True - # Keep the exact owner and request root enrolled until final unpin - # succeeds. A failed unpin is retryable without reconstructing which - # provider contribution or request allocation it belonged to. - if final_transfer and self.should_store_blocks: - self.kv_cache_manager.unpin_blocks_by_id( - transfer_metadata.pinned_block_ids) + 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 - if transfer_metadata.end_transfer(owner): + 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) - self._request_transfer_metadata.pop(req_id) - # We don't want to overwrite any error state. - if request.state != LlmRequestState.DISAGG_TRANS_ERROR: - request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - - return True + 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) - return False + 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.""" - if self._requests_in_transfer.get(request.py_request_id) is not request: - return False - transfer_metadata = self._request_transfer_metadata.get( - request.py_request_id) - return (transfer_metadata is not None - and transfer_metadata.has_owner(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.""" - return self._requests_in_transfer.get(request.py_request_id) is request + 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 + + 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``.""" - return { + requests = { req_id: request for req_id, request in self._requests_in_transfer.items() - if self._request_transfer_metadata[req_id].has_owner(owner) + 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 bool(self._requests_in_transfer - or self._transfer_admission_progress) + or self._transfer_admission_progress + or self._transfer_completion_progress) class PyExecutor: @@ -931,6 +1128,11 @@ 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 @@ -965,6 +1167,17 @@ def __init__( # 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. @@ -1332,6 +1545,179 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() + 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: @@ -1361,15 +1747,9 @@ def _release_transfer_owner(self, def _prepare_transfer_completion_response(self, request: LlmRequest) -> None: """Create and buffer the response needed before owner retirement.""" - response = request.create_response(False, self.dist.rank) - if response: - response.result.cached_tokens = request.cached_tokens - self._maybe_attach_ctx_usage(request, response) - # With ADP, _enqueue_responses does a tp_gather collective. Calling - # it from a completion callback would deadlock because only the - # owning DP rank reaches that callback. - self._pending_transfer_responses.append( - (request.py_request_id, response)) + self._prepare_transfer_terminal_response( + request, _TransferTerminalOutcome.SUCCEEDED, + _PYTHON_NATIVE_TRANSCEIVER_OWNER) def _end_transfer_and_maybe_terminate(self, request: LlmRequest, @@ -1379,6 +1759,10 @@ def _end_transfer_and_maybe_terminate(self, 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: # Keep the request active so the synchronized error path can @@ -1396,16 +1780,24 @@ def _end_transfer_and_maybe_terminate(self, self.active_requests.remove(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 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 # The PP=1 partial-reuse path may have released request resources # while the transfer lease remained live. The exact-request marker @@ -1417,9 +1809,13 @@ def _end_transfer_and_maybe_terminate(self, 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. @@ -1427,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. @@ -1736,30 +2225,194 @@ def cancel_request(self, id: int): """ self.executor_request_queue.enqueue_cancel_request(id) + 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 + 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 + # shutdown() and then immediately reads model_engine.model.model_config + # to compute kv_cache_max_memory. cleanup() would set + # model_engine.model = None, breaking that read with + # `'NoneType' object has no attribute 'model_config'`. + # The engine's __del__ still calls cleanup() at terminal teardown + # (when the executor's reference is dropped), which is sufficient for + # the GMS daemon registry eviction the cleanup hook was added for. + del self.model_engine + if self.draft_model_engine is not None: + del self.draft_model_engine + if self.virtual_memory_pools is not None: + keys = list(self.virtual_memory_pools.keys()) + for key in keys: + del self.virtual_memory_pools[key] + 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.") - 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 + 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 @@ -1767,19 +2420,66 @@ def shutdown(self): # 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() + 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. - if torch.cuda.is_available(): - torch.cuda.synchronize() + + 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: - transfer_manager.begin_shutdown() - is_python_native_transceiver = ( - self.kv_cache_transceiver is not None - and self._requires_physical_transfer_drain()) + 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 @@ -1789,119 +2489,243 @@ def shutdown(self): defer_transceiver_owners = getattr( self, "_defer_transceiver_owners_for_shutdown", None) if defer_transceiver_owners is not None: - defer_transceiver_owners() - transceiver_shutdown_result = self.kv_cache_transceiver.shutdown() - if transceiver_shutdown_result is False: - raise RuntimeError( - "KV cache transceiver still owns active transfer targets; " - "refusing to shut down resource managers") - if transceiver_shutdown_result is not True: - raise RuntimeError( - "Python-native KV cache transceiver did not prove physical " - "drain; refusing to shut down resource managers") + 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: - drain_deferred_terminations( - terminate_after_worker_shutdown=True) + 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: - drain_pending_cancellations() + 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): - raise RuntimeError( - "User cancellation ownership is still incomplete after " - "transceiver shutdown; refusing to shut down resource managers") + 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): - raise RuntimeError( - "KV transfer request ownership is still active after " - "transceiver shutdown; refusing to shut down resource managers") + 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._kv_connector_terminate_requests( - terminate_after_worker_shutdown=True) + 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: - drain_connector_completions( - terminate_after_worker_shutdown=True) + 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): - raise RuntimeError( - "Connector completion ownership is still incomplete after " - "executor shutdown; refusing to shut down resource managers") - 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") - 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") + 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: - 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") + 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: - self._do_terminate_request(request) + try: + self._do_terminate_request(request) + except Exception as error: + retain_post_connector_error(error) if termination_progress: - raise RuntimeError( - "Request resource termination is still incomplete; refusing " - "to shut down resource managers") + 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) - 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") + 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) - if (has_pending_releases is not None and has_pending_releases()): + 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( - "Request resource release is still pending; refusing to shut " - "down resource managers") - for manager in self.resource_manager.resource_managers.values(): - if manager: - manager.shutdown() - # 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 - # shutdown() and then immediately reads model_engine.model.model_config - # to compute kv_cache_max_memory. cleanup() would set - # model_engine.model = None, breaking that read with - # `'NoneType' object has no attribute 'model_config'`. - # The engine's __del__ still calls cleanup() at terminal teardown - # (when the executor's reference is dropped), which is sufficient for - # the GMS daemon registry eviction the cleanup hook was added for. - del self.model_engine - if self.draft_model_engine is not None: - del self.draft_model_engine - if self.virtual_memory_pools is not None: - 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 + "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: """ @@ -3828,6 +4652,20 @@ def _has_transceiver_transfer_owner(self, request: LlmRequest) -> bool: 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", {}) @@ -3875,6 +4713,8 @@ def _defer_transceiver_owners_for_shutdown(self) -> None: 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) @@ -3884,6 +4724,13 @@ def _defer_transceiver_owners_for_shutdown(self) -> None: # 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. @@ -4245,10 +5092,24 @@ def _advance_connector_completion( request = progress.request request_was_already_terminated = ( self._was_transfer_request_terminated(request)) - if not progress.response_prepared: - if progress.was_active: - PyExecutor._prepare_transfer_completion_response(self, request) - progress.response_prepared = True + 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 = ( @@ -4258,6 +5119,14 @@ def _advance_connector_completion( 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) @@ -4269,8 +5138,11 @@ def _advance_connector_completion( 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( @@ -4293,27 +5165,122 @@ def _drain_pending_connector_completions( 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: - pending = self._get_pending_connector_completions() - for request in self.kv_connector_manager.get_finished(): - 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: - pending[request_key] = _ConnectorCompletionProgress( - request=request, - was_active=bool(self.kv_cache_transceiver - and request in self.active_requests), - ) - + 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: self.kv_connector_manager.worker.wait_for_save( @@ -6728,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. @@ -6745,28 +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 self._requires_physical_transfer_drain() else None) + 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] - - if not self._finalize_deferred_transfer_termination(request): - self._end_transfer_and_maybe_terminate( - request, transfer_owner=transfer_owner) + 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( @@ -7102,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) @@ -7169,6 +8354,7 @@ 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() @@ -7185,16 +8371,15 @@ def _handle_errors(self, 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_cancellations = getattr( - self, "_pending_request_cancellations", {}) - cancellation_progress = pending_cancellations.get(request_key) - if (cancellation_progress is not None - and cancellation_progress.request is request): - # _handle_errors removes the request from active_requests, - # so the cancellation ledger must either perform final - # teardown or hand it to a remaining provider owner. - cancellation_progress.terminate_on_completion = True pending_ownership_request_keys.add(request_key) continue already_deferred = deferred_terminations.get(request_key) is request @@ -7238,26 +8423,54 @@ def _handle_errors(self, 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 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 - if requests is None: - self.active_requests.clear() - else: - self.active_requests = [ - request for request in self.active_requests - if request not in requests - ] + 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: + 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): @@ -7311,6 +8524,7 @@ def _do_terminate_request(self, request: LlmRequest): 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) @@ -7385,6 +8599,10 @@ 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 @@ -7421,7 +8639,10 @@ def _complete_request_cancellation_termination( 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 @@ -7459,6 +8680,11 @@ def _has_pending_ownership_transition(self, request: LlmRequest) -> bool: 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) @@ -7508,14 +8734,18 @@ def _finalize_deferred_transfer_termination( # retry; a remaining provider owns final teardown. should_terminate = force_terminate - if should_terminate and not request_was_already_terminated: - if terminate_after_worker_shutdown: - self._terminate_request_after_worker_shutdown(request) + 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._terminate_request(request) + 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, @@ -7542,7 +8772,8 @@ def _drain_deferred_transfer_terminations( except Exception as e: logger.error( "Failed to retry KV transfer cancellation for request " - f"{request.py_request_id}; retaining request resources: {e}") + f"{request.py_request_id}; retaining request resources: {e}" + ) continue if not is_drained: continue @@ -7582,6 +8813,17 @@ def _handle_canceled_requests(self): "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) @@ -7595,18 +8837,7 @@ def _handle_canceled_requests(self): f"Failed to cancel request {progress.cancel_id}; retaining " f"cancellation for retry: {e}") continue - if completed and progress.terminate_on_completion: - try: - self._complete_request_cancellation_termination(progress) - except Exception as e: - logger.error( - f"Failed to terminate cancelled request " - f"{progress.cancel_id}; retaining cancellation for " - f"retry: {e}") - continue - if (completed and (not progress.terminate_on_completion - or progress.termination_handoff_complete) - and pending.get(request_key) is progress): + if completed and pending.get(request_key) is progress: pending.pop(request_key) # Keep one ID while any exact child/request cancellation still has an @@ -7741,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. @@ -7767,6 +8999,15 @@ def _handle_responses(self, emit_first_iter: bool = True): new_active_requests.append(request) continue + terminal = self._get_transfer_terminal(request) + if (terminal is not None + and self._has_async_transfer_owner(request)): + # A provider already owns the exact logical response, but a + # sibling lease still prevents final teardown. Do not let the + # ordinary response pass create a second terminal result. + 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(): @@ -7842,14 +9083,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 @@ -7876,7 +9139,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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a2f0e15c8423..316a4317c742 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -339,6 +339,13 @@ def _teardown_kv_cache_managers_after_transceiver_shutdown( 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; " diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index b470525cb523..61fd9a7e4223 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2571,6 +2571,7 @@ class _ResourceReleaseProgress: release_plan: Tuple[BaseResourceManager, ...] next_manager: int = 0 in_doubt_manager: Optional[int] = None + manager_call_active: bool = False failure: Optional[str] = None @@ -2579,6 +2580,7 @@ class _TransferResourceReleaseProgress: request: LlmRequest completed_managers: Set[ResourceManagerType] in_doubt_manager: Optional[ResourceManagerType] = None + manager_call_active: bool = False failure: Optional[str] = None @@ -2636,68 +2638,102 @@ def update_resources( def free_resources(self, request: LlmRequest): request_key = id(request) - 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): - 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}") + 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}") + 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 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 - if progress.in_doubt_manager is not None: - 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}") - - while progress.next_manager < len(progress.release_plan): - # Mark the individual manager attempt ambiguous before calling - # it. An exception may be raised after that manager has already - # released some state, so blindly invoking it again could - # double-free. Only a normal return proves this step complete. - progress.in_doubt_manager = progress.next_manager - try: - progress.release_plan[progress.next_manager].free_resources( - request) - except Exception as error: - progress.failure = f"{type(error).__name__}: {error}" - raise + 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 - 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) - def release_resource_for_transfer( self, request: LlmRequest, resource_mgr_type: ResourceManagerType) -> None: @@ -2725,6 +2761,11 @@ def release_resource_for_transfer( 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 " @@ -2738,13 +2779,32 @@ def release_resource_for_transfer( return progress.in_doubt_manager = resource_mgr_type - try: - manager.free_resources(request) - except Exception as error: - progress.failure = f"{type(error).__name__}: {error}" - raise + 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: @@ -2752,8 +2812,10 @@ def has_in_doubt_resource_releases(self) -> bool: 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 for progress in + 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: diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/executor/test_async_transfer_manager.py index f9b40d04ae6a..bd73e88f2e8d 100644 --- a/tests/unittest/_torch/executor/test_async_transfer_manager.py +++ b/tests/unittest/_torch/executor/test_async_transfer_manager.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import threading from unittest.mock import MagicMock import pytest @@ -30,6 +31,32 @@ 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, @@ -111,6 +138,22 @@ def test_start_transfer_multiple_transfers_same_request(): 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] @@ -172,17 +215,23 @@ def test_end_transfer_rejects_different_request_with_same_id(): kv_cache_manager.unpin_blocks_by_id.assert_not_called() -def test_final_transfer_unpin_failure_preserves_owner_for_retry(): +def test_final_transfer_unpin_failure_is_retained_without_replay(): kv_cache_manager = MagicMock() kv_cache_manager.store_blocks_for_reuse.return_value = [100] - kv_cache_manager.unpin_blocks_by_id.side_effect = [RuntimeError("unpin failed"), None] + 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"): + with pytest.raises(RuntimeError, match="unpin failed after mutation"): manager.end_transfer(request, owner="python_native") metadata = manager._request_transfer_metadata[42] @@ -190,12 +239,112 @@ def test_final_transfer_unpin_failure_preserves_owner_for_retry(): 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 - assert kv_cache_manager.unpin_blocks_by_id.call_count == 2 + 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(): @@ -240,6 +389,70 @@ def test_successful_resource_release_clears_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() 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_py_executor_teardown.py b/tests/unittest/_torch/executor/test_py_executor_teardown.py index 10d0a6611453..eca36c23640f 100644 --- a/tests/unittest/_torch/executor/test_py_executor_teardown.py +++ b/tests/unittest/_torch/executor/test_py_executor_teardown.py @@ -28,6 +28,24 @@ class _ShutdownOwnershipStub: 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() @@ -39,6 +57,8 @@ def __init__(self, terminal_shutdown_result): 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() @@ -70,6 +90,9 @@ def test_native_shutdown_retries_before_releasing_resource_managers(monkeypatch) 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): @@ -87,18 +110,78 @@ def test_native_shutdown_rejects_legacy_none_result(monkeypatch): stub.manager.shutdown.assert_not_called() -@pytest.mark.parametrize("terminal_shutdown_result", [True, None]) -def test_estimation_teardown_retries_only_after_transceiver_drains( - terminal_shutdown_result, -): - """Manager teardown must wait for either current or legacy drain success.""" +@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.shutdown.side_effect = [False, terminal_shutdown_result] + 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="still owns active transfer targets"): + 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 ) @@ -112,9 +195,40 @@ def test_estimation_teardown_retries_only_after_transceiver_drains( 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 diff --git a/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py index 63f097ea2464..5f2f54848822 100644 --- a/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py +++ b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py @@ -22,8 +22,11 @@ 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, ) @@ -39,6 +42,7 @@ def __init__(self, request, *, transceiver=False, anonymous=0, events=None): 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) @@ -60,6 +64,15 @@ def _has_transfer_owner(self, request, owner): 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} @@ -86,8 +99,19 @@ def _end_transfer(self, request, owner=None): 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 = object.__new__(PyExecutor) + executor = _initialize_bare_executor_lifecycle_state(object.__new__(PyExecutor)) executor.active_requests = [request] executor._deferred_transfer_terminations = {} executor._error_budget = Mock(budget=1.0) @@ -307,7 +331,7 @@ def test_context_error_preserves_error_state_for_remaining_transfer_owner(): def _make_native_launch_executor(request): - executor = object.__new__(PyExecutor) + 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, @@ -342,6 +366,37 @@ def _make_native_launch_request(): 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() @@ -367,24 +422,328 @@ def test_transceiver_owner_mode_boundary(requires_physical_drain): @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) - executor = object.__new__(PyExecutor) + 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() - expected_owner = _PYTHON_NATIVE_TRANSCEIVER_OWNER if requires_physical_drain else None - executor._end_transfer_and_maybe_terminate.assert_called_once_with( - request, transfer_owner=expected_owner + 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(): @@ -473,7 +832,7 @@ def finish_by_reason(_reason): def _make_cancel_executor(requests): - executor = object.__new__(PyExecutor) + 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) @@ -609,7 +968,7 @@ def test_user_cancel_retries_transport_error_and_continues_siblings(): ) -def test_failed_cancellation_then_error_terminates_once_after_retry(): +def test_failed_cancellation_keeps_first_outcome_across_late_error(): request = _make_cancel_request( 7, LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, @@ -643,18 +1002,21 @@ def test_failed_cancellation_then_error_terminates_once_after_retry(): charge_budget=False, ) - assert progress.terminate_on_completion - assert executor.active_requests == [] + assert executor.active_requests == [request] executor._terminate_request.assert_not_called() executor._handle_canceled_requests() - executor._terminate_request.assert_called_once_with(request) 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( @@ -699,7 +1061,7 @@ def test_native_owner_release_retry_survives_response_processing(): def _make_connector_completion_executor(request): - executor = object.__new__(PyExecutor) + 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] @@ -741,44 +1103,296 @@ def _configure_finished_response(executor, request): def _configure_connector_poll(executor, request): - executor.kv_connector_manager = Mock() - executor.kv_connector_manager.get_finished.return_value = [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 -@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) +def _make_one_shot_connector_manager(request, *, empty_polls=0): + manager = Mock() + poll_count = 0 - executor._handle_canceled_requests() - finished_requests = executor._handle_responses() + def get_finished(): + nonlocal poll_count + poll_count += 1 + if poll_count == empty_polls + 1: + return [request] + return [] - assert finished_requests == [request] - assert executor.active_requests == [] - executor._terminate_request.assert_not_called() + 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._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( @@ -798,16 +1412,23 @@ def test_capability_false_cancel_keeps_anonymous_status_owner_after_connector(): executor._handle_canceled_requests() finished_requests = executor._handle_responses() - assert finished_requests == [request] - assert executor.active_requests == [] + assert finished_requests == [] + assert executor.active_requests == [request] executor.async_transfer_manager.end_transfer.assert_not_called() executor._terminate_request.assert_not_called() 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 == [ @@ -1007,11 +1628,11 @@ def test_user_cancel_without_native_context_owner_is_unchanged(state, generation def test_successful_request_termination_records_live_transfer_owner(): request = Mock(py_request_id=7) - executor = object.__new__(PyExecutor) + 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) + executor.dist = SimpleNamespace(rank=1, world_size=1) executor.kv_cache_transceiver = SimpleNamespace( requires_physical_drain_before_request_release=True ) @@ -1030,12 +1651,12 @@ def test_successful_request_termination_records_live_transfer_owner(): def test_failed_request_termination_is_not_recorded_as_complete(): request = Mock(py_request_id=7) - executor = object.__new__(PyExecutor) + 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) + executor.dist = SimpleNamespace(rank=1, world_size=1) executor.kv_cache_transceiver = SimpleNamespace( requires_physical_drain_before_request_release=True ) @@ -1054,7 +1675,7 @@ def test_failed_request_termination_is_not_recorded_as_complete(): def test_request_termination_retry_does_not_release_resources_twice(): request = Mock(py_request_id=7) - executor = object.__new__(PyExecutor) + 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 = [ @@ -1062,7 +1683,7 @@ def test_request_termination_retry_does_not_release_resources_twice(): None, ] executor.gather_all_responses = False - executor.dist = SimpleNamespace(rank=1) + executor.dist = SimpleNamespace(rank=1, world_size=1) executor.kv_cache_transceiver = None executor._request_resource_termination_progress = {} @@ -1079,7 +1700,7 @@ def test_request_termination_retry_does_not_release_resources_twice(): def test_request_termination_rejects_pending_transfer_admission(): request = Mock(py_request_id=7) - executor = object.__new__(PyExecutor) + 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() @@ -1091,13 +1712,13 @@ def test_request_termination_rejects_pending_transfer_admission(): def _make_shutdown_executor(events, request): - executor = object.__new__(PyExecutor) + 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) + 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() @@ -1111,6 +1732,7 @@ def _make_shutdown_executor(events, request): 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() @@ -1118,6 +1740,36 @@ def _make_shutdown_executor(events, request): 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) + 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", @@ -1128,6 +1780,7 @@ def test_capability_false_shutdown_does_not_call_legacy_transceiver(monkeypatch) 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) @@ -1135,10 +1788,11 @@ def test_capability_false_shutdown_does_not_call_legacy_transceiver(monkeypatch) 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_does_not_claim_anonymous_owner_drain(monkeypatch): +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, @@ -1150,14 +1804,16 @@ def test_capability_false_shutdown_does_not_claim_anonymous_owner_drain(monkeypa 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()) - with pytest.raises(RuntimeError, match="Asynchronous transfer ownership"): - executor.shutdown() + 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() - events.manager_shutdown.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): @@ -1165,6 +1821,7 @@ def test_shutdown_waits_for_final_connector_owner_after_native_cancel(monkeypatc "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, @@ -1183,8 +1840,7 @@ def test_shutdown_waits_for_final_connector_owner_after_native_cancel(monkeypatc executor.async_transfer_manager = _ExactTransferManager( request, transceiver=True, anonymous=1, events=events ) - executor.kv_connector_manager = Mock() - executor.kv_connector_manager.get_finished.return_value = [request] + executor.kv_connector_manager = _make_one_shot_connector_manager(request) executor._pending_connector_completions = {} executor._pending_transfer_responses = [] executor._maybe_attach_ctx_usage = Mock() @@ -1314,11 +1970,13 @@ def test_shutdown_retires_native_owner_but_vetoes_connector_owner(monkeypatch): events.manager_shutdown.assert_not_called() -def test_shutdown_ingests_finished_connector_owner_before_veto(monkeypatch): +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, @@ -1330,19 +1988,309 @@ def test_shutdown_ingests_finished_connector_owner_before_veto(monkeypatch): executor.active_requests = [] executor.force_terminate_ctx_for_partial_reuse = False events.transceiver_shutdown.return_value = True - executor.kv_connector_manager = Mock() - executor.kv_connector_manager.get_finished.return_value = [request] + executor.kv_connector_manager = _make_one_shot_connector_manager(request, empty_polls=1) executor.async_transfer_manager = _ExactTransferManager(request, anonymous=1) executor.shutdown() - executor.kv_connector_manager.get_finished.assert_called_once_with() + 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_once_with() + 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 = {} + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + 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 = {} + executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + 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", diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 90903997dc03..4b3eb854b55f 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -53,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) # --------------------------------------------------------------------------- # @@ -195,10 +202,10 @@ def test_make_kv_result_msg_uses_binary_frame(result_name): # --------------------------------------------------------------------------- # -# fan-in safety gate — equal split for uniform TP-by-head and even PP +# fan-in safety gate — equal split only within one PP stage # --------------------------------------------------------------------------- # def test_fanin_bounce_safe_gate(): - """Require uniform PP layers and no TP head duplication in either direction.""" + """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 @@ -209,22 +216,36 @@ def ov(dup, pp, peer_dup=1): 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), 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 - assert safe(ov(1, 1, peer_dup=2), ri([24])) is False # reciprocal duplication - # Equal PP layers are safe; reserve separately checks per-block byte sizes. - assert safe(ov(1, 4), ri([20, 20, 20, 20])) is True - # Uneven PP fan-in also falls 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 + 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 + + +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 # --------------------------------------------------------------------------- # @@ -253,6 +274,8 @@ def test_noop_behaviour(self): 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 @@ -386,7 +409,7 @@ 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, @@ -404,6 +427,22 @@ def _write_meta(): ) +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 TestTransportRollback: def _construct(self, agent, *, capacity=1024): @@ -1045,35 +1084,22 @@ def close(second=False): @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 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, writer_ranks=[7, 3]) is True + assert t.reserve(req, writer_ranks=[7]) is True assert req.bounce_dst_base == 0x100000 - # Exact rank order determines layout; per_writer = 100. assert t.writer_base((req.unique_rid, req.slice_id), 7) == 0x100000 - assert t.writer_base((req.unique_rid, req.slice_id), 3) == 0x100000 + 100 assert t.writer_base((req.unique_rid, req.slice_id), 9) is None - 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, writer_ranks=[7, 3]) is False - assert req.bounce_dst_base is None - - 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). - t = _make_transport(monkeypatch, block_bytes_per_group=[100, 200]) - req = _recv_req([2, 2]) # total = 2*100 + 2*200 = 600 + @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 - - 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]), writer_ranks=[7, 3]) is True + assert not t._recv_alloc.has_outstanding def test_reserve_heterogeneous_single_writer_ok(self, monkeypatch): # A single writer has no split, so heterogeneous slot bytes are fine. @@ -1085,6 +1111,34 @@ def test_reserve_single_writer_ok(self, monkeypatch): 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]) @@ -1096,6 +1150,21 @@ def test_reserve_binds_trusted_destination_intervals(self, monkeypatch): 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_destination_intervals_factory_runs_only_after_allocator_admission(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[8]) req = _recv_req([1]) @@ -1185,80 +1254,6 @@ 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]), writer_ranks=[3]) is False # total 2000 > cap 500 - def test_fanin_scatters_in_exact_rank_plan_order(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, writer_ranks=[3, 7]) is True - rid_slice = (req.unique_rid, req.slice_id) - # Writer for the higher planned base reports first; scatter follows the rank plan. - t.record_result( - rid_slice, - 7, - np.array([20], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=0x100000 + 100, - ) - 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=0x100000, - ) - ctx, descs = t._scatter_q.get_nowait() - # Each tail carries its own planned src_base in exact rank-plan order. - assert [t[0] for t in descs] == [0x100000, 0x100000 + 100] - 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, writer_ranks=[7, 3]) 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=0x100000 + 100, - ) # writer 1 bounced to base+100 - ctx, descs = t._scatter_q.get_nowait() - assert [t[0] for t in descs] == [0x100000 + 100] - 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, writer_ranks=[7, 3]) is True - rid_slice = (req.unique_rid, req.slice_id) - assert t.mark_writer_exposed(rid_slice, 7) - assert t.mark_writer_exposed(rid_slice, 3) - 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=0x100000 + 100, - ) - # 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 - 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. @@ -1271,7 +1266,7 @@ def test_on_done_fires_after_scatter_lands(self, monkeypatch): rid_slice, 3, np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), + np.array([200], dtype=np.int64), src_base=0x100000, on_done=lambda ok: calls.append(ok), ) @@ -1311,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, writer_ranks=[7, 3]) 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=0x100000) - t.record_result(rid_slice, 7, *arr, src_base=0x100000) # duplicate 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 = [] @@ -1456,7 +1439,7 @@ def fail_once(update, *, peer_rank=None): key, 7, np.array([0x2000], dtype=np.int64), - np.array([8], dtype=np.int64), + np.array([200], dtype=np.int64), src_base=0x100000, on_done=lambda succeeded: receiver._finish_bounce(key, succeeded, 7), ) @@ -1502,6 +1485,136 @@ def acknowledge(ok): 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) @@ -1510,20 +1623,21 @@ def test_scatter_failure_suppresses_later_unlaunched_queue_entries(self, monkeyp 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_calls, second_calls, first_logical_failures = [], [], [] t.record_result( first_key, 7, np.array([0x2000], dtype=np.int64), - np.array([8], 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([8], dtype=np.int64), + np.array([100], dtype=np.int64), src_base=0x100000, on_done=lambda ok: second_calls.append(ok), ) @@ -1544,10 +1658,82 @@ def test_scatter_failure_suppresses_later_unlaunched_queue_entries(self, monkeyp 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 == [] + def test_backend_quiescence_callback_fires_after_slot_settlement(self, monkeypatch): t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) @@ -1580,7 +1766,7 @@ def test_close_finishes_already_queued_scatter(self, monkeypatch): key, 3, np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), + np.array([200], dtype=np.int64), src_base=0x100000, ) @@ -1617,7 +1803,7 @@ def acknowledge(ok): key, 3, np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), + np.array([200], dtype=np.int64), src_base=0x100000, on_done=acknowledge, ) @@ -1786,7 +1972,10 @@ def _ctx( ) 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((7, 3, 11), per_writer_bytes=0x64, base_addr=0x1000) @@ -1813,7 +2002,7 @@ def test_fanin_holds_until_all_terminal(self): 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=0x1000 + 100, **self._dst(20)) + 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): @@ -1911,8 +2100,8 @@ def test_overlapping_destinations_within_writer_fail_closed(self): 3, succeeded=True, src_base=0x1000, - dst_ptrs=np.array([0x2000, 0x2004], dtype=np.int64), - sizes=np.array([8, 8], dtype=np.int64), + 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() @@ -1925,14 +2114,14 @@ def test_overlapping_destinations_across_writers_fail_closed(self): succeeded=True, src_base=0x1000, dst_ptrs=np.array([0x2000], dtype=np.int64), - sizes=np.array([8], 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([0x2004], dtype=np.int64), - sizes=np.array([8], dtype=np.int64), + 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() @@ -1949,13 +2138,13 @@ def test_scatter_destinations_must_stay_inside_trusted_intervals(self): succeeded=True, src_base=0x1000, dst_ptrs=np.array([0x20F0], dtype=np.int64), - sizes=np.array([0x20], 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_accept_complete_in_range_intervals(self): + def test_scatter_destinations_reject_truncated_in_range_intervals(self): c = self._ctx( per_writer_bytes=0x100, destination_intervals=[(0x2000, 0x100)], @@ -1967,6 +2156,22 @@ def test_scatter_destinations_accept_complete_in_range_intervals(self): 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( @@ -2007,6 +2212,7 @@ def test_scatter_failure_retains_ownership_without_a_positive_fence(self): 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() diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index dee5931b3358..58c8000a6666 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -55,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 @@ -645,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, @@ -655,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 = [ @@ -695,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 = [] diff --git a/tests/unittest/disaggregated/test_receive_ownership_wiring.py b/tests/unittest/disaggregated/test_receive_ownership_wiring.py index 77ebe01ad72d..5127aff3528a 100644 --- a/tests/unittest/disaggregated/test_receive_ownership_wiring.py +++ b/tests/unittest/disaggregated/test_receive_ownership_wiring.py @@ -31,6 +31,7 @@ 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, @@ -108,6 +109,53 @@ def _make_send_task_and_session(*, channel: str = "kv"): 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() @@ -170,7 +218,7 @@ def mark_protocol_conflict(self, key, on_done) -> None: def mark_backend_quiesced(self, key, on_done) -> None: self.backend_quiesced_callbacks.append((key, on_done)) - def retry_settlements(self) -> bool: + def retry_settlements(self, _scope=None) -> bool: return True @@ -240,6 +288,7 @@ def _make_receiver(*, bounce=None, writer_ranks=(WRITER_RANK,)) -> Receiver: 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), @@ -265,6 +314,10 @@ def _make_receiver(*, bounce=None, writer_ranks=(WRITER_RANK,)) -> Receiver: ) 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 @@ -842,12 +895,19 @@ def test_rx_session_cancel_retries_legacy_idle_release_after_task_failure() -> N 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() @@ -998,7 +1058,7 @@ def test_receiver_shutdown_waits_for_admitted_receive_dispatch() -> None: def get_sender_info(_params): discovery_started.set() - assert allow_discovery.wait(timeout=2) + allow_discovery.wait() return peer_info receiver._get_sender_info = get_sender_info @@ -1014,18 +1074,50 @@ def receive() -> None: receive_thread = threading.Thread( target=receive, name="blocked-receive", + daemon=True, ) + shutdown_thread = None receive_thread.start() - assert discovery_started.wait(timeout=1) - - assert receiver.shutdown() is 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()) + 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) - allow_discovery.set() - 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 @@ -1104,6 +1196,208 @@ def test_mamba_receive_uses_direct_path_until_bound_layout_includes_state_bytes( 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() @@ -1359,7 +1653,7 @@ class SourceOwner: source_owner = SourceOwner() source_owner_ref = weakref.ref(source_owner) - with pytest.raises(RuntimeError, match="peer lookup failed"): + with pytest.raises(RuntimeError, match="peer lookup failed") as exc_info: TxSession( request_id=REQUEST_ID, params=_params(), @@ -1371,9 +1665,114 @@ class SourceOwner: 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() @@ -1535,6 +1934,150 @@ def test_sender_failure_is_monotonic_while_sibling_source_access_drains(channel) 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) @@ -1551,6 +2094,7 @@ def test_local_cancel_installs_tombstone_and_retries_settlement() -> None: 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}} @@ -1624,6 +2168,7 @@ 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}} @@ -1664,6 +2209,7 @@ def test_external_backend_fence_coordinates_registry_and_bounce_settlement() -> 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 @@ -1824,6 +2370,7 @@ def invalidate(_agent_name: str) -> None: 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 @@ -1918,6 +2465,7 @@ def send() -> None: 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() @@ -1959,6 +2507,7 @@ def test_sender_retries_failed_worker_local_dealer_close() -> None: 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 diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index a317d2f632c0..2d093d0c34a5 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -49,6 +49,44 @@ class _FakeRequest: 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, shutdown_results: Optional[list[bool]] = None) -> None: self.sweep_count = 0 @@ -163,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, _cleanup_ready: ( - 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, ) ) @@ -288,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. @@ -328,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) @@ -336,6 +643,7 @@ 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: @@ -427,6 +735,166 @@ def test_gen_transfer_status_retains_cancelled_session_until_wait_is_terminal() 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) @@ -507,7 +975,7 @@ def test_async_send_rejects_a_distinct_owner_for_a_live_request_id() -> None: assert transceiver._send_sessions == {rid: session} -def test_context_status_poll_does_not_retire_replacement_owner() -> None: +def test_context_status_poll_retains_owner_and_blocks_replacement() -> None: rid = 21 original = _FakeRequest(request_id=rid) replacement = _FakeRequest(request_id=rid) @@ -529,8 +997,6 @@ def wait_complete(*, blocking: bool) -> WaitResult: return WaitResult.FAILED original_session.wait_complete.side_effect = wait_complete - replacement_session = Mock() - transceiver = object.__new__(KvCacheTransceiverV2) transceiver._session_admission_lock = threading.RLock() transceiver._shutdown_started = False @@ -563,23 +1029,25 @@ def poll_status() -> None: poll_thread.start() assert wait_started.wait(timeout=1) try: - assert transceiver.cancel_request(original) - with transceiver._session_admission_lock: - transceiver._send_sessions[rid] = replacement_session - transceiver._send_reqs[rid] = replacement + 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 == [([], [])] - assert transceiver._send_sessions == {rid: replacement_session} - assert transceiver._send_reqs == {rid: replacement} - replacement_session.close.assert_not_called() + 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_does_not_retire_replacement_owner() -> None: +def test_gen_status_poll_retains_owner_and_blocks_replacement() -> None: rid = 23 original = _FakeRequest(request_id=rid) replacement = _FakeRequest(request_id=rid) @@ -601,8 +1069,6 @@ def wait_complete(*, blocking: bool) -> WaitResult: return WaitResult.FAILED original_session.wait_complete.side_effect = wait_complete - replacement_session = Mock() - transceiver = object.__new__(KvCacheTransceiverV2) transceiver._session_admission_lock = threading.RLock() transceiver._shutdown_started = False @@ -633,20 +1099,22 @@ def poll_status() -> None: poll_thread.start() assert wait_started.wait(timeout=1) try: - assert transceiver.cancel_request(original) - with transceiver._session_admission_lock: - transceiver._recv_sessions[rid] = replacement_session - transceiver._recv_reqs[rid] = replacement + 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 == [([], [], [])] - assert transceiver._recv_sessions == {rid: replacement_session} - assert transceiver._recv_reqs == {rid: replacement} - replacement_session.close.assert_not_called() + 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: @@ -881,6 +1349,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: @@ -1136,6 +1605,130 @@ def test_transceiver_shutdown_retains_owner_when_cleanup_raises(failure_stage: s _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 From 7b4bc96be737abae903ccccb087d53000ee6a952 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:20:14 -0700 Subject: [PATCH 5/5] [None][fix] resolve KV ownership CI failures Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 54 +++++++++++++------ tensorrt_llm/_torch/pyexecutor/py_executor.py | 15 ++++-- .../test_disagg_inflight_cancel_gate.py | 6 +++ .../_torch/executor/test_py_executor.py | 26 +++++---- .../test_py_executor_transfer_termination.py | 13 ++++- tests/unittest/disaggregated/test_bounce.py | 9 ++-- .../test_receive_ownership_wiring.py | 41 ++++++++++++++ .../test_transceiver_bounded_polling.py | 1 + 8 files changed, 125 insertions(+), 40 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index f5533a33ceb5..e63e7f707ce6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -4069,7 +4069,10 @@ def __init__( self._aux_exposed_writer_ranks: set[int] = set() self._aux_publication_closed = not self._need_aux self._aux_result_conflict = False - self._aux_drained = not self._need_aux + # 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 @@ -4245,12 +4248,19 @@ 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._aux_exposed_writer_ranks.add(int(peer_rank)) + self._retain_aux_writer_exposure_locked(peer_rank) return True def _close_aux_publication_locked(self) -> None: @@ -4259,7 +4269,11 @@ def _close_aux_publication_locked(self) -> None: return self._aux_publication_closed = True if not self._kv_tasks: - self._aux_drained = True + # 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] @@ -4283,21 +4297,31 @@ def close_task_publication(self, task: KVRecvTask) -> None: 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: - 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" - ) + self._validate_receive_admission_locked() params = self._base_args.params slice_id = len(self._kv_tasks) task = KVRecvTask( @@ -4427,7 +4451,7 @@ def process_kv_protocol_conflict( # 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._aux_exposed_writer_ranks.add(int(peer_rank)) + self._retain_aux_writer_exposure_locked(peer_rank) self._close_aux_publication_locked() logger.error(detail) @@ -4517,7 +4541,7 @@ def process_aux_agent_result(self, peer_rank: int, status: AgentResult): # 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._aux_exposed_writer_ranks.add(peer_rank) + self._retain_aux_writer_exposure_locked(peer_rank) self._aux_results[peer_rank] = status if not identity_valid: self._aux_result_conflict = True @@ -4544,7 +4568,7 @@ def process_aux_agent_result(self, peer_rank: int, status: AgentResult): 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._aux_exposed_writer_ranks.add(peer_rank) + self._retain_aux_writer_exposure_locked(peer_rank) self._aux_result_conflict = True self._aux_status = TaskStatus.ERROR self._exception = RuntimeError( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f6c65990c046..529d393be981 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -9000,11 +9000,16 @@ def _handle_responses(self, emit_first_iter: bool = True): continue terminal = self._get_transfer_terminal(request) - if (terminal is not None - and self._has_async_transfer_owner(request)): - # A provider already owns the exact logical response, but a - # sibling lease still prevents final teardown. Do not let the - # ordinary response pass create a second terminal result. + 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 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_transfer_termination.py b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py index 5f2f54848822..40866fbb56fe 100644 --- a/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py +++ b/tests/unittest/_torch/executor/test_py_executor_transfer_termination.py @@ -1068,6 +1068,7 @@ def _make_connector_completion_executor(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() @@ -1380,6 +1381,8 @@ def test_connector_owned_cancel_waits_for_connector_completion(transceiver_kind) 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() @@ -1416,6 +1419,8 @@ def test_capability_false_cancel_keeps_anonymous_status_owner_after_connector(): 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() @@ -1585,6 +1590,8 @@ def test_cancelled_context_connector_owner_terminates_exactly_once(force_partial 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: @@ -1749,6 +1756,7 @@ def test_shutdown_discards_post_loop_native_response_without_collective(monkeypa 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) @@ -1832,6 +1840,7 @@ def test_shutdown_waits_for_final_connector_owner_after_native_cancel(monkeypatc 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] @@ -2177,7 +2186,7 @@ def test_shutdown_discards_ownerless_unbuffered_terminal_before_teardown( request = SimpleNamespace(py_request_id=7, is_dummy_request=False) executor = _make_shutdown_executor(events, request) executor._deferred_transfer_terminations = {} - executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + events.transceiver_shutdown.return_value = True executor.async_transfer_manager = _ExactTransferManager(request) executor.active_requests = [request] executor._claim_transfer_terminal( @@ -2248,7 +2257,7 @@ def test_shutdown_local_cleanup_failure_still_enters_connector_consensus(monkeyp request = Mock(py_request_id=7, is_dummy_request=False) executor = _make_shutdown_executor(events, request) executor._deferred_transfer_terminations = {} - executor.kv_cache_transceiver.requires_physical_drain_before_request_release = False + 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 = [] diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 4b3eb854b55f..1808ba934625 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -1181,7 +1181,7 @@ def test_destination_intervals_factory_runs_only_after_allocator_admission(self, 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 = {(0x2080, 0x80), (0x2000, 0x80)} + raw_intervals = {(0x2004, 4), (0x2000, 4)} calls = [] def build_intervals(): @@ -1196,10 +1196,10 @@ def 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, 0x2100),) + assert ctx._destination_intervals == ((0x2000, 0x2008),) raw_intervals.add((0x3000, 8)) - assert ctx._destination_intervals == ((0x2000, 0x2100),) + 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]) @@ -1383,7 +1383,7 @@ def acknowledge(ok): 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 + 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]) @@ -1420,6 +1420,7 @@ def test_settlement_retry_replays_registry_update_until_session_ack(self, monkey 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 diff --git a/tests/unittest/disaggregated/test_receive_ownership_wiring.py b/tests/unittest/disaggregated/test_receive_ownership_wiring.py index 5127aff3528a..16dea6e6fb4c 100644 --- a/tests/unittest/disaggregated/test_receive_ownership_wiring.py +++ b/tests/unittest/disaggregated/test_receive_ownership_wiring.py @@ -520,6 +520,47 @@ def test_aux_protocol_conflict_requires_full_exposed_ledger_or_backend_fence() - 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))) diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index 2d093d0c34a5..d024b3aff143 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -1079,6 +1079,7 @@ def wait_complete(*, blocking: bool) -> WaitResult: 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 = (