A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 02:15 UTC
[Rust Frontend] Add server router extension hook (#43774)
Signed-off-by: NolanHo <kujyo.eia.serias@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>
NolanHo committed
Jun 3, 2026 at 15:45 UTC
4aaed4ca225a3745aa1e18864dad0599d3ac7626
1 file changed
+17
-2
rust/src/server/src/lib.rs
+17
-2
@@ -12,7 +12,7 @@ mod utils;
12
use std::sync::{Arc, OnceLock};
13
14
use anyhow::{Context as _, Result};
15
-use axum::serve::ListenerExt as _;
15
+use axum::{Router, serve::ListenerExt as _};
16
pub use config::{Config, CoordinatorMode, HttpListenerMode};
17
use tokio::net::TcpListener;
18
use tokio::time::{Instant, sleep_until};
@@ -97,6 +97,21 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
97
/// The server owns one `vllm-chat` facade, which in turn owns the lower
98
/// `vllm-text` and `vllm-llm` layers, and shuts them down before returning.
99
pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> {
100
+ serve_with_router_extension(config, shutdown, |router| router).await
101
+}
102
+
103
+/// Run the OpenAI-compatible HTTP server with an opt-in router extension.
104
+///
105
+/// The extension receives the finalized vLLM router and can merge additional
106
+/// routes before the server starts accepting requests.
107
+pub async fn serve_with_router_extension<F>(
108
+ config: Config,
109
+ shutdown: CancellationToken,
110
+ extend_router: F,
111
+) -> Result<()>
112
+where
113
+ F: FnOnce(Router) -> Router,
114
+{
115
config.validate().context("invalid OpenAI frontend configuration")?;
116
117
// Also check shutdown during the (potentially long) startup handshake.
@@ -109,7 +124,7 @@ pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> {
124
.context("failed to bind listener for OpenAI server")?;
125
let bind_address = listener.local_addr()?;
126
let model = state.primary_model_name().to_owned();
112
- let app = build_router(state.clone());
127
+ let app = extend_router(build_router(state.clone()));
128
129
// Optionally bind the gRPC Generate server on a separate port. Bind
130
// synchronously here so bind errors (port in use, permission denied, ...)