fix(frontend): keep completion errors across full aggregation - #12786
fix(frontend): keep completion errors across full aggregation#12786KrishnanPrash wants to merge 1 commit into
Conversation
Signed-off-by: Krishnan Prashanth <kprashanth@nvidia.com>
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| async fn aggregate_completion_response( | ||
| stream: impl futures::Stream<Item = Annotated<NvCreateCompletionResponse>>, | ||
| parsing_options: ParsingOptions, | ||
| request_id: &str, | ||
| ) -> Result<NvCreateCompletionResponse, ErrorResponse> { | ||
| let backend_error = Arc::new(std::sync::OnceLock::new()); | ||
| let first_error = backend_error.clone(); | ||
| let stream = stream.inspect(move |response| { | ||
| if let Some(error) = extract_backend_error_if_present(response) { | ||
| // Keep the first error; it is the one that ended generation. | ||
| let _ = first_error.set(error); | ||
| } | ||
| }); | ||
|
|
||
| let aggregated = | ||
| NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options).await; | ||
|
|
||
| if let Some((message, status)) = backend_error.get() { | ||
| let error_response = backend_error_response(message.clone(), *status); | ||
| tracing::warn!(request_id, ?error_response, "Backend error detected"); | ||
| return Err(error_response); | ||
| } | ||
|
|
||
| aggregated.map_err(|e| { | ||
| tracing::error!(request_id, "Failed to fold completions stream: {e:?}"); | ||
| ErrorMessage::internal_server_error(&format!( | ||
| "Failed to fold completions stream for {request_id}" | ||
| )) | ||
| }) |
There was a problem hiding this comment.
🔍 Error frames no longer short-circuit; a backend that errors but keeps the stream open will hang the request
The removed preflight (check_for_backend_error) returned as soon as the first non-annotation event was an error frame, so the HTTP response was produced without waiting for the stream to terminate. aggregate_completion_response (lib/llm/src/http/service/openai.rs:965-993) only inspects the captured error after from_annotated_stream has folded the entire stream to completion. If a backend emits a typed error frame and then fails to close the stream (or is slow to do so), the client now blocks until the stream ends instead of getting an immediate 4xx. This is fine for backends where an error frame is terminal (which is what the new tests model), but it is worth confirming that all worker paths close the stream right after emitting event: "error".
Was this helpful? React with 👍 or 👎 to provide feedback.
WalkthroughCompletion aggregation now inspects all events and preserves typed backend HTTP statuses. Single and batch completions use the shared helper. HTTP tests cover invalid-argument errors before and after output for scalar and array prompts. ChangesCompletion error handling
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
lib/llm/tests/http-service.rs (1)
1628-1651: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a success-path test for the shared aggregation helper.
Both tests assert the error path.
aggregate_completion_responsenow inspects every event, so a frame that is misclassified as an error would fail a normal request. A test that serves an engine yielding only valid chunks and asserts HTTP 200 for scalar and array prompts would guard that behavior.🤖 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 `@lib/llm/tests/http-service.rs` around lines 1628 - 1651, Add a success-path test alongside test_completions_surface_backend_invalid_argument and test_completions_surface_backend_invalid_argument_after_first_chunk that serves an engine producing only valid completion chunks, invokes both scalar and array prompts, and asserts HTTP 200 responses. Exercise the shared aggregate_completion_response path and verify valid frames are not classified as errors.
🤖 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.
Nitpick comments:
In `@lib/llm/tests/http-service.rs`:
- Around line 1628-1651: Add a success-path test alongside
test_completions_surface_backend_invalid_argument and
test_completions_surface_backend_invalid_argument_after_first_chunk that serves
an engine producing only valid completion chunks, invokes both scalar and array
prompts, and asserts HTTP 200 responses. Exercise the shared
aggregate_completion_response path and verify valid frames are not classified as
errors.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 607351cb-2f6b-4349-b4a8-be6246c3537d
📒 Files selected for processing (2)
lib/llm/src/http/service/openai.rslib/llm/tests/http-service.rs
![dynamo-review-agent[bot]](https://avatars.githubusercontent.com/u/201626793?s=60&v=4){kind=small}
| }); | ||
|
|
||
| let aggregated = | ||
| NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options).await; |
There was a problem hiding this comment.
aggregate_completion_response records backend errors but waits for NvCreateCompletionResponse::from_annotated_stream to drain the entire merged stream before returning them, so a batch request with one immediate backend error keeps the other prompt streams generating and delays the 4xx response. Fix: short-circuit on the first backend error and drop the remaining stream instead of checking backend_error only after aggregation completes.
🤖 AI Fix
In lib/llm/src/http/service/openai.rs, rewrite aggregate_completion_response to poll the completion stream directly, return Err(backend_error_response(...)) immediately when a backend error frame is seen, and otherwise aggregate non-error Annotated<NvCreateCompletionResponse> items to the same final response.
| let backend_error = Arc::new(std::sync::OnceLock::new()); | ||
| let first_error = backend_error.clone(); | ||
| let stream = stream.inspect(move |response| { | ||
| if let Some(error) = extract_backend_error_if_present(response) { |
There was a problem hiding this comment.
Every non-streaming completion chunk now passes through extract_backend_error_if_present, whose normal-data path serializes NvCreateCompletionResponse to JSON even though that type cannot be a top-level {message, code} error payload, adding avoidable per-token CPU and allocations. Fix: use a completion-specific error extractor that skips generic data-payload JSON conversion for normal completion frames.
🤖 AI Fix
In lib/llm/src/http/service/openai.rs, add extract_completion_backend_error_if_present(event: &Annotated<NvCreateCompletionResponse>) that preserves typed event == "error" handling but does not run serde_json::to_value on event.data, and call it from aggregate_completion_response.
1 participant
Overview:
PR #12706 made
/v1/completionsreturn the backend's own HTTP status for a typed error such asBackend(InvalidArgument). It did this with a preflight that reads only the first stream event, so a backend that emits a normal chunk and then fails still returned HTTP 500 with the message replaced byFailed to fold completions stream.This PR removes the preflight and captures the typed error during aggregation instead, so an error at any position keeps its status and its message. The same change removes the batch barrier and the per-prompt boxing that #12706 introduced.
Repro
cargo test -p dynamo-llm --test http-service \ test_completions_surface_backend_invalid_argument_after_first_chunkThe endpoint shape that first showed the bug:
Before
After
Details:
lib/llm/src/http/service/openai.rsaggregate_completion_response. It taps the raw events on their way into the aggregator, keeps the first typed backend error in aOnceLock, and returns that error with its own status. The generic 500 now covers only a fold failure with no typed error. Bothcompletions_singleandcompletions_batchcall it, so the two handlers share one aggregation path.check_completion_batch_streams. Itstry_join_allheld every prompt stream until the slowest produced its first frame, so a fast prompt could not report metrics or releasehttp_queue_guarduntil the slow prompt answered.completions_batchmerges with a plainstream::select_allagain.BoxedCompletionResponseStream. It only gave the streaming and non-streaming arms the same type, at the cost of one allocation per prompt plus dynamic dispatch on every poll of the streaming path. Streaming is now identical to pre-fix(frontend): preserve completion backend error status #12706.warnwith a structurederror_responsefield instead of aterrorwith the response formatted into the message.lib/llm/tests/http-service.rsPOST /v1/completionsagainst a liveHttpService. One engine emits the typed error as its first frame, the other emits a chunk first. Each test asserts the single-prompt and the array-prompt shape, so both handlers are covered.Trade-off: without a preflight, a failing batch drains every prompt stream before it answers, as it did before #12706. In exchange, no sibling generation is abandoned mid-flight.
Related Issues
🚫 This PR is NOT linked to an issue:
Summary by CodeRabbit