diff --git a/Cargo.lock b/Cargo.lock index 21b517ae6cf6..8fc94a533d47 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2742,6 +2742,7 @@ dependencies = [ "tonic-build 0.13.1", "tower-http", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "url", "utoipa", diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index 5b5cd565cb34..f168997b8555 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -1788,6 +1788,7 @@ dependencies = [ "tonic-build 0.13.1", "tower-http", "tracing", + "tracing-opentelemetry", "url", "utoipa", "utoipa-swagger-ui", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 18ece605de13..75c70c7b0a98 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -1854,6 +1854,7 @@ dependencies = [ "tonic-build 0.13.1", "tower-http", "tracing", + "tracing-opentelemetry", "url", "utoipa", "utoipa-swagger-ui", diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index 970576a2e9cb..1d48ca29a04a 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -109,6 +109,7 @@ tokio = { workspace = true } tokio-stream = { workspace = true } tokio-util = { workspace = true } tracing = { workspace = true } +tracing-opentelemetry = { workspace = true } opentelemetry = { workspace = true } opentelemetry_sdk = { workspace = true } opentelemetry-otlp = { workspace = true } diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index cbec4df49e15..8931bf77937e 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -9,6 +9,10 @@ use dynamo_kv_router::{ }; use dynamo_runtime::{ error::{DynamoError, ErrorType, match_error_chain}, + logging::{ + DistributedTraceContext, get_distributed_tracing_context, + otel_parent_context_from_distributed, + }, metrics::frontend_perf::{STAGE_ROUTE, StageGuard}, pipeline::{ AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, Error, ManyOut, PushRouter, @@ -17,7 +21,8 @@ use dynamo_runtime::{ protocols::annotated::Annotated, }; use futures::stream::{self, StreamExt}; -use tracing::Instrument; +use tracing::{Instrument, Span}; +use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::{ kv_router::{ @@ -64,6 +69,35 @@ fn route_target(worker: WorkerWithDpRank) -> AffinityTarget { AffinityTarget::new(worker.worker_id, Some(worker.dp_rank)) } +fn route_request_span( + context_id: &str, + selection: &WorkerSelection, + phase: &RequestPhase, + trace_context: Option<&DistributedTraceContext>, +) -> Span { + let request_id = trace_context + .and_then(|context| context.request_id.as_deref()) + .unwrap_or(context_id); + let span = tracing::info_span!( + target: "request_span", + "kv_router.route_request", + request_id = %request_id, + worker_id = selection.worker.worker_id, + dp_rank = selection.worker.dp_rank, + overlap_blocks = selection.overlap_amount, + phase = ?phase, + trace_id = trace_context.map(|context| context.trace_id.as_str()), + parent_id = trace_context.map(|context| context.span_id.as_str()), + trace_flags = trace_context.map(|context| context.trace_flags.as_str()), + tracestate = trace_context.and_then(|context| context.tracestate.as_deref()), + x_request_id = trace_context.and_then(|context| context.x_request_id.as_deref()), + ); + if let Some(parent) = trace_context.and_then(otel_parent_context_from_distributed) { + let _ = span.set_parent(parent); + } + span +} + fn monitor_response_stream( mut response_stream: ManyOut>, context: Arc, @@ -411,19 +445,13 @@ where .await } }; - let dispatch_result = cancel_on_stop( - request_context.as_ref(), - dispatch.instrument(tracing::info_span!( - "kv_router.route_request", - request_id = %context_id, - worker_id = selection.worker.worker_id, - dp_rank = selection.worker.dp_rank, - overlap_blocks = selection.overlap_amount, - phase = ?phase, - )), - ) - .await - .and_then(|result| result); + let trace_context = get_distributed_tracing_context(); + let route_span = + route_request_span(&context_id, &selection, &phase, trace_context.as_ref()); + let dispatch_result = + cancel_on_stop(request_context.as_ref(), dispatch.instrument(route_span)) + .await + .and_then(|result| result); let response_stream = match dispatch_result { Ok(stream) => stream, Err(error) => { @@ -725,6 +753,10 @@ mod tests { discovery::EventTransportKind, distributed::{DiscoveryBackend, DistributedConfig, RequestPlaneMode}, error::{ErrorType, match_error_chain}, + logging::{ + DistributedTraceIdLayer, inject_trace_headers_into_map, + make_handle_payload_span_from_tcp_headers, + }, pipeline::{ AddressedRequest, AsyncEngineContext, Context, ManyIn, Operator, PushRouter, RouterMode, ServerStreamingEngine, StreamingDispatch, context::Controller, @@ -732,6 +764,7 @@ mod tests { storage::kv::Selector, }; use tokio::sync::watch; + use tracing_subscriber::layer::SubscriberExt; use super::*; use crate::{ @@ -774,6 +807,78 @@ mod tests { assert_send_sync::>(); } + #[test] + fn route_request_span_propagates_distributed_trace_context() { + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider}; + + const TRACE_ID: &str = "11111111111111111111111111111111"; + const PARENT_SPAN_ID: &str = "2222222222222222"; + const REQUEST_ID: &str = "33333333-3333-4333-8333-333333333333"; + + let provider = SdkTracerProvider::builder() + .with_sampler(Sampler::ParentBased(Box::new(Sampler::AlwaysOn))) + .build(); + let tracer = provider.tracer("kv-router-test"); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(tracer)) + .with(DistributedTraceIdLayer); + let _guard = tracing::subscriber::set_default(subscriber); + let mut inbound_headers = HashMap::from([ + ( + "traceparent".to_string(), + format!("00-{TRACE_ID}-{PARENT_SPAN_ID}-00"), + ), + ("tracestate".to_string(), "vendor=dynamo".to_string()), + ("x-request-id".to_string(), "external-1".to_string()), + ("request-id".to_string(), REQUEST_ID.to_string()), + ]); + let ingress_span = make_handle_payload_span_from_tcp_headers( + &inbound_headers, + "frontend", + "generate", + "test", + 1, + ); + let trace_context = ingress_span + .in_scope(get_distributed_tracing_context) + .expect("ingress span must expose its distributed trace context"); + let selection = WorkerSelection { + worker: WorkerWithDpRank::new(7, 3), + overlap_amount: 2, + effective_overlap_blocks: 2.0, + cached_tokens: 32, + routing_hashes: None, + router_hint: None, + }; + + let span = route_request_span( + "trace-propagation", + &selection, + &RequestPhase::Aggregated, + Some(&trace_context), + ); + assert_eq!( + span.metadata().map(|metadata| metadata.target()), + Some("request_span") + ); + + inbound_headers.clear(); + span.in_scope(|| inject_trace_headers_into_map(&mut inbound_headers)); + let traceparent = inbound_headers + .get("traceparent") + .expect("route span must provide an outbound trace context"); + let fields = traceparent.split('-').collect::>(); + + assert_eq!(fields.len(), 4); + assert_eq!(fields[1], TRACE_ID); + assert_ne!(fields[2], trace_context.span_id); + assert_eq!(fields[3], "00"); + assert_eq!(inbound_headers["tracestate"], "vendor=dynamo"); + assert_eq!(inbound_headers["x-request-id"], "external-1"); + assert_eq!(inbound_headers["request-id"], REQUEST_ID); + } + #[tokio::test] #[serial_test::serial] async fn terminal_item_does_not_skip_transport_eof() {