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>

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<String>,
 
+    /// 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::<ProfilerConfig>, value_name = "JSON")]
+    #[serde(default)]
+    pub profiler_config: Option<ProfilerConfig>,
+
     /// 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<String> {
+        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<String> {
+        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<String>,
+    /// Additional Python profiler config fields consumed by the engine layer.
+    #[serde(flatten)]
+    pub extra: serde_json::Map<String, Value>,
+}
+
 fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
     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,

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);
+}

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 {

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

@@ -77,6 +77,7 @@ impl ManagedEngineArgs {
         model: String,
         max_model_len: Option<u32>,
         max_logprobs: Option<i32>,
+        profiler_config: Option<String>,
         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());

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}");

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<String>,
 }
 
 impl Config {

rust/src/server/src/lib.rs

@@ -134,7 +134,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
             .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()),
     ))
 }
 

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

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<Arc<AppState>>) -> Result<StatusCode, ApiError> {
+    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<Arc<AppState>>) -> Result<StatusCode, ApiError> {
+    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)
+}

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<F>(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;
+}

rust/src/server/src/state.rs

@@ -47,6 +47,9 @@ pub struct AppState {
     model_path: Option<String>,
     /// Lazily initialized runtime for heavyweight request paths.
     request_runtime: OnceLock<BackgroundShutdownRuntime>,
+    /// Profiler mode that registers `/start_profile` and `/stop_profile`
+    /// routes when present.
+    pub profiler: Option<String>,
 }
 
 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<String>) -> 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);