diff --git a/docs/fern/pages/use-cases/reinforcement-learning/operations-and-simulation.md b/docs/fern/pages/use-cases/reinforcement-learning/operations-and-simulation.md index 9c08f21b238e..451a0482103c 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/operations-and-simulation.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/operations-and-simulation.md @@ -30,7 +30,7 @@ X-Dynamo-Session-ID: rollout-run42-sample7 X-Dynamo-Parent-Session-ID: rollout-run42-parent2 ``` -Use `x-request-id` for cross-component logs and distributed traces. Native SGLang `/generate` gives a body `rid` precedence over that header, so use one stable `rid` or omit it. +Use `x-request-id` for cross-component logs and distributed traces. Dynamo always generates its own internal request ID, so concurrent requests can reuse the same client value. Client-supplied IDs are correlation values only: native SGLang `/generate` returns the body `rid` (or `x-request-id` when `rid` is omitted) in the response, and `/inference/v1/generate` returns `x-request-id` (or the body `request_id`) as the response `request_id`. Native SGLang session requests are the exception: when the body has `session_params`, SGLang keys the session continuation by `rid` (or `x-request-id` when `rid` is omitted), so that value must stay unique and stable and must not be reused by concurrent requests. Dynamo does not define typed RL rollout, trainer-step, or policy-version fields. Keep those values in the framework and join them through opaque request or session identity. Session identity supports correlation; it does not enable affinity unless the router is configured for it and does not validate policy freshness. diff --git a/lib/llm/src/grpc/service/openai.rs b/lib/llm/src/grpc/service/openai.rs index 1e09369d3475..3f5007cf95b6 100644 --- a/lib/llm/src/grpc/service/openai.rs +++ b/lib/llm/src/grpc/service/openai.rs @@ -55,9 +55,8 @@ pub async fn completion_response_stream( ), Status, > { - // create the context for the request - // [WIP] from request id. - let request_id = get_or_create_request_id(request.inner.user.as_deref()); + // Client ids remain in the payload; only a server-minted id may key in-flight state. + let request_id = uuid::Uuid::new_v4().to_string(); let streaming = request.inner.stream.unwrap_or(false); let model_name = request.inner.model.clone(); let cancellation_labels = CancellationLabels { @@ -178,21 +177,6 @@ pub async fn completion_response_stream( Ok((stream, parsing_options)) } -/// Get the request ID from a primary source, or lastly create a new one if not present -// TODO: Similar function exists in lib/llm/src/http/service/openai.rs but with different signature and more complex logic (distributed tracing, headers) -fn get_or_create_request_id(primary: Option<&str>) -> String { - // Try to get the request ID from the primary source - if let Some(primary) = primary - && let Ok(uuid) = uuid::Uuid::parse_str(primary) - { - return uuid.to_string(); - } - - // Try to parse the request ID as a UUID, or generate a new one if missing/invalid - let uuid = uuid::Uuid::new_v4(); - uuid.to_string() -} - impl TryFrom for NvCreateCompletionRequest { type Error = Status; diff --git a/lib/llm/src/grpc/service/tensor.rs b/lib/llm/src/grpc/service/tensor.rs index 1ce74cefad75..230310669b04 100644 --- a/lib/llm/src/grpc/service/tensor.rs +++ b/lib/llm/src/grpc/service/tensor.rs @@ -62,8 +62,8 @@ pub async fn tensor_response_stream( streaming: bool, metadata: &MetadataMap, ) -> Result>, Status> { - // create the context for the request - let request_id = get_or_create_request_id(request.id.as_deref()); + // Client ids remain in the payload; only a server-minted id may key in-flight state. + let request_id = uuid::Uuid::new_v4().to_string(); let model_name = request.model.clone(); let cancellation_labels = CancellationLabels { model: model_name.clone(), @@ -177,20 +177,6 @@ pub async fn tensor_response_stream( Ok(stream) } -/// Get the request ID from a primary source, or lastly create a new one if not present -fn get_or_create_request_id(primary: Option<&str>) -> String { - // Try to get the request ID from the primary source - if let Some(primary) = primary - && let Ok(uuid) = uuid::Uuid::parse_str(primary) - { - return uuid.to_string(); - } - - // Try to parse the request ID as a UUID, or generate a new one if missing/invalid - let uuid = uuid::Uuid::new_v4(); - uuid.to_string() -} - /// Convert KServe InferParameter to Dynamo ParameterValue #[allow(clippy::result_large_err)] pub fn kserve_param_to_dynamo( diff --git a/lib/llm/src/http/service/generate.rs b/lib/llm/src/http/service/generate.rs index f279fefe7312..c064d2b0135e 100644 --- a/lib/llm/src/http/service/generate.rs +++ b/lib/llm/src/http/service/generate.rs @@ -76,7 +76,12 @@ fn generate_metric_error_type(status: StatusCode) -> ErrorType { #[derive(Debug)] struct GenerateRequestContext { + /// Server-minted id for the request context, routing, and logs. Never client-chosen: + /// in-flight request state is keyed on it. request_id: String, + /// Client-facing id echoed in the response body and the `vllm_tito` envelope, following + /// vLLM's `x-request-id` > body `request_id` precedence. Correlation only. + public_request_id: String, data_parallel_rank: Option, } @@ -130,12 +135,13 @@ fn resolve_generate_request_context( headers: &HeaderMap, body_request_id: Option<&str>, ) -> GenerateRequestContext { - let request_id = headers + let request_id = get_or_create_request_id(headers); + let public_request_id = headers .get(X_REQUEST_ID_HEADER) .and_then(|value| value.to_str().ok()) .map(ToOwned::to_owned) .or_else(|| body_request_id.map(ToOwned::to_owned)) - .unwrap_or_else(|| get_or_create_request_id(headers)); + .unwrap_or_else(|| request_id.clone()); let data_parallel_rank = headers .get(X_DATA_PARALLEL_RANK_HEADER) .and_then(|value| value.to_str().ok()) @@ -143,6 +149,7 @@ fn resolve_generate_request_context( GenerateRequestContext { request_id, + public_request_id, data_parallel_rank, } } @@ -152,11 +159,12 @@ fn dynamo_routing_priority(vllm_priority: i32) -> i32 { vllm_priority.saturating_neg() } -fn generate_dispatch_span(request_id: &str, model: &str) -> tracing::Span { +fn generate_dispatch_span(request_id: &str, public_request_id: &str, model: &str) -> tracing::Span { tracing::info_span!( target: "request_span", "generate", request_id = %request_id, + x_request_id = %public_request_id, model = %model, input_tokens = tracing::field::Empty, output_tokens = tracing::field::Empty, @@ -601,7 +609,7 @@ fn preprocessed_from_generate( request: GenerateRequest, model: &str, data_parallel_rank: Option, - request_id: &str, + public_request_id: &str, routing_metadata: GenerateRoutingMetadata, ) -> anyhow::Result { let tracker = Arc::new(RequestTracker::new()); @@ -610,7 +618,7 @@ fn preprocessed_from_generate( request, model, data_parallel_rank, - request_id, + public_request_id, routing_metadata, tracker, ) @@ -620,7 +628,7 @@ fn preprocessed_from_generate_with_tracker( request: GenerateRequest, model: &str, data_parallel_rank: Option, - request_id: &str, + public_request_id: &str, routing_metadata: GenerateRoutingMetadata, tracker: Arc, ) -> anyhow::Result { @@ -661,7 +669,7 @@ fn preprocessed_from_generate_with_tracker( } } }; - let vllm_tito = serde_json::to_value(VllmTitoEnvelope::new(&request, request_id))?; + let vllm_tito = serde_json::to_value(VllmTitoEnvelope::new(&request, public_request_id))?; let mut extra_args = serde_json::Map::new(); extra_args.insert("vllm_tito".to_string(), vllm_tito); if let Some(kv_transfer_params) = request.kv_transfer_params.as_ref() { @@ -893,7 +901,11 @@ async fn handler_generate( .to_string(); let tracker = Arc::new(RequestTracker::new()); tracker.record_isl(request.token_ids.len(), None); - let dispatch_span = generate_dispatch_span(&request_context.request_id, &metric_model); + let dispatch_span = generate_dispatch_span( + &request_context.request_id, + &request_context.public_request_id, + &metric_model, + ); let mut metric_lifecycle = dispatch_span.in_scope(|| { GenerateMetricLifecycle::new( state.as_ref(), @@ -981,7 +993,7 @@ async fn handler_generate( request, &model, request_context.data_parallel_rank, - &request_context.request_id, + &request_context.public_request_id, routing_metadata, tracker, ) { @@ -1026,6 +1038,7 @@ async fn handler_generate( engine, context, request_id, + request_context.public_request_id, state.clone(), response_options, metric_lifecycle, @@ -1049,6 +1062,7 @@ async fn generate_dispatch( engine: crate::types::openai::generate::GenerateStreamingEngine, context: Context, request_id: String, + public_request_id: String, state: Arc, response_options: GenerateResponseOptions, metric_lifecycle: GenerateMetricLifecycle, @@ -1134,7 +1148,7 @@ async fn generate_dispatch( request_context.as_ref(), GenerateResponse::from_annotated_stream_with_options( stream, - request_id.clone(), + public_request_id, response_options, ), ) @@ -1579,7 +1593,7 @@ pub(crate) mod tests { } #[derive(Clone)] - struct RequestIdCaptureLayer(Arc>>); + struct RequestIdCaptureLayer(Arc>); impl Layer for RequestIdCaptureLayer { fn on_new_span( @@ -1591,7 +1605,7 @@ pub(crate) mod tests { let mut visitor = RequestIdVisitor::default(); attrs.record(&mut visitor); if visitor.request_id.is_some() { - *self.0.lock().unwrap() = visitor.request_id; + *self.0.lock().unwrap() = visitor; } } } @@ -1599,18 +1613,23 @@ pub(crate) mod tests { #[derive(Default)] struct RequestIdVisitor { request_id: Option, + x_request_id: Option, } impl Visit for RequestIdVisitor { fn record_str(&mut self, field: &Field, value: &str) { if field.name() == "request_id" { self.request_id = Some(value.to_string()); + } else if field.name() == "x_request_id" { + self.x_request_id = Some(value.to_string()); } } fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { if field.name() == "request_id" { self.request_id = Some(format!("{value:?}")); + } else if field.name() == "x_request_id" { + self.x_request_id = Some(format!("{value:?}")); } } } @@ -1764,6 +1783,130 @@ pub(crate) mod tests { ); } + /// Echoes each request's tokens back, and holds every request until `barrier` releases them + /// together so both are in flight at once. Records what the worker saw for each request. + struct ConcurrentEchoEngine { + barrier: tokio::sync::Barrier, + seen: Mutex)>>, + } + + #[async_trait::async_trait] + impl AsyncEngine, ManyOut>, Error> + for ConcurrentEchoEngine + { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let envelope_request_id = request + .extra_args + .as_ref() + .and_then(|extra| extra["vllm_tito"]["request_id"].as_str()) + .map(str::to_string); + self.seen + .lock() + .unwrap() + .push((request.id().to_string(), envelope_request_id)); + self.barrier.wait().await; + let stream = futures::stream::iter([Annotated::from_data(LLMEngineOutput { + token_ids: request.token_ids.to_vec(), + index: Some(0), + finish_reason: Some(crate::protocols::common::FinishReason::Stop), + ..Default::default() + })]); + Ok(ResponseStream::new(Box::pin(stream), request.context())) + } + } + + /// Serve `engine` through a frontend source/sink pipeline, like the discovery watcher + /// does, so responses are matched to callers by request id. + pub(crate) fn engine_behind_frontend( + engine: crate::types::openai::generate::GenerateStreamingEngine, + ) -> crate::types::openai::generate::GenerateStreamingEngine { + use dynamo_runtime::pipeline::{ServiceBackend, Source}; + let frontend = dynamo_runtime::pipeline::ServiceFrontend::< + SingleIn, + ManyOut>, + >::new(); + frontend + .link(ServiceBackend::from_engine(engine)) + .unwrap() + .link_terminal(frontend) + .unwrap() + } + + async fn concurrent_generate_with_shared_client_id( + headers: HeaderMap, + body_request_id: Option<&str>, + expected_public_id: &str, + ) { + const MODEL: &str = "generate-concurrent-model"; + let engine = Arc::new(ConcurrentEchoEngine { + barrier: tokio::sync::Barrier::new(2), + seen: Mutex::new(Vec::new()), + }); + let service = HttpService::builder().build().unwrap(); + let state = service.state_clone(); + state + .manager() + .add_generate_model( + MODEL, + "generate-concurrent-test", + engine_behind_frontend(engine.clone()), + ) + .unwrap(); + + let send = |token_ids: Vec| { + let mut body = serde_json::json!({ + "token_ids": token_ids, + "sampling_params": {}, + "model": MODEL + }); + if let Some(id) = body_request_id { + body["request_id"] = serde_json::json!(id); + } + let request = serde_json::from_value(body).unwrap(); + handler_generate(State(state.clone()), headers.clone(), Json(request)) + }; + let (first, second) = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join(send(vec![1, 2, 3]), send(vec![7, 8])), + ) + .await + .expect("concurrent generate requests must not hang"); + + for (response, expected_tokens) in [(first, vec![1, 2, 3]), (second, vec![7, 8])] { + assert_eq!(response.status(), StatusCode::OK); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["request_id"], expected_public_id); + assert_eq!( + body["choices"][0]["token_ids"], + serde_json::json!(expected_tokens) + ); + } + + let seen = engine.seen.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert_ne!( + seen[0].0, seen[1].0, + "each request needs its own context id" + ); + for (context_id, envelope_request_id) in seen.iter() { + assert_ne!(context_id, expected_public_id); + assert_eq!(envelope_request_id.as_deref(), Some(expected_public_id)); + } + } + + #[tokio::test] + async fn concurrent_generate_requests_sharing_x_request_id_get_their_own_output() { + let mut headers = HeaderMap::new(); + headers.insert(X_REQUEST_ID_HEADER, "1".parse().unwrap()); + concurrent_generate_with_shared_client_id(headers, None, "1").await; + } + #[tokio::test] async fn generate_early_failures_are_counted() { let service = HttpService::builder().build().unwrap(); @@ -2509,7 +2652,8 @@ pub(crate) mod tests { let context = resolve_generate_request_context(&headers, Some("body-request")); - assert_eq!(context.request_id, "header-request"); + assert_eq!(context.public_request_id, "header-request"); + assert!(uuid::Uuid::parse_str(&context.request_id).is_ok()); assert_eq!(context.data_parallel_rank, Some(3)); } @@ -2520,23 +2664,36 @@ pub(crate) mod tests { let context = resolve_generate_request_context(&headers, Some("body-request")); - assert_eq!(context.request_id, "body-request"); + assert_eq!(context.public_request_id, "body-request"); + assert!(uuid::Uuid::parse_str(&context.request_id).is_ok()); assert_eq!(context.data_parallel_rank, None); + + let context = resolve_generate_request_context(&HeaderMap::new(), None); + assert_eq!(context.public_request_id, context.request_id); } #[test] fn generate_dispatch_span_uses_resolved_request_id() { - let captured_request_id = Arc::new(Mutex::new(None)); + let captured_request_id = Arc::new(Mutex::new(RequestIdVisitor::default())); let _interest_guard = crate::test_utils::tracing_interest_guard(); let _guard = tracing::subscriber::set_default( tracing_subscriber::registry().with(RequestIdCaptureLayer(captured_request_id.clone())), ); - let dispatch_span = generate_dispatch_span("header-request", "test-model"); + let context = resolve_generate_request_context(&HeaderMap::new(), Some("body-request")); + let dispatch_span = generate_dispatch_span( + &context.request_id, + &context.public_request_id, + "test-model", + ); assert_eq!( - captured_request_id.lock().unwrap().as_deref(), - Some("header-request") + captured_request_id.lock().unwrap().request_id.as_deref(), + Some(context.request_id.as_str()) + ); + assert_eq!( + captured_request_id.lock().unwrap().x_request_id.as_deref(), + Some("body-request") ); let fields = dispatch_span .metadata() @@ -2544,6 +2701,7 @@ pub(crate) mod tests { .fields(); for field in [ "request_id", + "x_request_id", "model", "input_tokens", "output_tokens", @@ -2593,6 +2751,7 @@ pub(crate) mod tests { generate_dispatch( engine, context, + request_id.clone(), request_id, state, response_options, diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 70cf631ed17b..1f6e2a55b818 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -1031,43 +1031,25 @@ pub async fn smart_json_error_middleware(request: Request, next: Next) -> /// Return the request ID for the current request. /// -/// The canonical request ID is set by `make_inference_request_span()` and stored +/// The canonical request ID is minted by `make_inference_request_span()` and stored /// in the `DistributedTraceContext` via `DistributedTraceIdLayer`. This function -/// retrieves it, falling back to a validated `x-dynamo-request-id` header value -/// (deprecated, DEP #7812) or a new UUID. +/// retrieves it, or generates a new UUID when no request span is active. /// -/// **Deprecation (DEP #7812):** The `x-dynamo-request-id` header is deprecated. -/// Clients should rely on server-generated request IDs instead of supplying their own. +/// Client-supplied ids never become the internal request id: in-flight request state +/// is keyed on this value, so two requests sharing a client id would collide. Clients +/// correlate requests with `x-request-id` instead. +/// +/// **Deprecation (DEP #7812):** the `x-dynamo-request-id` and `request-id` headers are +/// ignored on HTTP requests; a warning is logged when either is present. pub(super) fn get_or_create_request_id(headers: &HeaderMap) -> String { - // Validate x-dynamo-request-id header if present, warn on invalid values. - // DEP #7812: x-dynamo-request-id is deprecated — clients should rely on - // server-generated request IDs instead of supplying their own. - let validated_header = if let Some(raw) = headers.get(DYNAMO_REQUEST_ID_HEADER) { - tracing::warn!( - "{} header is deprecated (DEP #7812); server-generated request IDs should be used instead", - DYNAMO_REQUEST_ID_HEADER - ); - match raw.to_str() { - Err(_) => { - tracing::warn!( - "{} header must be a valid UTF-8 string", - DYNAMO_REQUEST_ID_HEADER - ); - None - } - Ok(s) if uuid::Uuid::parse_str(s).is_err() => { - tracing::warn!( - "{} header must be a valid UUID, got: {}", - DYNAMO_REQUEST_ID_HEADER, - s - ); - None - } - Ok(s) => Some(s.to_string()), + for header in [DYNAMO_REQUEST_ID_HEADER, "request-id"] { + if headers.contains_key(header) { + tracing::warn!( + header, + "deprecated request id header ignored (DEP #7812); the server generates the request ID, use x-request-id for client correlation" + ); } - } else { - None - }; + } // Prefer trace context (set by make_inference_request_span via DistributedTraceIdLayer) if let Some(trace_context) = get_distributed_tracing_context() @@ -1076,8 +1058,7 @@ pub(super) fn get_or_create_request_id(headers: &HeaderMap) -> String { return request_id; } - // Fallback: use validated header for backwards compat, or generate new UUID - validated_header.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()) + uuid::Uuid::new_v4().to_string() } pub(super) fn context_from_headers( diff --git a/lib/llm/src/http/service/sglang_generate.rs b/lib/llm/src/http/service/sglang_generate.rs index e663b7e8d849..acc0c42e1e40 100644 --- a/lib/llm/src/http/service/sglang_generate.rs +++ b/lib/llm/src/http/service/sglang_generate.rs @@ -63,7 +63,12 @@ fn canonical_generate_models( #[derive(Debug)] struct RequestContext { + /// Server-minted id for the request context, routing, and logs. Never client-chosen: + /// in-flight request state is keyed on it. request_id: String, + /// Client-facing id carried to the worker as the envelope `rid`, following SGLang's body + /// `rid` > `x-request-id` precedence. Correlation only. + public_request_id: String, data_parallel_rank: Option, } @@ -140,7 +145,8 @@ fn adapt_openai_error(response: super::openai::ErrorResponse) -> Response { } fn resolve_request_context(headers: &HeaderMap, body_request_id: Option<&str>) -> RequestContext { - let request_id = body_request_id + let request_id = get_or_create_request_id(headers); + let public_request_id = body_request_id .map(ToOwned::to_owned) .or_else(|| { headers @@ -148,7 +154,7 @@ fn resolve_request_context(headers: &HeaderMap, body_request_id: Option<&str>) - .and_then(|value| value.to_str().ok()) .map(ToOwned::to_owned) }) - .unwrap_or_else(|| get_or_create_request_id(headers)); + .unwrap_or_else(|| request_id.clone()); let data_parallel_rank = headers .get(X_DATA_PARALLEL_RANK_HEADER) .and_then(|value| value.to_str().ok()) @@ -156,6 +162,7 @@ fn resolve_request_context(headers: &HeaderMap, body_request_id: Option<&str>) - RequestContext { request_id, + public_request_id, data_parallel_rank, } } @@ -200,7 +207,7 @@ fn preprocessed_request( request: SglangGenerateRequest, model: &str, data_parallel_rank: Option, - request_id: &str, + public_request_id: &str, ) -> anyhow::Result { let max_tokens = request.max_new_tokens().map_err(anyhow::Error::msg)?; let min_tokens = request.min_new_tokens().map_err(anyhow::Error::msg)?; @@ -208,7 +215,7 @@ fn preprocessed_request( let routing_priority = request.priority.unwrap_or_default(); let cache_namespace = request.cache_namespace().map(str::to_owned); let streaming = request.stream; - let (input_ids, worker_envelope) = request.into_worker_envelope(request_id); + let (input_ids, worker_envelope) = request.into_worker_envelope(public_request_id); // The router records TTFT and ITL into the request's tracker. let tracker = Arc::new(RequestTracker::new()); tracker.record_isl(input_ids.len(), None); @@ -332,7 +339,7 @@ async fn handler( request, &model, request_context.data_parallel_rank, - &request_context.request_id, + &request_context.public_request_id, ) { Ok(preprocessed) => preprocessed, Err(error) => return error_response(StatusCode::BAD_REQUEST, error.to_string()), @@ -360,7 +367,8 @@ async fn handler( let dispatch_span = tracing::info_span!( target: "request_span", "generate", - request_id = %request_id + request_id = %request_id, + x_request_id = %request_context.public_request_id, ); let response = match tokio::spawn( dispatch( @@ -565,6 +573,151 @@ mod tests { } } + /// Echoes each request's input ids back, holding every request until `barrier` releases them + /// together so both are in flight at once. Records what the worker saw for each request. + struct ConcurrentEchoEngine { + barrier: tokio::sync::Barrier, + seen: std::sync::Mutex)>>, + } + + #[async_trait::async_trait] + impl + dynamo_runtime::engine::AsyncEngine< + dynamo_runtime::pipeline::SingleIn, + dynamo_runtime::pipeline::ManyOut>, + dynamo_runtime::pipeline::Error, + > for ConcurrentEchoEngine + { + async fn generate( + &self, + request: dynamo_runtime::pipeline::SingleIn, + ) -> Result< + dynamo_runtime::pipeline::ManyOut>, + dynamo_runtime::pipeline::Error, + > { + let rid = request + .extra_args + .as_ref() + .and_then(|extra| extra["sglang_tito"]["rid"].as_str()) + .map(str::to_string); + self.seen + .lock() + .unwrap() + .push((request.id().to_string(), rid)); + self.barrier.wait().await; + let output = crate::protocols::Annotated::from_data(LLMEngineOutput { + engine_data: Some(serde_json::json!({ + "sglang_response": {"output_ids": request.token_ids.to_vec()} + })), + ..Default::default() + }); + Ok(dynamo_runtime::engine::ResponseStream::new( + Box::pin(futures::stream::iter([output])), + request.context(), + )) + } + } + + async fn concurrent_generate_with_shared_client_id( + headers: HeaderMap, + body_rid: Option<&str>, + expected_public_id: &str, + ) { + use crate::local_model::runtime_config::SGLANG_GENERATE_CAPABILITY; + + let engine = Arc::new(ConcurrentEchoEngine { + barrier: tokio::sync::Barrier::new(2), + seen: std::sync::Mutex::new(Vec::new()), + }); + let state = service_v2::HttpService::builder() + .build() + .unwrap() + .state_clone(); + let mut card = crate::model_card::ModelDeploymentCard::default(); + card.worker_type = Some(crate::worker_type::WorkerType::Aggregated); + card.needs = Vec::new(); + card.runtime_config + .set_engine_specific(SGLANG_GENERATE_CAPABILITY, true) + .unwrap(); + let mut worker_set = + crate::discovery::WorkerSet::new("sglang-ns".to_string(), "sglang".to_string(), card); + worker_set.generate_engine = + Some(crate::http::service::generate::tests::engine_behind_frontend(engine.clone())); + assert!( + state + .manager() + .add_worker_set("sglang-concurrent-model", "sglang-ns", worker_set) + ); + + let send = |input_ids: Vec| { + let mut body = serde_json::json!({ + "input_ids": input_ids, + "sampling_params": {"max_new_tokens": 4}, + "stream": false + }); + if let Some(rid) = body_rid { + body["rid"] = serde_json::json!(rid); + } + let request = serde_json::from_value(body).unwrap(); + handler(State(state.clone()), headers.clone(), Ok(Json(request))) + }; + let (first, second) = tokio::time::timeout( + std::time::Duration::from_secs(10), + futures::future::join(send(vec![1, 2, 3]), send(vec![7, 8])), + ) + .await + .expect("concurrent SGLang generate requests must not hang"); + + for (response, expected_ids) in [(first, vec![1, 2, 3]), (second, vec![7, 8])] { + assert_eq!(response.status(), StatusCode::OK); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["output_ids"], serde_json::json!(expected_ids)); + } + + let seen = engine.seen.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert_ne!( + seen[0].0, seen[1].0, + "each request needs its own context id" + ); + for (context_id, rid) in seen.iter() { + assert_ne!(context_id, expected_public_id); + assert_eq!(rid.as_deref(), Some(expected_public_id)); + } + } + + #[tokio::test] + async fn concurrent_sglang_requests_sharing_rid_get_their_own_output() { + concurrent_generate_with_shared_client_id( + HeaderMap::new(), + Some("shared-rid"), + "shared-rid", + ) + .await; + } + + #[test] + fn request_context_keeps_client_ids_as_correlation_only() { + let mut headers = HeaderMap::new(); + headers.insert(X_REQUEST_ID_HEADER, "header-request".parse().unwrap()); + headers.insert(X_DATA_PARALLEL_RANK_HEADER, "2".parse().unwrap()); + + let context = resolve_request_context(&headers, Some("body-rid")); + assert_eq!(context.public_request_id, "body-rid"); + assert!(uuid::Uuid::parse_str(&context.request_id).is_ok()); + assert_eq!(context.data_parallel_rank, Some(2)); + + let context = resolve_request_context(&headers, None); + assert_eq!(context.public_request_id, "header-request"); + assert_ne!(context.request_id, "header-request"); + + let context = resolve_request_context(&HeaderMap::new(), None); + assert_eq!(context.public_request_id, context.request_id); + } + #[tokio::test] async fn worker_unavailable_dispatch_returns_503() { let engine: crate::types::openai::generate::GenerateStreamingEngine = diff --git a/lib/llm/tests/http-service.rs b/lib/llm/tests/http-service.rs index 962d4dcb652c..e79f238476f9 100644 --- a/lib/llm/tests/http-service.rs +++ b/lib/llm/tests/http-service.rs @@ -1620,19 +1620,169 @@ async fn test_request_id_annotation() { ); let received_uuid_str = received_request_id.unwrap(); - assert_eq!( + assert!( + uuid::Uuid::parse_str(&received_uuid_str).is_ok(), + "request_id annotation should be a UUID, got {received_uuid_str}" + ); + assert_ne!( received_uuid_str, request_uuid.to_string(), - "Received request ID should match the one we sent: expected {}, got {}", - request_uuid, - received_uuid_str + "client-supplied x-dynamo-request-id must not become the request id" ); - tracing::info!( - "✅ Request ID annotation test passed! Sent UUID: {}, Received: {}", - request_uuid, - received_uuid_str - ); + cancel_token.cancel(); + task.await.unwrap().unwrap(); +} + +/// Echoes each request's messages back and holds every request until `barrier` releases them +/// together, so concurrent requests overlap in the frontend. Records each request's context id. +struct ConcurrentEchoChatEngine { + barrier: tokio::sync::Barrier, + context_ids: std::sync::Mutex>, +} + +#[async_trait] +impl + AsyncEngine< + SingleIn, + ManyOut>, + Error, + > for ConcurrentEchoChatEngine +{ + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let (request, context) = request.transfer(()); + let ctx = context.context(); + self.context_ids.lock().unwrap().push(ctx.id().to_string()); + self.barrier.wait().await; + + let echo = serde_json::to_string(&request.inner.messages)?; + let mut generator = request.response_generator(ctx.id().to_string()); + let output = generator.create_choice( + 0, + Some(echo), + Some(dynamo_protocols::types::FinishReason::Stop), + None, + ); + Ok(ResponseStream::new( + Box::pin(futures::stream::iter([Annotated::from_data(output)])), + ctx, + )) + } +} + +/// Two concurrent requests on each OpenAI-compatible route share client-supplied +/// `x-dynamo-request-id` and `request-id` headers. Each must complete with its own output, and +/// the engine must see distinct server-minted context ids rather than the client's id. +#[tokio::test] +async fn test_concurrent_requests_with_shared_client_request_id() { + use dynamo_runtime::pipeline::{ServiceBackend, ServiceFrontend, Source}; + + const MODEL: &str = "concurrent-echo-model"; + let (listener, port) = bind_random_port().await; + let service = HttpService::builder() + .port(port) + .enable_chat_endpoints(true) + .enable_responses_endpoints(true) + .enable_anthropic_endpoints(true) + .build() + .unwrap(); + let state = service.state_clone(); + let token = CancellationToken::new(); + let cancel_token = token.clone(); + let task = tokio::spawn(async move { service.run_with_listener(token, listener).await }); + wait_for_service_ready(port).await; + + let engine = Arc::new(ConcurrentEchoChatEngine { + barrier: tokio::sync::Barrier::new(2), + context_ids: std::sync::Mutex::new(Vec::new()), + }); + // Discovery-built pipelines start with a frontend source that matches responses to callers + // by request id; serve the engine through one so the test exercises that map. + let frontend = ServiceFrontend::< + SingleIn, + ManyOut>, + >::new(); + let pipeline = frontend + .link(ServiceBackend::from_engine(engine.clone())) + .unwrap() + .link_terminal(frontend) + .unwrap(); + let card = ModelDeploymentCard::with_name_only(MODEL); + state + .manager() + .add_chat_completions_model(MODEL, card.mdcsum(), pipeline) + .unwrap(); + + let client_id = uuid::Uuid::new_v4().to_string(); + let client = reqwest::Client::new(); + type RequestBody = fn(&str) -> serde_json::Value; + let routes: [(&str, RequestBody); 3] = [ + ("/v1/chat/completions", |marker| { + serde_json::json!({ + "model": MODEL, + "messages": [{"role": "user", "content": marker}], + }) + }), + ( + "/v1/responses", + |marker| serde_json::json!({"model": MODEL, "input": marker}), + ), + ("/v1/messages", |marker| { + serde_json::json!({ + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + }) + }), + ]; + + for (path, body) in routes { + let send = |marker: &'static str| { + client + .post(format!("http://localhost:{port}{path}")) + .header("x-dynamo-request-id", &client_id) + .header("request-id", &client_id) + .json(&body(marker)) + .send() + }; + let (first, second) = timeout( + std::time::Duration::from_secs(10), + futures::future::join(send("alpha-marker"), send("bravo-marker")), + ) + .await + .unwrap_or_else(|_| panic!("{path}: concurrent requests must not hang")); + + for (response, own, other) in [ + (first, "alpha-marker", "bravo-marker"), + (second, "bravo-marker", "alpha-marker"), + ] { + let response = response.unwrap(); + let status = response.status(); + let text = response.text().await.unwrap(); + assert!(status.is_success(), "{path}: status {status}, body {text}"); + assert!(text.contains(own), "{path}: missing own output: {text}"); + assert!( + !text.contains(other), + "{path}: got the other request's output: {text}" + ); + } + + let context_ids = std::mem::take(&mut *engine.context_ids.lock().unwrap()); + assert_eq!(context_ids.len(), 2, "{path}"); + assert_ne!( + context_ids[0], context_ids[1], + "{path}: context ids must differ" + ); + for context_id in &context_ids { + assert_ne!( + context_id, &client_id, + "{path}: client id became the context id" + ); + } + } cancel_token.cancel(); task.await.unwrap().unwrap(); diff --git a/lib/llm/tests/kserve_service.rs b/lib/llm/tests/kserve_service.rs index dfe2b6c24948..0b5fbad88d18 100644 --- a/lib/llm/tests/kserve_service.rs +++ b/lib/llm/tests/kserve_service.rs @@ -105,6 +105,8 @@ pub mod kserve_test { ) -> Result>, Error> { let (request, context) = request.transfer(()); let ctx = context.context(); + assert!(uuid::Uuid::parse_str(ctx.id()).is_ok()); + assert_ne!(Some(ctx.id()), request.inner.user.as_deref()); // let generator = NvCreateChatCompletionStreamResponse::generator(request.model.clone()); let generator = request.response_generator(ctx.id().to_string()); @@ -226,6 +228,8 @@ pub mod kserve_test { // name "repeat", if so, send the same response as many time as the value of the tensor let (request, context) = request.transfer(()); let ctx = context.context(); + assert!(uuid::Uuid::parse_str(ctx.id()).is_ok()); + assert_ne!(Some(ctx.id()), request.id.as_deref()); let repeat_count = request .tensors @@ -504,16 +508,31 @@ pub mod kserve_test { let mut client = get_ready_client(TestPort::InferSuccess as u16, 5).await; let model_name = "split"; + let client_id = uuid::Uuid::new_v4().to_string(); let request = tonic::Request::new(ModelInferRequest { model_name: model_name.into(), model_version: "1".into(), - id: "1234".into(), + id: client_id.clone(), inputs: vec![text_input.clone()], ..Default::default() }); - let response = client.model_infer(request).await.unwrap(); - validate_infer_response(response, model_name); + let request = request.into_inner(); + let mut second_client = client.clone(); + let (first, second) = timeout( + Duration::from_secs(5), + futures::future::join( + client.model_infer(request.clone()), + second_client.model_infer(request), + ), + ) + .await + .expect("requests sharing a client id must not hang"); + for response in [first, second] { + let response = response.unwrap(); + assert_eq!(response.get_ref().id, client_id); + validate_infer_response(response, model_name); + } // Input data in raw_input_content let mut text_input = text_input.clone(); @@ -525,12 +544,13 @@ pub mod kserve_test { let request = tonic::Request::new(ModelInferRequest { model_name: model_name.into(), model_version: "1".into(), - id: "1234".into(), + id: client_id.clone(), inputs: vec![text_input], raw_input_contents: vec![serialized_input], ..Default::default() }); let response = client.model_infer(request).await.unwrap(); + assert_eq!(response.get_ref().id, client_id); validate_infer_response(response, model_name); } @@ -1249,6 +1269,7 @@ pub mod kserve_test { validate_tensor_response( response, model_name, + "1234", inputs, std::collections::HashMap::new(), ); @@ -1592,22 +1613,38 @@ pub mod kserve_test { } let model_name = "tensor"; + let client_id = uuid::Uuid::new_v4().to_string(); let inputs = vec![text_input.clone()]; let request = tonic::Request::new(ModelInferRequest { model_name: model_name.into(), model_version: "1".into(), - id: "1234".into(), + id: client_id.clone(), inputs: inputs.clone(), ..Default::default() }); - let response = client.model_infer(request).await.unwrap(); - validate_tensor_response( - response, - model_name, - inputs, - std::collections::HashMap::new(), - ); + let request = request.into_inner(); + let mut second_client = client.clone(); + let (first, second) = timeout( + Duration::from_secs(5), + futures::future::join( + client.model_infer(request.clone()), + second_client.model_infer(request), + ), + ) + .await + .expect("requests sharing a client id must not hang"); + for response in [first, second] { + let response = response.unwrap(); + assert_eq!(response.get_ref().id, client_id); + validate_tensor_response( + response, + model_name, + &client_id, + inputs.clone(), + std::collections::HashMap::new(), + ); + } // streaming response in model_infer(), expect failure let repeat = inference::model_infer_request::InferInputTensor { @@ -1702,6 +1739,7 @@ pub mod kserve_test { validate_tensor_response( Response::new(response.clone()), model_name, + "1234", inputs.clone(), std::collections::HashMap::from([ ("text_input".into(), serialized_text_input.clone()), @@ -1718,6 +1756,7 @@ pub mod kserve_test { fn validate_tensor_response( response: Response, model_name: &str, + request_id: &str, inputs: Vec, expected_raw_outputs: std::collections::HashMap>, ) { @@ -1733,7 +1772,7 @@ pub mod kserve_test { ); assert_eq!( response.get_ref().id, - "1234", + request_id, "Expected response of the same request ID" ); assert_eq!( diff --git a/lib/runtime/src/logging.rs b/lib/runtime/src/logging.rs index a927d65849ae..1da14ad14111 100644 --- a/lib/runtime/src/logging.rs +++ b/lib/runtime/src/logging.rs @@ -590,6 +590,8 @@ fn extract_trace_parent( (None, None, None, None) }; + // Transport hop only: workers read the frontend's forwarded `request-id`. + // HTTP ingress spans ignore it so a client cannot choose the internal request id. let request_id = headers .get("request-id") .or_else(|| headers.get("x-dynamo-request-id")) @@ -626,12 +628,9 @@ pub fn make_inference_request_span(req: &Request) -> Span { let version = format!("{:?}", req.version()); let (trace_parent, otel_context) = extract_trace_parent(req.headers()); - // Ensure every inference request has a request_id on the span. - // This is the single source of truth — workers and get_or_create_request_id - // read it back via DistributedTraceIdLayer. - let request_id = trace_parent - .request_id - .unwrap_or_else(|| Uuid::new_v4().to_string()); + // Always server-minted and the single source of truth; client `request-id` headers are + // ignored. Client correlation travels separately as `x_request_id`. + let request_id = Uuid::new_v4().to_string(); let span = tracing::info_span!( target: "request_span", @@ -676,10 +675,7 @@ pub fn make_system_request_span(req: &Request) -> Span { let version = format!("{:?}", req.version()); let (trace_parent, otel_context) = extract_trace_parent(req.headers()); - // Ensure every system request has a request_id on the span. - let request_id = trace_parent - .request_id - .unwrap_or_else(|| Uuid::new_v4().to_string()); + let request_id = Uuid::new_v4().to_string(); let span = tracing::debug_span!( target: "system_span", @@ -2527,6 +2523,48 @@ pub mod tests { ); } + #[test] + fn http_request_spans_mint_request_id_and_ignore_client_headers() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("test"); + // Core `set_default`: `SubscriberInitExt::set_default` installs the global `log` + // LogTracer, which would make a later `logging::init()` fail with SetLoggerError. + let _guard = tracing::subscriber::set_default( + tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(tracer)) + .with(DistributedTraceIdLayer), + ); + let client_id = "11111111-2222-3333-4444-555555555555"; + let make_req = || { + Request::builder() + .header("request-id", client_id) + .header("x-dynamo-request-id", client_id) + .header("x-request-id", "client-correlation") + .body(()) + .unwrap() + }; + + let mut seen = Vec::new(); + for make_span in [ + make_inference_request_span::<()>, + make_inference_request_span::<()>, + make_system_request_span::<()>, + ] { + let span = make_span(&make_req()); + let _enter = span.enter(); + let ctx = get_distributed_tracing_context().expect("span has trace context"); + let request_id = ctx.request_id.expect("span has a request_id"); + assert_ne!(request_id, client_id); + assert!(Uuid::parse_str(&request_id).is_ok()); + assert_eq!(ctx.x_request_id.as_deref(), Some("client-correlation")); + seen.push(request_id); + } + assert_ne!(seen[0], seen[1], "each HTTP request must get its own id"); + + let trace_parent = TraceParent::from_headers(make_req().headers()); + assert_eq!(trace_parent.request_id.as_deref(), Some(client_id)); + } + #[test] fn root_context_uses_otel_unsampled_decision() { let provider = SdkTracerProvider::builder() diff --git a/lib/runtime/src/pipeline/error.rs b/lib/runtime/src/pipeline/error.rs index 7d7fa3c9f1b1..d6259bf308ac 100644 --- a/lib/runtime/src/pipeline/error.rs +++ b/lib/runtime/src/pipeline/error.rs @@ -60,6 +60,11 @@ pub enum PipelineError { #[error("Unlinked response; response task was dropped or cancelled")] DetachedStreamSender, + /// A request was issued while another in-flight request through the same source already + /// uses this id. Responses are routed back by id, so the new request is refused. + #[error("Duplicate in-flight request id: {0}")] + DuplicateRequestId(String), + #[error("Serialzation Error: {0}")] SerializationError(String), diff --git a/lib/runtime/src/pipeline/network/egress/push_router.rs b/lib/runtime/src/pipeline/network/egress/push_router.rs index d20c398ac473..6ef445688774 100644 --- a/lib/runtime/src/pipeline/network/egress/push_router.rs +++ b/lib/runtime/src/pipeline/network/egress/push_router.rs @@ -91,6 +91,7 @@ fn is_inhibited(err: &(dyn std::error::Error + 'static)) -> bool { ErrorType::WorkerUnavailable, ]; // Pre-stream request rejections can carry an outer CannotConnect wrapper. + // A `Conflict` is a worker refusing a request id it is still serving. !is_worker_overloaded(err) && match_error_chain( err, @@ -99,6 +100,7 @@ fn is_inhibited(err: &(dyn std::error::Error + 'static)) -> bool { ErrorType::InvalidRequest, ErrorType::InvalidArgument, ErrorType::Backend(BackendError::InvalidArgument), + ErrorType::Conflict, ], ) } @@ -2773,6 +2775,7 @@ mod tests { (ErrorType::InvalidRequest, false), (ErrorType::InvalidArgument, false), (ErrorType::Backend(BackendError::InvalidArgument), false), + (ErrorType::Conflict, false), (ErrorType::CannotConnect, true), (ErrorType::Unavailable, true), ] { diff --git a/lib/runtime/src/pipeline/network/ingress/push_handler.rs b/lib/runtime/src/pipeline/network/ingress/push_handler.rs index 32b2ebc7c48b..2a41779de6b4 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_handler.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_handler.rs @@ -1043,11 +1043,27 @@ where /// and converting the enclosing `PipelineError` yields a bare /// `ErrorType::Unknown`, losing the worker error type. Any other variant carries /// no worker error and converts to `ErrorType::Unknown`. +/// +/// A [`PipelineError::DuplicateRequestId`] refusal is typed +/// [`crate::error::ErrorType::Conflict`]: the worker is healthy and still serving +/// the earlier request with that id, so fault detection must not quarantine it. pub(crate) fn typed_error_from_pipeline_error(e: &PipelineError) -> DynamoError { let source: &(dyn std::error::Error + 'static) = match e { PipelineError::GenerateError(inner) => inner.as_ref(), other => other, }; + let is_duplicate_id = std::iter::successors(Some(source), |e| e.source()).any(|e| { + matches!( + e.downcast_ref::(), + Some(PipelineError::DuplicateRequestId(_)) + ) + }); + if is_duplicate_id { + return DynamoError::builder() + .error_type(crate::error::ErrorType::Conflict) + .message(source.to_string()) + .build(); + } DynamoError::from(source) } @@ -1097,6 +1113,17 @@ mod tests { assert_eq!(error.reason().as_str(), "runtime.unclassified"); } + #[test] + fn duplicate_request_id_is_a_conflict() { + let e = PipelineError::GenerateError(anyhow::Error::new( + PipelineError::DuplicateRequestId("same-id".to_string()), + )); + let error = typed_error_from_pipeline_error(&e); + + assert_eq!(error.class(), ErrorType::Conflict); + assert_eq!(error.reason().as_str(), "request.conflict"); + } + #[derive(Clone, Default)] struct AdmissionCapture { started: Arc, diff --git a/lib/runtime/src/pipeline/nodes/sources/base.rs b/lib/runtime/src/pipeline/nodes/sources/base.rs index ee091a2ac269..6c025a8ac783 100644 --- a/lib/runtime/src/pipeline/nodes/sources/base.rs +++ b/lib/runtime/src/pipeline/nodes/sources/base.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::hash_map::Entry; + use crate::engine::AsyncEngineContextProvider; use super::*; @@ -59,15 +61,54 @@ impl Sink for impl AsyncEngine for Frontend { async fn generate(&self, request: In) -> Result { let (tx, rx) = oneshot::channel::(); + let id = request.id().to_string(); { + // The response is routed back by request id, so a second in-flight request with the + // same id would steal or orphan the first caller's stream. Refuse it instead. let mut sinks = self.sinks.lock().unwrap(); - sinks.insert(request.id().to_string(), tx); + match sinks.entry(id.clone()) { + Entry::Occupied(_) => return Err(PipelineError::DuplicateRequestId(id).into()), + Entry::Vacant(slot) => { + slot.insert(tx); + } + } } + // Removes the sender if `on_next` fails or this future is dropped before `on_next` + // returns; otherwise the stale entry would refuse every later request with this id. + let guard = PendingSinkGuard { + sinks: &self.sinks, + id: Some(id), + }; self.on_next(request, private::Token {}).await?; + // `on_data` has already taken the sender. Disarm before `rx.await` so the guard cannot + // remove a later request's entry for the same id. + guard.disarm(); Ok(rx.await.map_err(|_| PipelineError::DetachedStreamSender)?) } } +/// Owns a `Frontend` sink entry until the request has been handed downstream. +struct PendingSinkGuard<'a, Out> { + sinks: &'a Mutex>>, + id: Option, +} + +impl PendingSinkGuard<'_, Out> { + fn disarm(mut self) { + self.id = None; + } +} + +impl Drop for PendingSinkGuard<'_, Out> { + fn drop(&mut self) { + if let Some(id) = self.id.take() { + // A poisoned lock means another thread panicked; the entry is still ours to remove. + let mut sinks = self.sinks.lock().unwrap_or_else(|e| e.into_inner()); + sinks.remove(&id); + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/lib/runtime/tests/pipeline.rs b/lib/runtime/tests/pipeline.rs index 303143426ead..35645f4f9090 100644 --- a/lib/runtime/tests/pipeline.rs +++ b/lib/runtime/tests/pipeline.rs @@ -417,3 +417,148 @@ async fn distinct_response_context_retains_operator_graph() { assert!(backend_weak.upgrade().is_none()); assert!(operator_weak.upgrade().is_none()); } + +/// Holds every request inside `generate` until the test adds permits, so requests overlap in the +/// frontend's in-flight map. Echoes the request content back as its only output. +struct GatedEchoBackend { + gate: Arc, + arrived: tokio::sync::mpsc::UnboundedSender, +} + +#[async_trait::async_trait] +impl AsyncEngine, ManyOut>, Error> for GatedEchoBackend { + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let _ = self.arrived.send(request.content().clone()); + self.gate.acquire().await?.forget(); + let ctx = request.context(); + let content = request.content().clone(); + Ok(ResponseStream::new( + Box::pin(stream::iter([Annotated::Data(content)])), + ctx, + )) + } +} + +#[allow(clippy::type_complexity)] +fn make_gated_service() -> ( + ServiceEngine, ManyOut>>, + Arc, + tokio::sync::mpsc::UnboundedReceiver, +) { + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let (arrived, arrivals) = tokio::sync::mpsc::unbounded_channel(); + let frontend = ServiceFrontend::, ManyOut>>::new(); + let backend = ServiceBackend::from_engine(Arc::new(GatedEchoBackend { + gate: gate.clone(), + arrived, + })); + let service = frontend + .link(backend) + .unwrap() + .link_terminal(frontend) + .unwrap(); + (service, gate, arrivals) +} + +fn request_with_id(content: &str, id: &str) -> SingleIn { + Context::with_id_and_metadata(content.to_string(), id.to_string(), Default::default()) +} + +async fn collect_data(stream: ManyOut>) -> Vec { + stream + .filter_map(|item| async move { + match item { + Annotated::Data(data) => Some(data), + _ => None, + } + }) + .collect() + .await +} + +#[tokio::test] +async fn frontend_refuses_duplicate_in_flight_request_id() { + let (service, gate, mut arrivals) = make_gated_service(); + + let first_service = service.clone(); + let first = tokio::spawn(async move { + let stream = first_service + .generate(request_with_id("first", "same-id")) + .await?; + Ok::<_, Error>(collect_data(stream).await) + }); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), arrivals.recv()) + .await + .expect("first request must reach the backend") + .as_deref(), + Some("first") + ); + + // The first request is still parked, so its id is in flight. Keep the gate closed until + // the duplicate is answered, or the first id could be released first. + let second_service = service.clone(); + let second = tokio::spawn(async move { + second_service + .generate(request_with_id("second", "same-id")) + .await + .map(|_| ()) + }); + let second = tokio::time::timeout(Duration::from_secs(5), second).await; + gate.add_permits(2); + + let second_err = second + .expect("duplicate request must not wait for the first one") + .unwrap() + .expect_err("duplicate in-flight id must be refused") + .try_into_pipeline_error() + .expect("duplicate refusal is a PipelineError"); + assert!( + matches!(&second_err, PipelineError::DuplicateRequestId(id) if id == "same-id"), + "unexpected error: {second_err:?}" + ); + assert_eq!(first.await.unwrap().unwrap(), vec!["first".to_string()]); + + // Once the first request has completed, the id can be reused. + let stream = service + .generate(request_with_id("third", "same-id")) + .await + .unwrap(); + assert_eq!(collect_data(stream).await, vec!["third".to_string()]); +} + +#[tokio::test] +async fn frontend_releases_request_id_when_caller_is_dropped() { + let (service, gate, mut arrivals) = make_gated_service(); + + let parked_service = service.clone(); + let parked = tokio::spawn(async move { + let _ = parked_service + .generate(request_with_id("parked", "same-id")) + .await; + }); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), arrivals.recv()) + .await + .expect("parked request must reach the backend") + .as_deref(), + Some("parked") + ); + + // Drop the caller while its request is still being handed to the backend. + parked.abort(); + assert!(parked.await.unwrap_err().is_cancelled()); + + gate.add_permits(1); + let stream = tokio::time::timeout( + Duration::from_secs(5), + service.generate(request_with_id("retry", "same-id")), + ) + .await + .expect("retry must not wait") + .expect("a dropped caller must not keep its request id in flight"); + assert_eq!(collect_data(stream).await, vec!["retry".to_string()]); +} diff --git a/lib/sidecar/testkit/src/control.rs b/lib/sidecar/testkit/src/control.rs index fdc199c39bb2..cd121af04e16 100644 --- a/lib/sidecar/testkit/src/control.rs +++ b/lib/sidecar/testkit/src/control.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use futures::{StreamExt, stream::BoxStream}; @@ -155,12 +155,15 @@ impl RequestHandle

{ pub struct Controller { requests: Arc>>>, + /// Plans for requests whose worker id is not known in advance, claimed in arrival order. + unbound: Arc>>>, } impl Default for Controller

{ fn default() -> Self { Self { requests: Arc::default(), + unbound: Arc::default(), } } } @@ -169,6 +172,7 @@ impl Clone for Controller

{ fn clone(&self) -> Self { Self { requests: Arc::clone(&self.requests), + unbound: Arc::clone(&self.unbound), } } } @@ -185,13 +189,29 @@ impl Controller

{ handle } + /// Bind `plan` to the next request whose id was not registered with [`Self::request`]. + /// + /// HTTP ingress mints the worker request id on the server, so a test driving the frontend + /// cannot choose it; it queues the plan here before sending the request instead. + pub fn next_request(&self, plan: RequestPlan) -> RequestHandle

{ + let handle = RequestHandle::new(plan); + self.unbound.lock().unwrap().push_back(handle.clone()); + handle + } + pub async fn open(&self, request: &P::Request) -> Result, P::Error> { let handle = self .requests .lock() .unwrap() .entry(P::request_id(request).to_owned()) - .or_insert_with(|| RequestHandle::new(RequestPlan::default())) + .or_insert_with(|| { + self.unbound + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| RequestHandle::new(RequestPlan::default())) + }) .clone(); { let mut observed = handle.0.observed.lock().unwrap(); diff --git a/lib/sidecar/testkit/tests/router_sidecar_mocker_integration.rs b/lib/sidecar/testkit/tests/router_sidecar_mocker_integration.rs index 68bbc21bcc75..252678efab3c 100644 --- a/lib/sidecar/testkit/tests/router_sidecar_mocker_integration.rs +++ b/lib/sidecar/testkit/tests/router_sidecar_mocker_integration.rs @@ -242,29 +242,25 @@ async fn native_errors_reach_http_and_recover() { let mut child = env.spawn::(&peer.endpoint(), DisaggregationMode::Aggregated, 5); let router = env.ready("backend").await; let frontend = Frontend::start(&env).await; - for (index, is_streaming) in [false, true].into_iter().enumerate() { - let id = format!("00000000-0000-4000-8000-{:012}", index * 2 + 1); - let handle = control.request( - &id, - if is_streaming { - RequestPlan { - stream: Some(StreamFault { - at: StreamPoint::TokenResponse(1), - action: StreamAction::Fail, - pause: true, - }), - ..Default::default() - } - } else { - RequestPlan { - open: OpenAction::Reject, - ..Default::default() - } - }, - ); + for is_streaming in [false, true] { + let handle = control.next_request(if is_streaming { + RequestPlan { + stream: Some(StreamFault { + at: StreamPoint::TokenResponse(1), + action: StreamAction::Fail, + pause: true, + }), + ..Default::default() + } + } else { + RequestPlan { + open: OpenAction::Reject, + ..Default::default() + } + }); let mut response = bounded( "HTTP native failure response", - frontend.generate(&env.model, &id, is_streaming).send(), + frontend.generate(&env.model, is_streaming).send(), ) .await .unwrap(); @@ -329,13 +325,12 @@ async fn native_errors_reach_http_and_recover() { .await; peer.scheduler_idle().await; - let recovery_id = format!("00000000-0000-4000-8000-{:012}", index * 2 + 2); - let recovered = control.request(&recovery_id, RequestPlan::default()); + let recovered = control.next_request(RequestPlan::default()); let response = tokio::time::timeout(Duration::from_secs(20), async { loop { let response = bounded( "HTTP recovery response", - frontend.generate(&env.model, &recovery_id, false).send(), + frontend.generate(&env.model, false).send(), ) .await .unwrap(); diff --git a/lib/sidecar/testkit/tests/support/process.rs b/lib/sidecar/testkit/tests/support/process.rs index 8a9d929df14d..eb89d06acdc8 100644 --- a/lib/sidecar/testkit/tests/support/process.rs +++ b/lib/sidecar/testkit/tests/support/process.rs @@ -238,15 +238,12 @@ impl Frontend { } } - pub fn generate(&self, model: &str, id: &str, is_streaming: bool) -> reqwest::RequestBuilder { - self.client - .post(&self.url) - .header("x-dynamo-request-id", id) - .json(&serde_json::json!({ - "model": model, "prompt": "hello", "max_tokens": 3, - "stream": is_streaming, "temperature": 0, "ignore_eos": true, - "nvext": {"extra_fields": ["completion_token_ids"]}, - })) + pub fn generate(&self, model: &str, is_streaming: bool) -> reqwest::RequestBuilder { + self.client.post(&self.url).json(&serde_json::json!({ + "model": model, "prompt": "hello", "max_tokens": 3, + "stream": is_streaming, "temperature": 0, "ignore_eos": true, + "nvext": {"extra_fields": ["completion_token_ids"]}, + })) } pub async fn shutdown(mut self) { diff --git a/tests/frontend/test_request_tracing_logs.py b/tests/frontend/test_request_tracing_logs.py index a72cf06cd953..e1bb7e6a3454 100644 --- a/tests/frontend/test_request_tracing_logs.py +++ b/tests/frontend/test_request_tracing_logs.py @@ -381,11 +381,10 @@ def test_agg_404_error(tracing_services) -> None: def test_agg_invalid_uuid_warn(tracing_services) -> None: - """Invalid x-dynamo-request-id: WARN logged, request proceeds with generated ID.""" + """Deprecated x-dynamo-request-id is ignored: WARN logged, request uses a generated ID.""" port = tracing_services["frontend_port"] - # Send deprecated x-dynamo-request-id with invalid value to test deprecation warning - # Send both x-request-id (for log filtering) and deprecated x-dynamo-request-id (invalid) + # Send both x-request-id (for log filtering) and the deprecated x-dynamo-request-id rid = str(uuid.uuid4()) resp = requests.post( f"http://localhost:{port}/v1/chat/completions", @@ -406,15 +405,15 @@ def test_agg_invalid_uuid_warn(tracing_services) -> None: req_logs = get_request_logs(tracing_services["frontend"], rid) - # A WARN log should be emitted about the invalid UUID warn_logs = [ e for e in req_logs - if e.get("level") == "WARN" and "must be a valid UUID" in e.get("message", "") + if e.get("level") == "WARN" + and "deprecated request id header ignored" in e.get("message", "") + and e.get("header") == "x-dynamo-request-id" ] assert len(warn_logs) == 1 - # Request still gets a valid request_id (generated, not the invalid one) received = [e for e in req_logs if e.get("message") == "request received"] assert len(received) == 1 request_id = received[0].get("request_id", "")