From 08a8a4af3fb904ab809c6347b785407f180ccf22 Mon Sep 17 00:00:00 2001 From: Pranav Thakur <103683515+pranavthakur0-0@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:16:07 +0530 Subject: [PATCH] feat(rust): expose profiler control routes in Rust frontend (#46306) Co-authored-by: Bugen Zhao Signed-off-by: Bugen Zhao --- rust/src/cmd/src/cli.rs | 48 +++++- rust/src/cmd/src/cli/tests.rs | 112 ++++++++++++++ rust/src/engine-core-client/src/client.rs | 12 ++ rust/src/managed-engine/src/cli.rs | 5 + .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 3 + rust/src/server/src/lib.rs | 3 +- rust/src/server/src/routes.rs | 13 ++ rust/src/server/src/routes/profile.rs | 33 ++++ rust/src/server/src/routes/tests.rs | 143 ++++++++++++++++++ rust/src/server/src/state.rs | 10 ++ 11 files changed, 381 insertions(+), 2 deletions(-) create mode 100644 rust/src/server/src/routes/profile.rs diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index d45baadb01..3384fb9c67 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -13,8 +13,8 @@ use std::time::Duration; use clap::{Args, Parser, Subcommand}; use educe::Educe; -use serde::Deserialize; use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; use serde_json::Value; use serde_with::{DefaultOnNull, OneOrMany, serde_as}; use thiserror_ext::AsReport as _; @@ -290,6 +290,15 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub ssl_ciphers: Option, + /// Profiler configuration forwarded by the Python supervisor. + /// + /// When set with a non-null `profiler` type, the Rust frontend registers + /// the `/start_profile` and `/stop_profile` routes and forwards calls to + /// the engine via the `"profile"` utility RPC. + #[arg(long, value_parser = parse_json::, value_name = "JSON")] + #[serde(default)] + pub profiler_config: Option, + /// Unsupported Python vLLM frontend arguments recognized but not yet /// implemented in Rust. #[educe(Debug(ignore))] @@ -317,6 +326,20 @@ impl SharedRuntimeArgs { .map_or(DEFAULT_KEEP_ALIVE_TIMEOUT, Duration::from_secs) } + /// Return the configured profiler mode, when profiling is enabled. + pub fn profiler(&self) -> Option { + self.profiler_config.as_ref().and_then(|c| c.profiler.clone()) + } + + /// Return the profiler config JSON for managed Python engine forwarding. + pub fn profiler_config_json(&self) -> Option { + self.profiler_config + .as_ref() + .map(serde_json::to_string) + .transpose() + .expect("profiler config serialization should not fail") + } + /// Apply fallback logic for API key configuration from env variables. fn apply_env_api_key_fallback(&mut self) { if self.api_key.is_empty() @@ -345,6 +368,7 @@ impl SharedRuntimeArgs { let api_server_options = self.api_server_options(); let cors = self.cors_config(); let tls = self.tls_config(); + let profiler = self.profiler(); Config { transport_mode: TransportMode::Bootstrapped { @@ -377,6 +401,7 @@ impl SharedRuntimeArgs { grpc_port: self.grpc_port, shutdown_timeout, keep_alive_timeout, + profiler, } } @@ -397,6 +422,7 @@ impl SharedRuntimeArgs { let api_server_options = self.api_server_options(); let cors = self.cors_config(); let tls = self.tls_config(); + let profiler = self.profiler(); Config { transport_mode: TransportMode::HandshakeOwner { @@ -427,6 +453,7 @@ impl SharedRuntimeArgs { grpc_port: self.grpc_port, shutdown_timeout, keep_alive_timeout, + profiler, } } @@ -477,6 +504,23 @@ fn default_py_bootstrap_parser_selection() -> ParserSelection { ParserSelection::None } +/// Minimal profiler configuration parsed from `--profiler-config`. +/// +/// Only the `profiler` field is inspected by the Rust frontend to decide +/// whether to register the `/start_profile` and `/stop_profile` routes. +/// All other fields are accepted but ignored — they are consumed by the +/// Python engine layer. +#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize, Serialize)] +pub struct ProfilerConfig { + /// Profiler backend type (e.g. `"torch"`, `"cuda"`). When `null` or + /// absent, profiling is disabled. + #[serde(default)] + pub profiler: Option, + /// Additional Python profiler config fields consumed by the engine layer. + #[serde(flatten)] + pub extra: serde_json::Map, +} + fn parse_json(value: &str) -> Result { serde_json::from_str(value).map_err(|e| format!("invalid JSON object: {}", e.as_report())) } @@ -603,11 +647,13 @@ impl ServeArgs { pub fn to_managed_engine_config(&self, handshake_port: u16) -> ManagedEngineConfig { let reasoning_parser = effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model); + let profiler_config = self.runtime.profiler_config_json(); self.managed_engine.clone().into_config( self.runtime.model.clone(), self.runtime.max_model_len, self.runtime.max_logprobs, + profiler_config, reasoning_parser.as_deref(), self.runtime.language_model_only, self.runtime.disable_log_stats, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 551603fce0..8056ee9392 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -71,6 +71,7 @@ fn serve_args_forward_python_flags_with_separator() { ssl_ca_certs: None, ssl_cert_reqs: 0, ssl_ciphers: None, + profiler_config: None, }, managed_engine: ManagedEngineArgs { python: "../vllm/.venv/bin/python", @@ -180,6 +181,40 @@ fn serve_args_forward_disable_log_stats_to_managed_engine() { .assert_debug_eq(&config.python_args); } +#[test] +fn serve_args_forward_profiler_config_to_managed_engine() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--profiler-config", + r#"{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}"#, + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + assert_eq!(args.runtime.profiler().as_deref(), Some("torch")); + + let config = args.to_managed_engine_config(5555); + let profiler_flag_index = config + .python_args + .iter() + .position(|arg| arg == "--profiler-config") + .expect("profiler config flag"); + let profiler_config: serde_json::Value = + serde_json::from_str(&config.python_args[profiler_flag_index + 1]) + .expect("profiler config json"); + assert_eq!( + profiler_config, + serde_json::json!({ + "profiler": "torch", + "torch_profiler_dir": "/tmp/profile", + }) + ); +} + #[test] fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() { let cli = Cli::try_parse_from([ @@ -732,6 +767,7 @@ fn frontend_args_accept_json() { ssl_ca_certs: None, ssl_cert_reqs: 0, ssl_ciphers: None, + profiler_config: None, }, }, ), @@ -1253,6 +1289,7 @@ fn serve_args_accept_handshake_aliases() { ssl_ca_certs: None, ssl_cert_reqs: 0, ssl_ciphers: None, + profiler_config: None, }, managed_engine: ManagedEngineArgs { python: "python3", @@ -1392,6 +1429,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { grpc_port: None, shutdown_timeout: 0ns, keep_alive_timeout: 5s, + profiler: None, } "#]] .assert_debug_eq(&Config { @@ -1475,6 +1513,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { grpc_port: None, shutdown_timeout: 0ns, keep_alive_timeout: 5s, + profiler: None, } "#]] .assert_debug_eq(&config); @@ -1576,6 +1615,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present grpc_port: None, shutdown_timeout: 0ns, keep_alive_timeout: 5s, + profiler: None, } "#]] .assert_debug_eq(&config); @@ -1604,3 +1644,75 @@ fn serve_frontend_config_uses_unix_listener_when_uds_is_present() { } ); } + +#[test] +fn frontend_args_json_enables_profiling_when_profiler_config_set() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}}"#, + ]) + .unwrap(); + + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args"); + }; + assert_eq!(args.runtime.profiler().as_deref(), Some("torch")); + let config = args.into_config(); + assert_eq!(config.profiler.as_deref(), Some("torch")); +} + +#[test] +fn frontend_args_json_disables_profiling_when_profiler_config_absent() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, + ]) + .unwrap(); + + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args"); + }; + assert_eq!(args.runtime.profiler(), None); + let config = args.into_config(); + assert_eq!(config.profiler, None); +} + +#[test] +fn frontend_args_json_disables_profiling_when_profiler_type_is_null() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":null}}"#, + ]) + .unwrap(); + + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args"); + }; + assert_eq!(args.runtime.profiler(), None); + let config = args.into_config(); + assert_eq!(config.profiler, None); +} diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 5307aff2fa..42680706a3 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -738,6 +738,18 @@ impl EngineCoreClient { self.call_utility_consensus("is_scheduler_paused", ()).await } + /// Start profiling the engine. + pub async fn start_profile(&self, profile_prefix: Option<&str>) -> Result<()> { + self.call_utility::<(), _>("profile", (true, profile_prefix)).await?; + Ok(()) + } + + /// Stop profiling the engine. + pub async fn stop_profile(&self, profile_prefix: Option<&str>) -> Result<()> { + self.call_utility::<(), _>("profile", (false, profile_prefix)).await?; + Ok(()) + } + /// Shut down local client tasks and close transport state. pub async fn shutdown(self) -> Result<()> { let Self { diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index 36789fc7b4..9cf5e311ec 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -77,6 +77,7 @@ impl ManagedEngineArgs { model: String, max_model_len: Option, max_logprobs: Option, + profiler_config: Option, reasoning_parser: Option<&str>, language_model_only: bool, disable_log_stats: bool, @@ -93,6 +94,10 @@ impl ManagedEngineArgs { python_args.push("--max-logprobs".to_string()); python_args.push(max_logprobs.to_string()); } + if let Some(profiler_config) = profiler_config { + python_args.push("--profiler-config".to_string()); + python_args.push(profiler_config); + } if let Some(reasoning_parser) = reasoning_parser { python_args.push("--reasoning-parser".to_string()); python_args.push(reasoning_parser.to_string()); diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 03e474ad4c..d1c6c0fc31 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -77,6 +77,7 @@ async fn main() -> Result<()> { grpc_port: None, shutdown_timeout: Duration::ZERO, keep_alive_timeout: Duration::from_secs(5), + profiler: None, }; let bind_address = format!("127.0.0.1:{port}"); diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index ae8466b321..2ed3da2c1d 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -208,6 +208,9 @@ pub struct Config { /// Maximum idle time on a keep-alive HTTP connection before the server /// closes it (`VLLM_HTTP_TIMEOUT_KEEP_ALIVE`, default 5s). pub keep_alive_timeout: Duration, + /// Profiler mode that registers `/start_profile` and `/stop_profile` + /// routes when present. + pub profiler: Option, } impl Config { diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 78b18f8f47..008a29149e 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -134,7 +134,8 @@ async fn build_state(config: &Config) -> Result> { .with_api_server_options(config.api_server_options) .with_server_info(ServerInfoSnapshot::from_config(config)) .with_api_keys(config.api_keys.clone()) - .with_cors(config.cors.clone()), + .with_cors(config.cors.clone()) + .with_profiler(config.profiler.clone()), )) } diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 1f7dd03809..11918c4fbc 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -8,6 +8,7 @@ mod lora; mod metrics; pub(crate) mod openai; mod pause; +mod profile; mod server_info; mod sleep; mod tokenize; @@ -21,6 +22,7 @@ use axum::extract::DefaultBodyLimit; use axum::middleware::{from_fn, from_fn_with_state}; use axum::routing::{get, post}; use tower_http::trace::TraceLayer; +use tracing::warn; use crate::middleware; use crate::state::AppState; @@ -107,6 +109,17 @@ fn build_router_with_options( .route("/get_world_size", get(world_size::get_world_size)) } + if let Some(profiler) = &state.profiler { + warn!( + mode = profiler, + "profiler is enabled in the API server; \ + this should only be used for local development", + ); + router = router + .route("/start_profile", post(profile::start_profile)) + .route("/stop_profile", post(profile::stop_profile)); + } + let enable_request_id_headers = state.api_server_options.enable_request_id_headers; let enable_api_key_auth = state.has_api_keys(); let mut router = router diff --git a/rust/src/server/src/routes/profile.rs b/rust/src/server/src/routes/profile.rs new file mode 100644 index 0000000000..d70db30d73 --- /dev/null +++ b/rust/src/server/src/routes/profile.rs @@ -0,0 +1,33 @@ +use std::sync::Arc; + +use axum::extract::State; +use axum::http::StatusCode; +use tracing::info; + +use crate::error::ApiError; +use crate::state::AppState; +use crate::utils::utility_call_error; + +/// Start profiling the engine. +pub async fn start_profile(State(state): State>) -> Result { + info!("starting profiler"); + state + .engine_core_client() + .start_profile(None) + .await + .map_err(|error| utility_call_error("start_profile", error))?; + info!("profiler started"); + Ok(StatusCode::OK) +} + +/// Stop profiling the engine. +pub async fn stop_profile(State(state): State>) -> Result { + info!("stopping profiler"); + state + .engine_core_client() + .stop_profile(None) + .await + .map_err(|error| utility_call_error("stop_profile", error))?; + info!("profiler stopped"); + Ok(StatusCode::OK) +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 8e7953a263..6b70f51413 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -6183,3 +6183,146 @@ async fn world_size_excludes_data_parallelism_when_include_dp_false() { let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); assert_eq!(json, json!({"world_size": 2})); } + +// ========================= Profiler route tests ========================= + +async fn test_profiling_app_with_engine_script(script: F) -> (axum::Router, MockEngineTask) +where + F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static, +{ + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-openai-profiler".to_vec(); + + let engine_task = MockEngineTask::new(spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + move |dealer, push| script(dealer, push), + )); + + let client = EngineCoreClient::connect( + EngineCoreClientConfig::new_single(handshake_address) + .with_model_name("test-model") + .with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ), + ) + .await + .expect("connect client"); + + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); + ( + build_router(Arc::new( + AppState::new(vec!["test-model".to_string()], chat) + .with_profiler(Some("torch".to_string())), + )), + engine_task, + ) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn start_profile_route_sends_expected_utility_call() { + let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("profile")); + assert_eq!(array[3], Value::Array(vec![Value::from(true), Value::Nil])); + + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/start_profile") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + let status = response.status(); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body)); + assert!(body.is_empty()); + engine_task.await.expect("mock engine task"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stop_profile_route_sends_expected_utility_call() { + let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("profile")); + assert_eq!(array[3], Value::Array(vec![Value::from(false), Value::Nil])); + + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/stop_profile") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + let status = response.status(); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body)); + assert!(body.is_empty()); + engine_task.await.expect("mock engine task"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn profile_routes_are_hidden_when_profiling_is_disabled() { + let (chat, engine_task) = test_chat_with_engine_handle().await; + let app = build_router(Arc::new(AppState::new( + vec!["test-model".to_string()], + chat, + ))); + + for (method, uri) in [("POST", "/start_profile"), ("POST", "/stop_profile")] { + let response = app + .clone() + .call( + Request::builder() + .method(method) + .uri(uri) + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND, "{method} {uri}"); + } + + engine_task.abort_and_join().await; +} diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index b56ca59487..bdfd83e576 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -47,6 +47,9 @@ pub struct AppState { model_path: Option, /// Lazily initialized runtime for heavyweight request paths. request_runtime: OnceLock, + /// Profiler mode that registers `/start_profile` and `/stop_profile` + /// routes when present. + pub profiler: Option, } impl AppState { @@ -74,6 +77,7 @@ impl AppState { lora_manager: LoraManager::new(), model_path: None, request_runtime: OnceLock::new(), + profiler: None, } } @@ -95,6 +99,12 @@ impl AppState { self } + /// Set the profiler mode that enables `/start_profile` and `/stop_profile`. + pub fn with_profiler(mut self, profiler: Option) -> Self { + self.profiler = profiler; + self + } + /// Attach the runtime server information snapshot used by `/server_info`. pub(crate) fn with_server_info(mut self, server_info: ServerInfoSnapshot) -> Self { self.server_info = Some(server_info);