feat(pylon): handle Dynamo request-priority headers in Pylon - #673
feat(pylon): handle Dynamo request-priority headers in Pylon#673along-2017 wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughPylon now derives and sanitizes Dynamo request-priority headers. CLI configuration controls the behavior. mock-dynamo records priority headers in test snapshots. Unit and integration tests cover mapping, filtering, forwarding, and reset behavior. ChangesDynamo priority forwarding
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Pylon
participant UpstreamInferenceServer
participant MockDynamoTestControl
Client->>Pylon: Send x-priority and inbound Dynamo headers
Pylon->>Pylon: Parse, invert, and clamp priority
Pylon->>UpstreamInferenceServer: Forward derived x-dynamo-request-priority
UpstreamInferenceServer->>MockDynamoTestControl: Record priority headers
MockDynamoTestControl-->>Client: Return priority snapshot
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
🛡️ CodeQL Analysis🚨 Found 2 issue(s) Severity Breakdown:
📋 Top Issues🔗 View full details in Security tab 🕐 Last updated: 2026-08-04 22:29:46 UTC | Commit: 1c43f89 |
Pylon owns the engine-facing Dynamo header contract. On every tunneled inference request it now strips inbound x-dynamo-request-* headers so client-supplied values never reach the engine, and derives x-dynamo-request-priority from x-priority when that header is present (i32::MAX - min(x, i32::MAX); absent stays absent). Derivation sits behind the default-on --pylon-derive-dynamo-priority flag as a kill switch; the strip is unconditional. The emitted value is logged with the request id and recorded on the upstream request span. mock-dynamo records the priority headers seen on the latest request per endpoint and model in its /test-control snapshot so cluster QA can assert what actually reached the engine. The pylon and mock-dynamo crates gain rust_test Bazel targets; their in-crate tests previously did not run in CI. Refs: #620 Signed-off-by: along <along@nvidia.com>
Dynamo parses x-dynamo-request-strict-priority as a u32 queue tier, so the spoofed test value should look like the real contract. The strip assertion is name-based and unaffected. Signed-off-by: along <along@nvidia.com>
fcbc062 to
39efe42
Compare
FamousDirector
left a comment
FamousDirector
left a comment
There was a problem hiding this comment.
Reviewed against Dynamo's actual source (lib/llm/src/protocols/common/extensions.rs, lib/kv-router/src/scheduling/policy.rs, lib/llm/src/preprocessor.rs) and the priority-scheduling doc.
What matches. Polarity is correct: requestctx.go documents x-priority as "Lower value is higher priority, 0 is highest" and Dynamo is higher-wins i32, so the inversion runs the right direction. The gateway already rejects client-supplied X-Priority with 400, so the input to the mapping is trusted. HeaderValue::from(i32) is infallible, the strip runs before the insert so no duplicate header is possible, and the mapping never emits a negative, so Dynamo's negative-clamp path is never hit.
Three blocking issues, detailed inline:
- The claim that spoofed values never reach the engine does not hold. Dynamo resolves priority from the header and the request body, and
strict_priority(which Pylon never emits) sorts ahead of everything else in the router queue. x-dynamo-request-is the wrong prefix. The routing headers worth stripping (x-dynamo-worker-instance-id,x-dynamo-dp-rank,x-tenant-id, plus unprefixed aliases) all fall outside it.- Anchoring at
i32::MAXis a policy decision that is not stated anywhere: Dynamo treats the value as seconds of arrival-time bump, so every request carryingx-priorityoutranks every request without one, permanently.
One question outside the diff. rejectClientSuppliedPriority is registered on the LLM route group, "not globally". Any other route reaching stargate and then pylon would let a caller set x-priority directly, which now maps straight to engine priority. Worth confirming no such route exists before this merges.
| /// Client-supplied Dynamo request headers never reach the engine; Pylon is | ||
| /// the only writer of x-dynamo-request-* values. | ||
| pub(super) fn is_dynamo_request_header(name: &HeaderName) -> bool { | ||
| name.as_str().starts_with(DYNAMO_REQUEST_HEADER_PREFIX) |
There was a problem hiding this comment.
Blocking: the body path is unguarded, so "spoofed values never reach the engine" is not true.
Dynamo resolves priority from the header and the body. From resolve_request_priority in lib/llm/src/protocols/common/extensions.rs:
let priority = priority_header.and_then(|h| h.trim().parse::<i32>().ok())
.or_else(|| hints.and_then(|h| h.priority));
let strict_priority = strict_priority_header.and_then(|h| h.trim().parse::<u32>().ok())
.or_else(|| hints.and_then(|h| h.strict_priority));Three ways a client still controls engine scheduling:
strict_priorityis never emitted by Pylon and is the strongest lever. The router queue key is(strict_priority, policy_score)with the strict tier compared first (lib/kv-router/src/scheduling/policy.rs). A client sending{"nvext":{"agent_hints":{"strict_priority":4294967295}}}in the JSON body jumps ahead of every platform-prioritized request regardless of whatx-priorityresolved to. Stripping the header does nothing against it.- Absent
x-prioritymeans no header at all, so bodyprioritywins outright for any caller without a priority config. latency_sensitivity(deprecated, body-only, f64) feedspriority_jumpwhenever no priority exists. Same bypass.
The body reaches Dynamo verbatim: stargate does opaque body forwarding, and the gateway's proxyRequest passes request.Body through unmodified for the responses and embeddings paths. nvext and agent_hints appear nowhere in this repo, so nothing sanitizes them anywhere.
Three ways to close it, in rough order of preference:
- Sanitize
nvext.agent_hintsout of the body at the trust boundary (gateway or here). - Emit both
x-dynamo-request-priorityandx-dynamo-request-strict-priorityunconditionally, including whenx-priorityis absent, since a well-formed header beats the body per Dynamo's precedence rule. - Keep the current scope but drop the security claim from the PR description and the doc comment here, and file the body path as a known gap.
| /// Engine-facing priority header in the Dynamo contract; Pylon owns this | ||
| /// contract, so the constant stays out of the shared tunnel contract. | ||
| pub(super) const HEADER_DYNAMO_REQUEST_PRIORITY: &str = "x-dynamo-request-priority"; | ||
| const DYNAMO_REQUEST_HEADER_PREFIX: &str = "x-dynamo-request-"; |
There was a problem hiding this comment.
Blocking: x-dynamo-request- is the wrong boundary. The headers worth stripping are outside this prefix.
apply_header_routing_overrides in Dynamo honors all of these, and none of them match x-dynamo-request-:
| header | aliases also honored | effect |
|---|---|---|
x-dynamo-worker-instance-id |
x-worker-instance-id |
sets backend_instance_id and decode_worker_id |
x-dynamo-prefill-instance-id |
x-prefill-instance-id |
pins the prefill worker |
x-dynamo-dp-rank |
x-dp-rank, x-data-parallel-rank |
pins the DP rank |
x-dynamo-prefill-dp-rank |
x-prefill-dp-rank |
pins the prefill DP rank |
x-tenant-id |
none | sets cache_salt |
A client can pin itself to a chosen worker and DP rank, which bypasses the load balancer entirely and makes targeted worker overload trivial. x-tenant-id sets the KV cache namespace, which is a cross-tenant cache probing vector. Both are larger than the priority hole this PR closes, and the prefix match makes the strip look more complete than it is.
Suggest an explicit deny-list of the header names above plus their aliases, sourced from Dynamo's constants, rather than a prefix. A prefix match will keep silently missing whatever Dynamo adds next outside it.
| /// Dynamo schedules on an i32 where higher wins and silently drops values | ||
| /// that do not parse as i32, while x-priority is a u32 where lower wins, so |
There was a problem hiding this comment.
The stated failure mode is not what Dynamo does. An unparseable header does not leave the request with no priority; it falls back to the client-controlled body value. Dynamo's own test asserts it:
let r = resolve_request_priority(Some(&hints), Some("abc"), None);
assert_eq!(r.priority, Some(5)); // the body valueNot live today because this function always produces a well-formed value, but the comment is the thing a future author will reason from when they change the emitted format. Worth correcting to "a malformed header falls back to the client body value".
| pub(super) fn dynamo_request_priority(priority: u32) -> i32 { | ||
| i32::MAX - i32::try_from(priority).unwrap_or(i32::MAX) | ||
| } |
There was a problem hiding this comment.
Blocking: anchoring at i32::MAX is an unstated policy decision, and the unit is seconds.
Dynamo's FCFS queue key is:
OrderedFloat(ctx.request().priority_jump.max(0.0) - arrival_offset.as_secs_f64())with priority_jump = priority as f64. The priority is seconds of arrival-time bump. So x-priority: 0 becomes a bump of 2147483647 seconds, roughly 68 years.
Two consequences:
- Any request carrying
x-priority, even the least urgent one, outranks every request without one, permanently. The PR reasons carefully about "absent must not be promoted to maximum", but the mirror case is unaddressed: in a mixed fleet where some functions have priority config and some do not, unconfigured tenants sit behind all configured traffic forever. WSPT has the same shape,weight = 1.0 + priority_jump, so the1.0baseline vanishes and unconfigured traffic is always last there too. - Within configured traffic, one platform priority step equals one second of queue jump. That scale is a side effect of the anchor, not a chosen value.
The justification in the doc comment only argues for staying inside i32 range, which does not require the maximum. A bounded ceiling gives identical ordering, stays a valid i32, and leaves unconfigured traffic somewhere sane:
/// Platform priorities occupy [0, CEILING]; the derived value stays small so
/// requests without x-priority keep a meaningful position relative to them.
const PLATFORM_PRIORITY_CEILING: u32 = 1_000;
pub(super) fn dynamo_request_priority(priority: u32) -> i32 {
(PLATFORM_PRIORITY_CEILING.saturating_sub(priority.min(PLATFORM_PRIORITY_CEILING))) as i32
}If i32::MAX is deliberate policy ("platform traffic always beats non-platform traffic"), that is a defensible choice, but it needs to be stated here and in the PR description rather than falling out of the clamp argument.
There was a problem hiding this comment.
This kinda becomes an extension of the stargate/gateway contract, and then assumes that the api gateway understands the priority api of the underlying engines. And what happens if stargate is connected to a mix of engines supporting the same model (I assume this is theoretically possible), say a mix of dynamo and raw sglang instances. Each may have a different priority (or other) APIs
I think Pylon/Stargate HAS to present a single unified contract to the api gateway, and will have to internally consolidate the api differences between different backends.
The easiest way I can think of to normalize priority behavior is to have pylon set priority to the lowest-priority value if no priority is supplied. And this will be the behavior for all backends
| /// Absent stays absent: a request without x-priority must not be promoted to | ||
| /// maximum engine priority, so this reads the raw header instead of the | ||
| /// parsed default of 0. Malformed values were already rejected upstream. | ||
| pub(super) fn x_priority_header_value(headers: &HeaderMap) -> Option<u32> { | ||
| headers | ||
| .get(HEADER_PRIORITY)? |
There was a problem hiding this comment.
This is a second, independent parser for a header that validate_required_tunnel_headers in request_observer/headers.rs has already parsed on this exact path (non-health requests validate before reaching send_upstream_request, which is what makes the "rejected upstream" claim true).
The two agree today, both trim and both parse u32, but nothing enforces that they keep agreeing. The root cause is that RequiredTunnelHeaders.priority is a u32 built with unwrap_or_default(), which collapses absent into 0, exactly the distinction this PR needs.
Cleaner: make that field Option<u32> and thread the validated value into send_upstream_request. One parser, and the absent-versus-zero distinction lives in the type instead of in a comment.
| /// Derive x-dynamo-request-priority for the upstream engine from x-priority | ||
| #[arg( | ||
| long, | ||
| action = clap::ArgAction::Set, |
There was a problem hiding this comment.
With ArgAction::Set and default_value_t = true, a bare --pylon-derive-dynamo-priority is an error; only --pylon-derive-dynamo-priority=false works.
The PR description calls the kill switch --pylon-derive-dynamo-priority, which will send whoever is turning this off mid-incident down a wrong path. Worth spelling the =false form in the PR body and wherever this gets written down operationally. (Matches the surrounding flags, so no objection to the style itself.)
| long, | ||
| action = clap::ArgAction::Set, | ||
| default_value_t = true, | ||
| env = "PYLON_DERIVE_DYNAMO_PRIORITY" |
There was a problem hiding this comment.
The checklist marks documentation as up to date, but the diff contains no doc changes, and two existing docs specifically encode this class of rule:
src/libraries/rust/stargate/docs/api-gateway-contract.mdalready documents the analogous strip under "Internal header": "x-stargate-expected-queue-ms: Stargate-to-pylon only. Stargate strips caller values; pylon strips it before upstream forwarding." The newx-dynamo-request-*strip belongs alongside it, and the header table there should gain the derived header.docs/user/llm-request-router-load-balancing.mdcarries the ingressRequestHeaderModifierremove-list plus the warning that the stock HTTPRoute does not strip router-facing headers. Neither thex-dynamo-*family nor this flag and env var appear in it.
The flag and PYLON_DERIVE_DYNAMO_PRIORITY also need to land somewhere an operator will find them under pressure.
| rust_test( | ||
| name = "pylon_test", | ||
| crate = ":pylon", | ||
| deps = _WORKSPACE_DEPS + all_crate_deps(normal_dev = True), |
There was a problem hiding this comment.
Nit: this re-lists _WORKSPACE_DEPS while mock-dynamo's new rust_test lists only all_crate_deps(normal_dev = True). With crate = ":target" rules_rust inherits the crate's deps, so if mock-dynamo builds then the re-listing here is redundant. Worth picking one style so the next person adding a test target does not have to guess which one is load-bearing.
Good catch on these two crates having had no test target at all, by the way. That is the most valuable part of the PR independent of the feature.
| &HeaderName::from_bytes(b"X-Request-Id")?, | ||
| &retry | ||
| )); | ||
| for name in [b"X-Request-Id".as_slice(), b"X-Priority", b"X-Dynamo-Nvext"] { |
There was a problem hiding this comment.
This assertion is the clearest statement of the gap raised on DYNAMO_REQUEST_HEADER_PREFIX: any x-dynamo-* header that is not x-dynamo-request-* is forwarded to the engine. x-dynamo-nvext is harmless, but x-dynamo-worker-instance-id, x-dynamo-dp-rank and x-tenant-id are real Dynamo routing headers that pass this filter today.
Once the strip becomes a deny-list, this loop is the right place to pin which x-dynamo-* names are intentionally still forwarded.
| fn pylon_x_priority_header_value_distinguishes_absent_from_zero() { | ||
| let mut headers = HeaderMap::new(); |
There was a problem hiding this comment.
Minor: both the " 7 " and "not-a-priority" cases are unreachable in production. validate_required_tunnel_headers trims and returns 400 for an unparseable x-priority before send_upstream_request runs on any non-health path, and health requests never reach the derivation.
Fine as a unit-level contract for the function, but it is worth a line in the test saying so, otherwise it reads as coverage of a live path and hides the duplicate-parser issue raised on x_priority_header_value.
|
run this deslop prompt |
| request_quality_monitor: RequestQualityMonitorConfig::default(), | ||
| retry: PylonRetryConfig::default(), | ||
| queue_mismatch_retry: PylonQueueMismatchRetryConfig::default(), | ||
| derive_dynamo_priority: true, |
There was a problem hiding this comment.
pylon has to be backend agnostic. We shouldn't have multiple flags for each type of backend. imagine having derive_dynamo_priority, derive_sglang_priority, and derive_vllm_priority flags. Instead we need a backend_type, that contains the config block for that backend.
This can be a future PR
3 participants
TL;DR
Pylon now translates the platform's
x-priorityheader into Dynamo'sx-dynamo-request-prioritywhen forwarding to the upstream inference server, and strips any client-suppliedx-dynamo-request-*header so spoofed values never reach the engine.Additional Details
Dynamo schedules on
x-dynamo-request-priority(i32, higher wins) and silently ignores values that do not parse as i32. The platform resolves a caller priority and carries it asx-priority(u32, lower is more urgent, absent means unconfigured). Without translation the engine never sees the platform priority; without the strip a client could set engine priority directly.Pylon is the single contract point with the engine, so both pieces live only here; the gateway and Stargate stay Dynamo-agnostic.
i32::MAX - min(x, i32::MAX). Inverts polarity and always produces a valid i32, so Dynamo never drops it.x-priorityis present on the request. Absent stays absent, so unconfigured traffic is never promoted to maximum engine priority. Health requests are excluded.--pylon-derive-dynamo-priority(envPYLON_DERIVE_DYNAMO_PRIORITY). The strip is not gated, so disabling the feature cannot re-open header spoofing.GET /test-control.rust_testtarget here.For the Reviewer
The behavior change is concentrated in
send_upstream_requestandshould_forward_headerincrates/pylon-lib/src/quic_http_tunnel/core.rs. Everything else is plumbing (CLI flag to forwarding config), the mock-dynamo fixture, and tests.Two known follow-ups, deliberately not in this PR: the per-request info log may be demoted to debug once QA validation is done.
For QA
bazel test //src/libraries/rust/stargate/...passes (14 targets, includes the two new ones).x-dynamo-request-*headers do not, absentx-priorityemits nothing, and disabling the flag stops emission but not the strip.i32::MAX,u32::MAX) and the presence rule.derived dynamo request priorityor query mock-dynamo's/test-control.Issues
Closes #620
Checklist
Summary by CodeRabbit