diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 2dfb9fa1c2..1b4ccc7581 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/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::>(); + // 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}")) }) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index f7df2fd7bb..b1b9a1803a 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/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, abort_tx: mpsc::UnboundedSender, + /// 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(()) diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 1107e415d3..8f509be8e5 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/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, utility_reg: Mutex, @@ -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"), diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index 0493732b03..2172c98374 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/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:?}")] diff --git a/rust/src/engine-core-client/src/lib.rs b/rust/src/engine-core-client/src/lib.rs index e39e29c4e5..f4ae0e19ee 100644 --- a/rust/src/engine-core-client/src/lib.rs +++ b/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; diff --git a/rust/src/engine-core-client/src/runtime.rs b/rust/src/engine-core-client/src/runtime.rs new file mode 100644 index 0000000000..2015bff752 --- /dev/null +++ b/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); + +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 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 = 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::().ok()) + .filter(|value| *value > 0) + .unwrap_or(DEFAULT_ZMQ_WORKER_THREADS) +} diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 40f59675a6..c73da7a0a9 100644 --- a/rust/src/server/Cargo.toml +++ b/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 diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 5f135e0ed5..c22a06ddc3 100644 --- a/rust/src/server/src/lib.rs +++ b/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(), ); diff --git a/rust/src/server/src/middleware/mod.rs b/rust/src/server/src/middleware/mod.rs index 65d7b25b02..d61f8bc3cc 100644 --- a/rust/src/server/src/middleware/mod.rs +++ b/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; diff --git a/rust/src/server/src/middleware/offload.rs b/rust/src/server/src/middleware/offload.rs new file mode 100644 index 0000000000..cad560eb7a --- /dev/null +++ b/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( + state: Arc, +) -> impl tower::Layer> + Clone { + layer_fn(move |inner| RequestRuntimeService { + inner, + state: state.clone(), + }) +} + +/// Service produced by [`request_runtime_layer`]. +#[derive(Clone)] +pub(crate) struct RequestRuntimeService { + inner: S, + state: Arc, +} + +impl Service> for RequestRuntimeService +where + S: Service> + 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>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> 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 { + 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")); + } +} diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 1e83c42781..ce94e2ecab 100644 --- a/rust/src/server/src/routes.rs +++ b/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, diff --git a/rust/src/server/src/runtime.rs b/rust/src/server/src/runtime.rs new file mode 100644 index 0000000000..4a8711a6af --- /dev/null +++ b/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::() { + 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) +} diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 01b5b78962..eb37df40ea 100644 --- a/rust/src/server/src/state.rs +++ b/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, + /// Lazily initialized runtime for heavyweight request paths. + request_runtime: OnceLock, } 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,