Skip to content

fix(bindings): restore trace context across the Python chat-processor boundary (3 orphan traces → 1 waterfall)#2

Closed
renancloudwalk wants to merge 1 commit into
tmp-upstream-syncfrom
fix/python-processor-trace-context
Closed

fix(bindings): restore trace context across the Python chat-processor boundary (3 orphan traces → 1 waterfall)#2
renancloudwalk wants to merge 1 commit into
tmp-upstream-syncfrom
fix/python-processor-trace-context

Conversation

@renancloudwalk

@renancloudwalk renancloudwalk commented Jul 7, 2026

Copy link
Copy Markdown

Problem

With the Python chat processor (--dyn-chat-processor sglang), distributed tracing shatters: one request produces three unrelated OTel traces instead of one waterfall:

  1. http-request — alone, 1 span, orphan
  2. prefill tree — router.route_request → prefiller handle_payload → SGLang prefill_* spans, orphan root
  3. decode tree — router.route_request → decoder handle_payload → SGLang decode_* spans, orphan root

Cross-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:

Frontend config Traces per request
pure Rust (no --dyn-chat-processor) 1 — full waterfall: http-request → 2× router.route_request → prefiller + decoder + all SGLang spans (50 spans, one root)
--dyn-chat-processor sglang (stock, unpatched) 3 — orphan roots as above

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:

  1. Rust→Python (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). ✅
  2. Python→Rust (lib/bindings/python/rust/llm/routed_engine.rs): RoutedEngine::generate receives that same context back, uses it for cancellation linkage and metadata — but ignores trace_context — and runs inner.generate(...) inside pyo3_async_runtimes::tokio::future_into_py(async move {...}), a fresh tokio task where Span::current() is empty. ❌

Downstream, span parentage is propagated from Span::current() (inject_trace_headers_into_map at egress writes traceparent), so with no current span every router.route_request becomes a new root trace — one for the prefill dispatch, one for decode.

Fix

In RoutedEngine::generate, when the PyContext carries a DistributedTraceContext, create a routed_engine.generate span, set_parent it to the captured remote span context (same pattern as logging.rs make_handle_payload_span and the add_link in lib/backend-common/src/adapter.rs), and .instrument() the dispatch future with it.

Result: router.route_request (prefill + decode) parent to routed_engine.generate → which parents to http-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)

  • The Bootstrap arm of the PrefillRouter never sets migration_link (lib/llm/src/kv_router/prefill_router/mod.rs:287-292; spawn_prefill_task discards the PrefillCompletion.worker_link it 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 test test_disagg_decode_span_links_to_prefill_span can't see it (the sample worker never returns bootstrap_info, so the test only exercises the Completed arm).
  • Adjacent upstream PRs checked, neither covers this: feat(tracing): trace cancellation and migration attempts ai-dynamo/dynamo#11133 (cancellation/migration attempt tracing), fix(observability): idempotent logging init + disagg span-link test deadlock ai-dynamo/dynamo#11020 (logging-init idempotency + disagg span-link test deadlock fix).

Validation

  • cargo check clean, cargo fmt applied.
  • Runtime-validated on the live cluster (2026-07-07): built ai-dynamo-runtime wheel 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:
    • Before (stock wheel): 3 orphan traces.
    • After (patched wheel): 1 trace, 53 spans, rooted at http-requestrouted_engine.generate → 2× router.route_request → prefiller (prefill_forward, prefill_transfer_kv_cache) + decoder (decode_transferred, decode_loop) + all SGLang spans. Full PD waterfall.
    • Note: a wheel built from Jul-7 main (with or without this patch) hangs against the Jul-5 deployed Python frontend code — upstream API drift, unrelated to this change. Validation used the version-matched base.

Upstream

Intent is to validate on our cluster, then offer this upstream to ai-dynamo/dynamo.

@cloudwalk-review-agent cloudwalk-review-agent Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

@cloudwalk-review-agent cloudwalk-review-agent Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@renancloudwalk

Copy link
Copy Markdown
Author

Superseded by the personal-fork PR: renflowerz#1 (same patch, runtime-validated). Closing.

@renancloudwalk
renancloudwalk deleted the fix/python-processor-trace-context branch July 8, 2026 12:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant