[None][fix] PyExecutor: non-fatal KV-cache block-pool exhaustion guard

[Thachnh]

Problem

Under capacity_scheduler_policy: MAX_UTILIZATION + enable_block_reuse (e.g. Nemotron-3-Ultra-550B, FP8 KV, 256K), a saturated KV-cache block pool makes the C++ block manager raise a fatal assertion from inside prepare_resources:

• kvCacheManager.cpp:2480 — "Can't allocate new blocks for window size 262152. No free blocks left." (WindowBlockManager::allocateBlock)
• evictionPolicy.cpp:160 — "No free block found. This shouldn't happen!" (LRUEvictionPolicy::getFreeBlock)

The throw escapes the (previously unguarded) resource_manager.prepare_resources(scheduled_batch) call in the PyExecutor event loop and kills the worker thread on every rank: the engine stops stepping while the process stays up, so health/metrics go silent until the pod is force-killed → respawn → next burst kills it again.

This recurs because MAX_UTILIZATION's accounting deliberately over-subscribes (the three existing accounting patches 4504b7882 / 0ea3acc95 / a33649516 narrow but don't close it for this config). GUARANTEED_NO_EVICT would close it structurally but under-admits against the 262K window for a mostly-short workload, so it's not acceptable here.

Fix

Opt-in, mirrors TRTLLM_NAN_GUARD_NONFATAL: with TRTLLM_KV_EXHAUSTION_NONFATAL=1 (default off → unchanged fatal behavior), when the pool is exhausted we reject the request(s) that could not be allocated and keep stepping. Keeps MAX_UTILIZATION's concurrency for short-request traffic; degrades the rare over-subscription event from a pod-wide outage to a single graceful per-request reject. Python-only (no C++ rebuild).

• resource_manager.py: is_kv_cache_exhaustion_error (matches the assert signatures) + env gate, and a guard around add_sequence_batch/add_token in V1 KVCacheManager.prepare_resources. Offending requests are removed from the batch and collected (not freed here).
• scheduler.py: ScheduledRequests.kv_cache_rejected_requests channel.
• py_executor.py: collect after prepare_resources; drain at the top of the next iteration (fail via the request-scoped _handle_errors(charge_budget=False) path) in _executor_loop and _executor_loop_overlap.

Why it's safe

• removeSequence is idempotent + a clean no-op on an unknown/partially-onboarded sequence → free_resources is safe, no double-free.
• The alloc assert is a pre-allocation capacity check (hasFreeBlocks before mutating) → no half-allocated sequence → prepare_resources stays consistent after a caught throw.
• Free is deferred one iteration so no in-flight overlap forward can reference freed blocks.
• The drain runs in rank lockstep when enabled (process-wide env gate), so the attention-DP response gather inside _handle_errors is paired across ranks.
• PP loop intentionally not guarded (pipeline-depth in-flight micro-batches make a 1-iteration-deferred free unsafe); only the V1 KVCacheManager path (the crash path) is guarded.

Validation

• All changed files py_compile clean; classifier logic proven (failing→passing).
• In-image (pytrtllm:v1.3.0rc18-nemotron3ultra-kvguard): tensorrt_llm imports clean, classifier correct, ScheduledRequests field + PyExecutor methods present, gate defaults off.
• Unit test added: tests/unittest/_torch/executor/test_resource_manager.py::test_is_kv_cache_exhaustion_error_classification.
• Pending: multi-GPU soak on Nemotron-3-Ultra with TRTLLM_KV_EXHAUSTION_NONFATAL=1 under long-context burst (confirm graceful reject vs crash; confirm attention-DP pairing).

🤖 Generated with Claude Code