[BugFix] scheduler: Fix resuming of preempted requests after async load

[orozery]

The correct flow is for resuming preempted requests is:
RequestStatus.PREEMTPED -> resumed req in SchedulerOutput

However, for requests going through async loading after preemption this becomes:
RequestStatus.PREEMTPED -> RequestStatus.WAITING_FOR_REMOTE_KVS -> RequestStatus.WAITING
then, the scheduler incorrectly flags it as a new request in SchedulerOutput, instead of a resumed request.

This PR fixes the scheduler to correctly resume preempted requests after being async loaded.
RequestStatus.PREEMTPED -> RequestStatus.WAITING_FOR_REMOTE_KVS -> RequestStatus.PREEMTPED -> resumed

---

Note

Fixes incorrect classification of resumed requests after async KV loading.

• In scheduler.schedule(), when a request exits WAITING_FOR_REMOTE_KVS, set status to PREEMPTED if num_preemptions > 0 (otherwise WAITING), so resumed requests are emitted as resumed (not new) in SchedulerOutput.
• Test updates: parameterize KV preemption tests with is_async, add _step_until_kv_transfer_finished(...), and assert correct scheduling/output behavior for async KV transfers (including cached vs new reqs counts).

^{Written by Cursor Bugbot for commit 25476ad33e86c13746f894393ae7a2c9d129ec83. This will update automatically on new commits. Configure here.}

---

Note

Ensures preempted requests that underwent async KV loading are resumed correctly instead of being reclassified as new.

• In scheduler.schedule(), when leaving WAITING_FOR_REMOTE_KVS, set status=PREEMPTED if num_preemptions>0 (else WAITING), so resumed requests appear in scheduled_cached_reqs as resumed.
• Extend tests: parameterize KV preemption scenarios with is_async, add _step_until_kv_transfer_finished(...), and assert correct waiting/running transitions and cached vs new request counts.
• Minor doc tweak: clarify Request.num_preemptions meaning.

^{Written by Cursor Bugbot for commit 399053a497be3ccfc1bb17179a156c1bbb1bb8ea. This will update automatically on new commits. Configure here.}

---

Note

^{Cursor Bugbot is generating a summary for commit cb37857d209dbdca4f080484dcbd7a78aedd7dfc. Configure here.}

---

Note

Fixes misclassification of resumed requests after async KV loading.

• In scheduler.schedule(), when leaving WAITING_FOR_REMOTE_KVS, set status=PREEMPTED if num_preemptions>0 (else WAITING) so resumed requests appear in scheduled_cached_reqs.
• Tests: parameterize KV preemption with is_async, add _step_until_kv_transfer_finished(...), and assert correct waiting/running transitions and cached vs new request counts.
• Minor doc tweak: clarify meaning of Request.num_preemptions.

^{Written by Cursor Bugbot for commit cb37857d209dbdca4f080484dcbd7a78aedd7dfc. This will update automatically on new commits. Configure here.}

---

Note

Fixes misclassification of resumed requests following async KV transfers.

• In scheduler.schedule(), when leaving WAITING_FOR_REMOTE_KVS, set status=PREEMPTED if num_preemptions>0 (else WAITING), so resumed requests are emitted as resumed in scheduled_cached_reqs.
• Tests: parameterize KV preemption path with is_async, add _step_until_kv_transfer_finished(...), and assert correct waiting/running transitions and cached vs new request counts.
• Minor doc tweak: clarify Request.num_preemptions meaning.

^{Written by Cursor Bugbot for commit 3162ea59039b916a9628a22d6c45f4a4bd0106d1. This will update automatically on new commits. Configure here.}

---

Note

Fixes misclassification of resumed requests after async KV loading.

• In scheduler.schedule(), when leaving WAITING_FOR_REMOTE_KVS, set status to PREEMPTED if num_preemptions>0 (else WAITING), so resumed requests appear in scheduled_cached_reqs rather than scheduled_new_reqs.
• Extend tests: parameterize KV preemption path with is_async, add _step_until_kv_transfer_finished(...), and assert correct waiting/running transitions and cached vs new request counts.
• Clarify Request.num_preemptions comment to reflect "times preempted".

^{Written by Cursor Bugbot for commit a2e70f0. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Code Review

This pull request addresses a bug in the scheduler where preempted requests that undergo asynchronous KV cache loading were incorrectly treated as new requests instead of resumed requests. The fix introduces a new state tracking set, preempted_requests_being_async_loaded, to correctly manage the status of these requests. When a preempted request is scheduled for async loading, it's added to this set. Upon completion of the async load, its status is correctly restored to PREEMPTED. The changes also handle the cleanup of this tracking set when requests are aborted. The implementation is clear, well-contained, and effectively resolves the state transition issue. The logic appears robust and correct.

[NickLucche]

I am admittedly not too familiar with the full extent of the preemption lifecycle, perhaps @markmc can chime in on this?

In the meantime could we add a unit test to trace through this case? Thanks @orozery !

[orozery]

In the meantime could we add a unit test to trace through this case? Thanks @orozery !

@NickLucche I've added a unit test.

[njhill]

Thanks @orozery, good catch with this!

I think there is a simpler fix that can avoid introducing new state, pls see the inline comment.

[orozery]

I think there is a simpler fix that can avoid introducing new state, pls see the inline comment.

Thanks! I've updated to use your simplified fix.

[njhill]

Thanks @orozery!

[orozery]

@njhill I rebased and pushed a squashed commit