@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 08a8a4af3f

feat(rust): expose profiler control routes in Rust frontend (#46306)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Pranav Thakur committed Jul 2, 2026 at 14:16 UTC 08a8a4af3fb904ab809c6347b785407f180ccf22

11 files changed +381 -2

rust/src/cmd/src/cli.rs

+47 -1

@@ -13,8 +13,8 @@ use std::time::Duration;

14 use clap::{Args, Parser, Subcommand};

15 use educe::Educe;

16 -use serde::Deserialize;

16 use serde::de::DeserializeOwned;

17 +use serde::{Deserialize, Serialize};

18 use serde_json::Value;

19 use serde_with::{DefaultOnNull, OneOrMany, serde_as};

20 use thiserror_ext::AsReport as _;

@@ -290,6 +290,15 @@ pub struct SharedRuntimeArgs {

290 #[serde(default)]

291 pub ssl_ciphers: Option<String>,

292

293 + /// Profiler configuration forwarded by the Python supervisor.

294 + ///

295 + /// When set with a non-null `profiler` type, the Rust frontend registers

296 + /// the `/start_profile` and `/stop_profile` routes and forwards calls to

297 + /// the engine via the `"profile"` utility RPC.

298 + #[arg(long, value_parser = parse_json::<ProfilerConfig>, value_name = "JSON")]

299 + #[serde(default)]

300 + pub profiler_config: Option<ProfilerConfig>,

301 +

302 /// Unsupported Python vLLM frontend arguments recognized but not yet

303 /// implemented in Rust.

304 #[educe(Debug(ignore))]

@@ -317,6 +326,20 @@ impl SharedRuntimeArgs {

326 .map_or(DEFAULT_KEEP_ALIVE_TIMEOUT, Duration::from_secs)

327 }

328

329 + /// Return the configured profiler mode, when profiling is enabled.

330 + pub fn profiler(&self) -> Option<String> {

331 + self.profiler_config.as_ref().and_then(|c| c.profiler.clone())

332 + }

333 +

334 + /// Return the profiler config JSON for managed Python engine forwarding.

335 + pub fn profiler_config_json(&self) -> Option<String> {

336 + self.profiler_config

337 + .as_ref()

338 + .map(serde_json::to_string)

339 + .transpose()

340 + .expect("profiler config serialization should not fail")

341 + }

342 +

343 /// Apply fallback logic for API key configuration from env variables.

344 fn apply_env_api_key_fallback(&mut self) {

345 if self.api_key.is_empty()

@@ -345,6 +368,7 @@ impl SharedRuntimeArgs {

368 let api_server_options = self.api_server_options();

369 let cors = self.cors_config();

370 let tls = self.tls_config();

371 + let profiler = self.profiler();

372

373 Config {

374 transport_mode: TransportMode::Bootstrapped {

@@ -377,6 +401,7 @@ impl SharedRuntimeArgs {

401 grpc_port: self.grpc_port,

402 shutdown_timeout,

403 keep_alive_timeout,

404 + profiler,

405 }

406 }

407

@@ -397,6 +422,7 @@ impl SharedRuntimeArgs {

422 let api_server_options = self.api_server_options();

423 let cors = self.cors_config();

424 let tls = self.tls_config();

425 + let profiler = self.profiler();

426

427 Config {

428 transport_mode: TransportMode::HandshakeOwner {

@@ -427,6 +453,7 @@ impl SharedRuntimeArgs {

453 grpc_port: self.grpc_port,

454 shutdown_timeout,

455 keep_alive_timeout,

456 + profiler,

457 }

458 }

459

@@ -477,6 +504,23 @@ fn default_py_bootstrap_parser_selection() -> ParserSelection {

504 ParserSelection::None

505 }

506

507 +/// Minimal profiler configuration parsed from `--profiler-config`.

508 +///

509 +/// Only the `profiler` field is inspected by the Rust frontend to decide

510 +/// whether to register the `/start_profile` and `/stop_profile` routes.

511 +/// All other fields are accepted but ignored — they are consumed by the

512 +/// Python engine layer.

513 +#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize, Serialize)]

514 +pub struct ProfilerConfig {

515 + /// Profiler backend type (e.g. `"torch"`, `"cuda"`). When `null` or

516 + /// absent, profiling is disabled.

517 + #[serde(default)]

518 + pub profiler: Option<String>,

519 + /// Additional Python profiler config fields consumed by the engine layer.

520 + #[serde(flatten)]

521 + pub extra: serde_json::Map<String, Value>,

522 +}

523 +

524 fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {

525 serde_json::from_str(value).map_err(|e| format!("invalid JSON object: {}", e.as_report()))

526 }

@@ -603,11 +647,13 @@ impl ServeArgs {

647 pub fn to_managed_engine_config(&self, handshake_port: u16) -> ManagedEngineConfig {

648 let reasoning_parser =

649 effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model);

650 + let profiler_config = self.runtime.profiler_config_json();

651

652 self.managed_engine.clone().into_config(

653 self.runtime.model.clone(),

654 self.runtime.max_model_len,

655 self.runtime.max_logprobs,

656 + profiler_config,

657 reasoning_parser.as_deref(),

658 self.runtime.language_model_only,

659 self.runtime.disable_log_stats,

rust/src/cmd/src/cli/tests.rs

+112

@@ -71,6 +71,7 @@ fn serve_args_forward_python_flags_with_separator() {

71 ssl_ca_certs: None,

72 ssl_cert_reqs: 0,

73 ssl_ciphers: None,

74 + profiler_config: None,

75 },

76 managed_engine: ManagedEngineArgs {

77 python: "../vllm/.venv/bin/python",

@@ -180,6 +181,40 @@ fn serve_args_forward_disable_log_stats_to_managed_engine() {

181 .assert_debug_eq(&config.python_args);

182 }

183

184 +#[test]

185 +fn serve_args_forward_profiler_config_to_managed_engine() {

186 + let cli = Cli::try_parse_from([

187 + "vllm-rs",

188 + "serve",

189 + "Qwen/Qwen3-0.6B",

190 + "--profiler-config",

191 + r#"{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}"#,

192 + ])

193 + .unwrap();

194 +

195 + let Command::Serve(args) = cli.command else {

196 + panic!("expected serve args");

197 + };

198 + assert_eq!(args.runtime.profiler().as_deref(), Some("torch"));

199 +

200 + let config = args.to_managed_engine_config(5555);

201 + let profiler_flag_index = config

202 + .python_args

203 + .iter()

204 + .position(|arg| arg == "--profiler-config")

205 + .expect("profiler config flag");

206 + let profiler_config: serde_json::Value =

207 + serde_json::from_str(&config.python_args[profiler_flag_index + 1])

208 + .expect("profiler config json");

209 + assert_eq!(

210 + profiler_config,

211 + serde_json::json!({

212 + "profiler": "torch",

213 + "torch_profiler_dir": "/tmp/profile",

214 + })

215 + );

216 +}

217 +

218 #[test]

219 fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {

220 let cli = Cli::try_parse_from([

@@ -732,6 +767,7 @@ fn frontend_args_accept_json() {

767 ssl_ca_certs: None,

768 ssl_cert_reqs: 0,

769 ssl_ciphers: None,

770 + profiler_config: None,

771 },

772 },

773 ),

@@ -1253,6 +1289,7 @@ fn serve_args_accept_handshake_aliases() {

1289 ssl_ca_certs: None,

1290 ssl_cert_reqs: 0,

1291 ssl_ciphers: None,

1292 + profiler_config: None,

1293 },

1294 managed_engine: ManagedEngineArgs {

1295 python: "python3",

@@ -1392,6 +1429,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {

1429 grpc_port: None,

1430 shutdown_timeout: 0ns,

1431 keep_alive_timeout: 5s,

1432 + profiler: None,

1433 }

1434 "#]]

1435 .assert_debug_eq(&Config {

@@ -1475,6 +1513,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {

1513 grpc_port: None,

1514 shutdown_timeout: 0ns,

1515 keep_alive_timeout: 5s,

1516 + profiler: None,

1517 }

1518 "#]]

1519 .assert_debug_eq(&config);

@@ -1576,6 +1615,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present

1615 grpc_port: None,

1616 shutdown_timeout: 0ns,

1617 keep_alive_timeout: 5s,

1618 + profiler: None,

1619 }

1620 "#]]

1621 .assert_debug_eq(&config);

@@ -1604,3 +1644,75 @@ fn serve_frontend_config_uses_unix_listener_when_uds_is_present() {

1644 }

1645 );

1646 }

1647 +

1648 +#[test]

1649 +fn frontend_args_json_enables_profiling_when_profiler_config_set() {

1650 + let cli = Cli::try_parse_from([

1651 + "vllm-rs",

1652 + "frontend",

1653 + "--listen-fd",

1654 + "3",

1655 + "--input-address",

1656 + "ipc:///tmp/input.sock",

1657 + "--output-address",

1658 + "ipc:///tmp/output.sock",

1659 + "--args-json",

1660 + r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}}"#,

1661 + ])

1662 + .unwrap();

1663 +

1664 + let Command::Frontend(args) = cli.command else {

1665 + panic!("expected frontend args");

1666 + };

1667 + assert_eq!(args.runtime.profiler().as_deref(), Some("torch"));

1668 + let config = args.into_config();

1669 + assert_eq!(config.profiler.as_deref(), Some("torch"));

1670 +}

1671 +

1672 +#[test]

1673 +fn frontend_args_json_disables_profiling_when_profiler_config_absent() {

1674 + let cli = Cli::try_parse_from([

1675 + "vllm-rs",

1676 + "frontend",

1677 + "--listen-fd",

1678 + "3",

1679 + "--input-address",

1680 + "ipc:///tmp/input.sock",

1681 + "--output-address",

1682 + "ipc:///tmp/output.sock",

1683 + "--args-json",

1684 + r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,

1685 + ])

1686 + .unwrap();

1687 +

1688 + let Command::Frontend(args) = cli.command else {

1689 + panic!("expected frontend args");

1690 + };

1691 + assert_eq!(args.runtime.profiler(), None);

1692 + let config = args.into_config();

1693 + assert_eq!(config.profiler, None);

1694 +}

1695 +

1696 +#[test]

1697 +fn frontend_args_json_disables_profiling_when_profiler_type_is_null() {

1698 + let cli = Cli::try_parse_from([

1699 + "vllm-rs",

1700 + "frontend",

1701 + "--listen-fd",

1702 + "3",

1703 + "--input-address",

1704 + "ipc:///tmp/input.sock",

1705 + "--output-address",

1706 + "ipc:///tmp/output.sock",

1707 + "--args-json",

1708 + r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":null}}"#,

1709 + ])

1710 + .unwrap();

1711 +

1712 + let Command::Frontend(args) = cli.command else {

1713 + panic!("expected frontend args");

1714 + };

1715 + assert_eq!(args.runtime.profiler(), None);

1716 + let config = args.into_config();

1717 + assert_eq!(config.profiler, None);

1718 +}

rust/src/engine-core-client/src/client.rs

+12

@@ -738,6 +738,18 @@ impl EngineCoreClient {

738 self.call_utility_consensus("is_scheduler_paused", ()).await

739 }

740

741 + /// Start profiling the engine.

742 + pub async fn start_profile(&self, profile_prefix: Option<&str>) -> Result<()> {

743 + self.call_utility::<(), _>("profile", (true, profile_prefix)).await?;

744 + Ok(())

745 + }

746 +

747 + /// Stop profiling the engine.

748 + pub async fn stop_profile(&self, profile_prefix: Option<&str>) -> Result<()> {

749 + self.call_utility::<(), _>("profile", (false, profile_prefix)).await?;

750 + Ok(())

751 + }

752 +

753 /// Shut down local client tasks and close transport state.

754 pub async fn shutdown(self) -> Result<()> {

755 let Self {

rust/src/managed-engine/src/cli.rs

@@ -77,6 +77,7 @@ impl ManagedEngineArgs {

77 model: String,

78 max_model_len: Option<u32>,

79 max_logprobs: Option<i32>,

80 + profiler_config: Option<String>,

81 reasoning_parser: Option<&str>,

82 language_model_only: bool,

83 disable_log_stats: bool,

@@ -93,6 +94,10 @@ impl ManagedEngineArgs {

94 python_args.push("--max-logprobs".to_string());

95 python_args.push(max_logprobs.to_string());

96 }

97 + if let Some(profiler_config) = profiler_config {

98 + python_args.push("--profiler-config".to_string());

99 + python_args.push(profiler_config);

100 + }

101 if let Some(reasoning_parser) = reasoning_parser {

102 python_args.push("--reasoning-parser".to_string());

103 python_args.push(reasoning_parser.to_string());

rust/src/server/examples/external_engine_openai_qwen.rs

@@ -77,6 +77,7 @@ async fn main() -> Result<()> {

77 grpc_port: None,

78 shutdown_timeout: Duration::ZERO,

79 keep_alive_timeout: Duration::from_secs(5),

80 + profiler: None,

81 };

83 let bind_address = format!("127.0.0.1:{port}");

rust/src/server/src/config.rs

@@ -208,6 +208,9 @@ pub struct Config {

208 /// Maximum idle time on a keep-alive HTTP connection before the server

209 /// closes it (`VLLM_HTTP_TIMEOUT_KEEP_ALIVE`, default 5s).

210 pub keep_alive_timeout: Duration,

211 + /// Profiler mode that registers `/start_profile` and `/stop_profile`

212 + /// routes when present.

213 + pub profiler: Option<String>,

214 }

215

216 impl Config {

rust/src/server/src/lib.rs

+2 -1

@@ -134,7 +134,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {

134 .with_api_server_options(config.api_server_options)

135 .with_server_info(ServerInfoSnapshot::from_config(config))

136 .with_api_keys(config.api_keys.clone())

137 - .with_cors(config.cors.clone()),

137 + .with_cors(config.cors.clone())

138 + .with_profiler(config.profiler.clone()),

139 ))

140 }

141

rust/src/server/src/routes.rs

+13

@@ -8,6 +8,7 @@ mod lora;

8 mod metrics;

9 pub(crate) mod openai;

10 mod pause;

11 +mod profile;

12 mod server_info;

13 mod sleep;

14 mod tokenize;

@@ -21,6 +22,7 @@ use axum::extract::DefaultBodyLimit;

22 use axum::middleware::{from_fn, from_fn_with_state};

23 use axum::routing::{get, post};

24 use tower_http::trace::TraceLayer;

25 +use tracing::warn;

27 use crate::middleware;

28 use crate::state::AppState;

@@ -107,6 +109,17 @@ fn build_router_with_options(

109 .route("/get_world_size", get(world_size::get_world_size))

110 }

111

112 + if let Some(profiler) = &state.profiler {

113 + warn!(

114 + mode = profiler,

115 + "profiler is enabled in the API server; \

116 + this should only be used for local development",

117 + );

118 + router = router

119 + .route("/start_profile", post(profile::start_profile))

120 + .route("/stop_profile", post(profile::stop_profile));

121 + }

122 +

123 let enable_request_id_headers = state.api_server_options.enable_request_id_headers;

124 let enable_api_key_auth = state.has_api_keys();

125 let mut router = router

rust/src/server/src/routes/profile.rs new

+33

@@ -0,0 +1,33 @@

1 +use std::sync::Arc;

2 +

3 +use axum::extract::State;

4 +use axum::http::StatusCode;

5 +use tracing::info;

6 +

7 +use crate::error::ApiError;

8 +use crate::state::AppState;

9 +use crate::utils::utility_call_error;

10 +

11 +/// Start profiling the engine.

12 +pub async fn start_profile(State(state): State<Arc<AppState>>) -> Result<StatusCode, ApiError> {

13 + info!("starting profiler");

14 + state

15 + .engine_core_client()

16 + .start_profile(None)

17 + .await

18 + .map_err(|error| utility_call_error("start_profile", error))?;

19 + info!("profiler started");

20 + Ok(StatusCode::OK)

21 +}

22 +

23 +/// Stop profiling the engine.

24 +pub async fn stop_profile(State(state): State<Arc<AppState>>) -> Result<StatusCode, ApiError> {

25 + info!("stopping profiler");

26 + state

27 + .engine_core_client()

28 + .stop_profile(None)

29 + .await

30 + .map_err(|error| utility_call_error("stop_profile", error))?;

31 + info!("profiler stopped");

32 + Ok(StatusCode::OK)

33 +}

rust/src/server/src/routes/tests.rs

+143

@@ -6183,3 +6183,146 @@ async fn world_size_excludes_data_parallelism_when_include_dp_false() {

6183 let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");

6184 assert_eq!(json, json!({"world_size": 2}));

6185 }

6186 +

6187 +// ========================= Profiler route tests =========================

6188 +

6189 +async fn test_profiling_app_with_engine_script<F>(script: F) -> (axum::Router, MockEngineTask)

6190 +where

6191 + F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static,

6192 +{

6193 + let ipc = IpcNamespace::new().expect("create ipc namespace");

6194 + let handshake_address = ipc.handshake_endpoint();

6195 + let engine_id = b"engine-openai-profiler".to_vec();

6196 +

6197 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(

6198 + handshake_address.clone(),

6199 + engine_id.clone(),

6200 + move |dealer, push| script(dealer, push),

6201 + ));

6202 +

6203 + let client = EngineCoreClient::connect(

6204 + EngineCoreClientConfig::new_single(handshake_address)

6205 + .with_model_name("test-model")

6206 + .with_local_input_output_addresses(

6207 + Some(ipc.input_endpoint()),

6208 + Some(ipc.output_endpoint()),

6209 + ),

6210 + )

6211 + .await

6212 + .expect("connect client");

6213 +

6214 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));

6215 + (

6216 + build_router(Arc::new(

6217 + AppState::new(vec!["test-model".to_string()], chat)

6218 + .with_profiler(Some("torch".to_string())),

6219 + )),

6220 + engine_task,

6221 + )

6222 +}

6223 +

6224 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

6225 +#[serial]

6226 +async fn start_profile_route_sends_expected_utility_call() {

6227 + let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| {

6228 + boxed_test_future(async move {

6229 + let utility = recv_engine_message(dealer).await;

6230 + assert_eq!(utility[0].as_ref(), &[0x03]);

6231 +

6232 + let payload = decode_value(&utility[1]).expect("decode utility payload");

6233 + let array = payload.as_array().expect("utility payload array");

6234 + let call_id = array[1].as_u64().expect("call id");

6235 +

6236 + assert_eq!(array[2], Value::from("profile"));

6237 + assert_eq!(array[3], Value::Array(vec![Value::from(true), Value::Nil]));

6238 +

6239 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;

6240 + })

6241 + })

6242 + .await;

6243 +

6244 + let response = app

6245 + .clone()

6246 + .call(

6247 + Request::builder()

6248 + .method("POST")

6249 + .uri("/start_profile")

6250 + .body(Body::empty())

6251 + .expect("build request"),

6252 + )

6253 + .await

6254 + .expect("call app");

6255 +

6256 + let status = response.status();

6257 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

6258 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));

6259 + assert!(body.is_empty());

6260 + engine_task.await.expect("mock engine task");

6261 +}

6262 +

6263 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

6264 +#[serial]

6265 +async fn stop_profile_route_sends_expected_utility_call() {

6266 + let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| {

6267 + boxed_test_future(async move {

6268 + let utility = recv_engine_message(dealer).await;

6269 + assert_eq!(utility[0].as_ref(), &[0x03]);

6270 +

6271 + let payload = decode_value(&utility[1]).expect("decode utility payload");

6272 + let array = payload.as_array().expect("utility payload array");

6273 + let call_id = array[1].as_u64().expect("call id");

6274 +

6275 + assert_eq!(array[2], Value::from("profile"));

6276 + assert_eq!(array[3], Value::Array(vec![Value::from(false), Value::Nil]));

6277 +

6278 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;

6279 + })

6280 + })

6281 + .await;

6282 +

6283 + let response = app

6284 + .clone()

6285 + .call(

6286 + Request::builder()

6287 + .method("POST")

6288 + .uri("/stop_profile")

6289 + .body(Body::empty())

6290 + .expect("build request"),

6291 + )

6292 + .await

6293 + .expect("call app");

6294 +

6295 + let status = response.status();

6296 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

6297 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));

6298 + assert!(body.is_empty());

6299 + engine_task.await.expect("mock engine task");

6300 +}

6301 +

6302 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

6303 +#[serial]

6304 +async fn profile_routes_are_hidden_when_profiling_is_disabled() {

6305 + let (chat, engine_task) = test_chat_with_engine_handle().await;

6306 + let app = build_router(Arc::new(AppState::new(

6307 + vec!["test-model".to_string()],

6308 + chat,

6309 + )));

6310 +

6311 + for (method, uri) in [("POST", "/start_profile"), ("POST", "/stop_profile")] {

6312 + let response = app

6313 + .clone()

6314 + .call(

6315 + Request::builder()

6316 + .method(method)

6317 + .uri(uri)

6318 + .body(Body::empty())

6319 + .expect("build request"),

6320 + )

6321 + .await

6322 + .expect("call app");

6323 +

6324 + assert_eq!(response.status(), StatusCode::NOT_FOUND, "{method} {uri}");

6325 + }

6326 +

6327 + engine_task.abort_and_join().await;

6328 +}

rust/src/server/src/state.rs

+10

@@ -47,6 +47,9 @@ pub struct AppState {

47 model_path: Option<String>,

48 /// Lazily initialized runtime for heavyweight request paths.

49 request_runtime: OnceLock<BackgroundShutdownRuntime>,

50 + /// Profiler mode that registers `/start_profile` and `/stop_profile`

51 + /// routes when present.

52 + pub profiler: Option<String>,

53 }

55 impl AppState {

@@ -74,6 +77,7 @@ impl AppState {

77 lora_manager: LoraManager::new(),

78 model_path: None,

79 request_runtime: OnceLock::new(),

80 + profiler: None,

81 }

82 }

@@ -95,6 +99,12 @@ impl AppState {

99 self

100 }

101

102 + /// Set the profiler mode that enables `/start_profile` and `/stop_profile`.

103 + pub fn with_profiler(mut self, profiler: Option<String>) -> Self {

104 + self.profiler = profiler;

105 + self

106 + }

107 +

108 /// Attach the runtime server information snapshot used by `/server_info`.

109 pub(crate) fn with_server_info(mut self, server_info: ServerInfoSnapshot) -> Self {

110 self.server_info = Some(server_info);