diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs index 0efb17cf4..077ec9f46 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/bringup/upstream.rs @@ -94,7 +94,7 @@ pub(super) async fn send_completion_request( request_id: request_id.clone(), routing_key: None, model_id: model_id.to_string(), - priority: 0, + priority: None, input_tokens: u64::try_from(input_tokens).unwrap_or(u64::MAX), accepted_at: std::time::Instant::now(), }, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/lib.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/lib.rs index d6b9202e4..b32f5e681 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/lib.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/lib.rs @@ -45,9 +45,9 @@ pub use model_lifecycle::{ }; pub use queue_admission::PylonQueueMismatchRetryConfig; pub use quic_http_tunnel::{ - DEFAULT_MAX_SSE_BUFFER_BYTES, PylonRetryConfig, QuicHttpTunnelConfig, QuicHttpTunnelHandle, - ReverseQuicTunnelConfig, ReverseQuicTunnelHandle, TunnelError, TunnelForwardingConfig, - start_quic_http_tunnel, start_reverse_quic_tunnel, + DEFAULT_MAX_SSE_BUFFER_BYTES, DEFAULT_PRIORITY_CEILING, PylonRetryConfig, QuicHttpTunnelConfig, + QuicHttpTunnelHandle, ReverseQuicTunnelConfig, ReverseQuicTunnelHandle, TunnelError, + TunnelForwardingConfig, UpstreamBackend, start_quic_http_tunnel, start_reverse_quic_tunnel, }; pub use registration::{ ClientError, InferenceServerRegistrationClient, InferenceServerRegistrationConfig, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs index eb100fa3b..0a3697abe 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/queue_admission.rs @@ -304,7 +304,7 @@ impl LiveRequestState { }) .filter(|request| request.generation == *generation); model.queue_estimate_ms_for_priority_excluding( - required.priority, + required.queue_priority(), excluded_request, ) }) @@ -346,7 +346,7 @@ impl LiveRequestState { let request_id = required.request_id.clone(); let request = TrackedPromptRequest { generation, - priority: required.priority, + priority: required.queue_priority(), input_tokens: required.input_tokens, phase: TrackedPromptPhase::Pending, active_chat_output_tps: None, @@ -842,7 +842,7 @@ mod tests { request_id: request_id.to_string(), routing_key: None, model_id: model_id.to_string(), - priority, + priority: Some(priority), input_tokens, accepted_at: Instant::now(), } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel.rs index aafdaf23c..489b3174b 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel.rs @@ -13,6 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod backend; mod core; mod endpoint; mod http3; @@ -23,6 +24,7 @@ mod server; mod tests; mod webtransport; +pub use backend::{DEFAULT_PRIORITY_CEILING, UpstreamBackend}; pub use core::{DEFAULT_MAX_SSE_BUFFER_BYTES, PylonRetryConfig, TunnelForwardingConfig}; pub use endpoint::TunnelError; pub use reverse::{ReverseQuicTunnelConfig, ReverseQuicTunnelHandle, start_reverse_quic_tunnel}; diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/backend.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/backend.rs new file mode 100644 index 000000000..c797388e4 --- /dev/null +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/backend.rs @@ -0,0 +1,108 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Upstream inference-server dialects. +//! +//! Pylon speaks the platform tunnel contract upward and translates it into +//! the dialect of the engine it fronts. All engine-specific header names and +//! encodings live in this module. + +use std::fmt; +use std::str::FromStr; + +/// Which engine dialect pylon speaks to its local upstream. Future engines +/// add a variant and a submodule here, never a new CLI flag. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum UpstreamBackend { + /// Forward requests unchanged; derive nothing. Inbound engine priority + /// headers are still stripped. + Passthrough, + /// Derive the engine priority headers from `x-priority`. + #[default] + Dynamo, +} + +impl fmt::Display for UpstreamBackend { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Passthrough => "passthrough", + Self::Dynamo => "dynamo", + }) + } +} + +impl FromStr for UpstreamBackend { + type Err = String; + + fn from_str(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "passthrough" => Ok(Self::Passthrough), + "dynamo" => Ok(Self::Dynamo), + other => Err(format!( + "unknown upstream backend {other:?}; expected \"passthrough\" or \"dynamo\"" + )), + } + } +} + +/// Default priority band ceiling; see [`dynamo::request_priority`]. +pub const DEFAULT_PRIORITY_CEILING: u32 = 3600; + +pub(crate) mod dynamo { + use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; + + /// Engine priority headers pylon derives; the names stay out of the + /// shared tunnel contract because only pylon speaks them. + pub(crate) const HEADER_REQUEST_PRIORITY: &str = "x-dynamo-request-priority"; + pub(crate) const HEADER_REQUEST_STRICT_PRIORITY: &str = "x-dynamo-request-strict-priority"; + + /// Denylist of engine headers pylon owns: inbound values are stripped in + /// every backend mode so pylon stays their only writer. Scoped to the + /// priority headers for now; other engine headers are tracked separately. + const STRIPPED_REQUEST_HEADERS: [&str; 2] = + [HEADER_REQUEST_PRIORITY, HEADER_REQUEST_STRICT_PRIORITY]; + + pub(crate) fn is_stripped_engine_header(name: &HeaderName) -> bool { + STRIPPED_REQUEST_HEADERS.contains(&name.as_str()) + } + + /// Map the platform rank (lower wins, absent = unconfigured) to the + /// engine value (higher wins, read as seconds of queue head start): + /// `max(0, ceiling - rank)`, with absent as the lowest value. The head + /// start is bounded so prioritized traffic cannot starve the rest. + pub(crate) fn request_priority(priority: Option, ceiling: u32) -> i32 { + let ceiling = ceiling.min(i32::MAX as u32); + let rank = priority.unwrap_or(ceiling).min(ceiling); + (ceiling - rank) as i32 + } + + /// Emit both priority headers on every inference request. + pub(crate) fn apply_priority_headers( + priority: Option, + ceiling: u32, + upstream_headers: &mut HeaderMap, + ) -> i32 { + let dynamo_priority = request_priority(priority, ceiling); + upstream_headers.insert( + HeaderName::from_static(HEADER_REQUEST_PRIORITY), + HeaderValue::from(dynamo_priority), + ); + upstream_headers.insert( + HeaderName::from_static(HEADER_REQUEST_STRICT_PRIORITY), + HeaderValue::from_static("0"), + ); + dynamo_priority + } +} diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs index f5b89f5d6..22916aead 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/core.rs @@ -37,6 +37,7 @@ use tokio_util::{sync::CancellationToken, task::TaskTracker}; use tracing::{Instrument, Span, field}; use tracing_opentelemetry::OpenTelemetrySpanExt; +use super::backend::{self, DEFAULT_PRIORITY_CEILING, UpstreamBackend}; use crate::output_token_parser::{OutputTokenParser, OutputTokenProgress}; use crate::queue_admission::{ PylonQueueMismatchRetryConfig, QueueAdmissionDecision, QueueTrackedRequestGuard, @@ -105,6 +106,10 @@ pub struct TunnelForwardingConfig { pub request_quality_monitor: RequestQualityMonitorConfig, pub retry: PylonRetryConfig, pub queue_mismatch_retry: PylonQueueMismatchRetryConfig, + /// Engine dialect spoken to the local upstream; see [`UpstreamBackend`]. + pub upstream_backend: UpstreamBackend, + /// Priority band ceiling; see [`backend::dynamo::request_priority`]. + pub priority_ceiling: u32, pub metrics: Option>, #[cfg(test)] pub webtransport_stream_header_wait_tx: Option>, @@ -121,6 +126,8 @@ impl Default for TunnelForwardingConfig { request_quality_monitor: RequestQualityMonitorConfig::default(), retry: PylonRetryConfig::default(), queue_mismatch_retry: PylonQueueMismatchRetryConfig::default(), + upstream_backend: UpstreamBackend::default(), + priority_ceiling: DEFAULT_PRIORITY_CEILING, metrics: None, #[cfg(test)] webtransport_stream_header_wait_tx: None, @@ -141,6 +148,8 @@ pub(super) struct TunnelServerApp { pub(super) request_quality_monitor: RequestQualityMonitorConfig, pub(super) retry: PylonRetryConfig, pub(super) queue_mismatch_retry: PylonQueueMismatchRetryConfig, + pub(super) upstream_backend: UpstreamBackend, + pub(super) priority_ceiling: u32, pub(super) metrics: Option>, #[cfg(test)] pub(super) webtransport_stream_header_wait_tx: Option>, @@ -164,6 +173,8 @@ impl TunnelServerApp { request_quality_monitor: forwarding.request_quality_monitor, retry: forwarding.retry, queue_mismatch_retry: forwarding.queue_mismatch_retry, + upstream_backend: forwarding.upstream_backend, + priority_ceiling: forwarding.priority_ceiling, metrics: forwarding.metrics, #[cfg(test)] webtransport_stream_header_wait_tx: forwarding.webtransport_stream_header_wait_tx, @@ -654,13 +665,17 @@ pub(super) async fn forward_tunnel_request( } } + let priority = lifecycle + .as_ref() + .and_then(|lifecycle| lifecycle.required.priority); let response = match send_upstream_request( app, method, &path_and_query, &request_headers, body_bytes, - !health_request, + health_request, + priority, ) .await { @@ -699,9 +714,10 @@ async fn send_upstream_request( path_and_query: &str, request_headers: &HeaderMap, body_bytes: Vec, - traced: bool, + health_request: bool, + priority: Option, ) -> Result { - let span = if traced { + let span = if !health_request { let span = tracing::info_span!( "pylon_upstream_http_request", otel_parent = field::Empty, @@ -710,6 +726,8 @@ async fn send_upstream_request( inference_server.id = %app.inference_server_id, upstream.status = field::Empty, upstream.error = field::Empty, + priority = field::Empty, + dynamo.request_priority = field::Empty, ); let _ = span.set_parent(pylon_upstream_parent_context(request_headers)); if let Some(otel_parent) = otel_parent_from_headers(request_headers) { @@ -725,7 +743,18 @@ async fn send_upstream_request( upstream_headers.append(name, value.clone()); } } - if traced { + if !health_request { + if let Some(priority) = priority { + span.record("priority", priority); + } + if app.upstream_backend == UpstreamBackend::Dynamo { + let dynamo_priority = backend::dynamo::apply_priority_headers( + priority, + app.priority_ceiling, + &mut upstream_headers, + ); + span.record("dynamo.request_priority", dynamo_priority); + } inject_trace_context(&mut upstream_headers, &span.context()); } let send = async { @@ -1058,6 +1087,7 @@ pub(super) fn join_base_path(base: &str, path_and_query: &str) -> Result bool { !is_tunnel_control_header(name, retry) + && !backend::dynamo::is_stripped_engine_header(name) && !matches!( name.as_str(), "host" | "x-method" | "x-path" | HEADER_STARGATE_EXPECTED_QUEUE_MS diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/raw_quic.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/raw_quic.rs index df26999b8..e39d81d49 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/raw_quic.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/raw_quic.rs @@ -171,6 +171,8 @@ mod tests { request_quality_monitor: RequestQualityMonitorConfig::default(), retry: PylonRetryConfig::default(), queue_mismatch_retry: PylonQueueMismatchRetryConfig::default(), + upstream_backend: crate::quic_http_tunnel::UpstreamBackend::default(), + priority_ceiling: crate::quic_http_tunnel::DEFAULT_PRIORITY_CEILING, metrics: None, #[cfg(test)] webtransport_stream_header_wait_tx: None, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs index 6bdf49021..ae8ef3dab 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/quic_http_tunnel/tests.rs @@ -13,6 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use super::backend::{DEFAULT_PRIORITY_CEILING, UpstreamBackend, dynamo}; use super::core::{ MAX_SPECULATIVE_REQUEST_BODY_PREALLOC_BYTES, TunnelServerApp, extend_body_from_buf, is_health_request_path, otel_parent_from_headers, pylon_upstream_parent_context, @@ -477,6 +478,8 @@ fn pylon_request_header_filter_strips_tunnel_headers_case_insensitively() "X-Method", "X-Path", "X-Stargate-Expected-Queue-Ms", + "X-Dynamo-Request-Priority", + "X-Dynamo-Request-Strict-Priority", ] .into_iter() .chain(RETRY_CONTROL_REQUEST_HEADERS) @@ -493,6 +496,46 @@ fn pylon_request_header_filter_strips_tunnel_headers_case_insensitively() Ok(()) } +#[test] +fn pylon_dynamo_request_priority_inverts_within_bounded_ceiling() { + let ceiling = DEFAULT_PRIORITY_CEILING; + // Most urgent platform rank gets the full head start. + assert_eq!(dynamo::request_priority(Some(0), ceiling), ceiling as i32); + assert_eq!( + dynamo::request_priority(Some(7), ceiling), + (ceiling - 7) as i32 + ); + // Unconfigured and beyond-ceiling ranks both land at the lowest value. + assert_eq!(dynamo::request_priority(None, ceiling), 0); + assert_eq!(dynamo::request_priority(Some(ceiling), ceiling), 0); + assert_eq!(dynamo::request_priority(Some(u32::MAX), ceiling), 0); + // A ceiling beyond i32 is clamped so the emitted value stays a valid i32. + assert_eq!(dynamo::request_priority(Some(0), u32::MAX), i32::MAX); + assert_eq!(dynamo::request_priority(None, u32::MAX), 0); + // A ceiling of 0 collapses every rank to the lowest value. + assert_eq!(dynamo::request_priority(Some(0), 0), 0); + assert_eq!(dynamo::request_priority(None, 0), 0); +} + +#[test] +fn pylon_dynamo_priority_headers_are_always_emitted() { + let mut headers = HeaderMap::new(); + let emitted = dynamo::apply_priority_headers(Some(7), DEFAULT_PRIORITY_CEILING, &mut headers); + assert_eq!(emitted, (DEFAULT_PRIORITY_CEILING - 7) as i32); + assert_eq!( + headers["x-dynamo-request-priority"], + emitted.to_string().as_str() + ); + assert_eq!(headers["x-dynamo-request-strict-priority"], "0"); + + // Absent platform priority pins both headers to the lowest values. + let mut headers = HeaderMap::new(); + let emitted = dynamo::apply_priority_headers(None, DEFAULT_PRIORITY_CEILING, &mut headers); + assert_eq!(emitted, 0); + assert_eq!(headers["x-dynamo-request-priority"], "0"); + assert_eq!(headers["x-dynamo-request-strict-priority"], "0"); +} + #[test] fn pylon_trace_context_extracts_remote_parent() -> Result<()> { opentelemetry::global::set_text_map_propagator( @@ -765,7 +808,7 @@ async fn start_queue_mismatch_test_tunnel( request_id: "req-already-queued".to_string(), routing_key: Some("rk-1".to_string()), model_id: "model-a".to_string(), - priority: 0, + priority: None, input_tokens: 100, accepted_at: std::time::Instant::now(), }); @@ -1633,6 +1676,147 @@ async fn quic_tunnel_forwards_to_http_backend() { tunnel.shutdown().await; } +/// Echoes the Dynamo priority headers the backend received, so the tunnel +/// tests assert on what actually crossed the pylon-to-engine hop. +fn dynamo_priority_echo_router() -> Router { + Router::new() + .route( + "/health", + axum::routing::get(|req: Request| async move { + let dynamo_priority = req + .headers() + .get("x-dynamo-request-priority") + .and_then(|value| value.to_str().ok()) + .unwrap_or("absent") + .to_string(); + ([("x-echo-dynamo-priority", dynamo_priority)], "ok") + }), + ) + .route( + "/v1/chat/completions", + post(|req: Request| async move { + let echo_header = |name: &str| { + req.headers() + .get(name) + .and_then(|value| value.to_str().ok()) + .unwrap_or("absent") + .to_string() + }; + let dynamo_priority = echo_header("x-dynamo-request-priority"); + let dynamo_strict_priority = echo_header("x-dynamo-request-strict-priority"); + let mut sse = axum::response::Sse::new(async_stream::stream! { + yield Ok::<_, std::convert::Infallible>( + Event::default().data(r#"{"object":"chat.completion.chunk","choices":[{"delta":{"content":"ok"}}]}"#) + ); + yield Ok::<_, std::convert::Infallible>(Event::default().data("[DONE]")); + }) + .into_response(); + sse.headers_mut().insert( + HeaderName::from_static("x-echo-dynamo-priority"), + HeaderValue::from_str(&dynamo_priority).unwrap(), + ); + sse.headers_mut().insert( + HeaderName::from_static("x-echo-dynamo-strict-priority"), + HeaderValue::from_str(&dynamo_strict_priority).unwrap(), + ); + *sse.status_mut() = StatusCode::OK; + sse + }), + ) +} + +#[tokio::test] +async fn quic_tunnel_derives_dynamo_priority_from_x_priority() { + let (config, _metrics) = metered_test_tunnel_config_for(dynamo_priority_echo_router()).await; + let ceiling = config.forwarding.priority_ceiling; + let mut tunnel = RawTunnelTest::start(config).await; + + let mut headers = + tunnel_request_headers("/v1/chat/completions", "model-a", "req-dynamo-1", "11"); + headers.insert("x-priority", "7".parse().unwrap()); + // Spoofed engine headers must be replaced by pylon-derived values. + headers.insert("x-dynamo-request-priority", "42".parse().unwrap()); + headers.insert("x-dynamo-request-strict-priority", "1".parse().unwrap()); + tunnel + .send(headers, br#"{"messages":[],"stream":true}"#) + .await; + + let response_headers = tunnel.response_head(StatusCode::OK).await; + assert_eq!( + response_headers + .get("x-echo-dynamo-priority") + .unwrap() + .to_str() + .unwrap(), + (ceiling - 7).to_string() + ); + assert_eq!(response_headers["x-echo-dynamo-strict-priority"], "0"); + + tunnel.shutdown().await; +} + +#[tokio::test] +async fn quic_tunnel_emits_lowest_dynamo_priority_without_x_priority() { + let (config, _metrics) = metered_test_tunnel_config_for(dynamo_priority_echo_router()).await; + let mut tunnel = RawTunnelTest::start(config).await; + + let mut headers = + tunnel_request_headers("/v1/chat/completions", "model-a", "req-dynamo-2", "11"); + headers.insert("x-dynamo-request-priority", "42".parse().unwrap()); + tunnel + .send(headers, br#"{"messages":[],"stream":true}"#) + .await; + + // Unconfigured requests carry the lowest priority instead of no header. + let response_headers = tunnel.response_head(StatusCode::OK).await; + assert_eq!(response_headers["x-echo-dynamo-priority"], "0"); + assert_eq!(response_headers["x-echo-dynamo-strict-priority"], "0"); + + tunnel.shutdown().await; +} + +#[tokio::test] +async fn quic_tunnel_health_requests_carry_no_derived_priority() { + let (config, _metrics) = metered_test_tunnel_config_for(dynamo_priority_echo_router()).await; + let mut tunnel = RawTunnelTest::start(config).await; + + // Health requests skip validation, so no required tunnel headers. + let mut headers = HeaderMap::new(); + headers.insert("x-method", "GET".parse().unwrap()); + headers.insert("x-path", "/health".parse().unwrap()); + // A spoofed engine header is stripped on the health path too. + headers.insert("x-dynamo-request-priority", "42".parse().unwrap()); + tunnel.send(headers, b"").await; + + let response_headers = tunnel.response_head(StatusCode::OK).await; + assert_eq!(response_headers["x-echo-dynamo-priority"], "absent"); + + tunnel.shutdown().await; +} + +#[tokio::test] +async fn quic_tunnel_passthrough_backend_strips_but_derives_nothing() { + let (mut config, _metrics) = + metered_test_tunnel_config_for(dynamo_priority_echo_router()).await; + config.forwarding.upstream_backend = UpstreamBackend::Passthrough; + let mut tunnel = RawTunnelTest::start(config).await; + + let mut headers = + tunnel_request_headers("/v1/chat/completions", "model-a", "req-dynamo-3", "11"); + headers.insert("x-priority", "7".parse().unwrap()); + headers.insert("x-dynamo-request-strict-priority", "1".parse().unwrap()); + tunnel + .send(headers, br#"{"messages":[],"stream":true}"#) + .await; + + let response_headers = tunnel.response_head(StatusCode::OK).await; + assert_eq!(response_headers["x-echo-dynamo-priority"], "absent"); + // Stripping inbound engine priority headers is not gated by the backend. + assert_eq!(response_headers["x-echo-dynamo-strict-priority"], "absent"); + + tunnel.shutdown().await; +} + #[tokio::test] async fn quic_tunnel_rejects_pending_generation_before_upstream() { let upstream_hits = Arc::new(AtomicUsize::new(0)); diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs index 69d0ebc31..84c1192b3 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer.rs @@ -161,11 +161,12 @@ impl RequestObserver { generation: Option, runtime_state: PylonRuntimeState, ) -> Self { + let priority = required.queue_priority(); let RequiredTunnelHeaders { request_id, routing_key, model_id, - priority, + priority: _, input_tokens, accepted_at, } = required; @@ -605,14 +606,14 @@ mod tests { assert_eq!(required.routing_key.as_deref(), Some("rk-1")); assert_eq!(required.model_id, "model-a"); assert_eq!(required.input_tokens, 42); - assert_eq!(required.priority, 7); + assert_eq!(required.priority, Some(7)); } #[test] - fn validate_required_tunnel_headers_defaults_missing_priority_to_zero() { + fn validate_required_tunnel_headers_keeps_missing_priority_absent() { let required = validate_required_tunnel_headers(&request_headers("req-1", 42)).unwrap(); - assert_eq!(required.priority, 0); + assert_eq!(required.priority, None); } #[test] @@ -676,7 +677,7 @@ mod tests { request_id: "req-embeddings-terminal".to_string(), routing_key: Some("rk-1".to_string()), model_id: "model-embed".to_string(), - priority: 0, + priority: None, input_tokens: 12, accepted_at: Instant::now(), } diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs index e3e50e73c..2b11e8f4c 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/request_observer/headers.rs @@ -62,11 +62,19 @@ pub(crate) struct RequiredTunnelHeaders { pub request_id: String, pub routing_key: Option, pub model_id: String, - pub priority: u32, + pub priority: Option, pub input_tokens: u64, pub(crate) accepted_at: Instant, } +impl RequiredTunnelHeaders { + /// Priority for queue accounting and observation, where unconfigured + /// counts as 0. The engine derivation reads `priority` directly instead. + pub(crate) fn queue_priority(&self) -> u32 { + self.priority.unwrap_or_default() + } +} + pub(crate) fn validate_required_tunnel_headers( request_headers: &HeaderMap, ) -> Result { @@ -77,8 +85,7 @@ pub(crate) fn validate_required_tunnel_headers( .ok_or_else(|| MissingRequiredHeaderError::new(HEADER_MODEL))?; let input_tokens = parse_optional_numeric_header(request_headers, HEADER_INPUT_TOKENS)? .ok_or_else(|| MissingRequiredHeaderError::new(HEADER_INPUT_TOKENS))?; - let priority = - parse_optional_numeric_header(request_headers, HEADER_PRIORITY)?.unwrap_or_default(); + let priority = parse_optional_numeric_header(request_headers, HEADER_PRIORITY)?; Ok(RequiredTunnelHeaders { request_id, routing_key, diff --git a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs index 247db3288..1c77c6958 100644 --- a/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs +++ b/src/libraries/rust/stargate/crates/pylon-lib/src/stats/collector.rs @@ -2048,7 +2048,7 @@ mod tests { request_id: "req-queued-after-fallback-samples".to_string(), routing_key: None, model_id: "model-a".to_string(), - priority: 0, + priority: None, input_tokens: 50, accepted_at: std::time::Instant::now(), }); @@ -2479,7 +2479,7 @@ mod tests { request_id: "req-queued".to_string(), routing_key: None, model_id: "model-a".to_string(), - priority: 0, + priority: None, input_tokens: 32, accepted_at: std::time::Instant::now(), }, diff --git a/src/libraries/rust/stargate/crates/pylon/BUILD.bazel b/src/libraries/rust/stargate/crates/pylon/BUILD.bazel index cf18a04d6..e0238a026 100644 --- a/src/libraries/rust/stargate/crates/pylon/BUILD.bazel +++ b/src/libraries/rust/stargate/crates/pylon/BUILD.bazel @@ -3,7 +3,7 @@ load("@rules_shell//shell:sh_test.bzl", "sh_test") load("@stargate_crates//:defs.bzl", "aliases", "all_crate_deps") -load("@rules_rust//rust:defs.bzl", "rust_binary") +load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test") load("//rules/oci:defs.bzl", "rust_oci_image") # `all_crate_deps()` only resolves third-party crate_universe deps; @@ -28,6 +28,12 @@ rust_binary( deps = _WORKSPACE_DEPS + all_crate_deps(normal = True), ) +rust_test( + name = "pylon_test", + crate = ":pylon", + deps = _WORKSPACE_DEPS + all_crate_deps(normal_dev = True), +) + # Multi-arch OCI image. distroless/cc base, binary at /usr/local/bin/pylon. rust_oci_image( name = "image", diff --git a/src/libraries/rust/stargate/crates/pylon/src/main.rs b/src/libraries/rust/stargate/crates/pylon/src/main.rs index 335632df8..609a15a58 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/main.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/main.rs @@ -14,7 +14,9 @@ // limitations under the License. use anyhow::Result; -use pylon_lib::{EngineStatsStreamMode, ModelDiscoveryProvider, TunnelTransportProtocol}; +use pylon_lib::{ + EngineStatsStreamMode, ModelDiscoveryProvider, TunnelTransportProtocol, UpstreamBackend, +}; use stargate_protocol::BackendConnectivity; use stargate_protocol::tunnel_contract::HEADER_STARGATE_UPSTREAM_RETRYABLE; @@ -204,6 +206,25 @@ struct Args { /// Optional retry-after hint in milliseconds for local queue-mismatch retries #[arg(long, env = "PYLON_QUEUE_MISMATCH_RETRY_AFTER_MS", value_name = "MS")] pylon_queue_mismatch_retry_after_ms: Option, + /// Engine dialect spoken to the local upstream: "dynamo" derives the + /// engine priority headers from x-priority, "passthrough" derives nothing + #[arg( + long, + default_value = "dynamo", + env = "PYLON_UPSTREAM_BACKEND", + value_name = "BACKEND" + )] + pylon_upstream_backend: UpstreamBackend, + /// Priority band ceiling: x-priority rank 0 maps to this engine value and + /// ranks at or beyond it map to the lowest. Dynamo reads the derived + /// value as seconds of queue head start. + #[arg( + long, + default_value_t = pylon_lib::DEFAULT_PRIORITY_CEILING, + env = "PYLON_PRIORITY_CEILING", + value_name = "RANK" + )] + pylon_priority_ceiling: u32, /// Collect post-stream output quality metrics (gibberish checks) #[arg(long, default_value_t = false)] collect_quality_metrics: bool, @@ -243,7 +264,7 @@ async fn main() -> Result<()> { mod tests { use pylon_lib::{ EngineStatsStreamMode, ModelDiscoveryProvider, PylonQueueMismatchRetryConfig, - PylonRetryConfig, TunnelTransportProtocol, + PylonRetryConfig, TunnelForwardingConfig, TunnelTransportProtocol, }; use reqwest::header::HeaderName; @@ -413,6 +434,33 @@ mod tests { assert!(retry.retryable_upstream_status_codes.is_empty()); } + #[test] + fn pylon_upstream_backend_cli_defaults_match_runtime_defaults() { + let args = parse_args(""); + let defaults = TunnelForwardingConfig::default(); + + assert_eq!(args.pylon_upstream_backend, defaults.upstream_backend); + assert_eq!(args.pylon_priority_ceiling, defaults.priority_ceiling); + } + + #[test] + fn pylon_upstream_backend_cli_overrides_are_applied() { + let args = parse_argv(&[ + "--pylon-upstream-backend", + "passthrough", + "--pylon-priority-ceiling", + "600", + ]); + + assert_eq!(args.pylon_upstream_backend, UpstreamBackend::Passthrough); + assert_eq!(args.pylon_priority_ceiling, 600); + } + + #[test] + fn pylon_upstream_backend_cli_rejects_unknown_backend() { + assert!(try_parse_argv(&["--pylon-upstream-backend", "sglang"]).is_err()); + } + #[test] fn pylon_queue_mismatch_retry_cli_defaults_match_runtime_defaults() { let args = parse_args(""); diff --git a/src/libraries/rust/stargate/crates/pylon/src/startup.rs b/src/libraries/rust/stargate/crates/pylon/src/startup.rs index 979ef3a65..7338f84a9 100644 --- a/src/libraries/rust/stargate/crates/pylon/src/startup.rs +++ b/src/libraries/rust/stargate/crates/pylon/src/startup.rs @@ -28,8 +28,8 @@ use pylon_lib::{ ModelInitialization, ModelLifecycleConfig, ModelLifecycleHandle, ModelSource, PylonMetrics, PylonQueueMismatchRetryConfig, PylonRetryConfig, PylonRuntimeState, QuicHttpTunnelConfig, QuicHttpTunnelHandle, RequestQualityMonitorConfig, StatsCollectorConfig, StatsCollectorHandle, - TunnelForwardingConfig, start_engine_stats_stream, start_metrics_server, start_model_lifecycle, - start_quic_http_tunnel, start_stats_collector_with_engine_stats, + TunnelForwardingConfig, UpstreamBackend, start_engine_stats_stream, start_metrics_server, + start_model_lifecycle, start_quic_http_tunnel, start_stats_collector_with_engine_stats, stats_aggregator_update_channel, }; use reqwest::header::HeaderName; @@ -79,6 +79,8 @@ fn log_startup_complete( inference_server_id, cluster_id = %plan.cluster_id, upstream = %plan.upstream, + upstream_backend = %plan.upstream_backend, + priority_ceiling = plan.priority_ceiling, model_ids = ?model_ids, "pylon startup complete; stargate registration started (reverse tunnel mode)" ); @@ -89,6 +91,8 @@ fn log_startup_complete( cluster_id = %plan.cluster_id, inference_server_url = registration_inference_server_url, upstream = %plan.upstream, + upstream_backend = %plan.upstream_backend, + priority_ceiling = plan.priority_ceiling, model_ids = ?model_ids, "pylon startup complete; stargate registration started (direct tunnel mode)" ); @@ -101,6 +105,8 @@ pub(crate) struct PylonStartupPlan { model_source: ModelSource, pylon_retry: PylonRetryConfig, queue_mismatch_retry: PylonQueueMismatchRetryConfig, + upstream_backend: UpstreamBackend, + priority_ceiling: u32, model_initialization: ModelInitialization, bringup: BringupConfig, request_quality_monitor: RequestQualityMonitorConfig, @@ -146,6 +152,8 @@ impl PylonStartupPlan { model_source, pylon_retry: pylon_retry_config_from_args(args)?, queue_mismatch_retry: pylon_queue_mismatch_retry_config_from_args(args)?, + upstream_backend: args.pylon_upstream_backend, + priority_ceiling: args.pylon_priority_ceiling, model_initialization, bringup: BringupConfig { enabled: !args.disable_bringup, @@ -504,6 +512,8 @@ fn tunnel_forwarding_config_from_plan( metrics: Some(metrics), retry: plan.pylon_retry.clone(), queue_mismatch_retry: plan.queue_mismatch_retry.clone(), + upstream_backend: plan.upstream_backend, + priority_ceiling: plan.priority_ceiling, ..Default::default() } } @@ -1113,6 +1123,27 @@ mod tests { tunnel.shutdown().await; } + #[test] + fn upstream_backend_flows_from_args_to_forwarding_config() { + let (_, default_plan) = startup(&[]); + let forwarding = test_forwarding(&default_plan); + assert_eq!(forwarding.upstream_backend, UpstreamBackend::Dynamo); + assert_eq!( + forwarding.priority_ceiling, + pylon_lib::DEFAULT_PRIORITY_CEILING + ); + + let (_, passthrough_plan) = startup(&[ + "--pylon-upstream-backend", + "passthrough", + "--pylon-priority-ceiling", + "600", + ]); + let forwarding = test_forwarding(&passthrough_plan); + assert_eq!(forwarding.upstream_backend, UpstreamBackend::Passthrough); + assert_eq!(forwarding.priority_ceiling, 600); + } + #[test] fn direct_tunnel_config_from_plan_preserves_runtime_inputs() { let (args, plan) = startup(&[ diff --git a/src/libraries/rust/stargate/docs/api-gateway-contract.md b/src/libraries/rust/stargate/docs/api-gateway-contract.md index 48889ceb9..5c7323f6d 100644 --- a/src/libraries/rust/stargate/docs/api-gateway-contract.md +++ b/src/libraries/rust/stargate/docs/api-gateway-contract.md @@ -136,7 +136,7 @@ Optional trusted headers: | `x-routing-key` | Authenticated routing scope. Omit for unscoped. | | `x-routing-method` | Request-scoped load-balancer override, only for methods allowed by Stargate config. | | `x-cache-affinity-key` | Opaque cache/prefix identity. Required by some LB configs. | -| `x-priority` | Unsigned priority, default `0`. | +| `x-priority` | Unsigned priority rank; lower is more urgent. Omit when no priority is resolved. | | `x-request-slo-ms` | Per-request LB latency hint. | | `x-max-wait-ms` | Wait budget for temporarily infeasible candidates. | | `x-stargate-max-wait-ms` | Stargate internal retry budget. | @@ -144,10 +144,27 @@ Optional trusted headers: The gateway must synthesize or validate these headers. Do not pass public caller-supplied routing headers through blindly. -Internal header: +Internal headers: - `x-stargate-expected-queue-ms`: Stargate-to-pylon only. Stargate strips caller values; pylon strips it before upstream forwarding. +- `x-dynamo-request-priority` and `x-dynamo-request-strict-priority`: + pylon-to-engine only. Pylon strips inbound values in every backend mode, so + pylon is the only writer of these two headers. When pylon runs with + `--pylon-upstream-backend dynamo` (the default), it emits both headers on + every inference request: the priority is derived from `x-priority` as + `max(0, ceiling - x)` with a configurable ceiling + (`--pylon-priority-ceiling`, default 3600), requests without `x-priority` + carry the lowest value `0`, and the strict tier is always `0`. Always + emitting means the engine reads priority only from pylon, never from + client-supplied values. Other engine headers pass through unchanged; the + strip denylist is scoped to the priority headers. + + An absent `x-priority` and `x-priority: 0` are opposite ends of the range: + absence maps to the lowest engine priority, while rank `0` maps to the + highest. The gateway must not synthesize `x-priority: 0` for unconfigured + requests. Stargate treats an absent header as `0` for its own queue + accounting only; that default never reaches the engine. Body rules: