@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b817b23f7b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 06:00 UTC

[Rust Frontend] add --enable-request-id-headers flag support. (#43883)

Signed-off-by: Maria Guevara <kawaiiplush14@gmail.com>

Maria Guevara committed Jun 2, 2026 at 09:08 UTC b817b23f7b1dcb36543bcc9a431edc3c6c937a47
11 files changed +170 -12
rust/src/cmd/src/cli.rs
+11
@@ -165,6 +165,15 @@ pub struct SharedRuntimeArgs {
165 #[serde(default)]
166 pub enable_log_requests: bool,
167
168 + /// If specified, API server will add X-Request-Id header to responses.
169 + #[arg(
170 + long,
171 + default_missing_value = "true",
172 + num_args = 0..=1
173 + )]
174 + #[serde(default)]
175 + pub enable_request_id_headers: bool,
176 +
177 /// Disable periodic logging of engine statistics (throughput, queue depth,
178 /// cache usage).
179 #[arg(long)]
@@ -238,6 +247,7 @@ impl SharedRuntimeArgs {
247 default_chat_template_kwargs: self.default_chat_template_kwargs,
248 chat_template_content_format: self.chat_template_content_format,
249 enable_log_requests: self.enable_log_requests,
250 + enable_request_id_headers: self.enable_request_id_headers,
251 disable_log_stats: self.disable_log_stats,
252 grpc_port: self.grpc_port,
253 shutdown_timeout,
@@ -278,6 +288,7 @@ impl SharedRuntimeArgs {
288 default_chat_template_kwargs: self.default_chat_template_kwargs,
289 chat_template_content_format: self.chat_template_content_format,
290 enable_log_requests: self.enable_log_requests,
291 + enable_request_id_headers: self.enable_request_id_headers,
292 disable_log_stats: self.disable_log_stats,
293 grpc_port: self.grpc_port,
294 shutdown_timeout,
rust/src/cmd/src/cli/tests.rs
+46
@@ -43,6 +43,7 @@ fn serve_args_forward_python_flags_with_separator() {
43 default_chat_template_kwargs: None,
44 chat_template_content_format: Auto,
45 enable_log_requests: false,
46 + enable_request_id_headers: false,
47 disable_log_stats: false,
48 served_model_name: [],
49 },
@@ -116,6 +117,46 @@ fn serve_args_accept_explicit_deepseek_v32_renderer() {
117 assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
118 }
119
120 +#[test]
121 +fn serve_passes_enable_request_id_headers_into_config() {
122 + let cli = Cli::try_parse_from([
123 + "vllm-rs",
124 + "serve",
125 + "Qwen/Qwen3-0.6B",
126 + "--enable-request-id-headers",
127 + ])
128 + .unwrap();
129 +
130 + let Command::Serve(args) = cli.command else {
131 + panic!("expected serve args");
132 + };
133 + let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
134 + assert!(config.enable_request_id_headers);
135 +}
136 +
137 +#[test]
138 +fn frontend_args_json_passes_enable_request_id_headers_into_config() {
139 + let cli = Cli::try_parse_from([
140 + "vllm-rs",
141 + "frontend",
142 + "--listen-fd",
143 + "3",
144 + "--input-address",
145 + "ipc:///tmp/input.sock",
146 + "--output-address",
147 + "ipc:///tmp/output.sock",
148 + "--args-json",
149 + r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#,
150 + ])
151 + .unwrap();
152 +
153 + let Command::Frontend(args) = cli.command else {
154 + panic!("expected frontend args");
155 + };
156 + let config = args.into_config();
157 + assert!(config.enable_request_id_headers);
158 +}
159 +
160 #[test]
161 fn serve_args_reject_unknown_renderer_value() {
162 let error = Cli::try_parse_from([
@@ -218,6 +259,7 @@ fn frontend_args_accept_json() {
259 default_chat_template_kwargs: None,
260 chat_template_content_format: Auto,
261 enable_log_requests: false,
262 + enable_request_id_headers: false,
263 disable_log_stats: false,
264 served_model_name: [],
265 },
@@ -616,6 +658,7 @@ fn serve_args_accept_handshake_aliases() {
658 default_chat_template_kwargs: None,
659 chat_template_content_format: Auto,
660 enable_log_requests: false,
661 + enable_request_id_headers: false,
662 disable_log_stats: false,
663 served_model_name: [],
664 },
@@ -733,6 +776,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
776 default_chat_template_kwargs: None,
777 chat_template_content_format: Auto,
778 enable_log_requests: false,
779 + enable_request_id_headers: false,
780 disable_log_stats: false,
781 grpc_port: None,
782 shutdown_timeout: 0ns,
@@ -795,6 +839,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
839 default_chat_template_kwargs: None,
840 chat_template_content_format: Auto,
841 enable_log_requests: false,
842 + enable_request_id_headers: false,
843 disable_log_stats: false,
844 grpc_port: None,
845 shutdown_timeout: 0ns,
@@ -872,6 +917,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
917 default_chat_template_kwargs: None,
918 chat_template_content_format: Auto,
919 enable_log_requests: false,
920 + enable_request_id_headers: false,
921 disable_log_stats: false,
922 grpc_port: None,
923 shutdown_timeout: 0ns,
rust/src/cmd/src/cli/unsupported.rs
-9
@@ -620,15 +620,6 @@ pub struct ServerUnsupportedArgs {
620 #[arg(long)]
621 pub middleware: Option<Unsupported>,
622
623 - /// If specified, API server will add X-Request-Id header to responses.
624 - #[arg(
625 - long,
626 - visible_alias = "no-enable-request-id-headers",
627 - default_missing_value = "true",
628 - num_args = 0..=1
629 - )]
630 - pub enable_request_id_headers: Option<Unsupported>,
631 -
623 /// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint.
624 #[arg(
625 long,
rust/src/server/examples/external_engine_openai_qwen.rs
+1
@@ -68,6 +68,7 @@ async fn main() -> Result<()> {
68 default_chat_template_kwargs: None,
69 chat_template_content_format: ChatTemplateContentFormatOption::Auto,
70 enable_log_requests: false,
71 + enable_request_id_headers: false,
72 disable_log_stats: false,
73 grpc_port: None,
74 shutdown_timeout: Duration::ZERO,
rust/src/server/src/config.rs
+2
@@ -61,6 +61,8 @@ pub struct Config {
61 pub chat_template_content_format: ChatTemplateContentFormatOption,
62 /// Log a summary line for each completed request.
63 pub enable_log_requests: bool,
64 + /// When `true`, set `X-Request-Id` on every HTTP response.
65 + pub enable_request_id_headers: bool,
66 /// When `true`, suppress periodic stats logging (throughput, queue depth,
67 /// cache usage).
68 pub disable_log_stats: bool,
rust/src/server/src/lib.rs
+3 -1
@@ -85,7 +85,9 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
85 };
86
87 Ok(Arc::new(
88 - AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests),
88 + AppState::new(served_model_names, chat)
89 + .with_log_requests(config.enable_log_requests)
90 + .with_request_id_headers(config.enable_request_id_headers),
91 ))
92 }
93
rust/src/server/src/middleware/mod.rs
+2
@@ -1,5 +1,7 @@
1 mod load;
2 mod metrics;
3 +mod request_id;
4
5 pub use load::track_server_load;
6 pub use metrics::track_http_metrics;
7 +pub use request_id::set_request_id_header;
rust/src/server/src/middleware/request_id.rs new
+24
@@ -0,0 +1,24 @@
1 +use axum::extract::Request;
2 +use axum::http::HeaderValue;
3 +use axum::http::header::HeaderName;
4 +use axum::middleware::Next;
5 +use axum::response::Response;
6 +use uuid::Uuid;
7 +
8 +const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
9 +
10 +/// Echo the request's `X-Request-Id` on the response, or generate a fresh
11 +/// `uuid4` hex if the request did not provide one.
12 +///
13 +/// Original Python:
14 +/// `vllm.entrypoints.openai.server_utils.XRequestIdMiddleware`.
15 +pub async fn set_request_id_header(req: Request, next: Next) -> Response {
16 + let incoming = req.headers().get(&X_REQUEST_ID).cloned();
17 + let mut response = next.run(req).await;
18 + let value = incoming.unwrap_or_else(|| {
19 + HeaderValue::from_str(&Uuid::new_v4().simple().to_string())
20 + .expect("uuid hex is valid header value")
21 + });
22 + response.headers_mut().insert(X_REQUEST_ID, value);
23 + response
24 +}
rust/src/server/src/routes.rs
+9 -2
@@ -56,11 +56,18 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
56 .route("/is_sleeping", get(sleep::is_sleeping))
57 }
58
59 - router
59 + let enable_request_id_headers = state.enable_request_id_headers;
60 + let mut router = router
61 .with_state(state.clone())
62 .layer(from_fn_with_state(state, middleware::track_server_load))
63 .layer(from_fn(middleware::track_http_metrics))
63 - .layer(TraceLayer::new_for_http())
64 + .layer(TraceLayer::new_for_http());
65 +
66 + if enable_request_id_headers {
67 + router = router.layer(from_fn(middleware::set_request_id_header));
68 + }
69 +
70 + router
71 }
72
73 #[cfg(test)]
rust/src/server/src/routes/tests.rs
+63
@@ -746,6 +746,20 @@ async fn test_app() -> axum::Router {
746 )))
747 }
748
749 +async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
750 + let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
751 + b"engine-openai-request-id",
752 + default_stream_output_specs(),
753 + Arc::new(FakeChatBackend::new()),
754 + )
755 + .await;
756 + let app = build_router(Arc::new(
757 + AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
758 + .with_request_id_headers(true),
759 + ));
760 + (app, engine_task)
761 +}
762 +
763 async fn test_health_app_with_engine_script<F>(
764 script: F,
765 ) -> (axum::Router, Arc<AppState>, MockEngineTask)
@@ -947,6 +961,18 @@ async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) {
961 (status, body)
962 }
963
964 +async fn health_response(app: &axum::Router, request_id: Option<&str>) -> axum::response::Response {
965 + let mut builder = Request::builder().method("GET").uri("/health");
966 + if let Some(request_id) = request_id {
967 + builder = builder.header("X-Request-Id", request_id);
968 + }
969 +
970 + app.clone()
971 + .call(builder.body(Body::empty()).expect("build request"))
972 + .await
973 + .expect("call app")
974 +}
975 +
976 fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option<f64> {
977 rendered.lines().find_map(|line| {
978 let rest = line.strip_prefix(metric)?;
@@ -994,6 +1020,43 @@ async fn list_models_returns_configured_model() {
1020 assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
1021 }
1022
1023 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1024 +#[serial]
1025 +async fn request_id_header_is_absent_by_default() {
1026 + let app = test_app().await;
1027 + let response = health_response(&app, None).await;
1028 +
1029 + assert_eq!(response.status(), StatusCode::OK);
1030 + assert!(!response.headers().contains_key("x-request-id"));
1031 +}
1032 +
1033 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1034 +#[serial]
1035 +async fn request_id_header_generates_uuid_hex_when_enabled() {
1036 + let (app, _engine_task) = test_app_with_request_id_headers().await;
1037 + let response = health_response(&app, None).await;
1038 +
1039 + assert_eq!(response.status(), StatusCode::OK);
1040 + let request_id = response
1041 + .headers()
1042 + .get("x-request-id")
1043 + .expect("x-request-id header")
1044 + .to_str()
1045 + .expect("header is ascii");
1046 + assert_eq!(request_id.len(), 32);
1047 + assert!(request_id.chars().all(|ch| ch.is_ascii_hexdigit()));
1048 +}
1049 +
1050 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1051 +#[serial]
1052 +async fn request_id_header_echoes_incoming_header_when_enabled() {
1053 + let (app, _engine_task) = test_app_with_request_id_headers().await;
1054 + let response = health_response(&app, Some("req-123")).await;
1055 +
1056 + assert_eq!(response.status(), StatusCode::OK);
1057 + assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
1058 +}
1059 +
1060 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1061 #[serial]
1062 async fn version_returns_engine_vllm_version() {
rust/src/server/src/state.rs
+9
@@ -17,6 +17,8 @@ pub struct AppState {
17 pub chat: ChatLlm,
18 /// Whether to log a summary line for each completed request.
19 pub enable_log_requests: bool,
20 + /// Whether to set X-Request-Id on every HTTP response.
21 + pub enable_request_id_headers: bool,
22 /// Number of in-flight inference requests currently owned by this frontend.
23 server_load: AtomicU64,
24 }
@@ -39,6 +41,7 @@ impl AppState {
41 served_model_names,
42 chat,
43 enable_log_requests: false,
44 + enable_request_id_headers: false,
45 server_load: AtomicU64::new(0),
46 }
47 }
@@ -49,6 +52,12 @@ impl AppState {
52 self
53 }
54
55 + /// Enable X-Request-Id response headers.
56 + pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
57 + self.enable_request_id_headers = enabled;
58 + self
59 + }
60 +
61 /// The primary model name echoed back in API responses (the first served
62 /// name).
63 pub fn primary_model_name(&self) -> &str {