fix(scheduler): check host write-back capacity before mutating retrac…#688
Open
alexps9 wants to merge 9 commits into
Open
fix(scheduler): check host write-back capacity before mutating retrac…#688alexps9 wants to merge 9 commits into
alexps9 wants to merge 9 commits into
Conversation
…t state
Scheduler::scheduleRetract() called TakeFirstPages()/Insert<Device>() to
move a request's local KV pages into the shared prefix-cache tree
*before* checking EnsureCapacityByEvict<Host>() for the corresponding
write-back. When that host-capacity check failed, the function bailed
out via `return {}` — but the mutation had already happened, and the
ScheduleRetractEvent that would reconcile the request's
device_node_ref_ with the newly inserted tree node was never
constructed or applied. This leaves the request internally
inconsistent (its local allocator already shrunk, the tree already
owning pages the request's own bookkeeping doesn't know about), which
can later surface as an unrelated "OwnedPages::TakeFirst: count out of
range" crash or a permanently leaked device page.
Every other scheduleXxx() in this file (scheduleDecode,
scheduleDecodeFromRetracted, ...) already follows a strict
"check every capacity gate first, mutate only once all gates pass"
rule. Bring scheduleRetract() in line with that pattern: compute the
host match/capacity requirement first, and only take/insert the device
pages once EnsureCapacityByEvict<Host>() succeeds.
Added a regression test (RetractHostCapacityExhaustedSuite) that
forces a retract to hit an exhausted host pool and asserts the device
page ledger is fully consistent afterwards. Verified the new test
fails against the previous ordering (0 pages reclaimed instead of 1)
and passes with the fix; the full C++ suite (408 tests) passes.
Signed-off-by: alexps9 <alexps9@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 488e9b595e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Scheduler::scheduleRetract()moves a request's local KV pages into the shared device prefix cache before verifying that enough host capacity is available for the corresponding write-back:request->TakeFirstPages(alloc_count); kv_prefix_cache_.Insert<ResourceType::Device>(...);If
EnsureCapacityByEvict<ResourceType::Host>()subsequently fails, the function returns without creating or applying theScheduleRetractEventthat updates the request'sdevice_node_ref_.This leaves the request and prefix cache out of sync:
This can cause a later
OwnedPages::TakeFirst: count out of rangefailure or leave device pages permanently unreachable.Other scheduling paths in this file perform all capacity checks before mutating allocator or cache state.
scheduleRetract()was the exception.Fix
Reorder
scheduleRetract()so that it:EnsureCapacityByEvict<ResourceType::Host>().The host-side match is independent of the device insertion, since
Insert<ResourceType::Device>()only modifies the device resource state.Testing
Added:
The test exhausts the host page pool, forces the host-capacity check to fail, and verifies that all device pages are reclaimed after the requests finish or abort.
Test plan
cmake --build . --target tokenspeed_scheduler_testsmainand passes with this changetokenspeed_scheduler_testssuite passes