fix: ToT KV-pressure crashes — free pages on context drop + break collect_* on starvation#425
Merged
Merged
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
shsym
changed the title
shsym
force-pushed
the
task/679-tot-kv-drop-free
branch
from
June 16, 2026 22:52
6a30ff1 to
6f8c42c
Compare
The host WIT context `drop` handler only deleted the wasm resource-table entry and deferred all KV-page cleanup to the whole-instance DestroyAll on teardown. For a long-lived daemon inferlet (chat-apc serves every request on one instance) that leaks every intermediate context for the instance's lifetime. The tree-of-thought search forks many transient contexts per request — `breadth` children per frontier node, a second fork per node for value scoring, and spent parents. The Active scheduler can mask a slow leak by evicting droppable pages, but a heavy tree forks faster than eviction frees: a new fork's allocation is deferred in `when_allocated` waiting on pages that only the not-yet-freed leaked contexts hold, and the request deadlocks (the RemoteProtocol/ConnectError hang with a healthy driver reported downstream). Free the context eagerly when the guest drops its handle: each WIT Context resource maps 1:1 to a ContextId, so a dropped handle means the guest is done with it. This bounds the live footprint to the active beam and unblocks deferred forks immediately. Named snapshots are unaffected — save/snapshot persist a separate owner:None context id, so destroying the live owned context never touches them. The instance-teardown DestroyAll still cleans up anything left live.
shsym
force-pushed
the
task/679-tot-kv-drop-free
branch
from
June 16, 2026 22:56
6f8c42c to
d6eac10
Compare
The SDK Generator's collect_tokens/collect_text sugars folded an empty forward-pass output as "0 tokens accepted": pie's scheduler swallows a device failure / per-batch timeout / mid-stream KV eviction into an empty Output (Ok, not Err) rather than erroring. With zero accepted tokens the loop never advances tokens_generated and never hits max_tokens or a stop token, so it spins — and the next step issues an empty-input forward pass that hangs the driver. The HTTP layer then closes with no response. This is the tree-of-thought daemon crash under KV pressure (thinking-ON nodes are large, so they pressure KV and provoke mid-stream eviction). The chat-completions loop already guards this case in its own hand-rolled loop, but tree-of-thought generates via collect_text, which had no guard. Add the same starvation predicate to the collect_* sugars: a step whose raw host slots carry no Token slot terminates the loop with a clear error instead of spinning. A natural stop is not starvation — the host still samples Token(stop) (later truncated out of Output::tokens), so a Token slot is present. Unit-tested.
shsym
force-pushed
the
task/679-tot-kv-drop-free
branch
from
June 17, 2026 02:32
d6eac10 to
2714bd2
Compare
Contributor
|
Fix 1 is good. But I don't understand the "starvation" mechanism for 2. How come this can happen? If this is true, then we need to find and fix the root cause.
|
Contributor
|
@shsym Could you open a PR for fix 1? I think this should land in main |
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
2 participants
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.
Two complementary fixes for the tree-of-thought daemon crash under KV pressure.
1.
fix(context): free KV pages on guest context drop (RAII)— 528e8bdThe host WIT context
drophandler only deleted the wasm resource-table entry and deferred all KV-page cleanup to whole-instanceDestroyAll. A long-lived daemon inferlet (one instance serves every request), and ToT forks many transient contexts per request (breadthchildren + a value-scoring fork per node + spent parents). Those dropped forks' pages were never reclaimed → a heavy tree's fork allocations defer inwhen_allocatedwaiting on pages only the leaked contexts hold → deadlock.Free the context eagerly on drop (each WIT
Contextmaps 1:1 to aContextId). Snapshots unaffected (separateowner:Noneid). Verified before/after on real portable-Metal: unfixed deadlocks on a heavy tree, fixed completes 3/3.2.
fix(generation): break collect_* on forward-pass starvation— 2714bd2The direct fix. pie's scheduler folds a device failure / per-batch timeout / mid-stream KV eviction into an empty
Output(Ok, notErr). The SDKGenerator::collect_tokens/collect_textsugars folded that as "0 tokens accepted" — never advancingtokens_generated, never hittingmax_tokens/stop — so they spun, and the next empty-input forward pass hung the driver ("Server disconnected without sending a response", port dead, driver healthy).The chat-completions loop already guards this in its own hand-rolled loop, but ToT generates every node + value score via
collect_text, which had no guard. Add the same starvation predicate (a step whose raw host slots carry noTokenslot) to thecollect_*sugars — terminate with a clear error instead of spinning. A natural stop is not starvation (the host still samplesToken(stop), later truncated out ofOutput::tokens). Unit-tested (4 cases).Tests
cargo test --lib(inferlet SDK): 25 passed incl. 4 starvation-predicate tests.