diff --git a/deploy/tracing/docker-compose.yml b/deploy/tracing/docker-compose.yml index 06ea55050..67b597d46 100644 --- a/deploy/tracing/docker-compose.yml +++ b/deploy/tracing/docker-compose.yml @@ -12,6 +12,13 @@ # # to find requests whose prompt phase (first chunk on GPU → first token) # # blew past 500ms. Emitted spans: request → {queue, prefill, decode}. # +# Fronting openinfer with an OTel-enabled proxy (e.g. vllm-project/router with +# --enable-trace --otlp-traces-endpoint 127.0.0.1:4317) joins both sides into +# one trace — client → router http_request/http_client_request → openinfer +# request → queue/prefill/decode — as long as both export to the same Tempo. +# openinfer's HTTP layer stashes the incoming traceparent and the engine +# bridge joins it; see docs/subsystems/tracing/e2e-router-tracing.md. +# # Switching backend later (VictoriaTraces, etc.) is an endpoint/image change, # not a code change — everything downstream of OTLP is swappable. @@ -35,8 +42,10 @@ services: restart: unless-stopped environment: GF_AUTH_ANONYMOUS_ENABLED: "true" - # Viewer, not Admin: Explore/dashboards need no write access. - GF_AUTH_ANONYMOUS_ORG_ROLE: "Viewer" + # Editor, not Viewer: anonymous Viewers are denied `datasources:explore` + # in Grafana 11.3 (Access denied in the server log, Explore silently + # returns nothing). Loopback-only binding keeps the write surface local. + GF_AUTH_ANONYMOUS_ORG_ROLE: "Editor" GF_AUTH_DISABLE_LOGIN_FORM: "true" # Enable the TraceQL editor and search UX. GF_FEATURE_TOGGLES_ENABLE: "traceqlEditor traceqlSearch" diff --git a/docs/index.md b/docs/index.md index 7ac1f7066..04d668dd5 100644 --- a/docs/index.md +++ b/docs/index.md @@ -137,6 +137,12 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | --- | --- | | `subsystems/sampling/openinfer-sample.md` | `openinfer-sample` is the one crate every model routes through for batched token selection (`select_batch`) and host logprobs (`token_logprob_from_row`, generic over f32/bf16). Replaces `core::ops::select_batch_tokens_into` + three copies of the logprob math. Kimi keeps its sharded-vocab greedy argmax (a DP concern the whole-vocab `select_batch` can't express) but shares the non-greedy sampler and the logprob math. | +## subsystems / tracing + +| Path | TL;DR | +| --- | --- | +| `subsystems/tracing/e2e-router-tracing.md` | Single trace for client → vllm-router → openinfer → prefill/decode verified: the router injects traceparent with OTel on; openinfer's axum middleware stashes it and the bridge joins via `external_req_id` (incl. the `cmpl-`/`chatcmpl-` prefix pitfall). After upstream vllm#50370 merges, migrate per #790 and delete the middleware. | + ## subsystems / frontend | Path | TL;DR | diff --git a/docs/subsystems/tracing/e2e-router-tracing.md b/docs/subsystems/tracing/e2e-router-tracing.md new file mode 100644 index 000000000..a642f1b40 --- /dev/null +++ b/docs/subsystems/tracing/e2e-router-tracing.md @@ -0,0 +1,115 @@ +# Router → OpenInfer E2E Request Tracing + +> **TL;DR:** A single trace for client → vllm-router → openinfer → prefill/decode is working and verified on the local Tempo/Grafana stack. The router already emits spans and injects `traceparent` when OTel is on; on the openinfer side, `openinfer-vllm-frontend/src/trace_context.rs` stashes the incoming `traceparent`, correlates it via `X-Request-Id → external_req_id` (tolerating vllm-server's `cmpl-`/`chatcmpl-` prefixes), and the bridge uses it as the parent of the `request` root span. Once upstream PR vllm-project/vllm#50370 (HTTP-layer `trace_headers` population) merges, migrate per openinfer#790 and delete the middleware. +> +> **Last touched:** 2026-07 + +## Preparation + +- **Read**: + - `docs/index.md` — routing table; no existing tracing subsystem doc (prior tracing work lived only in deploy/tracing + code) + - `deploy/tracing/docker-compose.yml` / `tempo.yaml` — existing local Tempo+Grafana stack, OTLP gRPC on 4317; point `OPENINFER_TRACE_OTLP_ENDPOINT` at it and go. Spans already emitted: `request → {queue, prefill, decode}` + - `openinfer-core/src/tracing.rs` — fastrace → OTLP reporter init; zero cost when `OPENINFER_TRACE_OTLP_ENDPOINT` is unset + - `openinfer-vllm-frontend/src/bridge.rs` — `Span::root("request", SpanContext::random())` (bridge.rs:396): **always opened a fresh trace** — this was the break point; `EngineCoreRequest` arrives from vllm-server over ZMQ + - `openinfer-qwen3/src/scheduler/phase_trace.rs` — queue/prefill/decode spans are children of `request` (via `GenerateRequest.trace_parent`) + - `openinfer-vllm-frontend/src/lib.rs` — `vllm_server::serve_with_router_extension(config, shutdown, extend_router)` exposes a `FnOnce(Router) -> Router` hook, so middleware can be added without patching the git dependency + - vllm-server (verified on both the pinned rev 8e61b64 and upstream main) — `resolve_request_context` extracts only `X-Request-Id`/`X-data-parallel-rank`; the `EngineCoreRequest.trace_headers` protocol field exists but the HTTP layer never populates it; **bumping the pin would not fix this** + - vllm-router (local checkout /data/code/workspace-rustllm/router) — `--enable-trace` + `--otlp-traces-endpoint`; server span `http_request` (extracts the client's traceparent as parent) + client span `http_client_request`, injecting span context into `traceparent`/`tracestate` on worker-bound requests; with OTel off it still forwards client trace headers verbatim; all client headers (incl. `X-Request-Id`) forwarded as-is; `service.name=vllm-router`; a bare host:port endpoint gets `http://` prepended + - fastrace 0.7.17 — `SpanContext::decode_w3c_traceparent(&str) -> Option` (pub, `collector/id.rs:281`); `EngineCoreRequest.external_req_id` works as the correlation key + - `openinfer-engine/src/tracing_state.rs` — global AtomicBool gate; tests must not flip it (parallel-test races), so the middleware core was written as a pure function that never reads the global flag + +- **Relevant history**: + - `docs/roadmap/roadmap-2026-h2.md` — "observability wiring" is on the H2 plan; this task is part of it + - `docs/subsystems/router/kv-aware-routing.md` — earlier Dynamo-router multi-turn routing experiment (different router, but the e2e measurement approach carried over) + - No past tracing task docs found + +- **Plan**: + 1. Add W3C trace-context intake on the openinfer side (the only code change in this repo): + - New module `openinfer-vllm-frontend/src/trace_context.rs`: axum middleware reads `traceparent`; if present, take `X-Request-Id` (generate and inject a short id when absent — vllm-server resolves it into `external_req_id`) and stash into a bounded shared map (`external_req_id → traceparent`, TTL/capacity eviction) + - Mount the layer for all serving variants inside `serve_model_on_host_with_router_extension` + - `bridge.rs` add_request: when tracing is on, pop the map by `external_req_id`; on successful `decode_w3c_traceparent`, use it as the parent of the `request` root span; otherwise fall back to `SpanContext::random()` (current behavior) + - Unit tests: middleware injection/no-header paths + bridge parent resolution (TestReporter pattern from phase_trace.rs) + 2. Bring up the local stack: `docker compose -f deploy/tracing/docker-compose.yml up -d` (Tempo 4317 / Grafana 3000) + 3. Start openinfer: `OPENINFER_TRACE_OTLP_ENDPOINT=http://127.0.0.1:4317 cargo run --release -- --model-path models/Qwen3-4B --port 8000` (confirm weights and GPU first) + 4. Build and start the router (/data/code/workspace-rustllm/router): `cargo build --release`, then `vllm-router --worker-urls http://127.0.0.1:8000 --port 8090 --enable-trace --otlp-traces-endpoint 127.0.0.1:4317` + 5. Verify one chain: send `/v1/completions` through the router; assert via Tempo HTTP API (`:3200/api/search` + `/api/traces/`) that **one trace** contains router `http_request` → `http_client_request` → openinfer `request` → `queue`/`prefill`/`decode` with correct parenting + 6. Observe e2e behavior: small concurrent load (vllm-bench or multi-turn curl) — router forwarding overhead, queue wait, prefill/decode breakdown; also verify the client-supplied-traceparent case (the whole chain should hang under the client's span) + 7. Wrap up: update deploy/tracing comments and docs/index.md, write the Debrief + +- **Decision (2026-07-30 review)**: dual track — land the local MVP (middleware workaround) to validate the chain now, while opening an upstream PR against vllm-project/vllm (Rust server extracts traceparent into `trace_headers`, Python parity); track the "middleware → upstream" migration with an issue in this repo. After upstream merges and the pin is bumped, delete the middleware and read `EngineCoreRequest.trace_headers` in the bridge (the bridge change is shared by both tracks, so nothing is wasted). + +- **Risks / open questions**: + - Correlation relies on `X-Request-Id → external_req_id`: the middleware must inject the id before vllm-server reads headers (axum layer order) — verify empirically + - Router `--otlp-traces-endpoint` format (host:port vs URL) and its `service.name` — check at execution + - Single-machine single-worker means the router policy is irrelevant, but router retries/circuit-breaking may add extra spans under load — watch for them + - `models/Qwen3-4B` weights and GPU availability unconfirmed + +## Execution Log + +### Step 1: W3C trace-context intake on the openinfer side (MVP) +- Added `openinfer-vllm-frontend/src/trace_context.rs`: `TraceContextStash` (Arc>, TTL 120s / cap 4096, one-shot `take`) + `stash_trace_context` axum middleware; the core `stash_from_headers` avoids the global flag so unit tests stay race-free +- `lib.rs`: `mod trace_context`; the stash is created at the top of `serve_model_on_host_with_router_extension`, cloned into the engine task (bridge field) and into the extend_router wrapper (`from_fn_with_state` layer as the outermost wrapper, guaranteeing it runs before vllm-server reads headers) +- `bridge.rs`: `LocalEngineBridge` gains a `trace_stash` field; `start_request` destructures `external_req_id`; with tracing on, `take(id) → decode_w3c_traceparent → Span::root parent`, falling back to `SpanContext::random()` on miss/invalid (previous behavior) +- 3 unit tests (roundtrip + W3C decode, id injection, no-header no-op); `cargo test --release -p openinfer-vllm-frontend --lib` 30/30; clippy clean (fixed one single-pattern match → if let) +- Result: success + +### Step 2: Environment checks +- docker OK; Tempo/Grafana up via `deploy/tracing/docker-compose.yml up -d` (4317/3000) +- GPU: RTX 5070 Ti 16GB; model: `/data/models/Qwen3-4B` (no models/ dir in the repo on this machine) +- router and openinfer release builds running in parallel in the background + +### Step 3: Stack bring-up + first verification (fail → fix → pass) +- GPU was fully held by an old openinfer from pegainfer-2 (PID 725807, 13.5GB) → killed after user confirmation; server starts fine (`OPENINFER_TRACE_OTLP_ENDPOINT=http://127.0.0.1:4317`, :8000) +- Router on :8090: `vllm-router --worker-urls http://127.0.0.1:8000 --policy round_robin --enable-trace --otlp-traces-endpoint 127.0.0.1:4317` (built with zero changes, system libzmq present) +- First verification: the router's two spans chained correctly, but openinfer opened a separate random trace — **join failed** +- Diagnosis: direct-to-openinfer with `traceparent` + `X-Request-Id: dbg12345` still failed; the openinfer trace showed `request_id=cmpl-dbg12345-2b7dfbb4` → vllm-server prepends **`cmpl-`/`chatcmpl-`** to X-Request-Id before it becomes `external_req_id` (llm/request.rs: prepare() reuses the route-prefixed id), so the stash key (bare header value) never matched +- Fix: `TraceContextStash::take_for_external_req_id` — exact lookup first, then prefix-stripped; regression test `lookup_tolerates_vllm_api_prefixes`; 31/31 pass, clippy clean +- Re-verify (trace id `aaaa1111…`): **full chain in a single trace** — client(5555eeee) → router `http_request`(334.79ms) → router `http_client_request` → openinfer `request`(319.89ms) → `queue`(0.04ms) → `prefill`(72.06ms) → `decode`(247.67ms), parenting correct at every level +- Result: success + +### Unexpected +- The router's `http_client_request` span once showed a ~5.3s duration (vs the 334ms server span) — looks like a span-lifetime bug on its non-streaming path; a router-side instrumentation issue, doesn't affect chain verification; worth reporting upstream +- Tempo `/api/traces` returns span ids base64-encoded, and spans from multiple requests sharing one trace id get compacted together — use a fresh trace id per debug iteration + +### Step 4: Load + both endpoints (pass) +- 16 req / conc=8 / distinct prompts through the router: client p50 756ms (includes Python urllib per-request connection setup); Tempo shows 16/16 complete chains + - Phase p50: queue 0.02ms / prefill 25.1ms (max 101.7, batching) / decode 359.9ms (32 tok, bs≈8, ~11ms/tok @5070 Ti) / request 466.3ms + - Sampled single trace: router server span 388.50ms vs openinfer request 386.14ms, start offset 1.02ms → **router overhead is ~1ms**; another ~6.8ms sits in vllm-server HTTP/tokenize before the bridge opens its span (no spans there — shows as a gap) +- `/v1/chat/completions` (chatcmpl- prefix) verified the same way: client span → router's two spans → openinfer's four spans, parenting correct throughout +- No-client-traceparent case: the router's `http_request` becomes the trace root and openinfer attaches normally (the load test exercised exactly this) +- Result: success + +### Step 4.5: Tempo query caveats +- `/api/traces/` may return only the spans that have reached a block (ingester flushes on a ~5s window) — refetch and you get everything; not span loss + +### Step 5: Upstream PR + tracking issue +- Upstream PR: https://github.com/vllm-project/vllm/pull/50370 "[Rust Frontend] Propagate W3C trace headers to engine-core requests" (xiaguan fork, DCO signed, off upstream/main e5f48dfda) + - Prior-art check found **#44567** (same goal) closed unmerged after a maintainer objected to its "new protocol surface"; this PR deliberately takes the minimal route: no handshake, no gating, pure extraction into `trace_headers`, stated explicitly in the PR body + - 10 files changed: `resolve_request_context` extracts traceparent/tracestate → `ResolvedRequestContext.trace_headers` → completions/chat/generate converts → additive pass-through fields in vllm-text/vllm-chat into the llm `GenerateRequest` (llm/engine-core-client untouched); grpc/tokenize only got compile-required `None` fill-ins + - Gates: fmt/clippy (-D warnings) clean; nextest vllm-text+chat 321 passed, vllm-server 328 passed (incl. 12 new tests) +- Tracking issue: https://github.com/openinfer-project/openinfer/issues/790 — records the MVP state, the upstream PR, and the three migration steps (bump pin → bridge reads `trace_headers["traceparent"]` → delete trace_context.rs and its lib.rs wiring) +- Grafana access: the host's port 3000 collided with a Grafana on the user's mac, so the container was rebound to 4000 (repo compose file untouched); the user's browser initially landed on a different Grafana 12 instance (Sign in/Bookmarks visible) — `curl localhost:4000/api/health` returning 11.3.0 confirmed the right port-forward +- Anonymous Viewer is denied `datasources:explore` on Grafana 11.3 (Access denied in the server log; Explore renders but returns nothing) → recreated the container with `GF_AUTH_ANONYMOUS_ORG_ROLE=Editor`; the repo compose file is fixed in the same PR (fix(deploy)) +- PR follow-up: only bot placeholder reviews, no maintainer response; per the user's judgement (maintainers may want to build it themselves), posted a direction/timeline question on the PR (#issuecomment-5126068994) offering to reshape or close in favor of upstream's own implementation + +### Step 6: MVP packaged as a pegainfer PR +- Branch `feat/frontend-trace-context`, 3 Commitizen commits (feat(frontend) / fix(deploy) / docs(tracing)); the prek fmt hook reformatted on the first attempt, re-staged and passed +- Pushed to origin (GitHub); PR: https://github.com/openinfer-project/openinfer/pull/791; #790 commented with the link +- The Viewer→Editor compose fix rides along in fix(deploy) (anonymous Viewers lack `datasources:explore` on Grafana 11.3, contradicting the file's own usage comment) +- Review follow-ups: Codex bot iterated on the trace-context intake across 13 rounds; every substantive finding was fixed (each with a regression test), and round 13 on `5ffcafd` came back clean (CI 15/15 green). Convergence path: TTL on take (`fb20b34`) → invalidate on untraced reuse (`5ea8017`) → prefix-stripped lookup (`bcede4a`) → FIFO queue per id (`3dfba61`) → error-response discard (`907425c`) → token-exact discard + generation-route gating (`c7da3d4`) → **route-aware stash keys + untraced-slot markers** (`22645de`, the structural fix: keys computed deterministically at intake where the route is known, bridge lookup exact) → consume at engine intake (`e13343b`) → documented FIFO pairing ceiling for true duplicate ids (`7e9e038`) → middleware `InsertionGuard` (`6957687`) → guard disarmed on success responses for the streaming race (`e73b6f4`) → exactly one slot per bridge arrival (`5ffcafd`). Design after the loop: route-aware key + FIFO slots + token-exact discard + untraced markers + intake-time consumption + disarmable drop guard + TTL backstop. Also placated CI's newer clippy (`unchecked_time_subtraction`, `type_complexity`) and the repo's DCO/prek hooks; Codex also caught a `<4317>` doc placeholder. One explicitly documented non-goal: deterministic pairing for concurrent requests sharing one `X-Request-Id` is impossible by construction (identical correlation key) — noted in code and the finding thread. + +## Debrief + +- **Outcome**: the client → vllm-router → openinfer → prefill/decode chain renders as one trace, verified in Tempo (both endpoints, with and without a client traceparent, and under 16-way concurrency); router overhead is ~1ms. The local MVP (middleware + stash + bridge parent resolution) is implemented, tested, and up as PR #791; upstream PR vllm-project/vllm#50370 is open; migration is tracked by openinfer#790. +- **Pitfalls encountered**: + - Root cause of the first join failure: vllm-server prepends `cmpl-`/`chatcmpl-` to `X-Request-Id` before it becomes `external_req_id` — this protocol fact could only be pinned down from live trace attributes (`request_id=cmpl-dbg12345-…`); the route-layer prefix step was missed when reading code + - The direct-to-openinfer control experiment (bypassing the router) was the key bisection: it first narrowed the fault to the openinfer side, then span attributes pinned the prefix issue + - A same-goal upstream PR (#44567) was rejected for adding protocol surface/gating — always check rejection history before proposing upstream; minimal pass-through is the only acceptable shape + - Tempo `/api/traces` partial returns (ingester flush window) and same-trace-id block compaction each cost a debugging round +- **Lessons learned**: + - Verify trace chains with deterministic traceparents (carry your own trace id) + Tempo HTTP API assertions — faster and reproducible, unlike UI spelunking + - A missing feature in a git dependency doesn't force a fork: framework hooks like `serve_with_router_extension` + middleware can close the gap; but a workaround needs a tracking issue and an upstream PR, or it rots in the tree +- **Follow-ups**: + - openinfer#790: after the upstream PR merges and the pin is bumped, delete the middleware + - vllm-router `http_client_request` span occasionally ~5s inflated (non-streaming path) — suspected router instrumentation bug; candidate for a separate issue to vllm-project/router + - vllm-server's own OTel instrumentation (the HTTP/tokenize segment is currently a blank gap in the trace) — upstream has related PRs in flight (#39438, #39905); track them diff --git a/openinfer-vllm-frontend/src/bridge.rs b/openinfer-vllm-frontend/src/bridge.rs index 313d68ec9..575ac9d85 100644 --- a/openinfer-vllm-frontend/src/bridge.rs +++ b/openinfer-vllm-frontend/src/bridge.rs @@ -72,6 +72,7 @@ pub(crate) struct LocalEngineBridge { pub(crate) engine_index: u32, pub(crate) data_parallel_size: u32, pub(crate) load_watch: Option>, + pub(crate) trace_stash: crate::trace_context::TraceContextStash, } impl LocalEngineBridge { @@ -309,8 +310,22 @@ impl LocalEngineBridge { request_id, prompt_token_ids, sampling_params, + external_req_id, .. } = request; + // Consume any stashed trace context the moment the request arrives: + // the validation rejections below terminate the request without ever + // reaching span creation (and streaming clients still see a success + // status, so the middleware's error-path cleanup cannot cover them). + // Every terminal path must retire the entry, not leak it to a + // within-TTL reuse of the same external id. + let stashed_parent = if openinfer_engine::tracing_state::is_enabled() { + external_req_id + .as_deref() + .and_then(|id| self.trace_stash.take(id)) + } else { + None + }; let Some(prompt_tokens) = prompt_token_ids else { warn!("request {request_id} dropped: missing prompt_token_ids"); send_terminal_output( @@ -392,9 +407,15 @@ impl LocalEngineBridge { // keeps the default (tracing-off) path free of per-request span work, // and `from_span` on a noop span yields `None` so the scheduler skips // its span work too. + // Parent resolution: join the upstream trace when the HTTP layer + // stashed a traceparent for this request (e.g. from vllm-router); + // otherwise start a fresh trace. The stash is keyed by + // external_req_id already, so the lookup is an exact match. let trace_root = if openinfer_engine::tracing_state::is_enabled() { - Span::root("request", SpanContext::random()) - .with_property(|| ("request_id", tag.to_string())) + let parent = stashed_parent + .and_then(|traceparent| SpanContext::decode_w3c_traceparent(&traceparent)) + .unwrap_or_else(SpanContext::random); + Span::root("request", parent).with_property(|| ("request_id", tag.to_string())) } else { Span::noop() }; diff --git a/openinfer-vllm-frontend/src/lib.rs b/openinfer-vllm-frontend/src/lib.rs index d5929b8ae..7e4601f7a 100644 --- a/openinfer-vllm-frontend/src/lib.rs +++ b/openinfer-vllm-frontend/src/lib.rs @@ -26,6 +26,7 @@ use vllm_server::RendererSelection; mod bridge; mod lora; mod request_contract; +mod trace_context; mod wire; use bridge::LocalEngineBridge; @@ -220,11 +221,15 @@ where // in the registration wait. let server_shutdown = shutdown.child_token(); let bridge_shutdown = shutdown.child_token(); + // Shared between the axum layer (stash on HTTP intake) and every engine + // bridge below (pop on EngineCoreRequest) — see trace_context.rs. + let trace_stash = trace_context::TraceContextStash::default(); let engine_task = tokio::spawn({ let server_shutdown = server_shutdown.clone(); let bridge_shutdown = bridge_shutdown.clone(); let input_address = input_address.clone(); let output_address = output_address.clone(); + let trace_stash = trace_stash.clone(); async move { let handle = match engine.await { Ok(handle) => handle, @@ -253,6 +258,7 @@ where engine_index: engine_index as u32, data_parallel_size, load_watch: handle.load_watch_for(engine_index), + trace_stash: trace_stash.clone(), }; let shutdown = bridge_shutdown.clone(); bridges.spawn(async move { (engine_index, bridge.run(shutdown).await) }); @@ -334,6 +340,19 @@ where tls: None, }; + // Outermost layer: stash W3C trace context on intake so the bridge can + // join each request's root span to the upstream trace. Must run before + // vllm-server reads headers (it forwards X-Request-Id but drops + // traceparent), so it wraps the router after all other extensions. + let extend_router = { + let trace_stash = trace_stash.clone(); + move |router: Router| { + extend_router(router).layer(axum::middleware::from_fn_with_state( + trace_stash, + trace_context::stash_trace_context, + )) + } + }; let result = vllm_server::serve_with_router_extension(config, server_shutdown, extend_router).await; // Stop the bridge (no-op if the caller's shutdown already cancelled it), diff --git a/openinfer-vllm-frontend/src/trace_context.rs b/openinfer-vllm-frontend/src/trace_context.rs new file mode 100644 index 000000000..9a0168d64 --- /dev/null +++ b/openinfer-vllm-frontend/src/trace_context.rs @@ -0,0 +1,560 @@ +//! W3C trace-context intake for distributed request tracing. +//! +//! The pinned `vllm-server` HTTP layer forwards `X-Request-Id` but drops +//! `traceparent`, so an upstream trace (e.g. vllm-router with OTel enabled) +//! cannot reach the bridge's `request` root span. Until vllm-server populates +//! `EngineCoreRequest.trace_headers` itself (upstream PR tracked in +//! docs/subsystems/tracing/e2e-router-tracing.md), this module stashes the +//! incoming `traceparent` at the axum boundary, and the bridge pops it when +//! the matching `EngineCoreRequest` arrives over ZMQ. +//! +//! Correlation key: `EngineCoreRequest.external_req_id`, computed at intake +//! where the route is still known (route prefix + `X-Request-Id`, mirroring +//! vllm-server's completions/chat prefixing), so the bridge lookup is exact. +//! Requests without trace context reserve their slot with an untraced marker, +//! so overlapping attempts reusing one id each consume only their own slot. +//! Entries are popped on use; entries whose request never reaches the engine +//! are retired by a middleware drop guard (error responses, client +//! disconnects), or expire by TTL. + +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; + +use axum::extract::Request; +use axum::extract::State; +use axum::http::HeaderMap; +use axum::http::Method; +use axum::middleware::Next; +use axum::response::Response; + +/// Stale-entry lifetime and the stash's hard capacity. Entries are normally +/// popped by the bridge within milliseconds of insertion; both bounds exist +/// only as backstops for requests that never reach the engine. +const TTL: Duration = Duration::from_secs(120); +const CAPACITY: usize = 4096; + +/// Routes whose accepted requests produce an `EngineCoreRequest` (and thus +/// consume a stash entry at the bridge). Stashing anywhere else would leak +/// entries for requests that never reach the engine (e.g. traced `/metrics` +/// probes), eventually tripping the capacity clear and splitting live traces. +const GENERATION_PATHS: &[&str] = &[ + "/v1/completions", + "/v1/chat/completions", + "/inference/v1/generate", +]; + +/// `(token, traceparent, inserted-at)` entries queued under one id. A `None` +/// traceparent is an untraced attempt's marker: it reserves the attempt's +/// slot without supplying a parent to join. +type EntryQueue = VecDeque<(u64, Option, Instant)>; + +/// Per-insertion token: lets the HTTP layer discard exactly the entry it +/// stashed even when responses complete out of FIFO order. +static NEXT_TOKEN: AtomicU64 = AtomicU64::new(1); + +/// Stashed `traceparent` headers awaiting pickup by the engine bridge. +/// +/// Cheap to clone (inner `Arc`); one instance is shared between the axum +/// layer (insert) and every engine bridge task (take). Entries queue FIFO per +/// id: concurrent attempts reusing one `X-Request-Id` (e.g. hedged retries) +/// each keep their own slot instead of the latest insert overwriting the +/// rest, paired best-effort in intake order (see `take`'s pairing-ceiling +/// note). A per-attempt unique key is not available — `external_req_id` is +/// the only correlation key the bridge can derive downstream. +#[derive(Clone, Default)] +pub(crate) struct TraceContextStash { + inner: Arc>>, +} + +impl TraceContextStash { + /// Queue `traceparent` (`None` for an untraced marker) under `request_id`, + /// returning the insertion token the HTTP layer needs to discard exactly + /// this entry on the error path. + fn insert(&self, request_id: &str, traceparent: Option<&str>) -> u64 { + let token = NEXT_TOKEN.fetch_add(1, Ordering::Relaxed); + let mut inner = self.inner.lock().expect("trace context stash poisoned"); + let mut total: usize = inner.values().map(VecDeque::len).sum(); + if total >= CAPACITY { + let now = Instant::now(); + inner.retain(|_, queue| { + queue.retain(|(_, _, inserted)| now.duration_since(*inserted) < TTL); + !queue.is_empty() + }); + total = inner.values().map(VecDeque::len).sum(); + if total >= CAPACITY { + // Pathological: more live unreached requests than the stash + // holds. Dropping it splits some traces; the stash must never + // grow unbounded or stall serving. + inner.clear(); + } + } + inner.entry(request_id.to_owned()).or_default().push_back(( + token, + traceparent.map(str::to_owned), + Instant::now(), + )); + token + } + + /// Pop exactly one slot for `request_id`: the queue head. A fresh parent + /// comes back as `Some`; anything else — no entry, an expired one, or an + /// untraced marker — as `None`, which the bridge treats as "start a + /// fresh trace". + /// + /// One slot per bridge arrival, never a scan past an expired head: slots + /// pair attempts to arrivals in order, and an expired head belongs to + /// this arrival's own attempt (merely slow, e.g. a long body upload past + /// the TTL). Sliding into the next attempt's live parent would both + /// mis-attach this request and orphan the next one. + /// + /// Pairing ceiling: entries queue in middleware intake order, which the + /// bridge normally follows. Concurrent requests reusing one id (hedged + /// retries) can be reordered downstream — e.g. a LoRA request stalled in + /// body rewriting while a later request overtakes it — and would then + /// consume each other's parent. Exact pairing is impossible there: the + /// correlation key is identical by construction and the engine + /// `request_id`'s random suffix carries no intake information. Callers + /// needing deterministic pairing must use unique `X-Request-Id` values. + pub(crate) fn take(&self, request_id: &str) -> Option { + let mut inner = self.inner.lock().expect("trace context stash poisoned"); + let (result, queue_empty) = { + let queue = inner.get_mut(request_id)?; + let result = queue + .pop_front() + .and_then(|(_, traceparent, inserted)| { + (inserted.elapsed() < TTL).then_some(traceparent) + }) + .flatten(); + (result, queue.is_empty()) + }; + if queue_empty { + inner.remove(request_id); + } + result + } + + /// Drop exactly the entry tagged `token`, wherever it sits in the queue. + /// Responses can complete out of FIFO order, so the error path must not + /// assume the abandoned attempt is at the head. + fn discard_entry(&self, request_id: &str, token: u64) { + let mut inner = self.inner.lock().expect("trace context stash poisoned"); + let Some(queue) = inner.get_mut(request_id) else { + return; + }; + if let Some(pos) = queue.iter().position(|(t, _, _)| *t == token) { + queue.remove(pos); + } + if queue.is_empty() { + inner.remove(request_id); + } + } + + #[cfg(test)] + fn len(&self) -> usize { + self.inner + .lock() + .expect("trace context stash poisoned") + .values() + .map(VecDeque::len) + .sum() + } +} + +/// The id vllm-server will put into `EngineCoreRequest.external_req_id` for a +/// generation request on `path`: the route prefix plus the caller's +/// `X-Request-Id`, mirroring the pinned server's completions/chat prefixing; +/// other routes pass the id through verbatim. Computing it at intake — where +/// the route is still known — keeps the bridge lookup a plain exact match. +fn external_req_id_for(path: &str, request_id: &str) -> String { + match path { + "/v1/completions" => format!("cmpl-{request_id}"), + "/v1/chat/completions" => format!("chatcmpl-{request_id}"), + _ => request_id.to_owned(), + } +} + +/// Drop guard for one stashed insertion. The bridge consumes an entry the +/// moment its `EngineCoreRequest` arrives (accepted or rejected alike), so +/// when the middleware future ends without a successful response — an error +/// status, or cancellation when the client disconnects mid-pipeline — a +/// still-present entry can only belong to a request that never reached the +/// engine, and is discarded. For consumed entries the discard is a no-op. +struct InsertionGuard { + stash: TraceContextStash, + key: String, + token: u64, + armed: bool, +} + +impl InsertionGuard { + /// Stand down: the response succeeded, so the engine request was already + /// sent and the bridge will consume the entry when it processes it. + /// (Streaming response heads can go out while the bridge is still behind; + /// dropping armed here would delete the parent first.) + fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for InsertionGuard { + fn drop(&mut self) { + if self.armed { + self.stash.discard_entry(&self.key, self.token); + } + } +} + +/// Axum middleware: stash the request's `traceparent` under its external id. +/// +/// Only generation routes are eligible — other routes never produce an +/// `EngineCoreRequest`, so their entries would leak until the TTL. Does +/// nothing when request tracing is disabled. An [`InsertionGuard`] retires +/// the entry when the request cannot have reached the engine: an error +/// status, or a client disconnect dropping the future before the handler +/// finishes. On success the guard is disarmed and the bridge consumes the +/// entry itself. +pub(crate) async fn stash_trace_context( + State(stash): State, + mut request: Request, + next: Next, +) -> Response { + if !openinfer_engine::tracing_state::is_enabled() + || request.method() != Method::POST + || !GENERATION_PATHS.contains(&request.uri().path()) + { + return next.run(request).await; + } + let path = request.uri().path().to_owned(); + let guard = stash_from_headers(&stash, &path, request.headers_mut()).map(|(key, token)| { + InsertionGuard { + stash: stash.clone(), + key, + token, + armed: true, + } + }); + let response = next.run(request).await; + if response.status().is_success() { + if let Some(guard) = guard { + guard.disarm(); + } + } + response +} + +/// Read `traceparent` from `headers` and stash it under the request's +/// external id, generating and injecting `X-Request-Id` when absent so the +/// bridge and vllm-server agree on the correlation key. Returns the stash key +/// and insertion token when an entry was queued, `None` otherwise. +fn stash_from_headers( + stash: &TraceContextStash, + path: &str, + headers: &mut HeaderMap, +) -> Option<(String, u64)> { + let request_id = headers + .get("x-request-id") + .and_then(|v| v.to_str().ok()) + .map(str::to_owned); + let Some(traceparent) = headers + .get("traceparent") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned) + else { + // A request with no trace context of its own reserves its slot with + // an untraced marker, so an overlapping traced attempt reusing the id + // keeps its own parent and this attempt consumes nothing upstream. + // With no caller id at all, vllm-server mints a unique id downstream + // and no slot is needed. + let key = external_req_id_for(path, &request_id?); + let token = stash.insert(&key, None); + return Some((key, token)); + }; + let request_id = request_id.unwrap_or_else(|| { + // Mirror vllm-server's own id shape (8 hex chars) so logs read the + // same regardless of which side generated the id. + let mut id = uuid::Uuid::new_v4().simple().to_string(); + id.truncate(8); + if let Ok(value) = id.parse() { + headers.insert("x-request-id", value); + } + id + }); + let key = external_req_id_for(path, &request_id); + let token = stash.insert(&key, Some(&traceparent)); + Some((key, token)) +} + +#[cfg(test)] +mod tests { + use axum::http::HeaderValue; + + use super::*; + + const TRACEPARENT: &str = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + const TRACEPARENT_B: &str = "00-1af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + const COMPLETIONS: &str = "/v1/completions"; + const CHAT: &str = "/v1/chat/completions"; + const GENERATE: &str = "/inference/v1/generate"; + + fn traced_headers(id: &str, traceparent: &str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert("traceparent", HeaderValue::from_str(traceparent).unwrap()); + headers.insert("x-request-id", HeaderValue::from_str(id).unwrap()); + headers + } + + #[test] + fn stash_roundtrip_decodes_to_upstream_trace() { + let stash = TraceContextStash::default(); + let mut headers = traced_headers("req-1", TRACEPARENT); + + stash_from_headers(&stash, COMPLETIONS, &mut headers); + + let stashed = stash.take("cmpl-req-1").expect("traceparent stashed"); + let ctx = fastrace::collector::SpanContext::decode_w3c_traceparent(&stashed) + .expect("valid W3C traceparent"); + assert_eq!(ctx.encode_w3c_traceparent(), TRACEPARENT); + // One-shot: the bridge must not join a second request to the same span. + assert!(stash.take("cmpl-req-1").is_none()); + } + + #[test] + fn generates_and_injects_request_id_when_absent() { + let stash = TraceContextStash::default(); + let mut headers = HeaderMap::new(); + headers.insert("traceparent", HeaderValue::from_static(TRACEPARENT)); + + stash_from_headers(&stash, COMPLETIONS, &mut headers); + + let injected = headers + .get("x-request-id") + .expect("request id injected") + .to_str() + .expect("ascii request id"); + assert_eq!(injected.len(), 8); + assert!(injected.chars().all(|c| c.is_ascii_hexdigit())); + assert!(stash.take(&format!("cmpl-{injected}")).is_some()); + } + + #[test] + fn external_key_mirrors_route_prefixing() { + // Each route's stash key is exactly the external_req_id the bridge + // will see: completions/chat prepend, other routes pass through. + let stash = TraceContextStash::default(); + for (path, expected_key) in [ + (COMPLETIONS, "cmpl-dbg12345"), + (CHAT, "chatcmpl-dbg12345"), + (GENERATE, "dbg12345"), + ] { + let mut headers = traced_headers("dbg12345", TRACEPARENT); + stash_from_headers(&stash, path, &mut headers); + assert_eq!(stash.take(expected_key), Some(TRACEPARENT.to_owned())); + } + } + + #[test] + fn unprefixed_route_preserves_prefixed_header_ids() { + // /inference/v1/generate prepends nothing: a header literally named + // `cmpl-foo` must be looked up verbatim, not mistaken for a + // completions-generated prefix of `foo`. + let stash = TraceContextStash::default(); + let mut generated = traced_headers("cmpl-foo", TRACEPARENT_B); + stash_from_headers(&stash, GENERATE, &mut generated); + + assert_eq!(stash.take("cmpl-foo"), Some(TRACEPARENT_B.to_owned())); + } + + #[test] + fn prefixed_header_ids_collide_only_with_themselves() { + // Completions headers `foo` and `cmpl-foo` become `cmpl-foo` and + // `cmpl-cmpl-foo` at the bridge; each consumes its own traceparent. + let stash = TraceContextStash::default(); + for (id, tp) in [("foo", TRACEPARENT), ("cmpl-foo", TRACEPARENT_B)] { + let mut headers = traced_headers(id, tp); + stash_from_headers(&stash, COMPLETIONS, &mut headers); + } + + assert_eq!(stash.take("cmpl-foo"), Some(TRACEPARENT.to_owned())); + assert_eq!(stash.take("cmpl-cmpl-foo"), Some(TRACEPARENT_B.to_owned())); + } + + #[test] + fn take_drops_expired_entries() { + let stash = TraceContextStash::default(); + let expired = Instant::now() + .checked_sub(TTL + Duration::from_secs(1)) + .expect("TTL subtraction stays within Instant range"); + stash + .inner + .lock() + .expect("trace context stash poisoned") + .insert( + "old".to_owned(), + VecDeque::from([(1, Some(TRACEPARENT.to_owned()), expired)]), + ); + + assert!(stash.take("old").is_none()); + } + + #[test] + fn expired_head_slot_does_not_slide_into_next_attempt() { + // The first attempt spent longer than the TTL before reaching the + // bridge: its arrival consumes its own expired slot (fresh trace), + // leaving the second attempt's live parent untouched. + let stash = TraceContextStash::default(); + let expired = Instant::now() + .checked_sub(TTL + Duration::from_secs(1)) + .expect("TTL subtraction stays within Instant range"); + stash + .inner + .lock() + .expect("trace context stash poisoned") + .entry("cmpl-slow".to_owned()) + .or_default() + .push_back((1, Some(TRACEPARENT.to_owned()), expired)); + let mut second = traced_headers("slow", TRACEPARENT_B); + stash_from_headers(&stash, COMPLETIONS, &mut second); + + assert!(stash.take("cmpl-slow").is_none()); + assert_eq!(stash.take("cmpl-slow"), Some(TRACEPARENT_B.to_owned())); + assert_eq!(stash.len(), 0); + } + + #[test] + fn duplicate_ids_keep_separate_parents_fifo() { + // Hedged retries reusing one X-Request-Id concurrently: each attempt + // must consume a distinct parent, oldest first. + let stash = TraceContextStash::default(); + for tp in [TRACEPARENT, TRACEPARENT_B] { + let mut headers = traced_headers("hedged", tp); + stash_from_headers(&stash, COMPLETIONS, &mut headers); + } + assert_eq!(stash.len(), 2); + + assert_eq!(stash.take("cmpl-hedged"), Some(TRACEPARENT.to_owned())); + assert_eq!(stash.take("cmpl-hedged"), Some(TRACEPARENT_B.to_owned())); + assert!(stash.take("cmpl-hedged").is_none()); + assert_eq!(stash.len(), 0); + } + + #[test] + fn error_path_discards_the_rejected_attempts_own_entry() { + // Two overlapping traced requests share an id; the later one errors + // out first. Discarding must remove the later attempt's entry, not + // the queue head. + let stash = TraceContextStash::default(); + let mut first = traced_headers("dup", TRACEPARENT); + stash_from_headers(&stash, COMPLETIONS, &mut first); + let mut second = traced_headers("dup", TRACEPARENT_B); + let (_, token_b) = stash_from_headers(&stash, COMPLETIONS, &mut second).unwrap(); + + stash.discard_entry("cmpl-dup", token_b); + + assert_eq!(stash.take("cmpl-dup"), Some(TRACEPARENT.to_owned())); + assert!(stash.take("cmpl-dup").is_none()); + } + + #[test] + fn traced_retry_after_rejection_gets_fresh_parent() { + // First attempt is rejected before reaching the engine; the HTTP + // layer discards its entry on the error response, so the traced + // retry's own parent is what the bridge consumes. + let stash = TraceContextStash::default(); + let mut first = traced_headers("retry-2", TRACEPARENT); + let stashed = stash_from_headers(&stash, COMPLETIONS, &mut first); + assert_eq!( + stashed.as_ref().map(|(key, _)| key.as_str()), + Some("cmpl-retry-2") + ); + + let (key, token) = stashed.unwrap(); + stash.discard_entry(&key, token); + assert_eq!(stash.len(), 0); + + let mut retry = traced_headers("retry-2", TRACEPARENT_B); + stash_from_headers(&stash, COMPLETIONS, &mut retry); + assert_eq!(stash.take("cmpl-retry-2"), Some(TRACEPARENT_B.to_owned())); + assert_eq!(stash.len(), 0); + } + + #[test] + fn untraced_attempt_reserves_its_own_slot() { + // An overlapping untraced request must neither consume the traced + // attempt's parent nor delete it: it queues a marker and the bridge + // opens a fresh trace for exactly that attempt. + let stash = TraceContextStash::default(); + let mut traced = traced_headers("retry-1", TRACEPARENT); + stash_from_headers(&stash, COMPLETIONS, &mut traced); + assert_eq!(stash.len(), 1); + + let mut retry = HeaderMap::new(); + retry.insert("x-request-id", HeaderValue::from_static("retry-1")); + stash_from_headers(&stash, COMPLETIONS, &mut retry); + assert_eq!(stash.len(), 2); + + // The traced attempt still consumes its own parent first... + assert_eq!(stash.take("cmpl-retry-1"), Some(TRACEPARENT.to_owned())); + // ...and the untraced attempt pops its marker: no parent, one-shot. + assert!(stash.take("cmpl-retry-1").is_none()); + assert_eq!(stash.len(), 0); + } + + #[test] + fn dropped_guard_retires_only_its_own_entry() { + // Client disconnects mid-pipeline drop the middleware future; the + // guard must retire exactly the cancelled attempt's entry. + let stash = TraceContextStash::default(); + let mut first = traced_headers("g", TRACEPARENT); + let (key_a, token_a) = stash_from_headers(&stash, COMPLETIONS, &mut first).unwrap(); + let mut second = traced_headers("g", TRACEPARENT_B); + let (key_b, _) = stash_from_headers(&stash, COMPLETIONS, &mut second).unwrap(); + assert_eq!(key_a, key_b); + + drop(InsertionGuard { + stash: stash.clone(), + key: key_a, + token: token_a, + armed: true, + }); + + assert_eq!(stash.take(&key_b), Some(TRACEPARENT_B.to_owned())); + assert!(stash.take(&key_b).is_none()); + } + + #[test] + fn disarmed_guard_leaves_entry_for_bridge_consumption() { + // Successful (possibly streaming) responses disarm the guard: the + // entry must survive for the bridge, which may still be behind. + let stash = TraceContextStash::default(); + let mut headers = traced_headers("s", TRACEPARENT); + let (key, token) = stash_from_headers(&stash, COMPLETIONS, &mut headers).unwrap(); + + InsertionGuard { + stash: stash.clone(), + key: key.clone(), + token, + armed: true, + } + .disarm(); + + assert_eq!(stash.take(&key), Some(TRACEPARENT.to_owned())); + assert_eq!(stash.len(), 0); + } + + #[test] + fn ignores_requests_without_traceparent_or_id() { + let stash = TraceContextStash::default(); + let mut headers = HeaderMap::new(); + + stash_from_headers(&stash, COMPLETIONS, &mut headers); + + assert_eq!(stash.len(), 0); + assert!(headers.get("x-request-id").is_none()); + } +}