[Rust Frontend][Perf] Use dedicated runtime for HTTP/request-processing/ZMQ (#46051)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>

rust/src/chat/src/multimodal.rs

@@ -365,6 +365,7 @@ impl MultimodalModelInfo {
         let processor = self.image_processor.raw;
         let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();
 
+        // TODO: is it still necessary given that we've already in a dedicated runtime?
         tokio::task::spawn_blocking(move || {
             processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}"))
         })

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

@@ -15,6 +15,7 @@ use crate::protocol::handshake::EngineCoreReadyResponse;
 use crate::protocol::lora::LoraRequest;
 use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
 use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
+use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
 use crate::transport::{self, ConnectedEngine};
 
 pub(crate) mod imp;
@@ -201,6 +202,8 @@ pub struct EngineCoreClient {
     coordinator: Option<CoordinatorHandle>,
     abort_tx: mpsc::UnboundedSender<AbortRequest>,
 
+    /// Runtime used to send messages to the engine and drive all background tasks.
+    runtime: BackgroundShutdownRuntime,
     // Background tasks
     output_task: AbortOnDropHandle<()>,
     dispatcher_task: AbortOnDropHandle<()>,
@@ -280,21 +283,22 @@ impl EngineCoreClient {
         let (output_tx, output_rx) = mpsc::channel(64);
         let (abort_tx, abort_rx) = mpsc::unbounded_channel();
         let engines = connected.engines;
+        let runtime = build_zmq_runtime();
         let inner = Arc::new(ClientInner::new(
             connected.input_send,
+            runtime.handle().clone(),
             config.model_name.clone(),
             &engines,
         ));
-        let output_task = AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(
+        let output_task = AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(
             connected.output_socket,
             output_tx,
         )));
-        let dispatcher_task = AbortOnDropHandle::new(tokio::spawn(run_output_dispatcher_loop(
-            inner.clone(),
-            output_rx,
-        )));
+        let dispatcher_task = AbortOnDropHandle::new(
+            runtime.spawn(run_output_dispatcher_loop(inner.clone(), output_rx)),
+        );
         let abort_task =
-            AbortOnDropHandle::new(tokio::spawn(run_abort_loop(inner.clone(), abort_rx)));
+            AbortOnDropHandle::new(runtime.spawn(run_abort_loop(inner.clone(), abort_rx)));
 
         // If any engine reported a dp_stats_address in its ready response, use it
         // as the external coordinator address.
@@ -307,13 +311,13 @@ impl EngineCoreClient {
                     CoordinatorHandle::new_inproc(coordinator_transport.input_socket);
                 let (coordinator_output_tx, coordinator_output_rx) = mpsc::channel(64);
                 let coordinator_output_task =
-                    AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(
+                    AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(
                         coordinator_transport.output_socket,
                         coordinator_output_tx,
                     )));
-                let coordinator_task = AbortOnDropHandle::new(tokio::spawn(
-                    runner.run(coordinator_output_rx, inner.clone()),
-                ));
+                let coordinator_task = AbortOnDropHandle::new(
+                    runtime.spawn(runner.run(coordinator_output_rx, inner.clone())),
+                );
                 (
                     Some(handle),
                     Some(coordinator_output_task),
@@ -327,7 +331,7 @@ impl EngineCoreClient {
             {
                 let (handle, service) = CoordinatorHandle::connect_external(address).await?;
                 let coordinator_task =
-                    AbortOnDropHandle::new(tokio::spawn(service.run(inner.clone())));
+                    AbortOnDropHandle::new(runtime.spawn(service.run(inner.clone())));
                 (Some(handle), None, Some(coordinator_task))
             } else {
                 (None, None, None)
@@ -341,6 +345,7 @@ impl EngineCoreClient {
             inner,
             coordinator,
             abort_tx,
+            runtime,
             output_task,
             dispatcher_task,
             abort_task,
@@ -737,6 +742,7 @@ impl EngineCoreClient {
         let Self {
             inner,
             abort_tx,
+            runtime,
             output_task,
             dispatcher_task,
             abort_task,
@@ -757,6 +763,8 @@ impl EngineCoreClient {
 
         tasks.iter().for_each(|t| t.abort());
         join_all(tasks).await;
+        drop(inner);
+        drop(runtime);
 
         info!("engine-core client shut down");
         Ok(())

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

@@ -5,6 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use arc_swap::ArcSwapOption;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
+use tokio::runtime::Handle;
 use tokio::sync::mpsc;
 use tracing::{debug, info, trace, warn};
 use vllm_metrics::METRICS;
@@ -26,6 +27,8 @@ use crate::{Error, Result, transport};
 
 pub(crate) struct ClientInner {
     input_send: RouterSendHalf,
+    /// The runtime handle used for sending messages to the engine.
+    handle: Handle,
     model_name: String,
     request_reg: Mutex<RequestRegistry>,
     utility_reg: Mutex<UtilityRegistry>,
@@ -37,11 +40,13 @@ impl ClientInner {
     /// handshake completes.
     pub fn new(
         input_send: RouterSendHalf,
+        handle: Handle,
         model_name: String,
         engines: &[ConnectedEngine],
     ) -> Self {
         Self {
             input_send,
+            handle,
             model_name,
             request_reg: Mutex::new(RequestRegistry::new(engines)),
             utility_reg: Mutex::new(UtilityRegistry::default()),
@@ -213,9 +218,19 @@ impl ClientInner {
         // frames instead of always producing a single msgpack frame.
         let payload = encode_msgpack(payload)?;
         let mut input_send = self.input_send.clone();
-        transport::send_message(&mut input_send, engine_id, request_type.to_frame(), payload)
-            .await?;
-        Ok(())
+        let engine_id = engine_id.clone();
+
+        self.handle
+            .spawn(async move {
+                transport::send_message(
+                    &mut input_send,
+                    &engine_id,
+                    request_type.to_frame(),
+                    payload,
+                )
+                .await
+            })
+            .await?
     }
 
     /// Handle an abort request by sending the abort message to the engine.
@@ -434,6 +449,7 @@ mod tests {
         let (send, _) = socket.split();
         ClientInner::new(
             send,
+            Handle::current(),
             "test-model".to_string(),
             &[ConnectedEngine {
                 engine_id: EngineId::from(b"engine-0"),

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

@@ -29,6 +29,8 @@ pub enum Error {
     Io(#[from] std::io::Error),
     #[error("transport error")]
     Transport(#[from] zeromq::ZmqError),
+    #[error("ZMQ runtime task failed")]
+    ZmqRuntimeTask(#[from] tokio::task::JoinError),
     #[error("engine core reported fatal failure")]
     EngineCoreDead,
     #[error("startup handshake timed out while waiting for {stage} after {timeout:?}")]

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

@@ -4,6 +4,7 @@ mod error;
 mod metrics;
 pub mod mock_engine;
 pub mod protocol;
+pub mod runtime;
 #[cfg(any(test, feature = "test-util"))]
 pub mod test_utils;
 mod transport;

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

@@ -0,0 +1,75 @@
+use std::mem::ManuallyDrop;
+use std::ops::{Deref, DerefMut};
+use std::sync::OnceLock;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use tokio::runtime::Runtime;
+
+/// A wrapper around [`Runtime`] that shuts down the runtime in the background when dropped.
+///
+/// This can be useful in some cases, because sometimes we want to drop the runtime without
+/// blocking the current thread, for example, when it's nested inside another runtime.
+pub struct BackgroundShutdownRuntime(ManuallyDrop<Runtime>);
+
+impl Drop for BackgroundShutdownRuntime {
+    fn drop(&mut self) {
+        // Safety: The runtime is only dropped once here.
+        let runtime = unsafe { ManuallyDrop::take(&mut self.0) };
+        runtime.shutdown_background();
+    }
+}
+
+impl Deref for BackgroundShutdownRuntime {
+    type Target = Runtime;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for BackgroundShutdownRuntime {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl From<Runtime> for BackgroundShutdownRuntime {
+    fn from(runtime: Runtime) -> Self {
+        Self(ManuallyDrop::new(runtime))
+    }
+}
+
+const ZMQ_WORKER_THREADS_ENV: &str = "VLLM_RS_ZMQ_WORKER_THREADS";
+/// The number of tasks running on the ZMQ runtime is fixed and expected to remain
+/// small, and multiple engines share the same ZMQ socket. Therefore, based on
+/// benchmarks, a default value of 4 is generally sufficient.
+const DEFAULT_ZMQ_WORKER_THREADS: usize = 4;
+
+static ZMQ_RUNTIME_SEQUENCE: OnceLock<AtomicUsize> = OnceLock::new();
+
+/// Build a Tokio runtime for ZMQ tasks. Multiple calls to this function will
+/// return multiple runtimes with distinct thread name suffixes.
+pub(crate) fn build_zmq_runtime() -> BackgroundShutdownRuntime {
+    let sequence = ZMQ_RUNTIME_SEQUENCE
+        .get_or_init(|| AtomicUsize::new(0))
+        .fetch_add(1, Ordering::Relaxed);
+
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(zmq_worker_threads())
+        .thread_name_fn(move || format!("vllm-zmq-{sequence}"))
+        .enable_all()
+        .build()
+        .expect("failed to build vLLM ZMQ runtime")
+        .into()
+}
+
+/// Get the number of worker threads to use for the ZMQ runtime. If env var
+/// `VLLM_RS_ZMQ_WORKER_THREADS` is set and a valid positive integer, it will be used.
+/// Otherwise, the default value of `DEFAULT_ZMQ_WORKER_THREADS` will be used.
+fn zmq_worker_threads() -> usize {
+    std::env::var(ZMQ_WORKER_THREADS_ENV)
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(DEFAULT_ZMQ_WORKER_THREADS)
+}

rust/src/server/Cargo.toml

@@ -30,6 +30,7 @@ tokio-stream.workspace = true
 tokio-util.workspace = true
 tonic.workspace = true
 tonic-prost.workspace = true
+tower.workspace = true
 tower-http.workspace = true
 tracing.workspace = true
 tracing-futures.workspace = true

rust/src/server/src/lib.rs

@@ -7,6 +7,7 @@ mod listener;
 mod lora;
 mod middleware;
 mod routes;
+mod runtime;
 mod server_info;
 mod state;
 mod utils;
@@ -157,6 +158,9 @@ where
             .with_context(|| format!("failed to bind gRPC listener on {grpc_host}:{grpc_port}"))?;
         let addr = grpc_listener.local_addr()?;
         let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
+        let svc = TonicServer::builder()
+            .layer(middleware::request_runtime_layer(state.clone()))
+            .add_service(svc);
         info!(%addr, "starting gRPC server");
         Some((grpc_listener, svc))
     } else {
@@ -238,7 +242,7 @@ where
                 shutdown.cancelled().await;
                 return Ok(());
             };
-            let server = TonicServer::builder().add_service(svc).serve_with_incoming_shutdown(
+            let server = svc.serve_with_incoming_shutdown(
                 TcpListenerStream::new(grpc_listener),
                 shutdown.cancelled_owned(),
             );

rust/src/server/src/middleware/mod.rs

@@ -2,10 +2,12 @@ mod auth;
 mod cors;
 mod load;
 mod metrics;
+mod offload;
 mod request_id;
 
 pub use auth::authenticate_api_key;
 pub use cors::{cors_layer, strip_cors_on_no_origin};
 pub use load::track_server_load;
 pub use metrics::track_http_metrics;
+pub(crate) use offload::request_runtime_layer;
 pub use request_id::set_request_id_header;

rust/src/server/src/middleware/offload.rs

@@ -0,0 +1,134 @@
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use axum::http::Request;
+use axum::response::{IntoResponse, Response};
+use futures::future::BoxFuture;
+use tokio_util::task::AbortOnDropHandle;
+use tonic::Status;
+use tower::Service;
+use tower::layer::layer_fn;
+use tracing::error;
+
+use crate::error::{ApiError, server_error};
+use crate::state::AppState;
+
+/// Request paths that are run on the request runtime.
+///
+/// These routes can perform CPU-heavy request preparation, including JSON
+/// extraction, validation, chat-template rendering, tokenization, request
+/// lowering, and engine submission. Lightweight operational routes stay on the
+/// HTTP runtime.
+const OFFLOADED_PATHS: &[&str] = &[
+    // HTTP routes:
+    "/v1/chat/completions",
+    "/v1/completions",
+    "/tokenize",
+    "/detokenize",
+    "/inference/v1/generate",
+    // gRPC routes:
+    "/vllm.Generate/Generate",
+    "/vllm.Generate/GenerateStream",
+];
+
+/// Return a Tower layer that runs selected data-plane requests on the request runtime,
+/// so that we can offset heavy request parsing and preprocessing from the HTTP runtime.
+pub(crate) fn request_runtime_layer<S>(
+    state: Arc<AppState>,
+) -> impl tower::Layer<S, Service = RequestRuntimeService<S>> + Clone {
+    layer_fn(move |inner| RequestRuntimeService {
+        inner,
+        state: state.clone(),
+    })
+}
+
+/// Service produced by [`request_runtime_layer`].
+#[derive(Clone)]
+pub(crate) struct RequestRuntimeService<S> {
+    inner: S,
+    state: Arc<AppState>,
+}
+
+impl<S, B> Service<Request<B>> for RequestRuntimeService<S>
+where
+    S: Service<Request<B>> + Clone + Send + 'static,
+    S::Future: Send + 'static,
+    S::Response: RequestRuntimeErrorResponse + Send + 'static,
+    S::Error: Send + 'static,
+    B: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Request<B>) -> Self::Future {
+        if !should_offload(req.uri().path()) {
+            return Box::pin(self.inner.call(req));
+        }
+
+        // Axum extractors and route handlers execute inside the inner service,
+        // so offloading here moves request parsing and preprocessing off the
+        // HTTP runtime without wrapping each handler manually. For streaming
+        // HTTP responses, the response body is still polled on the HTTP runtime.
+        let clone = self.inner.clone();
+        let mut inner = std::mem::replace(&mut self.inner, clone);
+        let task = AbortOnDropHandle::new(self.state.request_runtime().spawn(inner.call(req)));
+
+        Box::pin(async move {
+            match task.await {
+                Ok(result) => result,
+                Err(error) => {
+                    error!(%error, "request runtime task failed");
+                    Ok(S::Response::request_runtime_error_response())
+                }
+            }
+        })
+    }
+}
+
+trait RequestRuntimeErrorResponse {
+    fn request_runtime_error_response() -> Self;
+}
+
+impl RequestRuntimeErrorResponse for Response {
+    fn request_runtime_error_response() -> Self {
+        server_error!("request runtime task failed").into_response()
+    }
+}
+
+impl RequestRuntimeErrorResponse for axum::http::Response<tonic::body::Body> {
+    fn request_runtime_error_response() -> Self {
+        Status::internal("request runtime task failed").into_http()
+    }
+}
+
+fn should_offload(path: &str) -> bool {
+    OFFLOADED_PATHS.contains(&path)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::should_offload;
+
+    #[test]
+    fn offloads_generation_and_tokenization_paths() {
+        assert!(should_offload("/v1/chat/completions"));
+        assert!(should_offload("/v1/completions"));
+        assert!(should_offload("/tokenize"));
+        assert!(should_offload("/detokenize"));
+        assert!(should_offload("/inference/v1/generate"));
+        assert!(should_offload("/vllm.Generate/Generate"));
+        assert!(should_offload("/vllm.Generate/GenerateStream"));
+    }
+
+    #[test]
+    fn passes_through_lightweight_paths() {
+        assert!(!should_offload("/health"));
+        assert!(!should_offload("/metrics"));
+        assert!(!should_offload("/v1/models"));
+    }
+}

rust/src/server/src/routes.rs

@@ -108,6 +108,7 @@ fn build_router_with_options(
     let enable_api_key_auth = state.has_api_keys();
     let mut router = router
         .with_state(state.clone())
+        .layer(middleware::request_runtime_layer(state.clone()))
         .layer(from_fn_with_state(
             state.clone(),
             middleware::track_server_load,

rust/src/server/src/runtime.rs

@@ -0,0 +1,54 @@
+use tokio::runtime::Builder;
+use tracing::{info, warn};
+use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;
+
+const REQUEST_WORKER_THREADS_ENV: &str = "VLLM_RS_REQUEST_WORKER_THREADS";
+const DEFAULT_MAX_REQUEST_WORKER_THREADS: usize = 32;
+
+/// Build a Tokio runtime for heavyweight request paths outside the HTTP runtime.
+///
+/// The server middleware uses this runtime for inference and tokenization
+/// routes so CPU-heavy request preparation does not monopolize the HTTP
+/// runtime's worker queue. Dropping the wrapper shuts the runtime down in the
+/// background.
+pub(crate) fn build_request_runtime() -> BackgroundShutdownRuntime {
+    Builder::new_multi_thread()
+        .enable_all()
+        .thread_name("vllm-request")
+        .worker_threads(request_worker_threads())
+        .build()
+        .expect("failed to build request runtime")
+        .into()
+}
+
+/// Get the number of worker threads to use for the request runtime.
+///
+/// If `VLLM_RS_REQUEST_WORKER_THREADS` is set to a valid positive integer, it is
+/// used directly. Otherwise, the runtime uses available parallelism capped by
+/// `DEFAULT_MAX_REQUEST_WORKER_THREADS`.
+fn request_worker_threads() -> usize {
+    if let Some(value) = std::env::var_os(REQUEST_WORKER_THREADS_ENV) {
+        match value.to_string_lossy().parse::<usize>() {
+            Ok(worker_threads) if worker_threads > 0 => return worker_threads,
+            _ => warn!(
+                value = %value.to_string_lossy(),
+                "ignoring invalid {REQUEST_WORKER_THREADS_ENV}"
+            ),
+        }
+    }
+
+    std::thread::available_parallelism()
+        .map(|parallelism| {
+            let available = parallelism.get();
+            let worker_threads = available.min(DEFAULT_MAX_REQUEST_WORKER_THREADS);
+            if worker_threads < available {
+                info!(
+                    available_parallelism = available,
+                    capped_worker_threads = worker_threads,
+                    "capping request runtime worker threads, set {REQUEST_WORKER_THREADS_ENV} to override"
+                );
+            }
+            worker_threads
+        })
+        .unwrap_or(DEFAULT_MAX_REQUEST_WORKER_THREADS)
+}

rust/src/server/src/state.rs

@@ -1,16 +1,20 @@
 use std::sync::Arc;
+use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};
 
 use serde_json::Value;
 use sha2::{Digest, Sha256};
+use tokio::runtime::Runtime;
 use tokio::time::{Duration, Instant, sleep_until};
 use tracing::warn;
 use vllm_chat::ChatLlm;
 use vllm_engine_core_client::EngineCoreClient;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
+use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;
 
 use crate::config::{ApiServerOptions, CorsConfig};
 use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
+use crate::runtime::build_request_runtime;
 use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
 
 const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
@@ -42,6 +46,8 @@ pub struct AppState {
     lora_manager: LoraManager,
     /// Backend model path reported as `root` for base-model cards.
     model_path: Option<String>,
+    /// Lazily initialized runtime for heavyweight request paths.
+    request_runtime: OnceLock<BackgroundShutdownRuntime>,
 }
 
 impl AppState {
@@ -68,6 +74,7 @@ impl AppState {
             server_load: AtomicU64::new(0),
             lora_manager: LoraManager::new(),
             model_path: None,
+            request_runtime: OnceLock::new(),
         }
     }
 
@@ -185,6 +192,12 @@ impl AppState {
         self.chat.engine_core_client()
     }
 
+    /// Runtime used by middleware to isolate heavyweight request handlers from
+    /// the HTTP reactor.
+    pub(crate) fn request_runtime(&self) -> &Runtime {
+        self.request_runtime.get_or_init(build_request_runtime)
+    }
+
     /// Return the current in-flight inference request count for the `/load`
     /// endpoint.
     pub fn server_load(&self) -> u64 {
@@ -214,6 +227,7 @@ impl AppState {
             match Arc::try_unwrap(self) {
                 Ok(state) => {
                     state.chat.shutdown().await?;
+                    drop(state.request_runtime); // shutdown in background
                     return Ok(());
                 }
                 Err(state) => self = state,