fix(bindings): restore trace context across the Python chat-processor boundary (3 orphan traces → 1 waterfall)#2
Conversation
There was a problem hiding this comment.
Looks good to merge.
This change is tightly scoped and addresses a concrete trace-context loss at the Python→Rust boundary in RoutedEngine::generate. The new dispatch_span re-parenting using trace_context + .instrument(dispatch_span) is consistent with the stated root cause, and the fallback behavior (invalid/missing IDs => no forced parent) remains safe and non-disruptive.
I did not find correctness, auth, money-movement, data-loss, concurrency, or deploy-safety regressions in the provided diff.
… boundary RoutedEngine::generate receives the PyContext that engine.rs populated with the caller's DistributedTraceContext (the http-request span identity), but ignored it and dispatched inside a fresh tokio task where Span::current() is empty. Every downstream router.route_request span therefore became an orphan trace root: one request = three unrelated traces (http-request / prefill / decode) whenever the Python chat processor (--dyn-chat-processor sglang) is in the request path. Re-parent a routed_engine.generate span to the captured trace context and instrument the dispatch future with it, so the full PD pipeline lands on the request's trace — matching the pure-Rust processor behavior (proven 1-trace vs 3-trace A/B on a live GLM-5.2 NIXL disagg cluster). No-op when the context carries no trace identity (bare test contexts, tracing disabled).
2f15573 to
56dde92
Compare
There was a problem hiding this comment.
Looks good to merge.
I don’t see a concrete correctness or concurrency regression in this change. The fix is scoped and aligns with the reported failure mode: creating a local routed_engine.generate span, attaching remote parent context from PyContext.trace_context(), and instrumenting the async future ensures downstream work gets a valid parent in the Python→Rust boundary case.
What I specifically checked:
- Fallback safety: missing/invalid trace IDs gracefully degrade to existing behavior (no forced parent).
- Async behavior: stream loop/cancellation semantics are unchanged;
stop_generating()paths remain intact. - API contract: method signature and payload handling are unchanged.
Minor note (non-blocking): TraceFlags::SAMPLED is hardcoded when rebuilding parent context, so unsampled upstream traces may appear sampled after this boundary. If preserving upstream sampling is important, consider carrying trace flags in DistributedTraceContext and restoring them here.
|
Superseded by the personal-fork PR: renflowerz#1 (same patch, runtime-validated). Closing. |
Problem
With the Python chat processor (
--dyn-chat-processor sglang), distributed tracing shatters: one request produces three unrelated OTel traces instead of one waterfall:http-request— alone, 1 span, orphanrouter.route_request→ prefillerhandle_payload→ SGLangprefill_*spans, orphan rootrouter.route_request→ decoderhandle_payload→ SGLangdecode_*spans, orphan rootCross-phase TTFT attribution (prefill compute vs NIXL handoff vs decode queueing) becomes impossible in Jaeger without manual timestamp correlation.
Proof (A/B on the same live cluster, same workers, minutes apart)
GLM-5.2 disagg P/D (2× NVFP4 prefillers + 1× FP8 decoder, NIXL bootstrap), ai-dynamo 1.3.0.dev20260705,
OTEL_EXPORT_ENABLED=true,DYN_LOGGING_JSONL=1, OTLP → Jaeger, SGLang--enable-trace:--dyn-chat-processor)http-request→ 2×router.route_request→ prefiller + decoder + all SGLang spans (50 spans, one root)--dyn-chat-processor sglang(stock, unpatched)Same etcd/NATS/workers. The only variable is the Python processor in the request path.
Root cause
The trace context survives the Rust→Python hop but is dropped on the Python→Rust hop:
lib/bindings/python/rust/engine.rs): the bridge captures the caller's trace identity into the PyContext —Context::new(ctx, get_distributed_tracing_context(), None, metadata). ✅lib/bindings/python/rust/llm/routed_engine.rs):RoutedEngine::generatereceives that same context back, uses it for cancellation linkage and metadata — but ignorestrace_context— and runsinner.generate(...)insidepyo3_async_runtimes::tokio::future_into_py(async move {...}), a fresh tokio task whereSpan::current()is empty. ❌Downstream, span parentage is propagated from
Span::current()(inject_trace_headers_into_mapat egress writestraceparent), so with no current span everyrouter.route_requestbecomes a new root trace — one for the prefill dispatch, one for decode.Fix
In
RoutedEngine::generate, when the PyContext carries aDistributedTraceContext, create arouted_engine.generatespan,set_parentit to the captured remote span context (same pattern aslogging.rsmake_handle_payload_spanand theadd_linkinlib/backend-common/src/adapter.rs), and.instrument()the dispatch future with it.Result:
router.route_request(prefill + decode) parent torouted_engine.generate→ which parents tohttp-request. One trace, full PD waterfall, matching the pure-Rust processor behavior.No behavior change when the context has no trace identity (bare test contexts, tracing disabled): falls back to
Span::current()(a no-op), i.e. exactly today's behavior.Related (not fixed here)
migration_link(lib/llm/src/kv_router/prefill_router/mod.rs:287-292;spawn_prefill_taskdiscards thePrefillCompletion.worker_linkit computes). On the Completed path (vLLM/TRT-LLM) the decode span gets an explicit OTel Link to the prefill worker span; SGLang (the only backend registering bootstrap endpoints) doesn't. With this PR the phases already share one trace, so the missing Link is cosmetic — but the asymmetry exists, and the upstream e2e testtest_disagg_decode_span_links_to_prefill_spancan't see it (the sample worker never returnsbootstrap_info, so the test only exercises the Completed arm).Validation
cargo checkclean,cargo fmtapplied.ai-dynamo-runtimewheel from this patch on the dev20260705 base (aeda23b4, matching the deployed Python code), installed into a test frontend with--dyn-chat-processor sglang, sent a request:http-request—routed_engine.generate→ 2×router.route_request→ prefiller (prefill_forward,prefill_transfer_kv_cache) + decoder (decode_transferred,decode_loop) + all SGLang spans. Full PD waterfall.Upstream
Intent is to validate on our cluster, then offer this upstream to ai-dynamo/dynamo.