[Rust Frontend][gRPC] Surface engine generation errors (#56405)
Signed-off-by: Alec Flowers <aflowers@nvidia.com> Co-authored-by: Codex <noreply@openai.com>
Alec committed
Sep 11, 2026 at 15:31 UTC
988d9b6777d077f843cd2164a222ac8535d36ed2
3 files changed
+93
-32
rust/src/server/src/grpc/convert.rs
+18
-17
@@ -391,13 +391,14 @@ pub fn to_sequence_output(
391
logprobs: Option<&DecodedLogprobs>,
392
finished: Option<&Finished>,
393
opts: &ResponseOpts,
394
-) -> pb::SequenceOutput {
394
+) -> Result<pb::SequenceOutput, Status> {
395
+ let finish_info = finished.map(|f| to_finish_info(f, token_ids)).transpose()?;
396
let (lp_values, rank_values, candidates) = match logprobs {
397
Some(lp) if opts.output_logprobs => output_logprobs_to_proto(lp),
398
_ => (vec![], vec![], vec![]),
399
};
400
400
- pb::SequenceOutput {
401
+ Ok(pb::SequenceOutput {
402
index: 0, // TODO: multi-sequence (n > 1) not supported
403
text: if opts.output_text {
404
delta.to_string()
@@ -413,11 +414,11 @@ pub fn to_sequence_output(
414
logprobs: lp_values,
415
ranks: rank_values,
416
candidate_tokens: candidates,
416
- finish_info: finished.map(|f| to_finish_info(f, token_ids)),
417
- }
417
+ finish_info,
418
+ })
419
}
420
420
-fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
421
+fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result<pb::FinishInfo, Status> {
422
use pb::finish_info::FinishReason as PbFinishReason;
423
424
let (finish_reason, stop_reason) = match &finished.finish_reason {
@@ -438,18 +439,17 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {
439
(PbFinishReason::Stop as i32, sr)
440
}
441
FinishReason::Length => (PbFinishReason::Length as i32, None),
441
- FinishReason::Abort | FinishReason::Error | FinishReason::Repetition(_) => {
442
- (PbFinishReason::Aborted as i32, None)
443
- }
442
+ FinishReason::Error => return Err(Status::internal("engine failed during generation")),
443
+ FinishReason::Abort | FinishReason::Repetition(_) => (PbFinishReason::Aborted as i32, None),
444
};
445
446
- pb::FinishInfo {
446
+ Ok(pb::FinishInfo {
447
num_output_tokens: finished.usage.output_token_count as u32,
448
finish_reason,
449
stop_reason,
450
kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
451
ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct),
452
- }
452
+ })
453
}
454
455
// ========================================================================================
@@ -709,7 +709,7 @@ mod tests {
709
let fin = finished(FinishReason::Stop(None));
710
let token_ids = [1_u32, 2, 3, 151643];
711
712
- let info = to_finish_info(&fin, &token_ids);
712
+ let info = to_finish_info(&fin, &token_ids).expect("finish info");
713
714
assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
715
assert_eq!(info.stop_reason, Some(PbStopReason::EosTokenId(151643)));
@@ -719,7 +719,7 @@ mod tests {
719
fn eos_stop_with_empty_token_ids_leaves_stop_reason_unset() {
720
let fin = finished(FinishReason::Stop(None));
721
722
- let info = to_finish_info(&fin, &[]);
722
+ let info = to_finish_info(&fin, &[]).expect("finish info");
723
724
assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
725
assert_eq!(info.stop_reason, None);
@@ -730,7 +730,7 @@ mod tests {
730
let fin = finished(FinishReason::Stop(Some(StopReason::TokenId(42))));
731
// Terminal token list should be ignored when an explicit stop reason is
732
// present.
733
- let info = to_finish_info(&fin, &[7, 42]);
733
+ let info = to_finish_info(&fin, &[7, 42]).expect("finish info");
734
735
assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
736
assert_eq!(info.stop_reason, Some(PbStopReason::StopTokenId(42)));
@@ -740,7 +740,7 @@ mod tests {
740
fn explicit_stop_string_is_preserved() {
741
let fin = finished(FinishReason::Stop(Some(StopReason::Text("</stop>".into()))));
742
743
- let info = to_finish_info(&fin, &[1, 2, 3]);
743
+ let info = to_finish_info(&fin, &[1, 2, 3]).expect("finish info");
744
745
assert_eq!(info.finish_reason, PbFinishReason::Stop as i32);
746
assert_eq!(
@@ -753,7 +753,7 @@ mod tests {
753
fn length_finish_has_no_stop_reason() {
754
let fin = finished(FinishReason::Length);
755
756
- let info = to_finish_info(&fin, &[1, 2, 3]);
756
+ let info = to_finish_info(&fin, &[1, 2, 3]).expect("finish info");
757
758
assert_eq!(info.finish_reason, PbFinishReason::Length as i32);
759
assert_eq!(info.stop_reason, None);
@@ -763,7 +763,7 @@ mod tests {
763
fn abort_finish_is_mapped_to_aborted() {
764
let fin = finished(FinishReason::Abort);
765
766
- let info = to_finish_info(&fin, &[]);
766
+ let info = to_finish_info(&fin, &[]).expect("finish info");
767
768
assert_eq!(info.finish_reason, PbFinishReason::Aborted as i32);
769
assert_eq!(info.stop_reason, None);
@@ -778,7 +778,8 @@ mod tests {
778
..Default::default()
779
};
780
781
- let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts);
781
+ let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts)
782
+ .expect("sequence output");
783
784
let finish = out.finish_info.expect("finish_info should be present");
785
assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32);
rust/src/server/src/grpc/inference.rs
+16
-15
@@ -241,12 +241,6 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
241
.instrument(request_span.clone())
242
.await
243
.map_err(|error| log_text_error(&request_span, started_at, "collection", error))?;
244
- info!(
245
- parent: &request_span,
246
- elapsed_ms = started_at.elapsed().as_millis() as u64,
247
- "gRPC inference request completed"
248
- );
249
-
244
// Build the single aggregated response.
245
let prompt_info = convert::to_prompt_info(
246
&collected.prompt_token_ids,
@@ -267,6 +261,11 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
261
collected.logprobs.as_ref(),
262
Some(&finish_info),
263
&response_opts,
264
+ )?;
265
+ info!(
266
+ parent: &request_span,
267
+ elapsed_ms = started_at.elapsed().as_millis() as u64,
268
+ "gRPC inference request completed"
269
);
270
271
Ok(Response::new(pb::GenerateResponse {
@@ -331,19 +330,21 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
330
logprobs,
331
},
332
finished,
334
- }) => Ok(pb::GenerateResponse {
333
+ }) => convert::to_sequence_output(
334
+ &decoded.text,
335
+ &token_ids,
336
+ logprobs.as_ref(),
337
+ finished.as_deref(),
338
+ &response_opts,
339
+ )
340
+ .map(|outputs| pb::GenerateResponse {
341
prompt_info: None,
336
- outputs: Some(convert::to_sequence_output(
337
- &decoded.text,
338
- &token_ids,
339
- logprobs.as_ref(),
340
- finished.as_deref(),
341
- &response_opts,
342
- )),
342
+ outputs: Some(outputs),
343
}),
344
};
345
346
- if tx.send(response).await.is_err() {
346
+ let failed = response.is_err();
347
+ if tx.send(response).await.is_err() || failed {
348
break;
349
}
350
}
rust/src/server/src/grpc/tests.rs
+59
@@ -680,6 +680,65 @@ async fn unary_generate_returns_collected_text() {
680
server_task.abort();
681
}
682
683
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
684
+#[serial]
685
+async fn engine_error_finish_returns_internal_for_unary_and_streaming() {
686
+ for streaming in [false, true] {
687
+ for prefix in [vec![], vec![b'h' as u32]] {
688
+ let mut output_specs = Vec::new();
689
+ if !prefix.is_empty() {
690
+ output_specs.push((prefix.clone(), None));
691
+ }
692
+ output_specs.push((vec![], Some(EngineCoreFinishReason::Error)));
693
+ let (mut client, server_task, engine_task) =
694
+ grpc_test_server(b"engine-grpc-error", output_specs).await;
695
+ let request = pb::GenerateRequest {
696
+ request_id: "test-engine-error".to_string(),
697
+ model: "test-model".to_string(),
698
+ prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
699
+ stopping: Some(pb::StoppingCriteria {
700
+ max_new_tokens: 10,
701
+ ..Default::default()
702
+ }),
703
+ response: Some(pb::ResponseOptions {
704
+ output_token_ids: true,
705
+ ..Default::default()
706
+ }),
707
+ ..Default::default()
708
+ };
709
+
710
+ let status = if streaming {
711
+ let mut stream =
712
+ client.generate_stream(request).await.expect("start stream").into_inner();
713
+ let mut received_tokens = Vec::new();
714
+ let status = loop {
715
+ match stream.message().await {
716
+ Ok(Some(response)) => {
717
+ if let Some(output) = response.outputs {
718
+ assert!(
719
+ output.finish_info.is_none(),
720
+ "error became a finish event"
721
+ );
722
+ received_tokens.extend(output.token_ids);
723
+ }
724
+ }
725
+ Ok(None) => panic!("engine error ended as a successful stream"),
726
+ Err(status) => break status,
727
+ }
728
+ };
729
+ assert_eq!(received_tokens, prefix);
730
+ status
731
+ } else {
732
+ client.generate(request).await.expect_err("engine error must fail the RPC")
733
+ };
734
+ assert_eq!(status.code(), tonic::Code::Internal);
735
+
736
+ engine_task.await.expect("mock engine task");
737
+ server_task.abort();
738
+ }
739
+ }
740
+}
741
+
742
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
743
#[serial]
744
async fn unary_generate_with_token_ids_prompt() {