fix(router): propagate trace context through worker dispatch

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 }

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<Sel>(
     mut response_stream: ManyOut<Annotated<LLMEngineOutput>>,
     context: Arc<dyn AsyncEngineContext>,
@@ -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,13 +753,18 @@ 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,
         },
         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::<KvPushRouter<WorkerSelectionPolicy>>();
     }
 
+    #[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::<Vec<_>>();
+
+        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() {