[Bugfix][KV Offload] Keep stepping until a reset_cache flush is delivered, and make reset_cache re-entrant

[sohom-cs]

Purpose

vLLM can clear its KV offload store on demand. RL training loops do this whenever they update the model weights (pause with cache clearing, or sleep), and there is also an HTTP endpoint for it. If the reset lands just as the last request's KV is still being copied to CPU memory, two things go wrong. First, the engine goes idle before it tells the workers to wait for the copies the reset just abandoned. Second, a second reset before any new request arrives crashes the engine on an assertion. Back-to-back resets are a normal pattern in RL loops, so this can take down the inference server in the middle of a training run. This PR keeps the engine stepping until that wait reaches the workers, and lets a second reset fold into the first instead of asserting.

Concretely, OffloadingConnectorScheduler.reset_cache() moves every in-flight job id into _current_batch_jobs_to_flush and deliberately leaves it set (offloading/scheduler.py:2026). Workers have to wait on those jobs before a store after the reset reuses their CPU chunks. Two things go wrong after that:

1. The flush set can sit undelivered. has_pending_push_work() (:1833) returns bool(self._jobs) or manager.has_pending_work(), and reset_cache() has just cleared _jobs. If the reset lands after the last request finished while its final store was still in flight, Scheduler.has_requests() and EngineCore.has_work() both go false. The engine then stops stepping, so build_connector_meta(), the only place jobs_to_flush is sent to the workers (:1815), is not called until some unrelated request shows up.
2. A second reset in that window kills the engine. reset_cache() starts with assert not self._current_batch_jobs_to_flush (:1995). Scheduler.reset_prefix_cache(reset_connector=True) calls reset_connector_cache() every time. It is reached from pause_generation(clear_cache=True) and sleep (EngineCore._reset_caches, vllm/v1/engine/core.py:858, the RL weight-update path), from POST /reset_prefix_cache?reset_external=true and from LLM.reset_prefix_cache(reset_connector=True). EngineCore drains all queued utility calls before it steps again. So two resets with no step in between are a normal sequence, and the second one raises AssertionError inside _handle_client_request.

Fix:

• has_pending_push_work() also returns true while _current_batch_jobs_to_flush is non-empty. The engine then takes one step, the flush set goes out in that step's metadata, and it is cleared as usual.
• reset_cache() no longer asserts that the flush set is empty. A second reset adds its ids to the pending set. The two other "not in the middle of a step" asserts stay.

Diff: +10 / −3 in scheduler.py.

Not a duplicate: two open PRs touch this code, and neither makes either change:

• [Core][KVConnector] Keep the offloaded KV tier across pause/sleep #57810 (keep the offloaded tier across pause/sleep) changes the preemption flush in build_connector_meta.
• [KV Offload] Expose OffloadingManager.lock()/unlock() and service the tiering control plane off the scheduler thread #58168 (draft) wraps manager.has_pending_work() in the new manager lock.

Both will textually conflict with this PR in the same functions. I'll rebase onto whichever lands first. The added check reads only scheduler-local state, so it needs no lock under #58168.

Test Plan

python -m pytest tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py -q -k reset_cache
python -m pytest tests/v1/kv_connector/unit/offloading_connector/ -q
pre-commit run --from-ref origin/main --to-ref HEAD

The new test, test_reset_cache_flush_is_delivered_when_idle_and_reset_is_reentrant, does four things:

1. Finishes a request with its store still in flight.
2. Resets twice with no step in between.
3. Checks that the engine still reports work.
4. Checks that the next schedule() carries exactly those job ids in jobs_to_flush and that nothing is pending afterwards.

Test Result

With the fix (macOS arm64, CPU, rebased on 32cc3f1ea):

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py -k reset_cache ... 5 passed
tests/v1/kv_connector/unit/offloading_connector/ ... 460 passed, 2 skipped

On main (fix reverted, new test kept):

>       assert not self._current_batch_jobs_to_flush
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py:1995: AssertionError
1 failed, 4 passed (-k reset_cache)

pre-commit (--from-ref origin/main --to-ref HEAD) and mypy-3.12 (manual stage): clean.

This is a scheduler control-path change; model outputs are unaffected, so no evals are needed.

Related: one of a few independent fixes from an audit of the KV transfer paths (CPU offload, NIXL, P2P): #59096, #59102, #59325, #59329. None depends on another; they can be reviewed and merged in any order.

AI assistance

I used an AI coding assistant (Claude) to audit this code path, write the fix and write the test. I reviewed every changed line and ran the tests above myself. The commit carries a Co-authored-by trailer, as AGENTS.md asks.

[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.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀