// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 use std::collections::BTreeMap; use std::error::Error as StdError; use std::sync::Arc; use std::time::Instant; use anyhow::{Error, Result}; use futures::{stream, stream::StreamExt}; use crate::{ http::service::metrics::Metrics, model_card::ModelDeploymentCard, protocols::{ TokenIdType, common::{ extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}, llm_backend::{BackendOutput, LLMEngineOutput, PreprocessedRequest}, preprocessor::MultimodalData, timing::RequestPhase, }, }, session_affinity::explicit_target, }; use dynamo_runtime::engine::Data; use dynamo_runtime::error::{self, BackendError, DynamoError, ErrorType}; use dynamo_runtime::metrics::prometheus_names::frontend_service; use dynamo_runtime::pipeline::{ AsyncEngineContext, AsyncEngineContextProvider, Context, ManyOut, Operator, PipelineOperator, ResponseStream, ServerStreamingEngine, SingleIn, async_trait, attach_first_response_guard, }; use dynamo_runtime::protocols::annotated::Annotated; /// Accessors the migration RetryManager needs from a response chunk. /// `token_ids` lets it replay already-delivered tokens; `worker_trace_link` /// lets it stamp the failed worker's span onto the next attempt's /// `migration_link`. pub(crate) trait HasTokenIds { fn token_ids(&self) -> &[TokenIdType]; fn worker_trace_link(&self) -> Option<&crate::protocols::common::preprocessor::TraceLink>; } impl HasTokenIds for BackendOutput { fn token_ids(&self) -> &[TokenIdType] { &self.token_ids } fn worker_trace_link(&self) -> Option<&crate::protocols::common::preprocessor::TraceLink> { self.worker_trace_link.as_ref() } } impl HasTokenIds for LLMEngineOutput { fn token_ids(&self) -> &[TokenIdType] { &self.token_ids } fn worker_trace_link(&self) -> Option<&crate::protocols::common::preprocessor::TraceLink> { self.worker_trace_link.as_ref() } } /// Check if an error chain indicates the request should be migrated. fn is_migratable(err: &(dyn StdError + 'static)) -> bool { const MIGRATABLE: &[ErrorType] = &[ ErrorType::CannotConnect, ErrorType::Disconnected, ErrorType::ConnectionTimeout, // a stalled/frozen worker's stream-inactivity timeout surfaces // as ResponseTimeout (push_router fault detection quarantines the worker // via the same signal); migrate instead of hanging to the stream timeout. ErrorType::ResponseTimeout, ErrorType::Backend(BackendError::EngineShutdown), // A truncated stream from a departed worker is recoverable by failover. ErrorType::Backend(BackendError::StreamIncomplete), // One overloaded worker: another may have room. Pool-wide exhaustion is // ResourceExhausted below and stays non-migratable. ErrorType::WorkerOverloaded, ]; const NON_MIGRATABLE: &[ErrorType] = &[ErrorType::Cancelled, ErrorType::ResourceExhausted]; error::match_error_chain(err, MIGRATABLE, NON_MIGRATABLE) } /// Whether a worker-scoped failure can be retried without violating an explicit route. /// /// The phase is read after the failed attempt because disaggregated routing updates the /// shared request tracker as it moves from prefill to decode. Session affinity does not /// populate these explicit routing fields, so an invalidated affinity binding remains /// eligible for migration. fn is_migratable_for_request( request: &PreprocessedRequest, err: &(dyn StdError + 'static), ) -> bool { if !is_migratable(err) { return false; } let allows_phase = |phase| match explicit_target(request, phase) { Ok(None) => true, Ok(Some(target)) => { tracing::debug!( ?phase, worker_id = target.worker_id, dp_rank = ?target.dp_rank, "Migration disabled for explicitly pinned worker" ); false } Err(error) => { tracing::warn!(?phase, %error, "Migration disabled for invalid explicit worker target"); false } }; let Some(tracker) = request.tracker.as_ref() else { // Without a tracker there is no authoritative current phase. Any explicit // phase target may be the hard pin that just failed, so do not guess. return [ RequestPhase::Prefill, RequestPhase::Decode, RequestPhase::Aggregated, ] .into_iter() .all(allows_phase); }; let phase = tracker.phase(); allows_phase(phase) } pub struct Migration { migration_limit: u32, max_seq_len: Option, model_name: Arc, metrics: Arc, } impl Migration { pub fn new( migration_limit: u32, max_seq_len: Option, model_name: String, metrics: Arc, ) -> Arc { tracing::debug!( "model {} migration limit {} max_seq_len {:?}", model_name, migration_limit, max_seq_len ); Arc::new(Self { migration_limit, max_seq_len, model_name: Arc::new(model_name), metrics, }) } pub fn from_mdc( mdc: &ModelDeploymentCard, migration_limit: u32, max_seq_len: Option, metrics: Arc, ) -> Arc { Self::new( migration_limit, max_seq_len, mdc.display_name.clone(), metrics, ) } /// Wrap as a `PipelineOperator` over the given response type to /// disambiguate between the `Operator` impls on `Migration` since /// the response type doesn't appear in the struct. #[allow(clippy::type_complexity)] pub(crate) fn into_operator_for( self: &Arc, ) -> Arc< PipelineOperator< SingleIn, ManyOut>, SingleIn, ManyOut>, >, > where Resp: Data + HasTokenIds, { Operator::into_operator(self) } } #[async_trait] impl Operator< SingleIn, ManyOut>, SingleIn, ManyOut>, > for Migration where Resp: Data + HasTokenIds, { async fn generate( &self, request: SingleIn, next: ServerStreamingEngine>, ) -> Result>> { // NOTE: Keep the migration operator in the request path at limit zero. The limit controls // replacement attempts; RetryManager continues to own the initial attempt uniformly. let (preprocessed_request, context) = request.transfer(()); let engine_ctx = context.context(); let engine_ctx_ = engine_ctx.clone(); let session_affinity = context .get_optional::(SESSION_AFFINITY_CONTEXT_KEY) .map_err(Error::msg)? .map(|session_id| session_id.as_ref().clone()); let retry_manager = RetryManager::build( engine_ctx, context.metadata().clone(), preprocessed_request, next, self.migration_limit, self.max_seq_len, self.model_name.clone(), self.metrics.clone(), session_affinity, ) .await?; let response_stream = stream::unfold(retry_manager, move |mut retry_manager| async move { retry_manager .next() .await .map(|response| (response, retry_manager)) }) .fuse(); Ok(ResponseStream::new(Box::pin(response_stream), engine_ctx_)) } } struct MigrationEvent { migration_type: &'static str, started_at: Instant, } impl MigrationEvent { fn new(migration_type: &'static str) -> Self { Self { migration_type, started_at: Instant::now(), } } } struct RetryManager where Resp: Data + HasTokenIds, { context: Arc, metadata: BTreeMap, request: PreprocessedRequest, session_affinity: Option, next_generate: ServerStreamingEngine>, next_stream: Option>>, retries_left: u64, max_seq_len: Option, model_name: Arc, metrics: Arc, /// Latest worker span pointer seen on the active stream; stamped as /// `migration_link` on the next retry. Populated by `track_response`. last_worker_link: Option, } impl RetryManager where Resp: Data + HasTokenIds, { #[allow(clippy::too_many_arguments)] pub async fn build( context: Arc, metadata: BTreeMap, mut preprocessed_request: PreprocessedRequest, next: ServerStreamingEngine>, mut retries_left: u32, max_seq_len: Option, model_name: Arc, metrics: Arc, session_affinity: Option, ) -> Result { // TODO: prompt_embeds take precedence over replayed token_ids. Disable migration for // embedding prompts until a retry can represent an embedding-based continuation. // TODO: Define a replay-capability contract for attempt-local decoder and sampler state. // Stop-string matching, generated-token penalties, and thinking-token budgets are not // currently checkpointed across workers. // Disable migration for structured-output (guided-decoding) requests. // Inference backends initialize the guided-decoding FSM (finite state machine) fresh // for every new request and only advance it on newly-generated tokens, not on // context/prompt tokens. Migrating a partial structured-output response would replay // already-generated tokens as context, causing the FSM to restart from the schema // root and producing duplicated or nested JSON. This applies to all backends // (vLLM, SGLang, TRT-LLM) equally. Propagate the error cleanly instead. if preprocessed_request .sampling_options .guided_decoding .is_some() { if retries_left > 0 { tracing::warn!( "Guided-decoding request: migration disabled — FSM state is not transferable (applies to all backends)" ); } retries_left = 0; } if preprocessed_request.sampling_options.n.unwrap_or(1) > 1 { if retries_left > 0 { tracing::warn!( "n>1 request: migration disabled - per-choice generation state is not transferable" ); } retries_left = 0; } if retries_left > 0 { preprocessed_request.migration_state = Some(Default::default()); } let mut slf = Self { context, metadata, request: preprocessed_request, session_affinity, next_generate: next, next_stream: None, retries_left: u64::from(retries_left) + 1, // +1 to account for the initial attempt max_seq_len, model_name, metrics, last_worker_link: None, }; slf.new_stream(None).await?; slf.exceed_max_seq_len(0); // disable migration if prompt len > max_seq_len Ok(slf) } pub async fn next(&mut self) -> Option> { loop { let response_stream = match self.next_stream.as_mut() { Some(stream) => stream, None => { tracing::error!("next() called with next_stream is None - should not happen"); return Some(Annotated::from_error("next_stream is None")); } }; if let Some(response) = response_stream.next().await { // Check if this is a migratable error that should trigger stream recreation. if let Some(err) = response.error.as_ref() && is_migratable_for_request(&self.request, err) { tracing::warn!(error = %err, "Stream disconnected, recreating stream"); self.metrics.inc_migration_ongoing_request(&self.model_name); let migration_event = MigrationEvent::new(frontend_service::migration_type::ONGOING_REQUEST); // NOTE: Delegate exhaustion to new_stream so retry accounting has one owner. // When no replacement is established, preserve the triggering stream error. if let Err(err) = self.new_stream(Some(migration_event)).await { tracing::warn!(error = ?err, "Cannot recreate stream"); } else { continue; } } self.track_response(&response); return Some(response); } return None; } } async fn new_stream(&mut self, mut migration_event: Option) -> Result<()> { let mut response_stream: Option>>> = None; while self.retries_left > 0 { self.retries_left -= 1; // Once any chunks have arrived from a previous attempt, stamp // that worker's span as `migration_link` so the next worker's // span renders an OTel Link back to it. Guarded so the initial // attempt doesn't clobber a `migration_link` set upstream. if let Some(link) = self.last_worker_link.as_ref() { self.request.migration_link = Some(link.clone()); } let mut request = Context::with_id_and_metadata( self.request.clone(), self.context.id().to_string(), self.metadata.clone(), ); if let Some(session_affinity) = self.session_affinity.as_ref() { request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_affinity.clone()); } self.context.link_child(request.context()); if self.context.is_stopped() || self.context.is_killed() { tracing::debug!("Abort creating new stream after context is stopped or killed"); self.record_migration_outcome( migration_event.as_ref(), frontend_service::migration_outcome::CANCELLED, ); return Err(DynamoError::builder() .error_type(ErrorType::Cancelled) .message(format!( "Context id {} is stopped or killed", self.context.id() )) .build() .into()); } let source_guards = self .request .multi_modal_data .as_ref() .into_iter() .flat_map(|media| media.values()) .flatten() .filter_map(|item| match item { MultimodalData::Decoded(descriptor) => descriptor.source_storage.clone(), _ => None, }) .collect::>(); if !source_guards.is_empty() { attach_first_response_guard(&mut request, Arc::new(source_guards)); } response_stream = Some(self.next_generate.generate(request).await); if let Some(err) = response_stream.as_ref().unwrap().as_ref().err() && is_migratable_for_request(&self.request, err.as_ref()) { tracing::warn!(error = %err, "Creating new stream, retrying"); if migration_event.is_none() { migration_event = Some(MigrationEvent::new( frontend_service::migration_type::NEW_REQUEST, )); } // Preserve the existing per-attempt metric contract. self.metrics.inc_migration_new_request(&self.model_name); continue; } break; } match response_stream { Some(Ok(next_stream)) => { self.record_migration_outcome( migration_event.as_ref(), frontend_service::migration_outcome::SUCCESS, ); self.next_stream = Some(next_stream); Ok(()) } Some(Err(err)) => { let outcome = if error::match_error_chain(err.as_ref(), &[ErrorType::Cancelled], &[]) { frontend_service::migration_outcome::CANCELLED } else { frontend_service::migration_outcome::FAILURE }; self.record_migration_outcome(migration_event.as_ref(), outcome); Err(err) // should propagate original error if any } None => { self.record_migration_outcome( migration_event.as_ref(), frontend_service::migration_outcome::FAILURE, ); Err(Error::msg("Migration limit exhausted")) } } } fn record_migration_outcome(&self, migration_event: Option<&MigrationEvent>, outcome: &str) { if let Some(event) = migration_event { self.metrics.observe_migration_duration( &self.model_name, event.migration_type, outcome, event.started_at.elapsed(), ); } } fn track_response(&mut self, response: &Annotated) { if self.retries_left == 0 { return; } let llm_engine_output = match response.data.as_ref() { Some(output) => output, None => return, }; // Capture the worker's engine.generate span pointer so a future // migration retry can render an OTel Link back to it. The adapter // stamps this on the first non-empty chunk; subsequent chunks may // also carry it. Keep the most-recently-seen value. if let Some(link) = llm_engine_output.worker_trace_link() { self.last_worker_link = Some(link.clone()); } let token_ids = llm_engine_output.token_ids(); let output_len = u32::try_from(token_ids.len()).unwrap_or(u32::MAX); if self.exceed_max_seq_len(output_len) { return; } // NOTE: A zero remaining token budget is not retry-budget exhaustion. Backends remain // authoritative for deciding how generation terminates at max_tokens. if let Some(max_tokens) = self.request.stop_conditions.max_tokens { self.request.stop_conditions.max_tokens = Some(max_tokens.saturating_sub(output_len)); } if let Some(min_tokens) = self.request.stop_conditions.min_tokens { self.request.stop_conditions.min_tokens = Some(min_tokens.saturating_sub(output_len)); } for token_id in token_ids.iter() { self.request.token_ids.push(*token_id); } } /// Returns `true` if the tracked request token length plus `new_output_len` /// exceeds the configured max_seq_len, in which case migration is disabled. fn exceed_max_seq_len(&mut self, new_output_len: u32) -> bool { if let Some(max_seq_len) = self.max_seq_len { let total_len = self.request.token_ids.len() as u32 + new_output_len; if total_len > max_seq_len { tracing::warn!( "Sequence length {} exceeds migration max_seq_len {}, \ disabling migration", total_len, max_seq_len ); self.metrics .inc_migration_max_seq_len_exceeded(&self.model_name); self.retries_left = 0; // disable migration return true; } } false } } #[cfg(test)] mod tests { use super::*; use crate::http::service::metrics::Metrics; use crate::protocols::common::{ GuidedDecodingOptions, OutputOptions, SamplingOptions, StopConditions, preprocessor::RoutingHints, timing::RequestTracker, }; use dynamo_runtime::error::{DynamoError, ErrorType}; use dynamo_runtime::pipeline::AsyncEngine; use dynamo_runtime::pipeline::context::Controller; use dynamo_runtime::protocols::maybe_error::MaybeError; use std::sync::atomic::{AtomicU32, Ordering}; use tokio::sync::mpsc; const TEST_MODEL: &str = "test-model"; fn migration_duration_count(metrics: &Metrics, migration_type: &str, outcome: &str) -> u64 { metrics.get_migration_duration_sample_count(TEST_MODEL, migration_type, outcome) } // a stalled/frozen worker's stream-inactivity timeout surfaces as // ErrorType::ResponseTimeout (push_router fault detection). It must be // migratable so the request fails over instead of hanging to the stream // timeout. A StreamIncomplete backend error (a truncated stream from a // departed worker) is likewise migratable. #[test] fn stall_and_incomplete_stream_errors_are_migratable() { let response_timeout = DynamoError::builder() .error_type(ErrorType::ResponseTimeout) .message("backend response inactivity timeout") .build(); assert!( is_migratable(&response_timeout), "ResponseTimeout (stalled worker) must be migratable" ); let stream_incomplete = DynamoError::builder() .error_type(ErrorType::Backend(BackendError::StreamIncomplete)) .message("stream ended before completion") .build(); assert!( is_migratable(&stream_incomplete), "StreamIncomplete (truncated stream from departed worker) must be migratable" ); } // Guard: genuinely non-migratable errors stay non-migratable. #[test] fn cancelled_and_exhausted_are_not_migratable() { for et in [ErrorType::Cancelled, ErrorType::ResourceExhausted] { let err = DynamoError::builder().error_type(et).message("x").build(); assert!(!is_migratable(&err), "{et:?} must not be migratable"); } } fn migratable_error(error_type: ErrorType) -> DynamoError { DynamoError::builder() .error_type(error_type) .message("worker failed") .build() } #[tokio::test] async fn explicit_worker_pin_blocks_migration_for_worker_failures() { let tracker = Arc::new(RequestTracker::new()); let mut request = create_mock_request(1); request.tracker = Some(tracker.clone()); request.routing = Some(RoutingHints { backend_instance_id: Some(7), ..Default::default() }); for phase in [ RequestPhase::Aggregated, RequestPhase::Prefill, RequestPhase::Decode, ] { let permit = tracker.set_phase(phase).await; for error_type in [ErrorType::Disconnected, ErrorType::WorkerOverloaded] { let error = migratable_error(error_type); assert!( !is_migratable_for_request(&request, &error), "backend pin must block {error_type:?} migration during {phase:?}" ); } drop(permit); } } #[tokio::test] async fn phase_specific_pin_only_blocks_its_matching_phase() { let error = migratable_error(ErrorType::WorkerOverloaded); let prefill_tracker = Arc::new(RequestTracker::new()); let mut prefill_pinned = create_mock_request(1); prefill_pinned.tracker = Some(prefill_tracker.clone()); prefill_pinned.routing = Some(RoutingHints { prefill_worker_id: Some(11), ..Default::default() }); let permit = prefill_tracker.set_phase(RequestPhase::Prefill).await; assert!(!is_migratable_for_request(&prefill_pinned, &error)); drop(permit); let permit = prefill_tracker.set_phase(RequestPhase::Decode).await; assert!(is_migratable_for_request(&prefill_pinned, &error)); drop(permit); let decode_tracker = Arc::new(RequestTracker::new()); let mut decode_pinned = create_mock_request(1); decode_pinned.tracker = Some(decode_tracker.clone()); decode_pinned.routing = Some(RoutingHints { decode_worker_id: Some(22), ..Default::default() }); let permit = decode_tracker.set_phase(RequestPhase::Prefill).await; assert!(is_migratable_for_request(&decode_pinned, &error)); drop(permit); let permit = decode_tracker.set_phase(RequestPhase::Decode).await; assert!(!is_migratable_for_request(&decode_pinned, &error)); drop(permit); } #[test] fn trackerless_request_treats_any_explicit_worker_as_pinned() { let error = migratable_error(ErrorType::WorkerOverloaded); for routing in [ RoutingHints { backend_instance_id: Some(7), ..Default::default() }, RoutingHints { prefill_worker_id: Some(11), ..Default::default() }, RoutingHints { decode_worker_id: Some(22), ..Default::default() }, ] { let mut request = create_mock_request(1); assert!(request.tracker.is_none()); request.routing = Some(routing); assert!(!is_migratable_for_request(&request, &error)); } let unpinned = create_mock_request(1); assert!(is_migratable_for_request(&unpinned, &error)); } // Helper to create a mock preprocessed request fn create_mock_request(max_tokens: u32) -> PreprocessedRequest { PreprocessedRequest::builder() .model("mock".to_string()) .token_ids(vec![1, 2, 3]) .stop_conditions(StopConditions { max_tokens: Some(max_tokens), ..Default::default() }) .sampling_options(SamplingOptions::default()) .output_options(OutputOptions::default()) .eos_token_ids(vec![]) .annotations(vec![]) .build() .unwrap() } // Helper to create mock LLM engine output fn create_mock_output(token_id: u32) -> Annotated { Annotated::from_data(BackendOutput { token_ids: vec![token_id], tokens: vec![], text: Some(format!("token_{token_id}")), cum_log_probs: None, log_probs: None, top_logprobs: None, finish_reason: None, stop_reason: None, index: None, disaggregated_params: None, encoder_result: None, worker_trace_link: None, completion_usage: None, engine_data: None, routing_data: None, }) } #[derive(Debug, Clone)] enum MockBehavior { /// Always succeeds with all responses Success, /// Fails on first call with NoResponders error, then succeeds on subsequent calls FailThenSuccess, FailThenSuccessWithAffinity, /// Fails on the first call and cancels the request before the retry FailThenCancel { context: Arc, }, /// Fails on the first call, then generate returns a cancellation error FailThenGenerateCancelled, /// One addressed worker rejects admission, then a replacement succeeds. WorkerOverloadSequence { worker_ids: Vec, }, /// Succeeds initially, fails mid-stream with specific error, then succeeds on retry MidStreamFail { fail_after: usize, }, /// Succeeds initially, fails mid-stream with specific error, then always fails on retry attempts MidStreamFailAlways { fail_after: usize, }, /// Succeeds initially, fails mid-stream, then always fails with stream error on retry attempts MidStreamFailAlwaysStreamError { fail_after: usize, }, /// Always fails with NoResponders error (same as FailThenSuccess first call) AlwaysFail, } // Unified mock server streaming engine that can simulate different scenarios struct MockEngine { behavior: MockBehavior, num_responses: usize, token_offset: u32, call_count: Arc, context_id: String, initial_min_tokens: Option, } impl MockEngine { fn new( behavior: MockBehavior, num_responses: usize, token_offset: u32, context_id: String, ) -> Self { Self { behavior, num_responses, token_offset, call_count: Arc::new(AtomicU32::new(0)), context_id, initial_min_tokens: None, } } fn with_min_tokens(mut self, min_tokens: u32) -> Self { self.initial_min_tokens = Some(min_tokens); self } } #[async_trait] impl AsyncEngine, ManyOut>, anyhow::Error> for MockEngine { async fn generate( &self, request: SingleIn, ) -> Result>> { let call_num = self.call_count.fetch_add(1, Ordering::SeqCst); if matches!(self.behavior, MockBehavior::FailThenSuccessWithAffinity) { let actual = request .get::(SESSION_AFFINITY_CONTEXT_KEY) .expect("session affinity context missing after migration wrapper"); assert_eq!(actual.as_str(), "session-123"); } let (preprocessed_request, context) = request.transfer(()); // Assert that the context_id matches the expected one assert_eq!( context.id().to_string(), self.context_id, "Context ID mismatch" ); // Calculate how many responses we've already generated based on request token_ids // Initial request has [1, 2, 3], so anything beyond that are generated responses let initial_tokens = 3; // [1, 2, 3] let responses_already_generated = preprocessed_request .token_ids .len() .saturating_sub(initial_tokens); // Assert that max_tokens reflects the expected remaining tokens let expected_max_tokens = self.num_responses .saturating_sub(responses_already_generated) as u32; assert_eq!( preprocessed_request.stop_conditions.max_tokens, Some(expected_max_tokens), "max_tokens should be {} but got {:?}", expected_max_tokens, preprocessed_request.stop_conditions.max_tokens ); if let Some(initial_min_tokens) = self.initial_min_tokens { let expected_min_tokens = initial_min_tokens.saturating_sub(responses_already_generated as u32); assert_eq!( preprocessed_request.stop_conditions.min_tokens, Some(expected_min_tokens), "min_tokens should be rebased for each replacement request" ); } match &self.behavior { MockBehavior::Success => { // Always succeed with remaining responses self.send_responses(responses_already_generated, self.num_responses) .await } MockBehavior::FailThenSuccess | MockBehavior::FailThenSuccessWithAffinity => { if call_num == 0 { // First call - return "No responders available" error to trigger retry return Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::CannotConnect) .message("no responders") .build() )); } else { // Subsequent calls - succeed with remaining responses self.send_responses(responses_already_generated, self.num_responses) .await } } MockBehavior::FailThenCancel { context } => { assert_eq!(call_num, 0, "cancelled retry must not reach the engine"); context.stop_generating(); Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::CannotConnect) .message("no responders") .build() )) } MockBehavior::FailThenGenerateCancelled => { let error_type = if call_num == 0 { ErrorType::CannotConnect } else { ErrorType::Cancelled }; Err(anyhow::anyhow!( DynamoError::builder() .error_type(error_type) .message("request cancelled") .build() )) } MockBehavior::WorkerOverloadSequence { worker_ids } => { let excluded = preprocessed_request .migration_state .as_ref() .expect("migration state missing") .excluded_worker_ids(); assert_eq!( excluded, worker_ids[..call_num as usize], "each retry must retain every previously rejected worker" ); if let Some(&worker_id) = worker_ids.get(call_num as usize) { let error = DynamoError::builder() .error_type(ErrorType::WorkerOverloaded) .message("selected worker is overloaded") .build(); preprocessed_request .migration_state .as_ref() .unwrap() .record_failure(worker_id, Some(error.clone())); return Err(anyhow::anyhow!(error)); } self.send_responses(responses_already_generated, self.num_responses) .await } MockBehavior::MidStreamFail { fail_after } => { let (tx, rx) = mpsc::channel(1); let token_offset = self.token_offset; let fail_after = *fail_after; let num_responses = self.num_responses; if call_num == 0 { // First call - send some responses then an error to simulate disconnection tokio::spawn(async move { // Send responses from current position to fail_after for i in responses_already_generated..fail_after.min(num_responses) { let response = create_mock_output(token_offset + 1 + i as u32); if tx.send(response).await.is_err() { break; } } // Send the specific error that triggers retry logic let error_response = Annotated::from_err( DynamoError::builder() .error_type(ErrorType::Disconnected) .message("Stream ended before generation completed") .build(), ); let _ = tx.send(error_response).await; }); } else { // Second call - send remaining responses from where we left off tokio::spawn(async move { for i in responses_already_generated..num_responses { let response = create_mock_output(token_offset + 1 + i as u32); if tx.send(response).await.is_err() { break; } } }); } let stream = tokio_stream::wrappers::ReceiverStream::new(rx); let ctx = Arc::new(Controller::new(self.context_id.clone())); Ok(dynamo_runtime::pipeline::ResponseStream::new( Box::pin(stream), ctx, )) } MockBehavior::MidStreamFailAlways { fail_after } => { if call_num == 0 { // First call - send some responses then an error to simulate disconnection let (tx, rx) = mpsc::channel(1); let token_offset = self.token_offset; let fail_after = *fail_after; let num_responses = self.num_responses; tokio::spawn(async move { // Send responses from current position to fail_after for i in responses_already_generated..fail_after.min(num_responses) { let response = create_mock_output(token_offset + 1 + i as u32); if tx.send(response).await.is_err() { break; } } // Send the specific error that triggers retry logic let error_response = Annotated::from_err( DynamoError::builder() .error_type(ErrorType::Disconnected) .message("Stream ended before generation completed") .build(), ); let _ = tx.send(error_response).await; }); let stream = tokio_stream::wrappers::ReceiverStream::new(rx); let ctx = Arc::new(Controller::new(self.context_id.clone())); Ok(dynamo_runtime::pipeline::ResponseStream::new( Box::pin(stream), ctx, )) } else { // Subsequent calls - always fail with NoResponders error (same as AlwaysFail) Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::CannotConnect) .message("no responders") .build() )) } } MockBehavior::MidStreamFailAlwaysStreamError { fail_after } => { let (tx, rx) = mpsc::channel(1); let token_offset = self.token_offset; let fail_after = *fail_after; let num_responses = self.num_responses; if call_num == 0 { // First call - send some responses then an error to simulate disconnection tokio::spawn(async move { // Send responses from current position to fail_after for i in responses_already_generated..fail_after.min(num_responses) { let response = create_mock_output(token_offset + 1 + i as u32); if tx.send(response).await.is_err() { break; } } // Send the specific error that triggers retry logic let error_response = Annotated::from_err( DynamoError::builder() .error_type(ErrorType::Disconnected) .message("Stream ended before generation completed") .build(), ); let _ = tx.send(error_response).await; }); let stream = tokio_stream::wrappers::ReceiverStream::new(rx); let ctx = Arc::new(Controller::new(self.context_id.clone())); Ok(dynamo_runtime::pipeline::ResponseStream::new( Box::pin(stream), ctx, )) } else { // Subsequent calls - immediately send stream error (no successful responses) tokio::spawn(async move { // Send the stream error immediately let error_response = Annotated::from_err( DynamoError::builder() .error_type(ErrorType::Disconnected) .message("Stream ended before generation completed") .build(), ); let _ = tx.send(error_response).await; }); let stream = tokio_stream::wrappers::ReceiverStream::new(rx); let ctx = Arc::new(Controller::new(self.context_id.clone())); Ok(dynamo_runtime::pipeline::ResponseStream::new( Box::pin(stream), ctx, )) } } MockBehavior::AlwaysFail => { // Always fail with NoResponders error (same as FailThenSuccess first call) Err(anyhow::anyhow!( DynamoError::builder() .error_type(ErrorType::CannotConnect) .message("no responders") .build() )) } } } } impl MockEngine { async fn send_responses( &self, start: usize, end: usize, ) -> Result>> { let (tx, rx) = mpsc::channel(1); let token_offset = self.token_offset; tokio::spawn(async move { for i in start..end { let response = create_mock_output(token_offset + 1 + i as u32); if tx.send(response).await.is_err() { break; } } }); let stream = tokio_stream::wrappers::ReceiverStream::new(rx); let ctx = Arc::new(Controller::new(self.context_id.clone())); Ok(dynamo_runtime::pipeline::ResponseStream::new( Box::pin(stream), ctx, )) } } /// Test case 1: No migration needed /// The two overload cases must migrate differently: one busy worker can be /// retried elsewhere, a pool with no free worker cannot. /// /// Collapsing them — as a single `ResourceExhausted` did — either strands a /// request that had a healthy worker available, or bounces a pool-wide /// rejection around until retries run out. #[test] fn worker_overload_migrates_but_pool_exhaustion_does_not() { let worker_busy = DynamoError::builder() .error_type(ErrorType::WorkerOverloaded) .message("Selected worker is overloaded, please retry later") .build(); assert!( is_migratable(&worker_busy), "one overloaded worker must fail over to another" ); let pool_exhausted = DynamoError::builder() .error_type(ErrorType::ResourceExhausted) .message("All workers are busy, please retry later") .build(); assert!( !is_migratable(&pool_exhausted), "pool-wide exhaustion must not migrate; no worker has room" ); } /// Tests the normal case where the RetryManager successfully processes all responses /// from a single stream without any failures or need for retries/migration. /// Expected behavior: All 10 responses should be received successfully. #[tokio::test] async fn test_retry_manager_no_migration() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::Success, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 0, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } assert_eq!(responses.len(), 10); for (i, response) in responses.iter().enumerate() { assert!(response.err().is_none()); if let Some(output) = &response.data { assert_eq!(output.token_ids, vec![101 + i as u32]); // 101, 102, 103, ..., 110 } } assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 0); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); } #[tokio::test] async fn maximum_migration_limit_does_not_overflow() { let context_id = uuid::Uuid::new_v4().to_string(); let mock_engine = Arc::new(MockEngine::new( MockBehavior::Success, 1, 100, context_id.clone(), )); let calls = mock_engine.call_count.clone(); let request = Context::with_id_and_metadata(create_mock_request(1), context_id, BTreeMap::new()); let migration = Migration::new( u32::MAX, None, TEST_MODEL.to_string(), Arc::new(Metrics::new()), ); let responses = migration .generate(request, mock_engine) .await .unwrap() .collect::>() .await; assert_eq!(calls.load(Ordering::SeqCst), 1); assert_eq!(responses.len(), 1); assert!(responses[0].error.is_none()); } #[tokio::test] async fn test_migration_preserves_session_affinity_across_retry() { let context_id = uuid::Uuid::new_v4().to_string(); let mock_engine = Arc::new(MockEngine::new( MockBehavior::FailThenSuccessWithAffinity, 1, 100, context_id.clone(), )); let calls = mock_engine.call_count.clone(); let mut request = Context::with_id_and_metadata(create_mock_request(1), context_id, BTreeMap::new()); request.insert( SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId::new("session-123"), ); let migration = Migration::new(1, None, TEST_MODEL.to_string(), Arc::new(Metrics::new())); let mut stream = migration.generate(request, mock_engine).await.unwrap(); while stream.next().await.is_some() {} assert_eq!(calls.load(Ordering::SeqCst), 2); } #[tokio::test] async fn explicit_backend_pin_does_not_retry_the_same_worker() { let context_id = uuid::Uuid::new_v4().to_string(); let mock_engine = Arc::new(MockEngine::new( MockBehavior::FailThenSuccess, 1, 100, context_id.clone(), )); let calls = mock_engine.call_count.clone(); let mut content = create_mock_request(1); content.routing = Some(RoutingHints { backend_instance_id: Some(7), ..Default::default() }); let request = Context::with_id_and_metadata(content, context_id, BTreeMap::new()); let migration = Migration::new(3, None, TEST_MODEL.to_string(), Arc::new(Metrics::new())); let result = migration.generate(request, mock_engine).await; assert!(result.is_err()); assert_eq!(calls.load(Ordering::SeqCst), 1); } #[tokio::test] async fn worker_overload_excludes_the_rejected_worker_on_retry() { let context_id = uuid::Uuid::new_v4().to_string(); let mock_engine = Arc::new(MockEngine::new( MockBehavior::WorkerOverloadSequence { worker_ids: vec![7, 8], }, 1, 100, context_id.clone(), )); let calls = mock_engine.call_count.clone(); let request = Context::with_id_and_metadata(create_mock_request(1), context_id, BTreeMap::new()); let migration = Migration::new(2, None, TEST_MODEL.to_string(), Arc::new(Metrics::new())); let mut stream = migration.generate(request, mock_engine).await.unwrap(); let responses = stream.by_ref().collect::>().await; assert_eq!(calls.load(Ordering::SeqCst), 3); assert_eq!(responses.len(), 1); assert!(responses[0].error.is_none()); } /// Test case 2: New request migration /// Tests the scenario where a worker becomes unreachable for new requests initially, /// triggering the RetryManager to retry the request. The MockEngine with FailThenSuccess /// fails on the first call with a "No responders available" error, then succeeds /// on subsequent calls, simulating a worker becoming available after initial failure. /// Expected behavior: All 10 responses should be received successfully after retry. #[tokio::test] async fn test_retry_manager_new_request_migration() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::FailThenSuccess, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } assert_eq!(responses.len(), 10); for (i, response) in responses.iter().enumerate() { assert!(response.err().is_none()); if let Some(output) = &response.data { assert_eq!(output.token_ids, vec![101 + i as u32]); // 101, 102, 103, ..., 110 } } assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 1); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::NEW_REQUEST, frontend_service::migration_outcome::SUCCESS, ), 1 ); } /// Test case 3: Ongoing request migration /// Tests the scenario where a worker fails mid-stream during an ongoing request. /// This simulates a connection being lost after partial response delivery, requiring /// the RetryManager to detect the failure (via "Stream ended before generation completed" error), /// create a new stream, and continue from where it left off. /// Expected behavior: 5 responses from first stream + 5 responses from retry stream = 10 total. #[tokio::test] async fn test_retry_manager_ongoing_request_migration() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFail { fail_after: 5 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } // Should have received all 10 responses (5 from first stream + 5 from second stream) assert_eq!(responses.len(), 10); // Check that we received responses from both streams for (i, response) in responses.iter().enumerate() { assert!(response.err().is_none()); if let Some(output) = &response.data { assert_eq!(output.token_ids, vec![101 + i as u32]); // 101, 102, 103, ..., 110 } } assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 0); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 1); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::ONGOING_REQUEST, frontend_service::migration_outcome::SUCCESS, ), 1 ); } #[tokio::test] async fn retry_rebases_min_tokens_from_the_delivered_prefix() { let context_id = uuid::Uuid::new_v4().to_string(); let mut request = create_mock_request(10); request.stop_conditions.min_tokens = Some(7); let mock_engine = Arc::new( MockEngine::new( MockBehavior::MidStreamFail { fail_after: 3 }, 10, 100, context_id.clone(), ) .with_min_tokens(7), ); let calls = mock_engine.call_count.clone(); let request = Context::with_id_and_metadata(request, context_id, BTreeMap::new()); let migration = Migration::new(1, None, TEST_MODEL.to_string(), Arc::new(Metrics::new())); let responses = migration .generate(request, mock_engine) .await .unwrap() .collect::>() .await; assert_eq!(calls.load(Ordering::SeqCst), 2); assert_eq!(responses.len(), 10); assert!(responses.iter().all(|response| response.error.is_none())); } /// Test case 4: New request migration - indefinite failure /// Tests the scenario where a worker becomes unreachable for new requests indefinitely. /// The RetryManager should exhaust all retries and return the original error from the first attempt. /// Expected behavior: Should receive an error after all retries are exhausted, with the original error. #[tokio::test] async fn test_retry_manager_new_request_migration_indefinite_failure() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(0); let mock_engine = Arc::new(MockEngine::new( MockBehavior::AlwaysFail, 0, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; // Should fail to build due to initial stream creation failure after exhausting all 3 retries let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let retry_manager_result = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await; assert!(retry_manager_result.is_err()); if let Err(error) = retry_manager_result { assert!(error.to_string().contains("no responders")); } assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 4); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::NEW_REQUEST, frontend_service::migration_outcome::FAILURE, ), 1 ); } /// Test case 5: Ongoing request migration - indefinite failure /// Tests the scenario where a worker fails mid-stream indefinitely during ongoing requests. /// The RetryManager should exhaust all retries and return the original stream disconnection error. /// Expected behavior: Should receive some responses from first stream, then error after retries exhausted. #[tokio::test] async fn test_retry_manager_ongoing_request_migration_indefinite_failure() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFailAlways { fail_after: 3 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) // 3 retries .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); // Collect all responses (both successful and error responses) while let Some(response) = retry_manager.next().await { responses.push(response); } // Should have received 4 total responses: 3 successful + 1 error assert_eq!(responses.len(), 4); // First 3 responses should be successful with tokens 101, 102, 103 for (i, response) in responses[0..3].iter().enumerate() { assert!(response.err().is_none()); if let Some(output) = &response.data { assert_eq!(output.token_ids, vec![101 + i as u32]); // 101, 102, 103 } } // 4th response should be a Disconnected error after retries are exhausted let error_response = &responses[3]; let err = error_response.err().expect("expected error response"); assert_eq!(err.error_type(), ErrorType::Disconnected); assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 3); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 1); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::ONGOING_REQUEST, frontend_service::migration_outcome::FAILURE, ), 1 ); } /// Test case 6: Ongoing request migration - indefinite failure with stream errors /// Tests the scenario where a worker fails mid-stream indefinitely during ongoing requests, /// and all retry attempts also fail with stream errors instead of NATS errors. /// Expected behavior: Should receive some responses from first stream, then error after retries exhausted. #[tokio::test] async fn test_retry_manager_ongoing_request_migration_indefinite_failure_stream_error() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFailAlwaysStreamError { fail_after: 3 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) // 3 retries .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); // Collect all responses (both successful and error responses) while let Some(response) = retry_manager.next().await { responses.push(response); } // Should have received 4 total responses: 3 successful + 1 error assert_eq!(responses.len(), 4); // First 3 responses should be successful with tokens 101, 102, 103 for (i, response) in responses[0..3].iter().enumerate() { assert!(response.err().is_none()); if let Some(output) = &response.data { assert_eq!(output.token_ids, vec![101 + i as u32]); // 101, 102, 103 } } // 4th response should be a Disconnected error after retries are exhausted let error_response = &responses[3]; let err = error_response.err().expect("expected error response"); assert_eq!(err.error_type(), ErrorType::Disconnected); assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 0); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 4); // 3 retries + 1 final failure assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::ONGOING_REQUEST, frontend_service::migration_outcome::SUCCESS, ), 3 ); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::ONGOING_REQUEST, frontend_service::migration_outcome::FAILURE, ), 1 ); } /// Test case 7: Request cancelled when creating new stream /// Tests the scenario where context.stop_generating() is called when creating a new stream. /// The RetryManager should detect that the context is stopped and abort creating new streams. /// Expected behavior: Should fail to build RetryManager with "Context is stopped or killed" error. #[tokio::test] async fn test_retry_manager_context_stopped_before_stream() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::Success, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); // Stop the context before building RetryManager ctx.stop_generating(); // Should fail to build due to stopped context let metrics = Arc::new(Metrics::new()); let retry_manager_result = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await; assert!(retry_manager_result.is_err()); if let Err(error) = retry_manager_result { assert!( error .to_string() .contains(&format!("Context id {} is stopped or killed", context_id)) ); // Verify the error is a typed DynamoError with Cancelled type let dynamo_err = error .downcast_ref::() .expect("Error should be a DynamoError"); assert_eq!( dynamo_err.error_type(), ErrorType::Cancelled, "Stopped/killed context should produce a Cancelled error" ); } assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 0); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 0); } #[tokio::test] async fn test_retry_manager_cancelled_during_migration() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let ctx = Arc::new(Controller::new(context_id.clone())); let mock_engine = Arc::new(MockEngine::new( MockBehavior::FailThenCancel { context: ctx.clone(), }, 10, 100, context_id, )); let next_generate: ServerStreamingEngine> = mock_engine; let metrics = Arc::new(Metrics::new()); let result = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await; let error = match result { Ok(_) => panic!("cancelled migration must fail"), Err(error) => error, }; let dynamo_error = error .downcast_ref::() .expect("error should be a DynamoError"); assert_eq!(dynamo_error.error_type(), ErrorType::Cancelled); assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 1); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::NEW_REQUEST, frontend_service::migration_outcome::CANCELLED, ), 1 ); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::NEW_REQUEST, frontend_service::migration_outcome::FAILURE, ), 0 ); } #[tokio::test] async fn test_retry_manager_generate_cancelled_during_migration() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::FailThenGenerateCancelled, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let metrics = Arc::new(Metrics::new()); let result = RetryManager::build( Arc::new(Controller::new(context_id)), BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await; let error = match result { Ok(_) => panic!("cancelled migration must fail"), Err(error) => error, }; let dynamo_error = error .downcast_ref::() .expect("error should be a DynamoError"); assert_eq!(dynamo_error.error_type(), ErrorType::Cancelled); assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 1); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::NEW_REQUEST, frontend_service::migration_outcome::CANCELLED, ), 1 ); assert_eq!( migration_duration_count( &metrics, frontend_service::migration_type::NEW_REQUEST, frontend_service::migration_outcome::FAILURE, ), 0 ); } /// Test case 8: No migration for guided-decoding (structured-output) requests /// /// Bug (#7634): When a worker crashes mid-stream during a structured-output /// (json_schema) request, migration appends already-generated token IDs back onto /// token_ids and replays the request to a new worker. However, backends initialize /// the guided-decoding FSM fresh for every new request and only advance it on newly- /// generated tokens — not on context/prompt tokens. This causes the FSM to restart /// from the schema root while treating already-generated tokens as context, producing /// duplicated or nested JSON in the final response. /// /// Fix: Disable migration for structured-output requests by zeroing retries_left in /// RetryManager::build() when guided_decoding is set, propagating the error cleanly. /// /// Expected behavior BEFORE fix: All 10 responses received (migration happened — wrong) /// Expected behavior AFTER fix: 3 successful + 1 error (migration blocked — correct) #[tokio::test] async fn test_retry_manager_no_migration_for_guided_decoding() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); let mut request = create_mock_request(10); // Set guided decoding (json_schema structured output) on the request request.sampling_options.guided_decoding = Some(GuidedDecodingOptions::new( Some(serde_json::json!({"type": "object", "properties": {"name": {"type": "string"}}})), None, None, None, None, None, None, )); // MidStreamFail after 3 tokens: without the fix, migration would succeed and // deliver all 10 responses; with the fix, migration is blocked and an error // is returned after the 3 partial responses. let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFail { fail_after: 3 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, // migration_limit=3 — should be ignored for guided-decoding requests None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } // Must receive 3 successful tokens + 1 Disconnected error, NOT all 10. // Before the fix this assertion fails because migration proceeds and returns 10. assert_eq!( responses.len(), 4, "Expected 3 successful + 1 error response (migration must be blocked for \ guided-decoding), but got {} responses", responses.len() ); // First 3 responses should be successful for (i, response) in responses[0..3].iter().enumerate() { assert!( response.err().is_none(), "Response {} should be successful", i ); } // Last response must be the stream-disconnection error let last = responses.last().unwrap(); let err = last .err() .expect("Last response should be a Disconnected error"); assert_eq!( err.error_type(), ErrorType::Disconnected, "Error type should be Disconnected" ); } /// Test case 9: max_seq_len exceeded limit + 1 disables migration /// /// Boundary test: prompt has 3 tokens, max_seq_len = 5. After 2 generated tokens the /// total is 5 (== max_seq_len) — still migratable. The 3rd generated token would push /// the total to 6 (> max_seq_len), which disables migration and stops caching. /// The failure is placed right at that point (fail_after: 3) so we see the error /// propagated instead of retried. #[tokio::test] async fn test_retry_manager_max_seq_len_exceeded() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); // Prompt = [1, 2, 3] (len 3). max_seq_len = 5. // Token 101 → total 4 ≤ 5: tracked. // Token 102 → total 5 ≤ 5: tracked. // Token 103 → would-be 6 > 5: NOT tracked, migration disabled. // Error follows immediately (fail_after: 3) → not retried. let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFail { fail_after: 3 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, Some(5), // prompt(3) + 3 generated = 6 > 5 → disables migration Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } // 3 successful tokens + 1 Disconnected error (migration disabled at token 103). assert_eq!( responses.len(), 4, "Expected 3 successful + 1 error (migration disabled by max_seq_len)" ); for (i, response) in responses[0..3].iter().enumerate() { assert!(response.err().is_none(), "Response {} should be OK", i); } let err = responses[3] .err() .expect("Last response should be Disconnected error"); assert_eq!(err.error_type(), ErrorType::Disconnected); // Migration was attempted but blocked because max_seq_len set retries_left to 0. // The ongoing metric is still incremented (it counts attempts, not successes). assert_eq!(metrics.get_migration_new_request_count(TEST_MODEL), 0); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 1); // max_seq_len limit was triggered once (at token 103). assert_eq!( metrics.get_migration_max_seq_len_exceeded_count(TEST_MODEL), 1 ); } /// Test case 10: Migration succeeds when sequence length is at max_seq_len /// /// Boundary test: prompt has 3 tokens, max_seq_len = 5. After 2 generated tokens /// the total is exactly 5 (== max_seq_len). The failure occurs at that point /// (fail_after: 2). Because we use strict inequality (> not >=), the request is /// still migratable and the retry succeeds. #[tokio::test] async fn test_retry_manager_max_seq_len_at_limit() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); // Prompt = [1, 2, 3] (len 3). max_seq_len = 5. // Token 101 → total 4 ≤ 5: tracked. // Token 102 → total 5 == 5: tracked (still migratable — strict >). // Error (fail_after: 2) → migration succeeds, retry delivers remaining tokens. let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFail { fail_after: 2 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, Some(5), // prompt(3) + 2 generated = 5 == max_seq_len → still migratable Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } // Migration succeeds — all 10 responses delivered assert_eq!(responses.len(), 10); for response in &responses { assert!(response.err().is_none()); } assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 1); // Tracked token_ids must equal exactly max_seq_len (5). // The 2 tokens from the first stream were tracked (prompt 3 + gen 2 = 5). // After migration the retry stream delivers remaining tokens, but the first // new token would push to 6 > 5, so tracking stops and no more are appended. assert_eq!( retry_manager.request.token_ids.len(), 5, "tracked token_ids should be exactly max_seq_len" ); // The limit was triggered once (first token of the retry stream exceeded 5). assert_eq!( metrics.get_migration_max_seq_len_exceeded_count(TEST_MODEL), 1 ); } /// Test case 11: Prompt length alone exceeds max_seq_len /// /// When the prompt tokens already exceed max_seq_len, migration is disabled /// in RetryManager::build before any tokens are generated. A mid-stream /// failure should propagate the error without attempting migration. #[tokio::test] async fn test_retry_manager_max_seq_len_exceeded_by_prompt() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); // Prompt = [1, 2, 3] (len 3). max_seq_len = 2, so prompt alone exceeds the limit. let request = create_mock_request(10); let mock_engine = Arc::new(MockEngine::new( MockBehavior::MidStreamFail { fail_after: 3 }, 10, 100, context_id.clone(), )); let next_generate: ServerStreamingEngine> = mock_engine; let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, Some(2), // prompt(3) > max_seq_len(2) → migration disabled at build time Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } // 3 successful tokens + 1 Disconnected error (migration disabled from the start). assert_eq!( responses.len(), 4, "Expected 3 successful + 1 error (migration disabled by prompt exceeding max_seq_len)" ); for (i, response) in responses[0..3].iter().enumerate() { assert!(response.err().is_none(), "Response {} should be OK", i); } let err = responses[3] .err() .expect("Last response should be Disconnected error"); assert_eq!(err.error_type(), ErrorType::Disconnected); // max_seq_len was exceeded at build time (prompt len 3 > 2). assert_eq!( metrics.get_migration_max_seq_len_exceeded_count(TEST_MODEL), 1 ); // Migration was attempted but blocked (retries_left was already 0). assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 1); } /// Smoke test for the byo-preprocessor response shape. #[tokio::test] async fn test_retry_manager_generic_over_llm_engine_output() { dynamo_runtime::logging::init(); let context_id = uuid::Uuid::new_v4().to_string(); struct LlmEngineMock(String); #[async_trait] impl AsyncEngine< SingleIn, ManyOut>, anyhow::Error, > for LlmEngineMock { async fn generate( &self, _request: SingleIn, ) -> Result>> { let responses = stream::iter((0..3u32).map(|i| { Annotated::from_data(LLMEngineOutput { token_ids: vec![200 + i], ..Default::default() }) })); let ctx = Arc::new(Controller::new(self.0.clone())); Ok(ResponseStream::new(Box::pin(responses), ctx)) } } let request = create_mock_request(3); let next_generate: ServerStreamingEngine> = Arc::new(LlmEngineMock(context_id.clone())); let ctx = Arc::new(Controller::new(context_id.clone())); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 1, None, Arc::new(TEST_MODEL.to_string()), metrics, None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(r) = retry_manager.next().await { responses.push(r); } assert_eq!(responses.len(), 3); assert_eq!( retry_manager.request.token_ids, vec![1, 2, 3, 200, 201, 202] ); } /// 2-hop migration: A → fail → B → fail → C. Each retry's /// `migration_link` must point at the *latest* failed worker, not the /// original. #[tokio::test] async fn test_retry_manager_propagates_migration_link_over_two_hops() { use crate::protocols::common::preprocessor::TraceLink; use std::sync::Mutex; dynamo_runtime::logging::init(); struct LinkingMockEngine { captured_links: Arc>>>, worker_links: Vec, context_id: String, call_count: Arc, } #[async_trait] impl AsyncEngine< SingleIn, ManyOut>, anyhow::Error, > for LinkingMockEngine { async fn generate( &self, request: SingleIn, ) -> Result>> { let call_num = self.call_count.fetch_add(1, Ordering::SeqCst) as usize; let (preprocessed_request, _ctx) = request.transfer(()); self.captured_links .lock() .unwrap() .push(preprocessed_request.migration_link.clone()); let (tx, rx) = mpsc::channel(1); let context_id = self.context_id.clone(); let fail_this_call = call_num < 2; let link = self.worker_links.get(call_num).cloned(); let responses_already_generated = preprocessed_request.token_ids.len().saturating_sub(3); let total_chunks: usize = 6; tokio::spawn(async move { let start = responses_already_generated; let end = if fail_this_call { (start + 2).min(total_chunks) } else { total_chunks }; for i in start..end { let mut out = create_mock_output(100 + 1 + i as u32); if i == start && let (Some(link), Some(data)) = (&link, out.data.as_mut()) { data.worker_trace_link = Some(link.clone()); } if tx.send(out).await.is_err() { return; } } if fail_this_call { let err = Annotated::from_err( DynamoError::builder() .error_type(ErrorType::Disconnected) .message("Stream ended before generation completed") .build(), ); let _ = tx.send(err).await; } }); let stream = tokio_stream::wrappers::ReceiverStream::new(rx); let ctx = Arc::new(Controller::new(context_id)); Ok(dynamo_runtime::pipeline::ResponseStream::new( Box::pin(stream), ctx, )) } } let context_id = uuid::Uuid::new_v4().to_string(); let request = create_mock_request(6); let captured = Arc::new(Mutex::new(Vec::>::new())); let link_a = TraceLink { trace_id: "0123456789abcdef0123456789abcdef".to_string(), span_id: "aaaaaaaaaaaaaaaa".to_string(), }; let link_b = TraceLink { trace_id: "0123456789abcdef0123456789abcdef".to_string(), span_id: "bbbbbbbbbbbbbbbb".to_string(), }; let engine = Arc::new(LinkingMockEngine { captured_links: captured.clone(), worker_links: vec![link_a.clone(), link_b.clone()], context_id: context_id.clone(), call_count: Arc::new(AtomicU32::new(0)), }); let next_generate: ServerStreamingEngine> = engine; let ctx = Arc::new(Controller::new(context_id)); let metrics = Arc::new(Metrics::new()); let mut retry_manager = RetryManager::build( ctx, BTreeMap::new(), request, next_generate, 3, None, Arc::new(TEST_MODEL.to_string()), metrics.clone(), None, ) .await .expect("Failed to build RetryManager"); let mut responses = Vec::new(); while let Some(response) = retry_manager.next().await { responses.push(response); } assert_eq!(responses.len(), 6, "expected all 6 chunks across 2 hops"); for response in &responses { assert!(response.err().is_none(), "no chunk should be an error"); } let links = captured.lock().unwrap(); assert_eq!( links.len(), 3, "engine.generate must be called 3 times for a 2-hop migration" ); assert!( links[0].is_none(), "first attempt has no predecessor — migration_link must be None" ); assert_eq!( links[1].as_ref(), Some(&link_a), "second attempt must link back to worker A" ); assert_eq!( links[2].as_ref(), Some(&link_b), "third attempt must link back to worker B (latest worker, not original)" ); assert_eq!(metrics.get_migration_ongoing_request_count(TEST_MODEL), 2); } }