[Bugfix][KV Offloading] Offload last block at request finish and prevent reuse race

[Alex-ai-future]

Purpose

Fix: last KV block that fills at request finish time is never offloaded, causing prefix cache misses for that block. Also fixes a data corruption race when the block is reused.

Bug principle: _build_store_jobs runs during schedule() — at that point the finishing token (EOS) hasn't been processed yet, so the last block is still partial and skipped. request_finished is called after EOS is appended, but it had no store kickoff logic — the last block was silently dropped.

Consequence chain: last block not offloaded → next request with the same prefix misses that block → extra prefill compute + no KV reuse for the final block.

Fix: Use orozery's simplified approach — request_finished calls update_offload_keys() to refresh block hashes and always keeps req_status alive. The next step's _build_store_jobs processes finished_req_ids (via itertools.chain) and creates store jobs for the now-complete blocks through the normal code path. This matches the approach described in #47107: build store jobs at request_finished, fence blocks, flush at next build_connector_meta.

Worker-side fence fix: Since store jobs for finished requests are deferred (created in the next schedule() but not submitted until the step after), when the block gets reused before submission, wait() in handle_preemptions finds the job in jobs_to_flush but it hasn't been submit_store'd yet → wait() is a no-op → block gets overwritten. Fix: in handle_preemptions, submit jobs_to_flush stores from store_jobs before calling wait(). This ensures wait() actually fences the blocks. Only block-reuse scenarios are affected; normal stores remain deferred to preserve token generation performance.

Related issue: #41704 (same class of bug in SimpleCPUOffloadScheduler, fixed by #41777 — different component, not a duplicate).

Test Plan

Test	Verifies
test_last_block_offloaded_at_request_finish	EOS fills last block → request_finished keeps req_status alive → next step's _build_store_jobs creates store job → block is offloaded
test_two_groups_full_and_sliding_window	Updated touch_calls expectations (2→4) due to _build_store_jobs processing finished requests

.venv/bin/python -m pytest tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py -v

Test Result

91 passed

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

[Alex-ai-future]

@orozery PTAL — this implements your request_finished TODO plus the handle_preemptions wait-before-submit fix.

[orozery]

Thanks @Alex-ai-future ! This looks pretty good.
However, I think we can make it simpler and thinner change:

In request_finished, just add req_status.update_offload_keys() and keep req_status alive when there are storable blocks remaining.

In _build_store_jobs - chain finished_req_ids into the existing loop:

for req_id in itertools.chain(
    scheduler_output.num_scheduled_tokens,
    scheduler_output.finished_req_ids or (),
):
    req_status = self._req_status.get(req_id)
    if req_status is None:
        continue
    req = req_status.req

    if req.is_finished():
        num_tokens_after_batch = req.num_tokens
    else:
        num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
        num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens

    # ... existing loop body unchanged ...

    # At end of loop body, after job creation:
    if req.is_finished():
        for bid in non_sliding_window_block_ids or ():
            self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
            if bid in self._current_batch_allocated_block_ids:
                self._current_batch_jobs_to_flush.add(job_id)
        if not req_status.transfer_jobs:
            del self._req_status[req_id]

[orozery]

Suggested change

	entry = kv_connector_metadata.store_jobs.get(job_id)
	if entry is not None:
	assert isinstance(entry.src_spec, GPULoadStoreSpec)
	success = self.worker.submit_store(
	job_id, entry.src_spec, entry.dst_spec
	)
	assert success
	del kv_connector_metadata.store_jobs[job_id]
	entry = kv_connector_metadata.store_jobs.pop(job_id, None)
	if entry is not None:
	self._unsubmitted_store_jobs.append(
	(job_id, entry.src_spec, entry.dst_spec)
	)

[orozery]

The changes to the test utils seems overly complicated.
If we can't test with the existing utils, I suggest we just skip adding tests for this.

[jonathanc-n]

I will close #47107 in favour of this

[Alex-ai-future]

Thanks @orozery, the simplified approach is much cleaner! All updated per your suggestions — finished_req_ids chaining in _build_store_jobs, .pop() into _unsubmitted_store_jobs, and test utils significantly trimmed. Kept one minimal test for the core fix, happy to drop it if you prefer.

[orozery]

@Alex-ai-future can we revert to using just _build_store_jobs without introducing _build_store_jobs_for_req?
I want to make the diff minimal.

[Alex-ai-future]

@orozery Updated per your feedback to minimize the diff:

• Inlined _build_store_jobs_for_req back into _build_store_jobs
• Removed store_jobs variable in build_connector_meta
• Added _maybe_cleanup_finished_req helper for duplicate cleanup

One simplification I made: removed the has_storable check in request_finished. Now _build_store_jobs handles all cleanup on the next step. The cleanup delay is one step (negligible) but saves ~15 lines.

Let me know if this looks OK or if I should restore the original check.

[orozery]

Can we please avoid changes in this file?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Alex-ai-future.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Alex-ai-future]

@orozery I removed the changes in test utils according to your request.

But to sum up, I made two changes without permission.

1. Remove the _req_status in the previous request_finished, which led to the change of the declaration cycle of all _req_status.

2. Reconstructing the logic of _build_store_jobs may increase the level of judgment but avoid duplicate code and make the logic clearer.

[orozery]

@Alex-ai-future I liked the previous diff with _maybe_cleanup_finished_req helper better.
Now you change a lot of indentations.

[Alex-ai-future]

Thank you @orozery, having worked on the offloading scheduler recently, I found its complexity quite challenging — particularly the interplay between _jobs, transfer_jobs, and _block_id_to_pending_jobs as manually-synchronized reverse indices. PR #48102 showed this can cause real bugs.

Would there be interest in extracting job state management into a dedicated registry? This way the scheduler file would focus on scheduling semantics rather than also owning data-structure bookkeeping — which might make it easier to reason about and extend.

[orozery]

This can be removed.
If we reached here, req_status.transfer_jobs is non-empty, and so self._maybe_cleanup_finished_req is a no-op.

[orozery]

Thank you @orozery, having worked on the offloading scheduler recently, I found its complexity quite challenging — particularly the interplay between _jobs, transfer_jobs, and _block_id_to_pending_jobs as manually-synchronized reverse indices. PR #48102 showed this can cause real bugs.

Would there be interest in extracting job state management into a dedicated registry? This way the scheduler file would focus on scheduling semantics rather than also owning data-structure bookkeeping — which might make it easier to reason about and extend.

I agree OffloadingConnectorScheduler is quite complex.
As a first refactoring step, I suggested to abstract the store job creation behind some API.
See #42050.
Unfortunately I haven't had time to review and iterate this design.
But it's definitely something I want to push for.

[mergify]

Hi @Alex-ai-future, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[Alex-ai-future]

Thank you @orozery, having worked on the offloading scheduler recently, I found its complexity quite challenging — particularly the interplay between _jobs, transfer_jobs, and _block_id_to_pending_jobs as manually-synchronized reverse indices. PR #48102 showed this can cause real bugs.
Would there be interest in extracting job state management into a dedicated registry? This way the scheduler file would focus on scheduling semantics rather than also owning data-structure bookkeeping — which might make it easier to reason about and extend.

I agree OffloadingConnectorScheduler is quite complex. As a first refactoring step, I suggested to abstract the store job creation behind some API. See #42050. Unfortunately I haven't had time to review and iterate this design. But it's definitely something I want to push for.

OK, I'm willing to spend some time dealing with this recently.

[orozery]

Thanks @Alex-ai-future !