[https://nvbugs/6104831][fix] Disagg KV transfer: non-blocking polling, structured cancel, bounded quarantine

[yifjiang]

Summary

Disaggregated KV transfer cancellation, timeout, and process-health
contract — a layered fix on top of the analysis in
docs/source/features/disagg-kv-transfer-hang-restart-analysis.md.
Earlier attempts (#13706, #13713, #13728) each fixed one slice of
the problem but still hung in the field. Concretely:

• PR [Draft] Fail closed on unquiesced disagg KV transfer buffers #13706 with Python changes restarted workers on every routine
  per-request timeout — the Python fail-close policy was too aggressive.
• PR [https://nvbugs/6104831][fix] Disaggregated KV transfer: lifecycle, cancellation, and quiescence hardening #13713 / PR [codex] Fail closed on unquiesced disagg KV transfer #13728 still hung even with Python cleanup removed,
  because the C++ status polling could call future.get() on an unready
  worker future, freezing the executor event loop while NIXL/UCX was
  wedged.

This PR separates the four concerns that the earlier PRs mixed:

Responsibility	Owner	Mechanism
Object lifetime	C++ transceiver	shared_ptr<LlmRequest> through the public API + tracked-future vectors.
Resource quiescence	Python executor +	Worker future stays pinned until ready / error / quarantine; Python's
	C++ transceiver	_can_terminate_request_now defers free_resources() until safe.
Request failure	Python executor	Surface a request-level error to the user only after C++ reports a final state.
Process health	C++ transceiver	Bounded quarantine + global progress deadline → isHealthy() == false.

Base: b9ce4b69d12fe5ba65d13893111b1a2ea29413ee. Two commits.

Commit 1 — non-blocking polling, structured cancel, bounded quarantine

C++ — status polling is non-blocking by default

checkContextTransferStatus and checkGenTransferStatus poll worker
futures with wait_for(0ms) and never call future.get() on an
unready entry. When atLeastRequestNum > 0 is requested, the polling
code admits additional ready futures but skips unready ones rather than
blocking the executor thread. New explicit blocking variants exist for
shutdown only:

• drainContextTransferStatus(bool markComplete)
• drainGenTransferStatus()

This is the actual hang fix — see the analysis doc's "Why the C++-only
fix is not sufficient" section.

C++ — structured TransferCancelResult enum

cancelRequestStructured(LlmRequest*) returns a structured outcome:

Value	Caller may release request resources?
NotFound	yes — C++ already reached a final state
AlreadyComplete	yes — worker future is ready
CancelledBeforeAdvertise	yes — no buffer was advertised to the peer
CancelRequestedInFlight	no — worker may still touch buffers
BackendUnhealthy	no — orchestration must restart
NotCancellable	no — retry later

The historical bool cancelRequest is preserved as a backward-compatible
wrapper that returns true only for the first three states.

C++ — bounded quarantine and explicit health signal

When a tracked transfer exceeds kvTransferTimeoutMs and the worker
future is still not ready, the entry is flipped to quarantined: an
error is surfaced to the caller, but the future stays pinned so that
NIXL/UCX cannot write into freed memory. A counter increments. If
quarantined entries exceed mQuarantineBudget or no worker has
reached a final state for longer than mGlobalProgressDeadlineMs,
isHealthy() flips false. getHealth() returns a TransceiverHealth
snapshot for orchestration / metrics. Defaults: budget 16, deadline 60s.

Python — per-request policy, no fail-close

PyExecutor adds _can_terminate_request_now and
_inflight_cancel_requested_ids. _do_terminate_request defers freeing
resources for any request that is still in disagg transfer state or
whose cancel is in flight. Per-request timeouts no longer clear
active_requests, the waiting queue, or set is_shutdown — that was
the #13706 restart-loop trigger and is explicitly forbidden by the
analysis doc.

The timeout/cancel path now uses cancel_request_structured and only
calls _handle_errors([request]) when the structured result is one of
CancelledBeforeAdvertise / AlreadyComplete / NotFound. For
CancelRequestedInFlight and BackendUnhealthy, the request stays in
active_requests, the cancel is recorded, and the next iteration
re-polls C++.

Python — ADP synchronized response flush (port from #13112)

The base commit predates #13112. Without it, calling _enqueue_responses
inside _end_transfer_and_maybe_terminate deadlocks with ADP because
only the owning DP rank reaches that point. Transfer-completion responses
are now buffered in _pending_transfer_responses and flushed at
synchronised executor-loop points where every DP rank participates in
the tp_gather collective. Flush is symmetric: empty + enable_attention_dp=True
still calls _enqueue_responses([]) so the collective completes.

Commit 2 — shared_ptr<LlmRequest> lifetime refactor

Equivalent to the lifetime hardening that #13713 applied. Pin the
LlmRequest C++ object through the entire transceiver path so it
outlives every worker thread access regardless of when Python's
_terminate_request drops its pybind reference.

• respondAndSendAsync, requestAndReceiveSync, requestAndReceiveAsync,
  cancelRequest, and cancelRequestStructured all take
  std::shared_ptr<LlmRequest>.
• TrackedFuture::request is now shared_ptr<LlmRequest>; the
  shared_ptr in mSenderFutures / mRequesterFutures keeps the
  LlmRequest pinned for as long as the transceiver tracks the worker
  future.
• trtGptModelInflightBatching.cpp passes the shared_ptr directly
  instead of .get()-stripping it.
• The nanobind trampoline (PyCacheTransceiver) signatures match.
  <nanobind/stl/shared_ptr.h> (already included in the binding file)
  provides a Python-aware shared_ptr that pins the wrapper alive while
  C++ holds it.

The Python _can_terminate_request_now guard stays in place — but its
job is no longer object lifetime (the shared_ptr handles that). It is
now strictly about resource quiescence: free_resources() releases
KV blocks back to the pool, and the transport may still be writing
into them after the LlmRequest C++ object exists. Termination must
still wait for the structured cancel result to signal safety.

Testing

• tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py
  — 14 unit tests covering the deferred-terminate guard, the structured-
  cancel decision tree (including the NotFound-after-final-state
  case), the unhealthy-backend pin, and ADP flush symmetry. They run
  without GPU or MPI so they fail fast in pre-merge CI:

  $ python3 -m pytest tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py -v
  ============================== 14 passed in 0.07s ==============================
  

• Existing cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
  exercises the public API only and is not affected by the
  mSenderFutures/mRequesterFutures storage refactor.

Test plan

• pytest tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py
• bot run --extra-stage "<cache-transceiver multi-GPU stages>" (C++ multi-GPU tests)
• 1P1D cancel-burst harness — confirm no hang, no restart loop, no
  buffer reuse (the analysis-doc regression scenarios).
• ADP gen-first regression test
  (accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first).
• Field repro of NVBug 6104831 with MALLOC_PERTURB_=85 to confirm
  the LlmRequest UAF class is closed by commit 2.

Documentation

• docs/source/features/disagg-kv-transfer-hang-restart-analysis.md
  — the analysis that anchors this fix direction.
• docs/source/features/disagg-kv-transfer-session-lifecycle.md
  — the contract this PR establishes (status polling, cancel result,
  quarantine budget, ADP flush, Python policy, object lifetime).

What this PR is not

• It does not add a BufferIndexHolder::poison() API (#13728's
  approach). Quarantine is tracked at the transceiver level via a
  bounded counter + global progress deadline, which is sufficient to
  flip isHealthy() and let orchestration restart cleanly without
  per-buffer poison plumbing.