[Bugfix][KV Offloading] Defer request finalization until final store

[Palaiologos1453]

Purpose

Fixes #49635.

OffloadingConnectorScheduler.request_finished() currently notifies the manager before the scheduler prepares a finished request's deferred final store. TieringOffloadingManager can then delete its request state, and the next prepare_store() crashes EngineCore with KeyError.

This change:

• builds all final store jobs before signaling manager.on_request_finished();
• tracks whether the finish signal was issued so same-step completions and cache reset clean up exactly once;
• updates the scheduler test runner to drain finished_req_ids and final transfer completions.

The manager assertion remains strict so future lifecycle regressions are surfaced rather than silently dropping an offload.

Test Plan

• Exercise final full and partial chunks under synchronous and asynchronous scheduling.
• Verify prepare_store() precedes on_request_finished() while complete_store() may follow it.
• Exercise pending stores and cache reset around the deferred notification.
• Run the complete offloading scheduler unit-test file.

Test Result

• python -m ruff check vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py tests/v1/kv_connector/unit/offloading_connector/utils.py: passed.
• python -m ruff format --check ...: passed.
• python -m compileall -q ...: passed.
• All 114 offloading scheduler tests collect locally. Full execution is unavailable in the local Windows environment because it lacks the vLLM native CUDA extension.
• The reporter independently ran 114/114 tests on the affected base and verified both the original reproducer and 16/100-request bursts without EngineCore restarts: [Bugfix][KV Offloading] Defer request finalization until final store #49671 (comment)

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[orozery]

Thanks for the quick fix @Palaiologos1453 !
I actually made a fix, I think in some places it may be a bit cleaner.
Can you take a look and adapt things you think are better?
orozery@fa07027

I actually had to adapt some tests as well.
Enabling CI to verify.

[nilig]

Thanks for the quick fix. Verified it against the original #49635 deployment and reproducer: redeployed gpt-oss-120b with offload_prompt_only:false, block size 64, a P2P secondary tier, nightly-4080263b, the generic_p2p overlay at 145a460c, and this PR's scheduler.py.

The single ~2,900-token completion that previously killed EngineCore ~35-60 seconds after completion produced no restart for 100 seconds after each of two independent runs, nor through the rest of the session (further bursts, ~15 minutes). The 16-way same-prefix burst that previously returned 16x EngineDeadError completed 16/16, p50 0.27s. A 100-request burst completed 100/100. Cross-pod prefix pulls continue to work with the fix (seed 1.20s versus pull 0.19s, external prefix-cache hits covering ~the full prefix), and there were no tracebacks across the tested engines.

Non-blocking: wonder if we could also harden TieringOffloadingManager.prepare_store() against a future ordering regression. A missing _req_state entry should degrade to a skipped offload with a warning rather than an engine-wide EngineDeadError. The state check needs to happen after _maybe_process_finished_jobs() but before primary_tier.prepare_store(), since checking afterward could leave primary-tier allocation or bookkeeping without a corresponding completion. The retrieved state could then be reused for the pending_primary_stores and request_level_tiers accesses.

[nilig]

@orozery verified your variant (fa07027) the same way as the PR head: the #49635 cluster reproducers stay clean (0 restarts through the crash windows, 16/16 and 100/100 bursts, cross-pod pulls intact), and its test file passes 114/114 on our older nightly-4080263b base - the harness adaptation in utils.py also resolves the two base-sensitive failures we saw with the PR head's tests. The non-blocking prepare_store() hardening note above applies to either variant.

[orozery]

Thanks @nilig !
I would not remove assertions. It is the mechanism that allows us to discover bugs.

[Palaiologos1453]

Thanks @orozery, I adapted the cleaner structure in ecf6bc52af: finalization now happens centrally after _build_store_jobs(), cleanup is gated by finished_signaled, and the test runner drains finished_req_ids plus final completions. I also updated the affected expectations and added the explicit prepare-store-before-finish ordering test.

The only part not copied from fa07027d is the FINISHED_ABORTED token-count special case, because current main already contains the newer allocated-block bound from #49146. I also kept the manager assertion strict, per your comment, rather than adding a missing-state fallback.

[orozery]

Thanks @Palaiologos1453 !

[Palaiologos1453]

I investigated the remaining failures in Buildkite #79938:

• v1-core-plus-kv-plus-metrics: 15 failures in v1/kv_offload/tiering/test_fs_tier.py, all caused by [Errno 22] Invalid argument while writing temporary files on agent h200-ci-2. The PR diff only contains the offloading scheduler and its unit-test harness; the same V1 job passed in #79927 after retry, and passed on main #79931 on h200-ci-4.
• distributed-compile-plus-comm-4-gpus and distributed-compile-plus-rpc-tests-2-gpus: both fail compile/fullgraph/test_basic_correctness.py::test_compile_correctness[test_setting3] because EngineCore reports Available KV cache memory: -0.45 GiB / no available cache blocks. This is also reproduced by the corresponding main CI run.
• multi-modal-processor-cpu-4: one external download timed out fetching Roboto-Regular.ttf from paddlenlp.bj.bcebos.com.
• pytorch-compilation-unit-tests-h100: Non-uniform page sizes: 2112 vs 18432 in compile/h100/test_startup.py, also reproduced on main #79931.
• language-models-tests-hybrid-2: the two failures for pfnet/plamo-2-1b (token mismatch and missing Plamo2ForCausalLM.generate) are reproduced verbatim by main #79931.

The KV-offload-specific LM Eval jobs (1xH200, 2xH100, 4xH100) passed in #79938, and the reporter's production reproducer remains clean. Could the six unrelated/transient failures be retried? Our account does not have Buildkite retry permission.