fix(openai): cancel timed-out generation before releasing its lane - #1150
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe change adds asynchronous cancellation state, propagates request contexts through OpenAI backends and routes, and stops Skippy generation during admission, execution, and token emission. It also separates admitted and active session batch-size queries. ChangesRequest cancellation flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Router
participant OpenAiRequestContext
participant SkippyBackend
participant GenerationTokenBudget
participant GenerationWorker
Router->>SkippyBackend: Start context-aware request
SkippyBackend->>GenerationTokenBudget: Acquire generation permit
GenerationTokenBudget-->>SkippyBackend: Return permit or cancellation
SkippyBackend->>GenerationWorker: Run generation with cancellation token
Router->>OpenAiRequestContext: Cancel on timeout or drop
OpenAiRequestContext-->>GenerationWorker: Notify cancellation
GenerationWorker-->>SkippyBackend: Return result or cancellation
SkippyBackend-->>Router: Return backend response or error
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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 |
|
This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review. |
e0f0fae to
787f8e6
Compare
787f8e6 to
cbcc9c2
Compare
|
@michaelneale @ndizazzo impactful enough bug that I'd like to do a release after we review and merge |
michaelneale
left a comment
michaelneale
left a comment
There was a problem hiding this comment.
Oh that is a nasty one. I bet it bites often and mysteriously.
ndizazzo
left a comment
ndizazzo
left a comment
There was a problem hiding this comment.
Checks out as clean to me
|
@coderabbitai full review |
✅ Action performedFull review finished. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
crates/skippy-server/src/frontend/linear_proposal.rs (1)
658-680: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winCancellation on the first token produces a misleading error and skips the repair path.
The new check at lines 659-662 breaks before
committed_tokens.push(token). If cancellation is observed on the first iteration,committed_tokensstays empty. Control then reaches line 676 and returns"linear proposal classifier committed no target prediction", which discards thecallback_erroryou just set and misreports the cause.The early return also skips
finish_linear_proposal_after_repairat line 689.verify_tokens_sampledhas already advanced the session toposition_after_verification, so the verified rows are never trimmed back on this path.Return the recorded
callback_errorwhen one exists, so cancellation surfaces as"request cancelled"and the empty-commit error stays reserved for the genuine classifier case.🐛 Proposed fix
if committed_tokens.is_empty() { + if let Some(error) = callback_error { + return Err(error); + } return Err(OpenAiError::backend( "linear proposal classifier committed no target prediction", )); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/frontend/linear_proposal.rs` around lines 658 - 680, Update the post-loop handling in the linear proposal flow around the committed_tokens check to return the recorded callback_error when present, including cancellation observed before the first token is committed. Preserve the existing empty-commit error for genuine classifier results, while ensuring the callback-error path continues through finish_linear_proposal_after_repair so verified session rows are trimmed correctly.crates/skippy-server/src/frontend/local_generation.rs (1)
648-658: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winCancellation is handled two different ways inside the same loop iteration.
The loop-top check at lines 648-654 treats cancellation as a graceful stop. It sets
receipt_cancelled = trueand breaks, sofinalize_generation_receiptmarks the observation cancelled and delivers a receipt.The new checks at lines 658 and 712 treat the same condition as a hard error. They return
Err("request cancelled"), sogeneration_succeededis false,receipt_cancelledstays false, and the sink receives anabortinstead of a cancelled receipt.These two paths are a few lines apart and can both observe the same cancellation. Make the proposal-path checks break out of the loop like the loop-top check, so receipt reporting stays consistent.
🐛 Proposed fix
if let Some(config) = self.linear_proposal_ingress.as_ref() && linear_proposals_enabled { - ensure_request_active(request.cancellation)?; + if request + .cancellation + .is_some_and(openai_frontend::CancellationToken::is_cancelled) + { + receipt_cancelled = true; + break; + }Apply the same change at line 712.
Also applies to: 712-712
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/frontend/local_generation.rs` around lines 648 - 658, Update the cancellation handling in the proposal paths guarded by linear_proposals_enabled, including both the visible check and the corresponding check near line 712, to follow the loop-top behavior: set receipt_cancelled and break rather than propagating ensure_request_active errors. Keep cancellation receipt finalization consistent across all checks within the generation loop.
🧹 Nitpick comments (3)
crates/skippy-server/src/frontend/local_generation.rs (2)
364-375: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument why the suffix prefill is chunked here.
RuntimeState::prefillalready chunks internally throughsession.prefill_chunked. The outer loop exists to create cancellation checkpoints between batches, not to satisfy a batch-size constraint. Without a comment, a later change may collapse this back into a single call and remove the cancellation granularity.The
.max(1)guard is correct and required, becauseslice::chunkspanics on a zero chunk size.📝 Proposed comment
let suffix = &prefill_tokens[restored_prefill_tokens..]; + // `RuntimeState::prefill` chunks internally. Chunk again at + // the session batch size so cancellation is observed between + // batches instead of only after the whole suffix completes. let batch_size = runtime .session_batch_size(&session_id) .map_err(openai_backend_error)? .max(1);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/frontend/local_generation.rs` around lines 364 - 375, Add a concise comment above the suffix chunking loop in the prefill restoration flow, near `RuntimeState::prefill`, explaining that the outer chunks provide cancellation checkpoints while `prefill` handles internal chunking; explicitly retain the `.max(1)` guard because `slice::chunks` rejects a zero size.
42-50: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe cancellation-check predicate is duplicated three times under
crates/skippy-server/src/frontend/. This PR introduces two identicalensure_request_activedefinitions plus two inline copies of the same body. The shared root cause is that no module owns this predicate. Define it once in thefrontendmodule and import it at each site.
crates/skippy-server/src/frontend/local_generation.rs#L42-L50: keep this definition but move it to the sharedfrontendmodule and re-import it here.crates/skippy-server/src/frontend/linear_proposal.rs#L772-L780: delete this identical definition and import the shared helper.crates/skippy-server/src/frontend/generation_flow.rs#L61-L63: replace the inline condition withensure_request_active(cancellation)?, and apply the same replacement to the second inline copy at lines 90-92.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/frontend/local_generation.rs` around lines 42 - 50, Centralize the cancellation predicate by moving ensure_request_active from local_generation.rs into the shared frontend module, then import and reuse it in crates/skippy-server/src/frontend/local_generation.rs:42-50, removing the local definition there. Delete the duplicate helper and import the shared one in crates/skippy-server/src/frontend/linear_proposal.rs:772-780. In crates/skippy-server/src/frontend/generation_flow.rs:61-63, replace both inline cancellation checks, including the second copy at lines 90-92, with ensure_request_active(cancellation)?.crates/skippy-server/src/frontend/backend.rs (1)
578-579: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider reusing
run_blocking_generation_workerin the streaming path.
run_generation_streamspawns the blocking worker directly and holds the permit withlet _permit = permit;. This repeats the pattern thatrun_blocking_generation_workernow encapsulates. The helper is generic over the return type, so a()closure fits.The current code is behaviorally correct. This is a consolidation only. Note that the helper awaits the join handle, while this path deliberately does not await, so the helper would need a spawn-only variant or the returned future would need to be spawned.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/frontend/backend.rs` around lines 578 - 579, Consolidate the streaming worker setup in run_generation_stream by reusing the permit-lifetime logic from run_blocking_generation_worker where possible. Preserve the current fire-and-forget behavior: do not await the blocking task directly; add or use a spawn-only variant, or spawn the helper future, while retaining the existing () closure behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/skippy-server/src/runtime_state.rs`:
- Around line 254-256: Rename session_batch_size to admit_session_batch_size and
update its existing call sites, distinguishing batch-size lookups that
intentionally admit a session from eviction-time queries. In
evict_resident_prefix_for_decode_batch, avoid invoking the admitting method
before prefill; use a non-admitting lookup or explicitly ensure the session is
active before querying its batch size so capacity errors are not misinterpreted
as a read failure.
---
Outside diff comments:
In `@crates/skippy-server/src/frontend/linear_proposal.rs`:
- Around line 658-680: Update the post-loop handling in the linear proposal flow
around the committed_tokens check to return the recorded callback_error when
present, including cancellation observed before the first token is committed.
Preserve the existing empty-commit error for genuine classifier results, while
ensuring the callback-error path continues through
finish_linear_proposal_after_repair so verified session rows are trimmed
correctly.
In `@crates/skippy-server/src/frontend/local_generation.rs`:
- Around line 648-658: Update the cancellation handling in the proposal paths
guarded by linear_proposals_enabled, including both the visible check and the
corresponding check near line 712, to follow the loop-top behavior: set
receipt_cancelled and break rather than propagating ensure_request_active
errors. Keep cancellation receipt finalization consistent across all checks
within the generation loop.
---
Nitpick comments:
In `@crates/skippy-server/src/frontend/backend.rs`:
- Around line 578-579: Consolidate the streaming worker setup in
run_generation_stream by reusing the permit-lifetime logic from
run_blocking_generation_worker where possible. Preserve the current
fire-and-forget behavior: do not await the blocking task directly; add or use a
spawn-only variant, or spawn the helper future, while retaining the existing ()
closure behavior.
In `@crates/skippy-server/src/frontend/local_generation.rs`:
- Around line 364-375: Add a concise comment above the suffix chunking loop in
the prefill restoration flow, near `RuntimeState::prefill`, explaining that the
outer chunks provide cancellation checkpoints while `prefill` handles internal
chunking; explicitly retain the `.max(1)` guard because `slice::chunks` rejects
a zero size.
- Around line 42-50: Centralize the cancellation predicate by moving
ensure_request_active from local_generation.rs into the shared frontend module,
then import and reuse it in
crates/skippy-server/src/frontend/local_generation.rs:42-50, removing the local
definition there. Delete the duplicate helper and import the shared one in
crates/skippy-server/src/frontend/linear_proposal.rs:772-780. In
crates/skippy-server/src/frontend/generation_flow.rs:61-63, replace both inline
cancellation checks, including the second copy at lines 90-92, with
ensure_request_active(cancellation)?.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 12b51c0c-13a7-42fe-b78e-ba0d75d25b0e
📒 Files selected for processing (12)
crates/openai-frontend/src/backend.rscrates/openai-frontend/src/guardrails/compact.rscrates/openai-frontend/src/guardrails/mod.rscrates/openai-frontend/src/hooks.rscrates/openai-frontend/src/router.rscrates/skippy-server/src/frontend/admission.rscrates/skippy-server/src/frontend/backend.rscrates/skippy-server/src/frontend/generation_flow.rscrates/skippy-server/src/frontend/linear_proposal.rscrates/skippy-server/src/frontend/local_generation.rscrates/skippy-server/src/frontend/tests/generation.rscrates/skippy-server/src/runtime_state.rs
f6a5b81 to
436ebe8
Compare
436ebe8 to
f119153
Compare
f119153 to
b536993
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/skippy-server/src/frontend/generation_flow/text_generation.rs`:
- Around line 74-77: Update the generation flow around reserve_cancellable and
generate_local_tokens to check cancellation after admission succeeds and before
prefill begins. Add a second cancellation guard in the lower generation layer
immediately before prompt_prefill_sample emission, ensuring cancelled requests
return without stale prefill work or token output while preserving normal
generation behavior.
In `@crates/skippy-server/src/frontend/linear_proposal/execution.rs`:
- Around line 200-203: Update the cancellation handling around the
token-processing loop and the no-committed-tokens branch to preserve the
existing “request cancelled” callback_error when committed_tokens is empty.
Ensure cancellation still runs the required finish_linear_proposal_after_repair
repair or retirement path before returning, rather than replacing the
cancellation error with the generic classifier error.
In `@crates/skippy-server/src/frontend/local_generation/linear_decode.rs`:
- Around line 90-95: Update the no-proposal handling in the linear decode flow
around query_linear_proposal so cancellation is checked again when it returns
NoProposal or DeadlineExceeded. Before returning NotUsed or proceeding to normal
decoding via run_decode_loop and decode_one_token, inspect request.cancellation
and return the existing cancellation error if cancelled.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 84ef8111-e9d8-496d-a225-0c8a00b18178
📒 Files selected for processing (15)
crates/openai-frontend/src/backend.rscrates/openai-frontend/src/guardrails/compact.rscrates/openai-frontend/src/guardrails/mod.rscrates/openai-frontend/src/hooks.rscrates/openai-frontend/src/router.rscrates/skippy-server/src/frontend/admission.rscrates/skippy-server/src/frontend/backend.rscrates/skippy-server/src/frontend/generation_flow/text_generation.rscrates/skippy-server/src/frontend/linear_proposal/execution.rscrates/skippy-server/src/frontend/local_generation/linear_decode.rscrates/skippy-server/src/frontend/local_generation/tests.rscrates/skippy-server/src/frontend/local_generation/token_generation.rscrates/skippy-server/src/frontend/tests/generation.rscrates/skippy-server/src/kv_integration/resident_prefix.rscrates/skippy-server/src/runtime_state/frame_operations.rs
🚧 Files skipped from review as they are similar to previous changes (8)
- crates/skippy-server/src/frontend/tests/generation.rs
- crates/openai-frontend/src/guardrails/compact.rs
- crates/openai-frontend/src/router.rs
- crates/skippy-server/src/frontend/backend.rs
- crates/skippy-server/src/frontend/admission.rs
- crates/openai-frontend/src/backend.rs
- crates/openai-frontend/src/guardrails/mod.rs
- crates/openai-frontend/src/hooks.rs
| let token_budget_reservation = self.generation_token_budget.reserve_cancellable( | ||
| GenerationTokenBudgetRequest::new(prompt_token_ids.len(), max_tokens), | ||
| GENERATION_ADMISSION_TIMEOUT, | ||
| cancellation, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Check cancellation before prefill starts.
reserve_cancellable can return successfully, then the request can be cancelled before generate_local_tokens starts. The local path runs prefill and prepare_decode_state before run_decode_loop checks cancellation. prepare_decode_state can emit prompt_prefill_sample during this window. A timed-out request can therefore perform stale work and emit a token after cancellation. Add cancellation checks before prefill and before prompt-prefill emission in the lower generation layer.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/skippy-server/src/frontend/generation_flow/text_generation.rs` around
lines 74 - 77, Update the generation flow around reserve_cancellable and
generate_local_tokens to check cancellation after admission succeeds and before
prefill begins. Add a second cancellation guard in the lower generation layer
immediately before prompt_prefill_sample emission, ensuring cancelled requests
return without stale prefill work or token output while preserving normal
generation behavior.
3 participants
Problem
When an OpenAI request timed out, dropping the async request future did not stop the blocking generation worker. The request could return and release its concurrency lane while that worker was still using the model runtime.
A subsequent request could then acquire the same lane and overlap with stale work from the timed-out request, breaking the lane's isolation and making timeout behavior dependent on worker timing.
Fix
Cancellation is propagated through the OpenAI and Skippy layers, and the timed-out worker is allowed to stop before its lane permit is released.
This is necessary because a blocking worker continues after its async wrapper is dropped; releasing the lane immediately would only hide the work, not cancel it. Tying lane release to confirmed worker termination preserves the existing concurrency contract.
Validation
OpenAI tests pass (154/154) and Skippy Server tests pass (353/353), including the single-lane timeout regression.
Summary by CodeRabbit
New Features
Bug Fixes
Tests